Training a model with Primus and JAX MaxText

Training a model with Primus and JAX MaxText#

Training performance validation with the ROCm JAX MaxText training Docker image on AMD Instinct accelerators.

Overview#

MaxText for ROCm is a specialized fork of upstream MaxText, designed to enable training of large language models (LLMs) on AMD GPUs. By leveraging AMD Instinct™ MI300X and MI355X GPUs, MaxText delivers scalability, performance, and resource utilization for AI workloads. See the GitHub repository at ROCm/maxtext.

AMD provides a ready-to-use Docker image for AMD Instinct MI300X and MI355X GPUs containing essential components, including JAX, XLA, ROCm libraries, and MaxText utilities.

For the full software stack of this image (ROCm, JAX, Transformer Engine, hipBLASLt, RCCL, TensorFlow, and the rest), see Release notes → rocm/jax-training:maxtext-v26.9. The release notes are the single source of truth for image contents.

Primus source: The v26.9 image bundles Primus at its build commit db1a4160 (/workspace/Primus). Use a release/v26.9 checkout all the same—primus-cli container mounts it into the image for you. See Release notes → Primus source for v26.9.

Important notes for v26.9#

Read this section before starting a training run. It collects the settings this release requires, the architecture-specific workarounds, and the known issues. The contents change from release to release, so re-read it when you move to a new image tag.

Required settings#

Use the release/v26.9 branch. It is the Primus branch matching the rocm/jax-training:maxtext-v26.9 image. The image bundles Primus again, at its build commit db1a4160 on that branch, with MaxText release/v26.9 (8d1d43c0) in /workspace/maxtext. primus-cli container mounts your checkout into the image and runs it, so a checkout keeps you current if later commits land on the branch. The branch carries the Primus-managed MaxText environment defaults: CUB radix sort is turned off for MoE top-k (--xla_gpu_enable_cub_radix_sort=false), stream priority is honored (DEBUG_HIP_IGNORE_STREAM_PRIORITY=0) so collectives overlap compute, and dynamic-slice fusion is turned off (--xla_gpu_enable_dynamic_slice_fusion=false) (#1228). A config’s own env: block or XLA_FLAGS_APPEND can still override them.

XLA autotuning now defaults to level 3. Level 4, the previous Primus default, adds the autotuner’s correctness checks, which intermittently aborted compilation of the train step on GPU 0 with Mismatched results with host and device comparison (redzone_allocator.cc) or HSA_STATUS_ERROR_MEMORY_APERTURE_VIOLATION. It shows up most under primus-cli container, where every run compiles from scratch; a warm JAX compilation cache in a long-lived container hides it. Level 3 keeps autotuning and drops only those checks (#1316). The change was checked for finite losses on the MI325X nanoo_fp8 MoE recipes; if a plain fp8 MoE run on MI355X produces NaN losses, set XLA_GPU_AUTOTUNE_LEVEL=4, which also restores the old default.

Keep Shardy on. Shardy is the partitioning system in JAX, and the JAX in this image uses it. Primus enables it by default (shardy: true in its MaxText module config), so you only need to act if you have turned it off. See the Shardy migration guide.

Changes you may need to act on#

MaxText moves to release/v26.9. Upstream drops the retry_when_tokens_dropped config key and renames the cosmos3-*-reasoner model configs to weaver-max / weaver-mini (#1316). Nothing in Primus uses either, but remove the key from local configs, since the MaxText config schema rejects unknown keys.

A preset HIP_VISIBLE_DEVICES is kept. The runner used to overwrite it with 0..GPUS_PER_NODE-1, so a job restricted to particular GPUs could land on another user’s GPU on a shared node. It now keeps the value you or your scheduler set, and warns when it lists a different number of GPUs than GPUS_PER_NODE (#1237).

If you are moving from v26.7 or earlier:

Drop pure_nnx_decoder from local configs. MaxText release/v26.8 removed the Linen modules and the pure_nnx_decoder key with them, and its config schema rejects unknown keys, so a config that still sets it fails validation. The three MI325X nanoo_fp8 recipes that set it (DeepSeek-V2-16B, Mixtral-8x7B, Qwen3-30B-A3B) were fixed in v26.8 (#1228).

XLA_FLAGS in a config now takes precedence over the Primus-managed defaults, as it always should have; use XLA_FLAGS_APPEND to add flags without replacing them (#1054). If you relied on the old behaviour, check the effective flags in the launch log.

Architecture-specific settings#

MI355X (gfx950) — disable RCCL WarpSpeed. RCCL’s WarpSpeed feature (RCCL_WARP_SPEED_AUTO) is a gfx950-only optimization that is enabled by default in gfx950 builds, and it can cause NaN losses during training. Primus automatically sets RCCL_WARP_SPEED_AUTO=0 when a gfx950 (MI355X) device is detected, so primus-cli handles it for you. If you launch training manually on MI355X outside of Primus, export it yourself:

export RCCL_WARP_SPEED_AUTO=0

This variable is a no-op on MI300X (gfx942).

Known issues#

The v26.9 QA runs found no MaxText issues. They covered the MaxText recipes on MI300X and MI355X; the MI325X recipes were not re-run on this image.

v26.9 runs on a ROCm nightly. The image takes ROCm 10.3.0a20261026 and a Transformer Engine 2.19 development build. Re-check convergence on your own hardware before relying on a recipe for a long run.

Loss curve discrepancy with packing=false. With packing=false the loss converges at a slightly higher value than in earlier images. To reproduce the earlier convergence, set NVTE_CK_USES_FWD_V3=0, which uses Flash Attention v2 for the forward pass instead of v3. This is being tracked for a future release.

Fixed in v26.9:

  • The MI355X Mixtral-8x7B fp8 recipe ran out of memory on this image; the recipe is fixed (#1338).

  • Runs that start grain worker processes (grain_worker_count above 0; MaxText’s default is 1) trained zero steps and still exited 0, because grain read absl flags that Primus had left unparsed. Synthetic-data runs were not affected (#1229).

  • Transformer Engine printed a return_max_logit warning on every backward attention call, filling most of a training log. A Primus patch passes the value TE was already falling back to, so the warning is gone and numerics are unchanged (#1316).

Supported features and models#

MaxText supports the following key features to train large language models efficiently:

  • Transformer Engine (TE)

  • Flash Attention (FA) 3, with or without input sequence packing

  • GEMM tuning

  • Multi-node support

  • NANOO FP8 (for MI300X) or FP8 (for MI355X)

The following models are pre-optimized for performance on the AMD Instinct MI300X, MI325X, and MI355X accelerators. Each cell lists the precisions that ship as experiment configs for that device; the config path is examples/maxtext/configs/<DEVICE>/<prefix>-<precision>-pretrain.yaml.

Model

Config prefix

MI300X / MI325X

MI355X

Llama 2 7B

llama2_7B

bf16, nanoo_fp8

bf16, fp8

Llama 2 70B

llama2_70B

bf16, nanoo_fp8

bf16, fp8

Llama 3/3.1 8B

llama3_8B

bf16, nanoo_fp8

bf16, fp8

Llama 3/3.1 70B

llama3_70B

bf16

bf16, fp8

Llama 3.1 405B

llama3.1_405B

—

bf16, fp8

Llama 3.3 70B

llama3.3_70B

bf16

bf16, fp8

DeepSeek-V2-lite (16B)

deepseek_v2_16B

bf16, nanoo_fp8

bf16, fp8

Gemma4 26B

gemma4_26B

bf16, nanoo_fp8

bf16, fp8

Gemma4 31B

gemma4_31B

bf16, nanoo_fp8

bf16, fp8

Grok-1

grok1

nanoo_fp8

—

Mixtral 8x7B

mixtral_8x7B

bf16, nanoo_fp8

bf16, fp8

Mixtral 8x22B

mixtral_8x22B

bf16

bf16

Qwen3 14B

qwen3_14B

bf16, nanoo_fp8

bf16, fp8

Qwen3 30B-A3B

qwen3_30B_A3B

bf16, nanoo_fp8

bf16, fp8

Browse examples/maxtext/configs/ for the authoritative list — it tracks whatever configs ship in your checkout.

Note: Some models, such as Llama 3, require an external license agreement through a third party (for example, Meta). The only models supported in this workflow are those listed above.

System validation#

If you have already validated your system, skip this step. Otherwise, complete the system validation and optimization steps to set up your system before starting training.

Environment setup#

This Docker image is optimized for specific model configurations outlined below. Performance can vary for other training workloads, as AMD doesn’t validate configurations and run conditions outside those described.

This container should not be expected to provide generalized performance across all training workloads. Expect the container to perform in the model configurations described below; other configurations and run conditions are not validated by AMD.

Multi-node prerequisites (RDMA only)#

For multi-node runs, make sure all packages are installed based on the network device you use. You only need the setup below if you are using multi-node with RDMA—otherwise skip this part.

Install the packages below for building and installing the RDMA driver:

apt install iproute2 -y
apt install -y linux-headers-"$(uname -r)" libelf-dev
apt install -y gcc make libtool autoconf librdmacm-dev rdmacm-utils infiniband-diags ibverbs-utils perftest ethtool libibverbs-dev rdma-core strace libibmad5 libibnetdisc5 ibverbs-providers libibumad-dev libibumad3 libibverbs1 libnl-3-dev libnl-route-3-dev

Refer to your NIC manufacturer’s webpage for further steps about compiling and installing the RoCE driver. For Broadcom, see the section Compiling Broadcom NIC Software from Source in the Ethernet Networking Guide for AMD Instinct MI300X GPU Clusters.

Multi-node environment variables#

Set the following environment variables.

Master address—change localhost to the master node’s hostname:

export MASTER_ADDR="${MASTER_ADDR:-localhost}"

Number of nodes—set the number of nodes you want to train on (for example, 2, 4, 8):

export NNODES="${NNODES:-1}"

Node rank—set the rank of each node (0 for master, 1 for the first worker node, and so on):

export NODE_RANK="${NODE_RANK:-0}"

Network interface—update the network interface in the script to match your system’s network interface. To find your network interface, run this outside the container:

ip a

Then update the following variable in the script:

export NCCL_SOCKET_IFNAME=ens50f0np0

RDMA interface—first make sure the packages above are installed on all the nodes. Then set the RDMA interfaces to use for communication:

# If using Broadcom NIC
export NCCL_IB_HCA=rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7

# If using Mellanox NIC
export NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_8,mlx5_9

For Primus-specific networking guidance, see Multi-node networking.

Training and benchmarking options#

Use the following instructions to set up the environment, configure the script to train models, and reproduce the benchmark results on the MI300X, MI325X, MI350X, and MI355X accelerators with the Docker image.

All training runs through primus-cli. There are two ways to reach it:

Method

Use it when

primus-cli (recommended)

Any new work. One CLI covers direct, container, and Slurm launches, and the same YAML configurations work across every Primus backend.

Standalone benchmarking

You want to launch the container yourself and reproduce a single benchmark configuration by hand.

JAX MaxText is integrated into Primus, which supports multiple backends including Megatron-LM, TorchTitan, and JAX MaxText alongside ROCm-optimized components. The unified primus-cli runs training jobs with the JAX MaxText backend and is the path we recommend.

Standalone benchmarking#

Use the following command to pull the Docker image from Docker Hub:

docker pull rocm/jax-training:maxtext-v26.9

Single-node training#

Setup#

Note: Adjust the following variables based on your environment.

Export variables:

  • HF_TOKEN is your Hugging Face token to access models, tokenizers, and data. See User access tokens for more information.

  • HF_HOME is where huggingface_hub will store local data. Refer to the Hugging Face CLI documentation on how to download the data. If you already have downloaded or cached Hugging Face artifacts, set this variable to that path. Downloaded files typically get cached to a place like ~/.cache/huggingface.

export HF_TOKEN=<Your HuggingFace token>
export HF_HOME=<Location of saved/cached HuggingFace models>

Launch the Docker container:

docker run -it \
  --device /dev/dri --device /dev/kfd \
  --network host --ipc host --group-add video \
  --cap-add SYS_PTRACE --security-opt seccomp=unconfined --privileged \
  -v $HOME:$HOME -v $HOME/.ssh:/root/.ssh \
  -v $HF_HOME:/hf_cache -e HF_HOME=/hf_cache \
  -e HF_TOKEN=$HF_TOKEN \
  --shm-size 64G --name training_env \
  rocm/jax-training:maxtext-v26.9

Execute the training_env container (optional if you are already in the container):

docker start training_env
docker exec -it training_env bash

The image carries Primus at its build commit in /workspace/Primus, but run from your release/v26.9 checkout. The preceding docker run command mounts $HOME, so a checkout under your home directory is visible at the same path. Run training with primus-cli in direct mode; configs live under examples/maxtext/configs/<DEVICE>/ where <DEVICE> is MI300X, MI325X, or MI355X.

cd ~/Primus   # your release/v26.9 checkout

# Unquantized (bf16), e.g. Llama 2 7B on MI300X
# Note: RCCL_WARP_SPEED_AUTO=0 is auto-set by Primus on MI355X (gfx950).
./primus-cli direct -- train pretrain \
  --config examples/maxtext/configs/MI300X/llama2_7B-bf16-pretrain.yaml

For quantized training, replace -bf16- in the config name with -fp8- (MI355X) or -nanoo_fp8- (MI300X):

# nanoo_fp8 on MI300X
./primus-cli direct -- train pretrain \
  --config examples/maxtext/configs/MI300X/llama2_7B-nanoo_fp8-pretrain.yaml

# fp8 on MI355X
./primus-cli direct -- train pretrain \
  --config examples/maxtext/configs/MI355X/llama2_7B-fp8-pretrain.yaml

Benchmarking examples#

Every listed model has a bf16 variant. Quantization is device-specific: MI300X uses NANOO FP8 (-nanoo_fp8-) and MI355X uses FP8 (-fp8-). The columns below show which quantized variant is also available:

Model

MI300X (bf16 + …)

MI355X (bf16 + …)

Llama 2 7B

-nanoo_fp8-

-fp8-

Llama 2 70B

-nanoo_fp8-

-fp8-

Llama 3/3.1 8B

-nanoo_fp8-

-fp8-

Llama 3/3.1 70B

bf16 only

-fp8-

Llama 3.3 70B

bf16 only

-fp8-

DeepSeek-V2-lite

-nanoo_fp8-

-fp8-

Gemma4 26B

-nanoo_fp8-

-fp8-

Gemma4 31B

-nanoo_fp8-

-fp8-

Mixtral-8x7B

-nanoo_fp8-

-fp8-

Qwen3 14B

-nanoo_fp8-

-fp8-

Qwen3 30B-A3B

-nanoo_fp8-

-fp8-

Multi-node training#

Note: These scripts will launch the Docker container and execute the benchmark, so run them outside of any Docker container.

Multi-node training is launched through the unified primus-cli in Slurm mode. The general form for a multi-node run is:

# From your release/v26.9 Primus checkout
# RCCL_WARP_SPEED_AUTO=0 is auto-set by Primus on MI355X (gfx950).
./primus-cli --config my_maxtext_config.yaml slurm srun -N <NUM_NODES> \
  -- train pretrain --config examples/maxtext/configs/<DEVICE>/<model>-<precision>-pretrain.yaml

where <DEVICE> is MI300X, MI325X or MI355X, <model> is one of the MaxText configs (for example, llama2_7B, llama2_70B, llama3_8B, llama3_70B, gemma4_26B, gemma4_31B, mixtral_8x7B, qwen3_14B, qwen3_30B_A3B), and <precision> is bf16, fp8 (MI355X), or nanoo_fp8 (MI300X/MI325X).

Example commands#

  1. Multi-node training with the Llama 2 7B model on 2 nodes:

./primus-cli --config my_maxtext_config.yaml slurm srun -N 2 \
  -- train pretrain --config examples/maxtext/configs/MI300X/llama2_7B-bf16-pretrain.yaml

  1. Multi-node training with the Llama 2 70B model on 4 nodes:

./primus-cli --config my_maxtext_config.yaml slurm srun -N 4 \
  -- train pretrain --config examples/maxtext/configs/MI300X/llama2_70B-bf16-pretrain.yaml

  1. Multi-node training with the Llama 3 8B model on 2 nodes:

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

  1. Multi-node training with the Llama 3 70B model on 8 nodes:

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

  1. Multi-node training with the Llama 3.1 405B model on MI355X (gfx950) with 8 nodes:

./primus-cli --config my_maxtext_config.yaml slurm srun -N 8 \
  -- train pretrain --config examples/maxtext/configs/MI355X/llama3.1_405B-bf16-pretrain.yaml

Profiling with JAX XPlane Profiler#

MaxText has built-in XPlane profiling support via JAX’s profiler. Traces capture GPU kernel timelines, RCCL collectives, HLO graphs, and more. The output can be viewed in TensorBoard’s Trace Viewer or analyzed with TraceLens.

Key MaxText profiler flags#

The following MaxText config keys control profiling:

profiler=xplane                    # Use xplane format (produces .xplane.pb files)
skip_first_n_steps_for_profiler=2  # Skip compilation/warmup steps
profiler_steps=5                   # Number of steps to profile
upload_all_profiler_results=True   # Save all GPU profiles (not just GPU0)

Choosing step counts:

  • steps should be greater than skip_first_n_steps_for_profiler + profiler_steps (for example, steps=12 with skip=2 and profile=5 gives 5 warmup + 5 profiled + 2 cooldown)

  • skip_first_n_steps_for_profiler=2 skips step 0 (compilation) and step 1 (warmup)

  • profiler_steps=5 is typically enough; more steps mean larger .xplane.pb files

Enabling the profiler in an experiment config#

The Primus MaxText experiment configs (examples/maxtext/configs/<DEVICE>/<model>-<precision>-pretrain.yaml in your Primus checkout) already include a profiler key under overrides (set to "" by default). To profile a config without passing flags on every launch, edit its overrides block and set the profiler fields:

profiler: "xplane"
skip_first_n_steps_for_profiler: 2
profiler_steps: 5
upload_all_profiler_results: True
steps: 12

Then run the config as usual:

./primus-cli direct -- train pretrain \
  --config examples/maxtext/configs/MI300X/llama3_8B-bf16-pretrain.yaml

Profile output will be written under the base_output_directory specified in the YAML (see Output structure below).

Example: profile a model standalone in Docker#

#!/bin/bash
set -e

IMAGE="$1"       # Docker image, for example rocm/jax-training:maxtext-v26.9
TAG="$2"         # Short tag for output folder, for example v26.9_llama2_7b
PROFILE_DIR="/path/to/profiles/${TAG}"

mkdir -p "${PROFILE_DIR}"

docker run --rm --privileged --network=host \
  --device=/dev/dri --device=/dev/kfd --ipc=host \
  -v "${PROFILE_DIR}:/mnt/profile" \
  "${IMAGE}" bash -c '
export XLA_PYTHON_CLIENT_MEM_FRACTION=.97
export LD_LIBRARY_PATH=/usr/local/lib/:/opt/rocm/lib:$LD_LIBRARY_PATH
export XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=True --xla_gpu_enable_command_buffer= <your other XLA flags>"
export GPU_MAX_HW_QUEUES=2
# On MI355X (gfx950), disable RCCL WarpSpeed to avoid NaN losses (no-op on MI300X)
export RCCL_WARP_SPEED_AUTO=0

cd /workspace/maxtext

python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
  run_name=profile \
  base_output_directory=/mnt/profile \
  hardware=gpu \
  steps=12 \
  model_name=<your-model> \
  dataset_type=synthetic \
  enable_checkpointing=False \
  enable_goodput_recording=False \
  monitor_goodput=False \
  <your model-specific flags> \
  profiler=xplane \
  skip_first_n_steps_for_profiler=2 \
  profiler_steps=5 \
  upload_all_profiler_results=True
' 2>&1 | tee "${PROFILE_DIR}/run.log"

echo "Profile files:"
find "${PROFILE_DIR}" -name "*.xplane.pb" -o -name "*.trace.json.gz" 2>/dev/null

Output structure#

MaxText writes profiles in TensorBoard format:

<base_output_directory>/
└── profile/
    └── tensorboard/
        └── plugins/
            └── profile/
                └── <YYYY_MM_DD_HH_MM_SS>/
                    ├── <hostname>.xplane.pb      # Raw XPlane proto (GPU timelines)
                    ├── <hostname>.trace.json.gz  # Trace viewer data
                    └── *.hlo_proto.pb            # HLO graphs for each compiled module

Viewing traces in TensorBoard#

pip install tensorboard tensorboard-plugin-profile

# Point --logdir at the directory containing the tensorboard/ folder
tensorboard --logdir /path/to/profiles/<TAG>/profile --port 6006

Navigate to Profile > Trace Viewer in the TensorBoard UI.

Tips:

  • Zoom into a single training step (skip the first profiled step as it may have residual warmup)

  • Look at individual GPU streams to see compute/RCCL overlap

Keeping profile files small#

  • Use profiler_steps=5 (not more) to keep .xplane.pb under approximately 100 MB

  • Too many steps can produce files over 500 MB that TensorBoard struggles to load

  • enable_checkpointing=False avoids checkpoint I/O noise in the trace

  • dataset_type=synthetic eliminates data loading variability

Profiling with rocprofv3#

If you need to collect a trace and the JAX profiler isn’t working, you can use rocprofv3 as a temporary workaround:

rocprofv3 --hip-trace --kernel-trace --memory-copy-trace --rccl-trace --output-format pftrace -d ./v3_traces -- python3 app.py

  • Replace python3 app.py with any command line command that you want to run, such as ./primus-cli direct -- train pretrain --config examples/maxtext/configs/MI300X/llama2_7B-bf16-pretrain.yaml (run from your Primus checkout).

  • You can set the directory where you want the .json traces to be saved using -d <TRACE_DIRECTORY>.

  • The resulting traces can be opened in Perfetto.

Contents