Training a model with Primus and JAX MaxText

Training a model with Primus and JAX MaxText#

The JAX MaxText for ROCm training Docker image provides a prebuilt environment for training on AMD Instinct MI355X, MI350X, MI325X, and MI300X GPUs, with essential components such as JAX, XLA, ROCm libraries, and MaxText utilities.

The image also integrates with Primus, a high-level training framework that supports multiple backends. You can use the unified primus-cli to run training jobs using the JAX MaxText backend.

It includes the following software components:

Software component

Version

ROCm

7.14.0

JAX

0.10.0

Python

3.12.3

Transformer Engine

2.15.0.dev0+rocm7.15.0a20260707.72d01a0

hipBLASLt

1.4.1+cd957402

MaxText with on ROCm provides the following key features to train large language models efficiently:

  • Transformer Engine (TE)

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

  • GEMM tuning

  • Multi-node support

  • NANOO FP8 (for MI300X series GPUs) and FP8 (for MI355X and MI350X) quantization support

Supported models#

The following models are pre-optimized for performance on AMD Instinct GPUs. Some instructions, commands, and available training configurations in this documentation might vary by model – select one to get started.

Model
Variant
Variant
Variant
Variant

Note

Some models, such as Llama 3, require an external license agreement through a third party (for example, Meta).

System validation#

Before running AI workloads, it’s important to validate that your AMD hardware is configured correctly and performing optimally.

If you have already validated your system settings, including aspects like NUMA auto-balancing, you can skip this step. Otherwise, complete the procedures in the System validation and optimization guide to properly configure your system settings before starting training.

To test for optimal performance, consult the recommended System health benchmarks. This suite of tests will help you verify and fine-tune your system’s configuration.

Environment setup#

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

Pull the Docker image#

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

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

Multi-node configuration#

See Multi-node setup for AI workloads to configure your environment for multi-node training.

Benchmarking#

Once the setup is complete, choose between two options to reproduce the benchmark results:

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)

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 sufficient; more steps produce larger .xplane.pb files.

Profiling with MAD or madengine#

The model YAML configs under scripts/jax-maxtext/env_scripts/ include a profiler key (set to "" by default). To enable profiling when running through MAD or madengine, edit the YAML config for your model 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 benchmark as usual:

export MAD_SECRETS_HFTOKEN="your personal Hugging Face token to access gated models"
madengine run \
    --tags jax_maxtext_train_llama-3.1-8b \
    --keep-model-dir \
    --live-output \
    --timeout 28800

Use --keep-model-dir so the container’s output directory is preserved after the run. Profile output is written under the base_output_directory specified in the YAML.

Example: Profile a model standalone in Docker#

#!/bin/bash
set -e

IMAGE="$1"       # Docker image, e.g. rocm/jax-training:maxtext-v26.5
TAG="$2"         # Short tag for output folder, e.g. v26.5_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.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. Zoom into a single training step (skip the first profiled step as it may have residual warmup) and look at individual GPU streams to see compute/RCCL overlap.

To keep profile files small, use profiler_steps=5 to keep .xplane.pb files under approximately 100 MB. Too many steps can produce files over 500 MB that TensorBoard struggles to load. Use enable_checkpointing=False to avoid checkpoint I/O noise in the trace, and dataset_type=synthetic to eliminate 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 -- <command>

Replace <command> with the command you want to profile, such as ./jax-maxtext_benchmark_report.sh -m Llama-2-7B. Use -d <TRACE_DIRECTORY> to specify where the .json traces are saved. The resulting traces can be opened in Perfetto.

Known issues#

  • There is a known performance regression for Mixtral-8x7B in v26.5. This is being tracked and will be addressed in a future release.

  • There is a discrepancy in the loss curve when setting packing=False. It converges at a slightly higher value than previous Docker images. To achieve the same convergence as past Docker images, set NVTE_CK_USES_FWD_V3=0 (using FA v2 for forward instead of FA v3). This is being tracked and will be addressed in a future release.

  • On MI355X (gfx950), RCCL’s WarpSpeed feature (RCCL_WARP_SPEED_AUTO) — a gfx950-only optimization enabled by default in gfx950 builds — can cause NaN losses during training. To avoid this, set RCCL_WARP_SPEED_AUTO=0. For MAD-integrated benchmarking, this is already applied automatically in the gfx950 environment scripts under scripts/jax-maxtext/env_scripts/ (for example, gfx950_llama3_8b_env.sh). If you launch training manually on MI355X, export RCCL_WARP_SPEED_AUTO=0 yourself. This variable is a no-op on MI300X (gfx942).

  • The v26.5 Docker image ships JAX 0.10.0, which requires Shardy (the JAX partitioning system) to be enabled. Set shardy=True during the training run. You might get related errors if it’s not configured correctly. See the Shardy migration guide for more details.

Further reading#

Previous versions#

See primus-jax-maxtext-history to find documentation for previous releases of the Primus with MaxText training recipe.

Model
Contents