Training a model with Primus and PyTorch TorchTitan

Primus (AMD-AGI/Primus) is a unified training framework designed to enable efficient training of large-scale foundation models on AMD GPUs. It supports multiple backends including TorchTitan with full support for container-based and bare-metal execution.

The following sections explain how to run Primus training through containerized ROCm/MAD workflows and the ROCm/madengine interface for automated benchmarking.

Training flow

The following diagram illustrates the end-to-end flow of how a model is trained through Primus and madengine, from user CLI commands through the internal container invocation chain.

        flowchart TD
    A["<b>Clone MAD repo</b><br/><i>git clone https://github.com/ROCm/MAD</i>"] --> B["<b>Install madengine</b><br/><i>pip install madengine</i>"]
    A --> C["<b>Initialize Primus submodule</b><br/><i>git submodule update --init --recursive</i>"]
    A --> D["<b>Docker login</b><br/><i>Registry access</i>"]
    B --> F["<b>Build Docker Image</b><br/><i>madengine build --tags</i>"]
    B --> E["<b>Discover configs</b><br/><i>madengine discover --tags</i>"]
    C --> F
    D --> F
    F --> G["<b>Run training</b><br/><i>madengine run --tags</i>"]
    

Environment setup

1. Clone the MAD repository:

   git clone https://github.com/ROCm/MAD
   cd MAD
   

2. Install madengine:

   # Recommended: create and activate a Python virtual environment
   python3 -m venv .venv
   source .venv/bin/activate
   
   # Install madengine
   pip install -r requirements.txt
   

3. Initialize the Primus submodule:

   git submodule update --init --recursive scripts/Primus
   

4. Log in to Docker for image registry access:

   docker login
   

Supported models

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

Model

Meta Llama

DeepSeek

Variant

Llama 3.1 8BLlama 3.1 70B

Variant

DeepSeek V3 16B

Tip

To discover all supported model configurations, try madengine’s model discovery feature:

# List all Primus model configs
madengine discover --tags primus

# List all MI300X model configs
madengine discover --tags MI300X

# List all TorchTitan configs
madengine discover --tags torchtitan

Or, browse the AMD-AGI/Primus repository for available configs by AMD device architecture.

Single node training

Training with madengine involves a two-step process: build the Docker image, then run the model. Complete the preceding environment setup steps and navigate to the MAD repository root before getting started.

Note

Primus training tags used in the next sections follow the naming convention primus_train/<backend>_<GPU_ARCH>_<MODEL_CONFIG>, where:

• <backend> is the Primus backend: megatron or torchtitan

• <GPU_ARCH> is the target accelerator (for example, MI300X, MI355X)

• <MODEL_CONFIG> matches the YAML filename in the Primus repository under examples/<backend>/configs/. See AMD-AGI/Primus for available configs.

Build the Docker image

Build the Docker image with the desired model configuration. For instance:

MI300X

Use the following command to build the Docker image for training Llama 3.1 8B with the BF16 precision configuration file using Primus torchtitan:

madengine build \
    --tags primus_train/torchtitan_MI300X_llama3.1_8B-BF16-pretrain \
    --additional-context '{"gpu_vendor": "AMD", "guest_os": "UBUNTU"}'

To build for Llama 3.1 8B with FP8 precision, use the following command:

madengine build \
    --tags primus_train/torchtitan_MI300X_llama3.1_8B-FP8-pretrain \
    --additional-context '{"gpu_vendor": "AMD", "guest_os": "UBUNTU"}'

MI355X

Use the following command to build the Docker image for training Llama 3.1 8B with the BF16 precision configuration file using Primus torchtitan:

madengine build \
    --tags primus_train/torchtitan_MI355X_llama3.1_8B-BF16-pretrain \
    --additional-context '{"gpu_vendor": "AMD", "guest_os": "UBUNTU"}'

To build for Llama 3.1 8B with FP8 precision, use the following command:

madengine build \
    --tags primus_train/torchtitan_MI355X_llama3.1_8B-FP8-pretrain \
    --additional-context '{"gpu_vendor": "AMD", "guest_os": "UBUNTU"}'

MI300X

Use the following command to build the Docker image for training Llama 3.1 70B with the BF16 precision configuration file using Primus torchtitan:

madengine build \
    --tags primus_train/torchtitan_MI300X_llama3.1_70B-BF16-pretrain \
    --additional-context '{"gpu_vendor": "AMD", "guest_os": "UBUNTU"}'

To build for Llama 3.1 70B with FP8 precision, use the following command:

madengine build \
    --tags primus_train/torchtitan_MI300X_llama3.1_70B-FP8-pretrain \
    --additional-context '{"gpu_vendor": "AMD", "guest_os": "UBUNTU"}'

MI355X

Use the following command to build the Docker image for training Llama 3.1 70B with the BF16 precision configuration file using Primus torchtitan:

madengine build \
    --tags primus_train/torchtitan_MI355X_llama3.1_70B-BF16-pretrain \
    --additional-context '{"gpu_vendor": "AMD", "guest_os": "UBUNTU"}'

To build for Llama 3.1 70B with FP8 precision, use the following command:

madengine build \
    --tags primus_train/torchtitan_MI355X_llama3.1_70B-FP8-pretrain \
    --additional-context '{"gpu_vendor": "AMD", "guest_os": "UBUNTU"}'

MI300X

Use the following command to build the Docker image for training DeepSeek V3 16B using Primus torchtitan:

madengine build \
    --tags primus_train/torchtitan_MI300X_deepseek_v3_16b-pretrain \
    --additional-context '{"gpu_vendor": "AMD", "guest_os": "UBUNTU"}'

MI355X

Use the following command to build the Docker image for training DeepSeek V3 16B using Primus torchtitan:

madengine build \
    --tags primus_train/torchtitan_MI355X_deepseek_v3_16b-pretrain \
    --additional-context '{"gpu_vendor": "AMD", "guest_os": "UBUNTU"}'

The base Docker image is defined in docker/primus.ubuntu.amd.Dockerfile. To use a different base image (for example, a newer Primus release), edit the BASE_DOCKER argument at the top of that file:

ARG BASE_DOCKER=docker.io/rocm/primus:v26.5

Note

MAD_SYSTEM_GPU_ARCHITECTURE is automatically detected at runtime via rocminfo. You do not need to provide it during the build step.

Run the model

Run the model with the built image:

MI300X

Use the following command to train Llama 3.1 8B with the BF16 precision configuration file using Primus torchtitan:

madengine run --tags primus_train/torchtitan_MI300X_llama3.1_8B-BF16-pretrain --live-output

To train Llama 3.1 8B with FP8 precision, use the following command:

madengine run --tags primus_train/torchtitan_MI300X_llama3.1_8B-FP8-pretrain --live-output

MI355X

Use the following command to train Llama 3.1 8B with the BF16 precision configuration file using Primus torchtitan:

madengine run --tags primus_train/torchtitan_MI355X_llama3.1_8B-BF16-pretrain --live-output

To train Llama 3.1 8B with FP8 precision, use the following command:

madengine run --tags primus_train/torchtitan_MI355X_llama3.1_8B-FP8-pretrain --live-output

MI300X

Use the following command to train Llama 3.1 70B with the BF16 precision configuration file using Primus torchtitan:

madengine run --tags primus_train/torchtitan_MI300X_llama3.1_70B-BF16-pretrain --live-output

To train Llama 3.1 70B with FP8 precision, use the following command:

madengine run --tags primus_train/torchtitan_MI300X_llama3.1_70B-FP8-pretrain --live-output

MI355X

Use the following command to train Llama 3.1 70B with the BF16 precision configuration file using Primus torchtitan:

madengine run --tags primus_train/torchtitan_MI355X_llama3.1_70B-BF16-pretrain --live-output

To train Llama 3.1 70B with FP8 precision, use the following command:

madengine run --tags primus_train/torchtitan_MI355X_llama3.1_70B-FP8-pretrain --live-output

MI300X

Use the following command to train DeepSeek V3 16B using Primus torchtitan:

madengine run --tags primus_train/torchtitan_MI300X_deepseek_v3_16b-pretrain --live-output

MI355X

Use the following command to train DeepSeek V3 16B using Primus torchtitan:

madengine run --tags primus_train/torchtitan_MI355X_deepseek_v3_16b-pretrain --live-output

Note

--live-output is optional. It streams the training logs to your terminal in real time.

Passing environment variables to the container

To pass environment variables into the running container, use the docker_env_vars field in the --additional-context mapping:

madengine run \
    --tags <tag> \
    --live-output \
    --additional-context '{"docker_env_vars": {"MAD_SECRET_HFTOKEN": "<your_hf_token>", "HSA_NO_SCRATCH_RECLAIM": "1"}}'

Note

The MAD_SECRET_HFTOKEN environment variable is only required when training with real data (that is, mock_data: false in the config). The default configs use mock data and do not require a token. Inside the container, this is automatically mapped to HF_TOKEN.

Multi-node training

Multi-node training via madengine is not yet available. Multi-node support is planned for a future release.

Further reading

• To learn more about the Primus framework, see the AMD Primus documentation.

• 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 Customer Acceptance Guide.

Previous versions

See Primus TorchTitan training performance testing version history.