2026-05-19
55 min read time
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 |
7.2.1 |
JAX |
0.8.2 |
Python |
3.12 |
Transformer Engine |
2.8.0.dev0+9b312832 |
hipBLASLt |
1.3.0+bfcf25fa18 |
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
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.
Note
Some models, such as Llama 3, require an external license agreement through a third party (for example, Meta).
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.
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.
Use the following command to pull the Docker image from Docker Hub.
See Multi-node setup for AI workloads to configure your environment for multi-node training.
Once the setup is complete, choose between two options to reproduce the benchmark results:
The following run commands are tailored to Llama 2 7B. See Supported models to switch to another available model.
Download the Docker image and required packages
Pull the rocm/jax-training:maxtext-v26.3 Docker image from Docker Hub.
Run the Docker container.
Use these commands if you exit the training_env container and need to return to it.
Clone the Primus repository.
Run the training job with primus-cli
For detailed usage instructions for primus-cli, see the Primus CLI User Guide.
Use the following examples to run training with primus-cli:
Direct mode: run directly on the current host or within an existing Docker container
Container mode: run in Docker containers
Slurm mode: run distributed training on a Slurm cluster
The following run command is tailored to Llama 2 7B. 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.
Use this command to run the performance benchmark test on the Llama 2 7B model using one GPU with the bf16 data type on the host machine.
MAD launches a Docker container with the name container_ci-jax_maxtext_train_llama-2-7b. The latency and throughput reports of the model are collected in the following path: ~/MAD/perf.csv/.
The following commands are optimized for Llama 2 7B. 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.
Single node training
Set up environment variables.
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.
In the Docker container, clone the ROCm MAD repository and navigate to the benchmark scripts directory at MAD/scripts/jax-maxtext.
Run the setup scripts to install libraries and datasets needed for benchmarking.
To run the training benchmark without quantization, use the following command:
For quantized training, run the script with the appropriate option for your Instinct GPU.
For fp8 quantized training on MI355X and MI350X GPUs, use the following command:
For nanoo_fp8 quantized training on MI300X series GPUs, use the following command:
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.
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.3.
For example, to run a multi-node training benchmark on Llama 2 7B:
The following run commands are tailored to Llama 2 70B. See Supported models to switch to another available model.
Download the Docker image and required packages
Pull the rocm/jax-training:maxtext-v26.3 Docker image from Docker Hub.
Run the Docker container.
Use these commands if you exit the training_env container and need to return to it.
Clone the Primus repository.
Run the training job with primus-cli
For detailed usage instructions for primus-cli, see the Primus CLI User Guide.
Use the following examples to run training with primus-cli:
Direct mode: run directly on the current host or within an existing Docker container
Container mode: run in Docker containers
Slurm mode: run distributed training on a Slurm cluster
The following run command is tailored to Llama 2 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.
Use this command to run the performance benchmark test on the Llama 2 70B model using one GPU with the bf16 data type on the host machine.
MAD launches a Docker container with the name container_ci-jax_maxtext_train_llama-2-70b. The latency and throughput reports of the model are collected in the following path: ~/MAD/perf.csv/.
The following commands are optimized for Llama 2 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.
Single node training
Set up environment variables.
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.
In the Docker container, clone the ROCm MAD repository and navigate to the benchmark scripts directory at MAD/scripts/jax-maxtext.
Run the setup scripts to install libraries and datasets needed for benchmarking.
To run the training benchmark without quantization, use the following command:
For quantized training, run the script with the appropriate option for your Instinct GPU.
For fp8 quantized training on MI355X and MI350X GPUs, use the following command:
For nanoo_fp8 quantized training on MI300X series GPUs, use the following command:
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.
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.3.
For example, to run a multi-node training benchmark on Llama 2 70B:
The following run commands are tailored to Llama 3 8B. See Supported models to switch to another available model.
Download the Docker image and required packages
Pull the rocm/jax-training:maxtext-v26.3 Docker image from Docker Hub.
Run the Docker container.
Use these commands if you exit the training_env container and need to return to it.
Clone the Primus repository.
Run the training job with primus-cli
For detailed usage instructions for primus-cli, see the Primus CLI User Guide.
Use the following examples to run training with primus-cli:
Direct mode: run directly on the current host or within an existing Docker container
Container mode: run in Docker containers
Slurm mode: run distributed training on a Slurm cluster
The following commands are optimized for Llama 3 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.
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.
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.3.
For example, to run a multi-node training benchmark on Llama 3 8B:
The following run commands are tailored to Llama 3 70B. See Supported models to switch to another available model.
Download the Docker image and required packages
Pull the rocm/jax-training:maxtext-v26.3 Docker image from Docker Hub.
Run the Docker container.
Use these commands if you exit the training_env container and need to return to it.
Clone the Primus repository.
Run the training job with primus-cli
For detailed usage instructions for primus-cli, see the Primus CLI User Guide.
Use the following examples to run training with primus-cli:
Direct mode: run directly on the current host or within an existing Docker container
Container mode: run in Docker containers
Slurm mode: run distributed training on a Slurm cluster
The following commands are optimized for Llama 3 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.
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.
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.3.
For example, to run a multi-node training benchmark on Llama 3 70B:
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.
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.
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/.
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.
Single node training
Set up environment variables.
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.
In the Docker container, clone the ROCm MAD repository and navigate to the benchmark scripts directory at MAD/scripts/jax-maxtext.
Run the setup scripts to install libraries and datasets needed for benchmarking.
To run the training benchmark without quantization, use the following command:
For quantized training, run the script with the appropriate option for your Instinct GPU.
For fp8 quantized training on MI355X and MI350X GPUs, use the following command:
For nanoo_fp8 quantized training on MI300X series GPUs, use the following command:
Multi-node training
For multi-node training examples, choose a model from Supported models with an available multi-node training script.
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.
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.
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/.
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.
Single node training
Set up environment variables.
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.
In the Docker container, clone the ROCm MAD repository and navigate to the benchmark scripts directory at MAD/scripts/jax-maxtext.
Run the setup scripts to install libraries and datasets needed for benchmarking.
To run the training benchmark without quantization, use the following command:
For quantized training, run the script with the appropriate option for your Instinct GPU.
For fp8 quantized training on MI355X and MI350X GPUs, use the following command:
Multi-node training
For multi-node training examples, choose a model from Supported models with an available multi-node training script.
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.
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.
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.3.
For example, to run a multi-node training benchmark on Llama 3.1 405B (multi-node):
The following run commands are tailored to Llama 3.3 70B. See Supported models to switch to another available model.
Download the Docker image and required packages
Pull the rocm/jax-training:maxtext-v26.3 Docker image from Docker Hub.
Run the Docker container.
Use these commands if you exit the training_env container and need to return to it.
Clone the Primus repository.
Run the training job with primus-cli
For detailed usage instructions for primus-cli, see the Primus CLI User Guide.
Use the following examples to run training with primus-cli:
Direct mode: run directly on the current host or within an existing Docker container
Container mode: run in Docker containers
Slurm mode: run distributed training on a Slurm cluster
The following run command is tailored to Llama 3.3 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.
Use this command to run the performance benchmark test on the Llama 3.3 70B model using one GPU with the bf16 data type on the host machine.
MAD launches a Docker container with the name container_ci-jax_maxtext_train_llama-3.3-70b. The latency and throughput reports of the model are collected in the following path: ~/MAD/perf.csv/.
The following commands are optimized for Llama 3.3 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.
Single node training
Set up environment variables.
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.
In the Docker container, clone the ROCm MAD repository and navigate to the benchmark scripts directory at MAD/scripts/jax-maxtext.
Run the setup scripts to install libraries and datasets needed for benchmarking.
To run the training benchmark without quantization, use the following command:
For quantized training, run the script with the appropriate option for your Instinct GPU.
For fp8 quantized training on MI355X and MI350X GPUs, use the following command:
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.
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.3.
For example, to run a multi-node training benchmark on Llama 3.3 70B:
The following run commands are tailored to DeepSeek-V2-Lite (16B). See Supported models to switch to another available model.
Download the Docker image and required packages
Pull the rocm/jax-training:maxtext-v26.3 Docker image from Docker Hub.
Run the Docker container.
Use these commands if you exit the training_env container and need to return to it.
Clone the Primus repository.
Run the training job with primus-cli
For detailed usage instructions for primus-cli, see the Primus CLI User Guide.
Use the following examples to run training with primus-cli:
Direct mode: run directly on the current host or within an existing Docker container
Container mode: run in Docker containers
Slurm mode: run distributed training on a Slurm cluster
The following run command is tailored to DeepSeek-V2-Lite (16B). 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.
Use this command to run the performance benchmark test on the DeepSeek-V2-Lite (16B) model using one GPU with the bf16 data type on the host machine.
MAD launches a Docker container with the name container_ci-jax_maxtext_train_deepseek-v2-lite-16b. The latency and throughput reports of the model are collected in the following path: ~/MAD/perf.csv/.
The following commands are optimized for DeepSeek-V2-Lite (16B). 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.
Single node training
Set up environment variables.
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.
In the Docker container, clone the ROCm MAD repository and navigate to the benchmark scripts directory at MAD/scripts/jax-maxtext.
Run the setup scripts to install libraries and datasets needed for benchmarking.
To run the training benchmark without quantization, use the following command:
For quantized training, run the script with the appropriate option for your Instinct GPU.
For fp8 quantized training on MI355X and MI350X GPUs, use the following command:
For nanoo_fp8 quantized training on MI300X series GPUs, use the following command:
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.
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.3.
For example, to run a multi-node training benchmark on DeepSeek-V2-Lite (16B):
The following run commands are tailored to Mixtral 8x7B. See Supported models to switch to another available model.
Download the Docker image and required packages
Pull the rocm/jax-training:maxtext-v26.3 Docker image from Docker Hub.
Run the Docker container.
Use these commands if you exit the training_env container and need to return to it.
Clone the Primus repository.
Run the training job with primus-cli
For detailed usage instructions for primus-cli, see the Primus CLI User Guide.
Use the following examples to run training with primus-cli:
Direct mode: run directly on the current host or within an existing Docker container
Container mode: run in Docker containers
Slurm mode: run distributed training on a Slurm cluster
The following run command is tailored to Mixtral 8x7B. 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.
Use this command to run the performance benchmark test on the Mixtral 8x7B model using one GPU with the bf16 data type on the host machine.
MAD launches a Docker container with the name container_ci-jax_maxtext_train_mixtral-8x7b. The latency and throughput reports of the model are collected in the following path: ~/MAD/perf.csv/.
The following commands are optimized for Mixtral 8x7B. 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.
Single node training
Set up environment variables.
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.
In the Docker container, clone the ROCm MAD repository and navigate to the benchmark scripts directory at MAD/scripts/jax-maxtext.
Run the setup scripts to install libraries and datasets needed for benchmarking.
To run the training benchmark without quantization, use the following command:
For quantized training, run the script with the appropriate option for your Instinct GPU.
For fp8 quantized training on MI355X and MI350X GPUs, use the following command:
For nanoo_fp8 quantized training on MI300X series GPUs, use the following command:
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.
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.3.
For example, to run a multi-node training benchmark on Mixtral 8x7B:
The following run commands are tailored to Qwen 14B. See Supported models to switch to another available model.
Download the Docker image and required packages
Pull the rocm/jax-training:maxtext-v26.3 Docker image from Docker Hub.
Run the Docker container.
Use these commands if you exit the training_env container and need to return to it.
Clone the Primus repository.
Run the training job with primus-cli
For detailed usage instructions for primus-cli, see the Primus CLI User Guide.
Use the following examples to run training with primus-cli:
Direct mode: run directly on the current host or within an existing Docker container
Container mode: run in Docker containers
Slurm mode: run distributed training on a Slurm cluster
The following run command is tailored to Qwen 14B. 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.
Use this command to run the performance benchmark test on the Qwen 14B model using one GPU with the bf16 data type on the host machine.
MAD launches a Docker container with the name container_ci-jax_maxtext_train_qwen3-14b. The latency and throughput reports of the model are collected in the following path: ~/MAD/perf.csv/.
The following commands are optimized for Qwen 14B. 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.
Single node training
Set up environment variables.
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.
In the Docker container, clone the ROCm MAD repository and navigate to the benchmark scripts directory at MAD/scripts/jax-maxtext.
Run the setup scripts to install libraries and datasets needed for benchmarking.
To run the training benchmark without quantization, use the following command:
For quantized training, run the script with the appropriate option for your Instinct GPU.
For fp8 quantized training on MI355X and MI350X GPUs, use the following command:
For nanoo_fp8 quantized training on MI300X series GPUs, use the following command:
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.
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.3.
For example, to run a multi-node training benchmark on Qwen 14B:
The following run commands are tailored to Qwen 30B A3B. See Supported models to switch to another available model.
Download the Docker image and required packages
Pull the rocm/jax-training:maxtext-v26.3 Docker image from Docker Hub.
Run the Docker container.
Use these commands if you exit the training_env container and need to return to it.
Clone the Primus repository.
Run the training job with primus-cli
For detailed usage instructions for primus-cli, see the Primus CLI User Guide.
Use the following examples to run training with primus-cli:
Direct mode: run directly on the current host or within an existing Docker container
Container mode: run in Docker containers
Slurm mode: run distributed training on a Slurm cluster
The following run command is tailored to Qwen 30B A3B. 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.
Use this command to run the performance benchmark test on the Qwen 30B A3B model using one GPU with the bf16 data type on the host machine.
MAD launches a Docker container with the name container_ci-jax_maxtext_train_qwen3-30b-a3b. The latency and throughput reports of the model are collected in the following path: ~/MAD/perf.csv/.
The following commands are optimized for Qwen 30B A3B. 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.
Single node training
Set up environment variables.
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.
In the Docker container, clone the ROCm MAD repository and navigate to the benchmark scripts directory at MAD/scripts/jax-maxtext.
Run the setup scripts to install libraries and datasets needed for benchmarking.
To run the training benchmark without quantization, use the following command:
For quantized training, run the script with the appropriate option for your Instinct GPU.
For fp8 quantized training on MI355X and MI350X GPUs, use the following command:
For nanoo_fp8 quantized training on MI300X series GPUs, use the following command:
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.
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.3.
For example, to run a multi-node training benchmark on Qwen 30B A3B:
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.
The following MaxText config keys control profiling:
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.
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:
Then run the benchmark as usual:
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.
MaxText writes profiles in TensorBoard format:
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.
If you need to collect a trace without the JAX profiler, use rocprofv3:
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.
You might see NaNs in the losses while using real data (not synthetic data) when setting packing=True and NVTE_CK_IS_V3_ATOMIC_FP32=0. Set NVTE_CK_IS_V3_ATOMIC_FP32=1 for production training when using real data and input sequence packing (packing=True).
There is a known slight performance regression for DeepSeek-V2-lite (16B) in v26.3. This is being tracked and will be addressed in a future release.
JAX 0.9.1 Early Access known issues:
There is a known performance regression for MoE models (DeepSeek-V2-lite and Mixtral-8x7B).
The trace viewer in profiling may be missing some information in the flame graph.
Shardy is a new config in JAX 0.6.0. You might get related errors if it’s not configured correctly. To disable it, set shardy=False during the training run. See the Shardy migration guide to enable it.
To learn more about MAD and the madengine CLI, see the MAD usage guide.
To learn more about system settings and management practices to configure your system for AMD Instinct MI300X Series GPUs, see AMD Instinct MI300X system optimization.
For a list of other ready-made Docker images for AI with ROCm, see AMD Infinity Hub.
See JAX MaxText training performance testing version history to find documentation for previous releases of the ROCm/jax-training Docker image.