Train LLMs on AMD GPU: ROCm, MI300X/MI355X, ZAYA1 (2026)

Zyphra trained ZAYA1 entirely on AMD Instinct hardware, released it under Apache 2.0, and published enough about the stack to make the "can AMD do pretraining" conversation shift from "theoretically possible" to "someone already did it." CUDA has been the default for foundation-model training not because ROCm was categorically incapable, but because of ecosystem depth: FlashAttention 3, NCCL maturity, Megatron's kernel library, and years of battle-tested tooling. ZAYA1 proved you don't need all of that to train a production model.

This guide covers what the AMD pretraining stack actually looks like in 2026: hardware specs for MI300X and MI355X, the ROCm software stack with CUDA equivalents, how to port a training script, where the gotchas still live, cost math against NVIDIA at the same token budget, and how to launch AMD training on Spheron.

One clarification up front: ZAYA1 refers specifically to ZAYA1-8B, a mixture-of-experts model with 8.4B total parameters and 760M active. Zyphra trained it on 1,024 AMD Instinct MI300X GPUs in a cluster co-built with IBM, using AMD Pensando Pollara networking. Per-GPU throughput figures are not publicly reported, so the throughput estimates in this guide are based on community benchmarks.

TL;DR: Training LLMs on AMD GPUs with ROCm in 2026

• Proof: Zyphra trained ZAYA1-8B on 1,024 AMD Instinct MI300X GPUs and released it under Apache 2.0.
• Stack: ROCm 6.2 PyTorch images, RCCL in place of NCCL, DeepSpeed ZeRO-3 or Megatron-Core; FlashAttention 3 does not run on ROCm, so use Flash Attention 2.
• Memory: AMD MI300X carries 192GB HBM3 per GPU, so 70B BF16 weights (140GB) fit on one card; optimizer states still need sharding.
• On Spheron: MI300X rents on-demand at $3.59/hr per GPU as of 17 Sep 2026, in 1x or 2x VMs. Compare live GPU rates.
• MI355X: 288GB HBM3e, not on Spheron as of 17 Sep 2026; this guide's MI355X prices are 25 Jun 2026 market estimates.

Pricing fluctuates based on GPU availability. Spheron rates above are live as of 17 Sep 2026; other providers reflect their most recent published rates and may have changed. Check current GPU pricing → for live rates.

Why CUDA Lock-In Was Real (and What ZAYA1 Changed)

The CUDA moat in pretraining came from a stack of compounding dependencies. FlashAttention 3 is Hopper-specific and delivers 1.5-2x attention throughput over its predecessor. NCCL has years of tuning for distributed collectives. Megatron-Core's fused attention kernels and optimized data loaders are CUDA-native. The NVIDIA container registry ships ready-to-run images with everything pre-compiled.

On ROCm in 2023-2024, the equivalent pieces existed but with friction: Flash Attention 2 (ROCm/CK build) worked but needed manual configuration, RCCL lagged NCCL in throughput on some topologies, and torch.compile() on ROCm had more INDUCTOR backend misses. Teams with tight throughput requirements or custom CUDA kernels had real reasons to stay on NVIDIA.

What ZAYA1 changed is the risk perception. When a company trains a foundation model on AMD hardware, ships it under a permissive license, and it works, it removes the "unproven" label from the stack. Teams evaluating AMD for cost reasons now have a concrete proof point. For a broader comparison of what the ROCm stack looks like for inference workloads, see ROCm vs CUDA: AMD vs NVIDIA AI Stack Compared.

The 2026 AMD Pretraining Stack

The table below maps every major training stack component to its CUDA equivalent and its ROCm equivalent:

The PyTorch device string interception is the key insight: tensor.cuda() and .to("cuda") work without changes on ROCm because HIP intercepts those calls. You get a ROCm GPU without touching your model code. The friction is at the edges: custom CUDA kernels, FlashAttention 3 dependencies, and any code that calls into CUDA-specific C extensions.

For the parallelism strategy decisions (FSDP2 vs ZeRO-3 vs Megatron 3D), see the distributed LLM training guide, which covers all three in depth with multi-node torchrun configs.

MI300X vs MI355X: Which Hardware for Pretraining

The 192 GB memory on MI300X is the most practically significant number. A 70B model in BF16 weighs 140 GB. On an NVIDIA H100 with 80 GB, you cannot fit a 70B model on a single GPU without quantization or model parallelism. On MI300X, it fits on one GPU with 52 GB to spare. That eliminates tensor parallelism all-reduces for many 70B runs, which is a meaningful throughput advantage. Spheron's on-demand AMD MI300X rental comes as 1x or 2x VMs on a ROCm image, enough to port a training script and validate loss curves on real 192 GB hardware before you commit to a cluster. For the full memory-versus-cost breakdown against H100 on inference workloads specifically, see our AMD MI300X vs NVIDIA H100 guide. For what that memory headroom is worth on a full pretraining run rather than a single fine-tuning job, see our cost to train a 70B parameter LLM from scratch, which breaks down the GPU-hour math behind both the lean and overtrained cases.

When to use MI300X:

• 7B to 70B pretraining where the model fits in 192 GB
• Cost-optimized runs: on 25 Jun 2026, third-party MI300X market estimates ran roughly 55-60% below Spheron's H100 SXM5 on-demand rate that day; check the live rates in the cost table below before counting on that gap
• Teams migrating from multi-GPU NVIDIA setups and looking to reduce per-node count

When to use MI355X:

• 70B+ models in FP16 that you want to run without FSDP sharding across nodes (140 GB weights + optimizer states exceed MI300X capacity)
• Long-sequence pretraining (128K+ tokens) where the 8 TB/s bandwidth reduces attention memory pressure
• Pre-production scale runs where you want to minimize inter-node communication overhead

For a more detailed hardware spec comparison between AMD and NVIDIA silicon, the AMD MI300X vs NVIDIA H200 guide covers the hardware in depth.

Porting a Training Run from CUDA to ROCm

Container and device flags

CUDA:

docker run --gpus all \
  -it nvcr.io/nvidia/pytorch:24.01-py3

ROCm:

docker run \
  --device=/dev/kfd \
  --device=/dev/dri \
  --security-opt seccomp=unconfined \
  -it rocm/pytorch:rocm6.2_ubuntu22.04_py3.10_pytorch_2.4.0

The --device=/dev/kfd and --device=/dev/dri flags expose AMD GPU kernel driver and DRM devices instead of the NVIDIA device nodes. --security-opt seccomp=unconfined is required for ROCm GPU access.

Environment variables

# CUDA version
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
export NCCL_SOCKET_IFNAME=eth0
export NCCL_IB_DISABLE=0

# ROCm equivalent
export AMD_VISIBLE_DEVICES=0,1,2,3,4,5,6,7   # replaces CUDA_VISIBLE_DEVICES
export NCCL_SOCKET_IFNAME=eth0                 # same variable, RCCL reads it
export RCCL_DEBUG=INFO                         # ROCm equivalent of NCCL_DEBUG

RCCL reads the same NCCL_* environment variables as NCCL. Your torchrun launch script needs no changes beyond swapping CUDA_VISIBLE_DEVICES for AMD_VISIBLE_DEVICES.

Custom CUDA kernel porting

# Convert a custom CUDA kernel to HIP
hipify-perl my_cuda_kernel.cu > my_hip_kernel.hip

# Rebuild with hipcc instead of nvcc
hipcc -O3 -o my_hip_kernel.so my_hip_kernel.hip

hipify-perl handles most CUDA API calls automatically. What it cannot convert: PTX assembly (requires full rewrite as GCN/CDNA assembly), CUDA intrinsics with no HIP equivalent, and any code that calls into cuBLAS/cuDNN directly (use rocBLAS/MIOpen instead). That's the same wall HipKittens is built to remove for GEMM and attention kernels: since landing as an official AITER backend in March 2026, it matches hand-tuned CDNA assembly on those specific operators without requiring a from-scratch rewrite.

DeepSpeed ROCm install

# Standard install (no ROCm ops compilation)
pip install deepspeed

# With ROCm-compiled ops (faster, requires ROCm dev headers)
DS_BUILD_OPS=1 pip install deepspeed

Gotchas that will burn time

1. FlashAttention 3 is Hopper-specific. If your training script sets attn_implementation="flash_attention_3", it will fail on ROCm. Use attn_implementation="flash_attention_2" instead. FA2 (ROCm/CK build) is available and delivers most of the throughput benefit of FA3 for sequence lengths under 8K.

2. bitsandbytes ROCm support is partial. BnB's 4-bit quantization ops rely on CUDA kernels. On ROCm, use GPTQ or AWQ quantization instead if you need quantized pretraining or QLoRA. Full-precision and BF16 pretraining on ROCm does not need BnB.

3. torch.compile() on ROCm has more INDUCTOR backend misses. The compiler works but compilation time is longer and more ops fall back to eager mode. This is a ROCm-specific limitation as of 2026. Benchmark your throughput with and without torch.compile() before committing to a long pretraining run.

4. Custom PTX assembly requires manual rewrite. hipify-perl will not help with PTX. If your pretraining stack uses PTX-level attention kernels or memory access patterns, those need full rewrites as CDNA assembly.

5. FP8 pretraining numerics may differ. MI300X and MI355X support FP8 training via ROCm, but the FP8 number representation and rounding behavior may differ from H100 FP8. If you are porting an FP8 pretraining recipe from H100, validate loss curves carefully for the first 500-1000 steps before assuming it is equivalent.

This is not theoretical, and it bites well below the memory ceiling. Our timed H100 vs MI300X fine-tuning run hit an out-of-memory error at batch size 4 on a Llama 3.1 8B LoRA job that fits comfortably inside 192GB, because the failure was in ROCm's activation memory handling during the backward pass rather than in raw capacity. The fix was batch size 2 with gradient accumulation 8, keeping the same effective batch of 16 as the H100 leg.

For context on custom kernel development in CUDA that would need porting, see the CUDA 13 tile programming guide.

Cost Comparison: AMD vs NVIDIA at the Same Token Budget

Spheron has rented the AMD MI300X on-demand since 10 Sep 2026, so its row below is a live Spheron rate, as are the NVIDIA rows. MI355X is still not on Spheron, so its row reflects market-rate estimates from providers offering MI-series capacity, and a second MI300X row keeps the market estimate for comparison. Both market estimates date from 25 Jun 2026:

Pricing fluctuates based on GPU availability. Spheron rates above are live as of 17 Sep 2026; the MI300X and MI355X market estimates are from 25 Jun 2026 and may have changed. Check current GPU pricing → for live rates.

Cost-per-1T-tokens math:

The worked example below keeps its 25 Jun 2026 inputs frozen: Spheron's H100 SXM5 on-demand rate that day and the midpoint of that day's MI300X market estimate. Swap in the live rates from the table above to rerun it.

The formula: (1T / (throughput * 3600)) * $/hr gives cost per GPU for 1T tokens.

For H100 SXM5 at 2,750 tokens/sec and $5.01/hr (Spheron on-demand, 25 Jun 2026):

• (1,000,000,000,000 / (2750 * 3600)) * 5.01 = ~$506,000 per GPU. For an 8-GPU node training the same 1T token budget, total node cost is also ~$506,000: 8 GPUs finish ~8× faster at 8× the hourly rate, so GPU-hours stay constant.

For MI300X at 2,500 tokens/sec and $2.05/hr (midpoint market estimate, 25 Jun 2026):

• (1,000,000,000,000 / (2500 * 3600)) * 2.05 = ~$228,000 per GPU. An 8-GPU node training the same 1T token budget also costs ~$228,000 by the same logic.

On those 25 Jun 2026 inputs, the AMD advantage came from two directions: a lower market-estimate $/hr at comparable throughput, and fewer GPUs needed for large models due to 192+ GB memory. Rerun the formula with Spheron's live MI300X and H100 SXM5 rates before assuming the same gap. The memory side holds regardless of price: a 70B model that needs 2 H100 nodes for FSDP sharding may fit on 1 MI300X node, so for that workload you are pricing one MI300X node-hour against two H100 node-hours.

The caveats: these MI300X throughput estimates are based on community benchmarks, not controlled Spheron measurements. Real throughput depends on framework version, attention implementation, and your specific model architecture. The token budget context for why 1T tokens matters for pretraining is covered in detail in the AI pretraining data curation guide. AMD is not the only NVIDIA alternative teams are pricing out either; our AWS Trainium 3 vs H200 and B200 comparison covers the same tradeoff for custom silicon, including what the Neuron SDK migration actually costs against the raw hourly savings.

Launching AMD Training on Spheron

Step 1: Provision AMD GPU instances

Go to app.spheron.ai and pick the MI300X in the compute marketplace. As of 17 Sep 2026, the self-serve MI300X listing is on-demand 1x or 2x VMs on a fixed Ubuntu ROCm image in one US region (Michigan), billed per minute after a 20-minute minimum on 1x and 60 minutes on 2x. That covers porting, loss-curve validation and short single-node runs like the 2-GPU launch in Step 4. Larger MI300X configurations, including 8-GPU nodes, are available from Spheron as reserved capacity: contact the team for a quote instead of launching them from the marketplace. The MI355X is not on Spheron. An MI300X instance can only be destroyed, not stopped, so copy checkpoints off it before you tear it down.

Step 2: Pull the ROCm container

docker pull rocm/pytorch:rocm6.2_ubuntu22.04_py3.10_pytorch_2.4.0

AMD publishes versioned images on Docker Hub under rocm/pytorch. Always pin to a specific version for reproducible training environments.

Step 3: Install training dependencies

# Inside the ROCm container
pip install transformers datasets accelerate

# DeepSpeed with ROCm ops
DS_BUILD_OPS=1 pip install deepspeed

# Megatron-Core (from source for ROCm)
git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM
pip install -e .

Step 4: torchrun launch on a 2x MI300X VM

# Single node: a 2x MI300X VM on Spheron
export AMD_VISIBLE_DEVICES=0,1
export RCCL_DEBUG=INFO

torchrun \
  --standalone \
  --nproc_per_node=2 \
  pretrain.py \
  --model-name my-7b-config \
  --bf16

RCCL handles the collective communication between the two GPUs automatically. To scale the same script past one node, you need 8-GPU AMD nodes. Spheron's self-serve MI300X listing is 1x or 2x VMs (as of 17 Sep 2026), and larger configurations come through its reserved GPU option on request. There you swap --standalone for --nnodes, --node_rank, --master_addr and --master_port, set --nproc_per_node=8 and NCCL_SOCKET_IFNAME to the node interface, and the flags are identical to CUDA-based setups.

Step 5: SIGTERM checkpoint handler

import signal
import sys

checkpoint_dir = "/checkpoints"

def handle_sigterm(signum, frame):
    print("SIGTERM received, saving emergency checkpoint...")
    save_checkpoint(
        step=current_step,
        model=model,
        optimizer=optimizer,
        path=f"{checkpoint_dir}/emergency_step_{current_step}.pt"
    )
    sys.exit(1)

signal.signal(signal.SIGTERM, handle_sigterm)

Spheron listed no MI300X spot tier as of 17 Sep 2026, so on Spheron this handler matters for NVIDIA spot instances, where the preemption signal fires 30-120 seconds before reclamation; the same pattern works on spot AMD capacity from other providers. It writes an emergency checkpoint before the instance exits. On restart, the training loop polls checkpoint_dir for the latest checkpoint and resumes from there.

Step 6: Checkpoint configuration

Save incremental checkpoints (optimizer state + step counter) every 100-200 steps. Save full checkpoints every 500-1000 steps. Store checkpoints to a persistent NFS mount or S3-compatible store accessible from all nodes.

if step % 100 == 0:
    # Incremental: fast, saves optimizer state
    save_incremental_checkpoint(step, optimizer_state, lr_scheduler_state)

if step % 500 == 0:
    # Full: slower, saves everything needed for cold restart
    save_full_checkpoint(step, model, optimizer, lr_scheduler)

For complete checkpointing patterns including distributed checkpoint formats and NFS/S3 configuration for spot training, see the spot GPU training resilience guide.

Spheron rents the AMD MI300X and NVIDIA H100 SXM5 in one marketplace, so you can run the same training script on both and pick based on real cost per token, not vendor claims.

Rent MI300X → | H100 SXM5 on Spheron →

Get started on Spheron →