Running PyTorch on AMD ROCm: A Portability Playbook for Cheaper GPUs

Every ML team that priced GPUs in the last year has had the same conversation: NVIDIA capacity is expensive, lead times are long, and the cloud quotes for AMD MI300X and MI355X instances are materially cheaper per GB of HBM and often per token served. The question that follows is always the same — how much of our PyTorch stack actually has to change?

The honest 2026 answer is: less than you fear, more than the marketing says. PyTorch has had a first-class ROCm build for years, torch.compile works, FSDP works, vLLM ships ROCm wheels. But the failure modes are concentrated in a handful of predictable places — custom CUDA kernels, quantization libraries, NCCL-specific flags, and the long tail of pip install dependencies that hard-code cu124 in their wheel names.

This tutorial is a portability playbook: how to assess a codebase for vendor lock-in, how to port it, what to benchmark before you sign anything, and how to keep one codebase running on both vendors afterwards.

Step 0: decide what you are actually buying

Before any porting work, be clear which of the three benefits you want, because they demand different amounts of effort:

Goal	Typical effort	Risk
Serve an open-weights model on cheaper hardware	Days	Low — inference stacks are well covered
Train/fine-tune with FSDP on rented capacity	1–3 weeks	Medium — collectives, checkpoint portability
Run bespoke research code with custom kernels	4+ weeks	High — kernel rewrites dominate

Most commercial wins are in the first row. A 70B-class model served on MI300X's 192 GB of HBM per GPU often fits on fewer devices than on an 80 GB part, which removes a tensor-parallel hop and simplifies the whole deployment. That structural win is frequently larger than any raw FLOPS comparison.

Step 1: inventory your CUDA surface area

Run this audit before you provision anything. Grep is genuinely the right tool here.

# Hard-coded device strings and CUDA-only APIs
grep -rn "\.cuda()\|device='cuda'\|device=\"cuda\"" --include="*.py" .
grep -rn "torch.cuda\." --include="*.py" . | grep -v "is_available\|current_device\|empty_cache"

# Custom kernels and extension builds
grep -rln "load_inline\|CUDAExtension\|__global__\|cutlass\|cub::" .

# Vendor-pinned requirements
grep -rn "cu12\|nvidia-\|flash-attn\|bitsandbytes\|xformers\|apex\|transformer-engine" requirements*.txt pyproject.toml 2>/dev/null

Classify each hit:

• Benign. .cuda() and device='cuda' work unchanged on ROCm — the ROCm build of PyTorch deliberately reports torch.cuda.is_available() == True and uses the cuda device string. This surprises people. You do not need a mass rename; you need to stop assuming what the hardware is elsewhere in the code.
• Needs a shim. torch.cuda.nvtx, nvidia-smi parsing in your metrics code, nvmlDeviceGetPowerUsage, pynvml — all have ROCm equivalents (rocm-smi, amdsmi) but not the same API.
• Needs replacement. Handwritten CUDA C++, CUTLASS templates, Transformer Engine FP8 paths, bitsandbytes kernels, apex fused optimizers.

The third category is your real project plan. Everything else is an afternoon.

Step 2: write device code that is vendor-neutral

The single highest-value refactor is to stop branching on vendor and start branching on capability. A small device module pays for itself immediately:

# devices.py
import torch

def get_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")          # true on both CUDA and ROCm builds
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")

def is_rocm() -> bool:
    return torch.version.hip is not None

def vendor() -> str:
    if not torch.cuda.is_available():
        return "cpu"
    return "amd" if is_rocm() else "nvidia"

def supports_bf16() -> bool:
    return torch.cuda.is_available() and torch.cuda.is_bf16_supported()

def supports_fp8() -> bool:
    if not torch.cuda.is_available():
        return False
    name = torch.cuda.get_device_name().lower()
    if is_rocm():
        return any(a in name for a in ("mi300", "mi325", "mi355"))
    major, _ = torch.cuda.get_device_capability()
    return major >= 9  # Hopper and newer

def attention_backend() -> str:
    # Prefer the built-in SDPA dispatch; it has vendor-specific kernels behind it.
    return "sdpa"

torch.version.hip is the canonical ROCm check — it is None on CUDA builds and a version string like 6.3.x on ROCm builds. Use it for capability gating, never for device selection.

The second refactor: use torch.nn.functional.scaled_dot_product_attention instead of importing a vendor-specific attention package. SDPA dispatches to efficient backends on both vendors, composes with torch.compile, and removes your single most common porting blocker. If you genuinely need the FlashAttention API surface, AMD ships a ROCm-compatible fork, but SDPA is the portable default and in 2026 it is fast enough that most teams never look further.

Step 3: the dependency minefield

This is where ports actually stall. A practical substitution table:

CUDA-only package	ROCm-portable approach
flash-attn	F.scaled_dot_product_attention, or AMD's ROCm FlashAttention build
bitsandbytes (4/8-bit)	torchao quantization, or GPTQ/AWQ kernels with ROCm support
apex fused optimizers	torch.optim.AdamW(fused=True) — upstream and portable
xformers	SDPA; memory-efficient attention is upstream now
Transformer Engine FP8	torchao.float8 on supported parts, else BF16
nvidia-dali	torchdata / native DataLoader with more workers
NCCL env tuning	RCCL — same variable names mostly, different tuning values
pynvml	amdsmi (or shell out to rocm-smi --showuse)

torchao is the important line in that table. Moving quantization off bitsandbytes and onto torchao is worth doing even if you stay on NVIDIA, because it removes a binary dependency that historically breaks on every PyTorch minor release. Vendor-portability is a side benefit of an upgrade you wanted anyway — this is the general pattern of a good ROCm port.

Step 4: container and environment setup

Do not fight local installs. Start from AMD's published PyTorch container, which pins a matched ROCm/PyTorch pair:

docker run -it --rm \
  --device=/dev/kfd --device=/dev/dri \
  --group-add video --ipc=host \
  --security-opt seccomp=unconfined \
  --shm-size 16G \
  -v "$PWD":/workspace -w /workspace \
  rocm/pytorch:latest

--device=/dev/kfd and /dev/dri are the ROCm equivalent of the NVIDIA container runtime; --shm-size matters as much as it does on NVIDIA, and forgetting it produces the same baffling DataLoader hangs.

Two environment variables you will meet immediately:

• HSA_OVERRIDE_GFX_VERSION — forces a gfx architecture string for parts that are supported-but-unlisted (common on workstation cards; you should not need it on MI300-class datacenter parts).
• PYTORCH_ROCM_ARCH — the build-time equivalent of TORCH_CUDA_ARCH_LIST, needed if you compile extensions.

Smoke-test in one line before you trust anything else:

import torch
print(torch.__version__, torch.version.hip, torch.cuda.device_count())
print(torch.cuda.get_device_name(0))
x = torch.randn(8192, 8192, device="cuda", dtype=torch.bfloat16)
print((x @ x).float().norm().item())

Step 5: porting custom kernels

If your audit found handwritten CUDA, you have three routes, in increasing order of effort and decreasing order of regret:

1. Delete it. A surprising fraction of hand-rolled fusions from 2022 are now slower than torch.compile on the eager PyTorch equivalent. Benchmark the naive version first — we have removed several hundred lines of CUDA on client projects and gained throughput.
2. Rewrite in Triton. Triton has an AMD backend, so one Triton kernel serves both vendors, and it composes with torch.compile instead of fighting it. This is the right target for anything genuinely custom.
3. HIPify the C++. hipify-torch / hipify-perl mechanically translate CUDA C++ to HIP, and CUDAExtension in torch.utils.cpp_extension will build HIP sources on a ROCm install. Works well for straightforward kernels; warp-level intrinsics need attention because AMD CDNA parts have a wavefront of 64, not 32. Any kernel with a hard-coded 32 in a shuffle or a warp-reduction is a bug waiting to happen.

That wavefront difference is the single most common source of "it runs but the numbers are wrong" on a port. Search for __shfl_, warpSize, and literal 32s in reductions before you trust a HIPified kernel.

Step 6: distributed training

FSDP2 and DDP both work on ROCm via RCCL, and torch.distributed code needs no changes — the backend name stays "nccl". What changes is tuning:

import torch.distributed as dist
dist.init_process_group(backend="nccl")  # maps to RCCL on ROCm builds

Checklist for multi-node runs:

• Validate interconnect bandwidth before the first training job: rccl-tests' all_reduce_perf is the counterpart to nccl-tests, and a busbw number an order of magnitude below spec means a fabric misconfiguration, not a PyTorch problem.
• NCCL_DEBUG=INFO still works and RCCL honours it; the topology dump is your first stop on a hang.
• Checkpoints are portable. torch.distributed.checkpoint files written on one vendor load on the other — this is what makes a hybrid fleet practical, and it is worth testing explicitly on day one.
• Expect to re-tune batch size and activation-checkpointing policy. With 192 GB per device you can often drop activation checkpointing entirely, which is a real speedup that never shows up in a spec-sheet comparison.

Step 7: benchmark like a buyer, not a fan

Never port on faith. Run a benchmark that mirrors your production shape, and quote results in cost, not FLOPS:

• For serving: tokens/sec at your real input/output length distribution, at your real concurrency, plus p95 TTFT and p95 inter-token latency. Then divide by instance price to get cost per million tokens. That is the only number that settles the argument.
• For training: time per optimizer step at fixed global batch size, and model FLOPs utilisation (MFU). Then cost per epoch.
• Always: an accuracy or eval-suite check. Numerics differ slightly between vendors; BF16 accumulation order is not identical. Your eval harness should confirm parity, not your intuition.

A minimal serving A/B:

vllm serve meta-llama/Llama-3.3-70B-Instruct --max-model-len 8192 &
python -m vllm.entrypoints.cli.benchmark.main serve \
  --model meta-llama/Llama-3.3-70B-Instruct \
  --dataset-name sharegpt --num-prompts 500 --request-rate 8

Run the identical command on both fleets, record throughput, p95 latencies and instance hourly rate, and put the three numbers in a table. Decisions made this way survive contact with a CFO.

What still hurts

Being straight about the rough edges, as of 2026:

• Bleeding-edge research code. Repos published the week of a conference often assume CUDA. Give them a few weeks, or budget for the port.
• FP8 and exotic quantization. Improving fast, but support matrices vary by part and by library; verify on your exact SKU rather than trusting a blog post (including this one).
• Profiling ergonomics. The PyTorch profiler works, and rocprofiler is capable, but the tooling ecosystem around Nsight is deeper. Budget extra time for your first deep performance investigation.
• Long-tail libraries. Graph learning, 3D vision and simulation packages with custom kernels are the usual stragglers.

None of these block the mainstream case — serving and fine-tuning open-weights transformers — which is exactly where the money is.

A sane adoption path

1. Audit CUDA surface area (an afternoon) and get a real effort estimate.
2. Refactor to capability-based device code and SDPA — do this regardless of vendor plans; it is good hygiene.
3. Stand up one inference workload on rented ROCm capacity and benchmark cost per million tokens honestly.
4. Keep CI running on both vendors once anything ships, or the port silently rots within two releases.
5. Only then consider moving training, and only with rccl-tests numbers in hand.

The strategic prize is not that one vendor is better. It is optionality: a codebase that runs on whatever silicon is available and cheap this quarter is worth more than one that is 10% faster on hardware you cannot get a quota for.

If you are weighing a ROCm migration, need a vendor-neutral refactor of a PyTorch codebase, or want an independent benchmark of your own workload across fleets before you commit to a contract, get in touch — a scoped one-week assessment usually tells you the cost per million tokens and the real porting effort, with no procurement agenda attached.