Training a model with Primus and JAX MaxText

Training a model with Primus and JAX MaxText#

2026-02-24

41 min read time

Applies to Linux

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:

rocm/jax-training:maxtext-v26.2

Software component	Version
ROCm	7.1.1
JAX	0.8.2
Python	3.12
Transformer Engine	2.8.0.dev0+aec00a7f
hipBLASLt	1.2.x

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

Meta Llama

DeepSeek

Mistral AI

Variant

Llama 2 7B

Llama 2 70B

Llama 3 8B

Llama 3 70B

Llama 3.1 8B

Llama 3.1 70B

Llama 3.1 405B (multi-node)

Llama 3.3 70B

DeepSeek-V2-Lite (16B)

Mixtral 8x7B

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.2

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:

MAD-integrated benchmarking

The following run command is tailored to Llama 3.1 8B. See Supported models to switch to another available model.

Clone the ROCm Model Automation and Dashboarding (ROCm/MAD) repository to a local directory and install the required packages on the host machine.
```
git clone https://github.com/ROCm/MAD
cd MAD
pip install -r requirements.txt
```

Use this command to run the performance benchmark test on the Llama 3.1 8B model using one GPU with the bf16 data type on the host machine.

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

MAD launches a Docker container with the name container_ci-jax_maxtext_train_llama-3.1-8b. The latency and throughput reports of the model are collected in the following path: ~/MAD/perf.csv/.

Standalone benchmarking

The following commands are optimized for Llama 3.1 8B. See Supported models to switch to another available model. Some instructions and resources might not be available for all models and configurations.

Download the Docker image and required scripts

Run the JAX MaxText benchmark tool independently by starting the Docker container as shown in the following snippet.

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

Single node training

Set up environment variables.
```
export MAD_SECRETS_HFTOKEN=<Your Hugging Face token>
export HF_HOME=<Location of saved/cached Hugging Face models>
```
MAD_SECRETS_HFTOKEN is your Hugging Face access token to access models, tokenizers, and data. See User access tokens.

HF_HOME is where huggingface_hub will store local data. See huggingface_hub CLI. If you already have downloaded or cached Hugging Face artifacts, set this variable to that path. Downloaded files typically get cached to ~/.cache/huggingface.

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 MAD_SECRETS_HFTOKEN=$MAD_SECRETS_HFTOKEN
    --shm-size 64G \
    --name training_env \
    rocm/jax-training:maxtext-v26.2

In the Docker container, clone the ROCm MAD repository and navigate to the benchmark scripts directory at MAD/scripts/jax-maxtext.
```
git clone https://github.com/ROCm/MAD
cd MAD/scripts/jax-maxtext
```
Run the setup scripts to install libraries and datasets needed for benchmarking.
```
./jax-maxtext_benchmark_setup.sh -m Llama-3.1-8B
```
To run the training benchmark without quantization, use the following command:
```
./jax-maxtext_benchmark_report.sh -m Llama-3.1-8B
```
For quantized training, run the script with the appropriate option for your Instinct GPU.
MI355X and MI350X
For fp8 quantized training on MI355X and MI350X GPUs, use the following command:
./jax-maxtext_benchmark_report.sh -m Llama-3.1-8B -q fp8
MI325X and MI300X
For nanoo_fp8 quantized training on MI300X series GPUs, use the following command:
./jax-maxtext_benchmark_report.sh -m Llama-3.1-8B -q nanoo_fp8

Multi-node training

For multi-node training examples, choose a model from Supported models with an available multi-node training script.

MAD-integrated benchmarking

The following run command is tailored to Llama 3.1 70B. See Supported models to switch to another available model.

Clone the ROCm Model Automation and Dashboarding (ROCm/MAD) repository to a local directory and install the required packages on the host machine.
```
git clone https://github.com/ROCm/MAD
cd MAD
pip install -r requirements.txt
```

Use this command to run the performance benchmark test on the Llama 3.1 70B model using one GPU with the bf16 data type on the host machine.

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

MAD launches a Docker container with the name container_ci-jax_maxtext_train_llama-3.1-70b. The latency and throughput reports of the model are collected in the following path: ~/MAD/perf.csv/.

Standalone benchmarking

The following commands are optimized for Llama 3.1 70B. See Supported models to switch to another available model. Some instructions and resources might not be available for all models and configurations.

Download the Docker image and required scripts

Run the JAX MaxText benchmark tool independently by starting the Docker container as shown in the following snippet.

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

Single node training

Set up environment variables.
```
export MAD_SECRETS_HFTOKEN=<Your Hugging Face token>
export HF_HOME=<Location of saved/cached Hugging Face models>
```
MAD_SECRETS_HFTOKEN is your Hugging Face access token to access models, tokenizers, and data. See User access tokens.

HF_HOME is where huggingface_hub will store local data. See huggingface_hub CLI. If you already have downloaded or cached Hugging Face artifacts, set this variable to that path. Downloaded files typically get cached to ~/.cache/huggingface.

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 MAD_SECRETS_HFTOKEN=$MAD_SECRETS_HFTOKEN
    --shm-size 64G \
    --name training_env \
    rocm/jax-training:maxtext-v26.2

In the Docker container, clone the ROCm MAD repository and navigate to the benchmark scripts directory at MAD/scripts/jax-maxtext.
```
git clone https://github.com/ROCm/MAD
cd MAD/scripts/jax-maxtext
```
Run the setup scripts to install libraries and datasets needed for benchmarking.
```
./jax-maxtext_benchmark_setup.sh -m Llama-3.1-70B
```
To run the training benchmark without quantization, use the following command:
```
./jax-maxtext_benchmark_report.sh -m Llama-3.1-70B
```
For quantized training, run the script with the appropriate option for your Instinct GPU.
MI355X and MI350X
For fp8 quantized training on MI355X and MI350X GPUs, use the following command:
./jax-maxtext_benchmark_report.sh -m Llama-3.1-70B -q fp8

Multi-node training

For multi-node training examples, choose a model from Supported models with an available multi-node training script.

Standalone benchmarking

The following commands are optimized for Llama 3.1 405B (multi-node). See Supported models to switch to another available model. Some instructions and resources might not be available for all models and configurations.

Download the Docker image and required scripts

Run the JAX MaxText benchmark tool independently by starting the Docker container as shown in the following snippet.

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

Multi-node training

The following SLURM scripts will launch the Docker container and run the benchmark. Run them outside of any Docker container. The unified multi-node benchmark script accepts a configuration file that specifies the model and training parameters.

sbatch -N <NUM_NODES> jax_maxtext_multinode_benchmark.sh <config_file.yml> [docker_image]

<NUM_NODES>: The number of nodes to use for training (for example, 2, 4, 8).
<config_file.yml>: Path to the YAML configuration file containing model and training parameters. Configuration files are available in the scripts/jax-maxtext/env_scripts/ directory for different models and GPU architectures.
[docker_image] (optional): The Docker image to use. If not specified, it defaults to rocm/jax-training:maxtext-v26.2.

For example, to run a multi-node training benchmark on Llama 3.1 405B (multi-node):

MI355X and MI350X (gfx950)

sbatch -N 4 jax_maxtext_multinode_benchmark.sh env_scripts/gfx950_llama3_405b.yml

Known issues#

You might see NaNs in the losses when setting packing=True. As a workaround, turn off input sequence packing (packing=False). This will be fixed in a future release.

Previous versions#

See JAX MaxText training performance testing version history to find documentation for previous releases of the ROCm/jax-training Docker image.

Training a model with Primus and JAX MaxText

Contents

Training a model with Primus and JAX MaxText#

Supported models#

System validation#

Environment setup#

Pull the Docker image#

Multi-node configuration#

Benchmarking#

Known issues#

Further reading#

Previous versions#