AMD, ROCM, PyTorch, and AI on Ubuntu: The Rules of the Jungle

Are you the awkward lanky kid trying to join the other cool athletes on the field? Or perhaps you went out to buy an NVIDIA video card, only to realize the 2GB VRAM won’t get you very far? And better yet “upgraded” to an AMD? And then when you got it home and then started hacking away on spacy and Stable Diffusion you realize that it’s a nightmare?

Yeah, that’s me. As I’m writing this there’s still a side project I’m trying to get working with Pytorch + Cupy. But let me untangle this mess. In this guide I’ll walk through what the stack looks like and the various components, in hopes that it can shed light on what you might be up against.

AMD ≠ Nvidia

Duh.

Men and women are different too.

But what I mean is that the underlying drivers are not the same.

In short, Nvidia uses uses CUDA, and AMD uses ROCM.

The current tech industry relies heavily on CUDA. Nvidia isn’t sharing their tech with AMD, so AMD is essentially creating a software layer that says 😉…