End-to-end training recipes

Task-oriented, copy-paste commands for launching pretraining runs with each Primus backend on AMD Instinct™ GPUs.

This page covers what is common to every backend — image, architecture folders, environment, shared setup — and gives one worked example each. For the complete per-model command set, follow the backend recipe page in the table below. For the concepts behind the workflow (how backends work, YAML structure and inheritance, parallelism vocabulary, the full configuration inventory), see Pretraining.

Choose your recipe

Backend	Image family	Configurations	Full recipe
Megatron-LM	rocm/primus	examples/megatron/configs/<ARCH>/	Megatron-LM training
TorchTitan (PyTorch)	rocm/primus	examples/torchtitan/configs/<ARCH>/	TorchTitan training
JAX MaxText	rocm/jax-training:maxtext-…	examples/maxtext/configs/<ARCH>/	JAX MaxText training
Megatron Bridge (post-training)	rocm/primus	examples/megatron_bridge/configs/<ARCH>/	Post-training

Each backend recipe page opens with an Important notes section listing the settings that release requires, the architecture-specific tuning, and any known issues. Read it before your first run on a new image tag.

Image contents. The exact ROCm, PyTorch/JAX, Transformer Engine, and RCCL versions in every published tag are in Release notes, which is the single source of truth for image contents.

---

How recipes are structured

Every recipe follows the same four-step pattern:

1. Clone the Primus branch matching your image, on the host.

2. Set the GPU-architecture environment (performance environment variable settings differ by GPU).

3. Pick the configuration YAML for your GPU architecture under examples/<backend>/configs/<ARCH>/ in the Primus repository.

4. Launch with runner/primus-cli — in container mode from the host, which starts the container for you.

Choosing a launch mode

Mode	Run it from	What it does
container (recommended)	The host, in your Primus checkout	Starts the container, mounts your checkout into it at the same path, and runs the training inside. Nothing to set up by hand.
direct	Inside a container, or a bare-metal install	Runs training in the current environment. Use it when you already have a shell inside the container, or when Primus is installed directly on the host.
slurm	The host, in your Primus checkout	Allocates nodes and runs container mode on each of them.

Container mode mounts your checkout at the same absolute path inside the container and runs from there, so the branch you cloned is the code that executes — the /workspace/Primus copy baked into the image is not used. It also forwards a list of environment variables from the host (including HF_TOKEN, the gfx942 tuning variables, and the NCCL_* networking variables), so exports you make on the host take effect inside. The forwarded list is container.options.env in runner/.primus.yaml.

Only the Primus tree is mounted automatically. Mount datasets, checkpoints, and output directories with --volume /host/path (or --volume /host/path:/container/path).

GPU-architecture config folders

Configuration YAMLs are organized by GPU architecture. Always pick the folder that matches your hardware:

Backend	MI300X	MI325X	MI355X / MI350X
examples/megatron/configs/	yes	yes	yes
examples/torchtitan/configs/	yes	yes	yes
examples/maxtext/configs/	yes	—	yes
examples/megatron_bridge/configs/	yes	—	yes

MI350X uses the same configurations as MI355X because both are based on the gfx950 architecture. If a configuration for your model is not available in the architecture-specific folder, use the closest match from the same generation as a starting point.

GPU-architecture environment variables

Megatron-LM and TorchTitan on MI300X/MI325X (gfx942) benefit from the following settings. They are not needed on MI355X/MI350X (gfx950):

# MI300X / MI325X only -- improves performance
export HSA_NO_SCRATCH_RECLAIM=1
export PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32=1
export NVTE_CK_IS_V3_ATOMIC_FP32=1

JAX MaxText does not use these. Its backend adapter applies the correct architecture environment automatically, and the one variable you may need to export yourself is RCCL_WARP_SPEED_AUTO=0 on MI355X. See JAX MaxText → Architecture-specific settings.

Choosing the Docker image

For container and Slurm modes (direct mode runs in whatever environment you launched it from), the default image is rocm/primus:v26.5, set in runner/.primus.yaml. JAX MaxText has its own separate image family, rocm/jax-training:maxtext-…, which is not the default — pass it explicitly in container and Slurm modes.

The image is picked in priority order: DOCKER_IMAGE environment variable > --image CLI argument > config file. See Selecting the container image for a full explanation, and Configuration system for configuration loading.

---

Shared setup for all backends

These apply across all backends. Set them up before running the recipes below.

Get the Primus source

Clone the branch matching your image, on the host. Every command on this page runs from this directory:

git clone --recurse-submodules https://github.com/AMD-AGI/Primus.git
cd Primus
git checkout release/v26.5
git submodule update --init --recursive

Container mode mounts this checkout into the container, so this is the code that runs — you do not need the /workspace/Primus copy baked into the image, which lags the release branch. See Release notes → Primus source for v26.5.

Pull the image (optional)

primus-cli container pulls the image on first use, so this is only needed if you want to warm the cache:

docker pull rocm/primus:v26.5                  # Megatron-LM, TorchTitan, Megatron Bridge
docker pull rocm/jax-training:maxtext-v26.5    # JAX MaxText

Starting a container by hand instead

If you want an interactive shell in the container — for debugging, or to run primus-cli direct yourself — start one manually and bind your Primus checkout:

docker run -it \
    --device /dev/dri --device /dev/kfd --device /dev/infiniband \
    --network host --ipc host \
    --group-add video --cap-add SYS_PTRACE \
    --security-opt seccomp=unconfined --privileged \
    -v $PWD:$PWD -w $PWD --shm-size 128G \
    --name primus_training_env \
    rocm/primus:v26.5

Re-enter it later with docker start primus_training_env && docker exec -it primus_training_env bash. Inside, use primus-cli direct. Remember to re-export HF_TOKEN and any architecture or NCCL_* variables, since a manual docker run does not forward them.

Bind only the directories you need rather than your whole home directory.

Hugging Face token (for gated models or real data)

export HF_TOKEN=<your_hf_token>

runner/.primus.yaml forwards HF_TOKEN into the container automatically. MaxText configurations might also read ${HF_TOKEN:""} directly.

Mock vs. real data

• Mock/synthetic data (default for most examples): validates the stack without datasets. Megatron and TorchTitan set mock_data: true; MaxText sets dataset_type: "synthetic".

• Real data: set mock_data: false and point train_data_path (for Megatron) or the backend’s dataset fields at paths visible inside your container mounts.

Multi-node networking checklist

The primus-cli launcher sets sensible NCCL_* defaults, but auto-detection can pick the wrong device on multi-NIC nodes. Before multi-node, confirm and export if needed:

export NCCL_IB_HCA=<your_rdma_interfaces>      # from `ibv_devices`
export NCCL_SOCKET_IFNAME=<your_net_interface> # from `ip a`
export GLOO_SOCKET_IFNAME=<same_as_NCCL_SOCKET_IFNAME>
export NCCL_IB_GID_INDEX=3                     # 3 for RoCE (1 for AMD AINIC)

For AMD AINIC clusters also set USING_AINIC=1, NCCL_PXN_DISABLE=0, NCCL_IB_GID_INDEX=1. See Multi-node networking for the full reference.

---

Megatron-LM

Configurations: examples/megatron/configs/<ARCH>/ | Precisions: BF16, FP8 (all architectures); MXFP8 and MXFP4 for Llama 3.1 8B on MI355X

➜ Full recipe with every model and precision: Megatron-LM training

Pretrain Llama 3.1 8B BF16 on MI355X / MI350X, from your Primus checkout on the host:

./runner/primus-cli container \
  --log_file /tmp/primus_llama3.1_8B.log \
  -- train pretrain \
  --config examples/megatron/configs/MI355X/llama3.1_8B-BF16-pretrain.yaml

The same model on MI300X / MI325X — export the gfx942 performance variables first, and container mode forwards them into the container:

export HSA_NO_SCRATCH_RECLAIM=1
export PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32=1
export NVTE_CK_IS_V3_ATOMIC_FP32=1

./runner/primus-cli container \
  --log_file /tmp/primus_llama3.1_8B.log \
  -- train pretrain \
  --config examples/megatron/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml

If you already have a shell inside the container, swap container for direct in any of these commands.

Switch model or precision by changing the config filename (for example llama3.1_70B-FP8-pretrain.yaml, mixtral_8x7B_v0.1-BF16-pretrain.yaml). See the parallelism table in Pretraining.

Multi-node with Slurm:

./runner/primus-cli slurm srun -N 8 -p <partition> -- train pretrain \
  --config examples/megatron/configs/MI300X/llama3.1_8B-FP8-pretrain.yaml \
  --micro_batch_size 4 --global_batch_size 1024

Scale batch size with node count and align tensor_model_parallel_size, pipeline_model_parallel_size, and expert_model_parallel_size to your topology. The Megatron-LM recipe lists per-model, per-node-count batch sizes.

Model-specific notes:

• Zebra-Llama (hybrid Mamba+MLA) pretrain presets ship at examples/megatron/configs/<ARCH>/zebra_llama_1B-pretrain.yaml (and _3B, _8B), and run via the legacy runtime — prefix the command with PRIMUS_TRAIN_RUNTIME=legacy. Megatron Bridge SFT variants live under examples/megatron_bridge/configs/<ARCH>/.

• MoE models (DeepSeek-V2-Lite, Mixtral, Qwen3-A3B, GPT-OSS) may need extra grouped-GEMM or router flags; the Megatron-LM recipe gives the exact command per model.

---

TorchTitan (PyTorch)

Configurations: examples/torchtitan/configs/<ARCH>/ | Precisions: BF16, FP8

➜ Full recipe with every model and precision: TorchTitan training

Uses the same rocm/primus container as Megatron-LM. TorchTitan parameters use a dotted namespace (for example --training.local_batch_size).

Pretrain Llama 3.1 8B BF16 on MI355X / MI350X, from your Primus checkout on the host:

./runner/primus-cli container \
  --log_file /tmp/primus_llama3.1_8B.log \
  -- train pretrain \
  --config examples/torchtitan/configs/MI355X/llama3.1_8B-BF16-pretrain.yaml

On MI300X / MI325X, export the gfx942 performance variables first and use the MI300X config path.

Multi-node with Slurm:

./runner/primus-cli slurm srun -N 4 -- train pretrain \
  --config examples/torchtitan/configs/MI355X/llama3.1_70B-FP8-pretrain.yaml \
  --training.local_batch_size 6 \
  --training.global_batch_size 192 \
  --training.mock_data True

Available models include Llama 3.1 (8B/70B/405B), Llama 4 (17Bx16E/17Bx128E), DeepSeek V3 (16B/236B/671B), Qwen 3 (0.6B–32B), and GPT-OSS (20B/120B). See examples/torchtitan/configs/<ARCH>/.

---

JAX MaxText

Configurations: examples/maxtext/configs/<ARCH>/ | Precisions: BF16

➜ Full recipe with every model and precision: JAX MaxText training

MaxText uses a different Docker image than the PyTorch backends and it is not the default in runner/.primus.yaml, so pass it explicitly with --image in container and Slurm modes.

On MI355X, export RCCL_WARP_SPEED_AUTO=0 before launching or training can produce NaN losses. It is a no-op on MI300X. See Important notes.

Pretrain Llama 3 8B on MI355X, from your Primus checkout on the host:

export RCCL_WARP_SPEED_AUTO=0
./runner/primus-cli container --image rocm/jax-training:maxtext-v26.5 \
  -- train pretrain \
  --config examples/maxtext/configs/MI355X/llama3_8B-pretrain.yaml

If you already have a shell inside the MaxText container, use direct instead — no --image needed:

export RCCL_WARP_SPEED_AUTO=0
./runner/primus-cli direct \
  -- train pretrain \
  --config examples/maxtext/configs/MI355X/llama3_8B-pretrain.yaml

Slurm mode — supply the image (and any environment variables) via a config file:

./runner/primus-cli --config my_maxtext_config.yaml slurm srun -N 8 \
  -- train pretrain \
  --config examples/maxtext/configs/MI300X/llama3_8B-pretrain.yaml

MaxText parallelism is set with ici_* (intra-node) and dcn_* (inter-node) fields — see the MaxText config table and MaxText parameters.

Quantized MaxText runs. The examples/maxtext/configs/ YAMLs are BF16 only, so there is no FP8 config to select by path. The image does support FP8 (gfx950) and NANOO FP8 (gfx942) — reach them through the -q fp8 / -q nanoo_fp8 flags of the standalone benchmark scripts, described in JAX MaxText → Standalone benchmarking.

---

Megatron Bridge (post-training)

Megatron Bridge configurations are under examples/megatron_bridge/configs/<ARCH>/ and are primarily SFT and LoRA post-training recipes (for example qwen3_32b_sft_posttrain.yaml, llama31_70b_lora_posttrain.yaml). Launch with train posttrain:

./runner/primus-cli container \
  --log_file /tmp/primus_qwen3_32b_sft.log \
  -- train posttrain \
  --config examples/megatron_bridge/configs/MI355X/qwen3_32b_sft_posttrain.yaml

See Post-training for the full SFT/LoRA workflow.

---

Related documentation

• Megatron-LM training, TorchTitan training, JAX MaxText training: full per-model recipes.

• Release notes: what is inside each published image tag.

• Pretraining: backend concepts, configuration walkthroughs, parallelism vocabulary, full configuration inventories.

• Post-training: SFT and LoRA via Megatron Bridge.

• CLI reference: direct / container / slurm modes and flags.

• Configuration system: YAML inheritance, overrides, image/env precedence.

• Performance tuning: HipBLASLt autotuning, Primus-Turbo, FP8, MoE.