Pretraining workflows

August 12, 2026 · View on GitHub

Primus is a YAML-driven training stack for AMD GPUs. You select a backend (Megatron-LM, TorchTitan, JAX MaxText, Megatron Bridge), point train pretrain at a configuration YAML, and launch Primus with the unified CLI (runner/primus-cli) in direct, container, or Slurm mode. See CLI reference and Configuration system.

This section helps you understand concepts related to the Primus workflow: how backends work, YAML structure and inheritance, parallelism vocabulary, the full per-backend configuration inventory, and so on. If you already understand the concepts and just need the specific commands to run your training with Primus, see End-to-end training recipes.

Overview

The following table describes the four backend types supported by Primus and their typical uses.

BackendFrameworkTypical use
Megatron-LMframework: megatronLarge-scale transformer pretraining with Megatron-style parallelism (TP/PP/EP).
TorchTitanframework: torchtitanPyTorch-native scaled training (FSDP / tensor / pipeline / expert parallelism per config).
MaxText (JAX)framework: maxtextJAX/MaxText single- and multi-node runs; parallelism via MaxText ici_* / dcn_* settings.
Megatron Bridgeframework: megatron_bridgeBridge-oriented workflows (configure like other backends; see parameter reference).

Several setup steps apply to all backends (mock vs. real data, Hugging Face tokens, scaling to multiple nodes, and HipBLASLt autotuning). After you read the backend section that applies to you, see Common patterns below.

Megatron-LM pretraining

Quick start (container mode)

From the root of the clone of the Primus repository, with Docker or Podman available, the following command starts the training in container mode:

./runner/primus-cli container -- train pretrain \
  --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml

This uses the default image from runner/.primus.yaml (rocm/primus:v26.5 unless overridden). The project tree is mounted into the container automatically by runner/primus-cli-container.sh.

Example configurations under examples/megatron/configs/MI300X/

The following files ship in the repository (sorted by name). Parallelism columns are taken from tensor_model_parallel_size / pipeline_model_parallel_size / expert_model_parallel_size in each file (literals or ${PRIMUS_TP:…} defaults).

ConfigTPPPEP
deepseek_v2-BF16-pretrain.yaml${PRIMUS_TP:1}${PRIMUS_PP:4}${PRIMUS_EP:8}
deepseek_v2-FP8-pretrain.yaml${PRIMUS_TP:1}${PRIMUS_PP:4}${PRIMUS_EP:8}
deepseek_v2_lite-BF16-pretrain.yaml${PRIMUS_TP:1}${PRIMUS_PP:1}${PRIMUS_EP:8}
deepseek_v2_lite-FP8-pretrain.yaml${PRIMUS_TP:1}${PRIMUS_PP:1}${PRIMUS_EP:8}
deepseek_v3-BF16-pretrain.yaml${PRIMUS_TP:1}${PRIMUS_PP:1}${PRIMUS_EP:8}
deepseek_v3-FP8-pretrain.yaml${PRIMUS_TP:1}${PRIMUS_PP:1}${PRIMUS_EP:8}
gpt_oss_20B-BF16-pretrain.yaml${PRIMUS_TP:1}${PRIMUS_PP:1}${PRIMUS_EP:8}
gpt_oss_20B-FP8-pretrain.yaml${PRIMUS_TP:1}${PRIMUS_PP:1}${PRIMUS_EP:8}
grok1-BF16-pretrain.yaml148
grok1-FP8-pretrain.yaml148
grok2-BF16-pretrain.yaml148
grok2-FP8-pretrain.yaml148
llama2_13B-BF16-pretrain.yaml111
llama2_13B-FP8-pretrain.yaml111
llama2_70B-BF16-pretrain.yaml111
llama2_70B-FP8-pretrain.yaml111
llama2_7B-BF16-pretrain.yaml111
llama2_7B-FP8-pretrain.yaml111
llama3.1_405B-BF16-pretrain.yaml881
llama3.1_405B-FP8-pretrain.yaml881
llama3.1_70B-BF16-pretrain.yaml111
llama3.1_70B-FP8-pretrain.yaml111
llama3.1_8B-BF16-pretrain.yaml111
llama3.1_8B-FP8-pretrain.yaml111
llama3.2_1B-BF16-pretrain.yaml111
llama3.2_1B-FP8-pretrain.yaml111
llama3.2_3B-BF16-pretrain.yaml111
llama3.2_3B-FP8-pretrain.yaml111
llama3.3_70B-BF16-pretrain.yaml111
llama3.3_70B-FP8-pretrain.yaml111
llama3_70B-BF16-pretrain.yaml111
llama3_70B-FP8-pretrain.yaml111
llama3_8B-BF16-pretrain.yaml111
llama3_8B-FP8-pretrain.yaml111
llama4_17B128E-BF16-pretrain.yaml${PRIMUS_TP:1}${PRIMUS_PP:1}${PRIMUS_EP:8}
llama4_17B128E-FP8-pretrain.yaml${PRIMUS_TP:1}${PRIMUS_PP:1}${PRIMUS_EP:8}
llama4_17B16E-BF16-pretrain.yaml${PRIMUS_TP:1}${PRIMUS_PP:1}${PRIMUS_EP:8}
llama4_17B16E-FP8-pretrain.yaml${PRIMUS_TP:1}${PRIMUS_PP:1}${PRIMUS_EP:8}
mamba_370M-pretrain.yaml111
mixtral_8x22B_v0.1-BF16-pretrain.yaml148
mixtral_8x22B_v0.1-FP8-pretrain.yaml148
mixtral_8x7B_v0.1-BF16-pretrain.yaml118
mixtral_8x7B_v0.1-FP8-pretrain.yaml118
qwen2.5_14B-BF16-pretrain.yaml111
qwen2.5_14B-FP8-pretrain.yaml111
qwen2.5_32B-BF16-pretrain.yaml111
qwen2.5_32B-FP8-pretrain.yaml111
qwen2.5_3B-BF16-pretrain.yaml111
qwen2.5_3B-FP8-pretrain.yaml111
qwen2.5_72B-BF16-pretrain.yaml111
qwen2.5_72B-FP8-pretrain.yaml111
qwen2.5_7B-BF16-pretrain.yaml111
qwen2.5_7B-FP8-pretrain.yaml111
qwen3_14B-BF16-pretrain.yaml111
qwen3_14B-FP8-pretrain.yaml111
qwen3_235B_A22B-BF16-pretrain.yaml118
qwen3_235B_A22B-FP8-pretrain.yaml118
qwen3_30B_A3B-BF16-pretrain.yaml${PRIMUS_TP:1}${PRIMUS_PP:1}${PRIMUS_EP:8}
qwen3_30B_A3B-FP8-pretrain.yaml${PRIMUS_TP:1}${PRIMUS_PP:1}${PRIMUS_EP:8}
qwen3_32B-BF16-pretrain.yaml111
qwen3_32B-FP8-pretrain.yaml111
qwen3_4B-BF16-pretrain.yaml111
qwen3_4B-FP8-pretrain.yaml111
qwen3_5_35B_A3B-BF16-pretrain.yaml${PRIMUS_TP:1}${PRIMUS_PP:1}${PRIMUS_EP:8}
qwen3_5_35B_A3B-FP8-pretrain.yaml${PRIMUS_TP:1}${PRIMUS_PP:1}${PRIMUS_EP:8}
qwen3_8B-BF16-pretrain.yaml111
qwen3_8B-FP8-pretrain.yaml111
zebra_llama_1B-pretrain.yaml111
zebra_llama_3B-pretrain.yaml111
zebra_llama_8B-pretrain.yaml111

Sample YAML file (llama2_7B-BF16-pretrain.yaml) explained

Path: examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml

SectionRole
work_group, user_name, exp_name, workspaceRun identity and output root (supports ${VAR:default} substitution).
modules.pre_trainer.frameworkmegatron selects Megatron-LM integration.
config: pre_trainer.yamlModule preset under primus/configs/modules/megatron/.
model: llama2_7B.yamlModel preset under primus/configs/models/megatron/ (extends llama2_base.yaml → …).
overridesRun-specific training knobs: iterations, batching, LR, parallelism (tensor_model_parallel_size, pipeline_model_parallel_size, expert_model_parallel_size), data paths, checkpoints, Primus Turbo flags, etc.

The sample sets mock_data: true and train_data_path: null so you can validate the stack without real corpora.

Mock data versus real data

  • Mock data: Set mock_data: true and leave train_data_path / valid_data_path empty (as in llama2_7B-BF16-pretrain.yaml).
  • Real data: Set mock_data: false and populate Megatron-compatible data paths (and tokenizer assets) in overrides. Use paths visible inside your container mounts.

Multi-node training with Slurm

./runner/primus-cli slurm srun -N 4 -p <partition> -- train pretrain \
  --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml

runner/primus-cli-slurm-entry.sh derives MASTER_ADDR, NNODES, and NODE_RANK from Slurm and forwards them into the container. Align tensor_model_parallel_size, pipeline_model_parallel_size, and expert_model_parallel_size with your cluster width and job size.

TorchTitan pretraining

Quick start

./runner/primus-cli container -- train pretrain \
  --config examples/torchtitan/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml

Example configurations under examples/torchtitan/configs/MI300X/

File
deepseek_v3_16b-BF16-pretrain.yaml
deepseek_v3_16b-FP8-pretrain.yaml
deepseek_v3_236b-BF16-pretrain.yaml
deepseek_v3_236b-FP8-pretrain.yaml
deepseek_v3_671b-pretrain.yaml
llama3.1_405B-BF16-pretrain.yaml
llama3.1_405B-FP8-pretrain.yaml
llama3.1_70B-BF16-pretrain.yaml
llama3.1_70B-FP8-pretrain.yaml
llama3.1_8B-BF16-pretrain.yaml
llama3.1_8B-FP8-pretrain.yaml
llama4_17Bx128E-BF16-pretrain.yaml
llama4_17Bx128E-FP8-pretrain.yaml
llama4_17Bx16E-BF16-pretrain.yaml
llama4_17Bx16E-FP8-pretrain.yaml
qwen3_0.6B-pretrain.yaml
qwen3_1.7B-pretrain.yaml
qwen3_14B-pretrain.yaml
qwen3_32B-pretrain.yaml
qwen3_4B-pretrain.yaml
qwen3_8B-pretrain.yaml

Sample YAML file (llama3.1_8B-BF16-pretrain.yaml) explained

Path: examples/torchtitan/configs/MI300X/llama3.1_8B-BF16-pretrain.yaml

SectionRole
framework: torchtitanSelects the TorchTitan integration.
config: pre_trainer.yamlModule preset under primus/configs/modules/torchtitan/.
model: llama3.1_8B.yamlModel preset under primus/configs/models/torchtitan/.
overrides.training, lr_scheduler, activation_checkpoint, primus_turboRun-specific batching, steps, checkpointing, and Turbo options.

Some configurations omit an explicit parallelism: block; in that case the default values come from the module and model presets (primus/configs/modules/torchtitan/pre_trainer.yaml and the chosen model YAML). Other examples (for example DeepSeek and Qwen) set parallelism: inline with tensor_parallel_degree, pipeline_parallel_degree, expert_parallel_degree, etc.

MaxText (JAX) pretraining

Quick start

./runner/primus-cli container -- train pretrain \
  --config examples/maxtext/configs/MI300X/llama2_7B-pretrain.yaml

JAX-specific requirements

Install JAX/MaxText dependencies from the repository root:

pip install -r requirements-jax.txt

Example configurations under examples/maxtext/configs/MI300X/

FileKey parallelism (ici_* intra-node, dcn_* inter-node)
deepseek_v2_16B-pretrain.yamlici_fsdp_parallelism: 1, ici_data_parallelism: 1, dcn_fsdp_parallelism: 1, dcn_data_parallelism: -1
grok1-pretrain.yamlici_fsdp_parallelism: 1, ici_data_parallelism: 1, dcn_fsdp_parallelism: 1, dcn_data_parallelism: -1
llama2_70B-pretrain.yamlici_fsdp_parallelism: 8, ici_data_parallelism: 1, dcn_fsdp_parallelism: 1, dcn_data_parallelism: -1
llama2_7B-pretrain.yamlici_fsdp_parallelism: 8, ici_data_parallelism: 1, dcn_fsdp_parallelism: 1, dcn_data_parallelism: -1
llama3.3_70B-pretrain.yamlici_fsdp_parallelism: 8, ici_data_parallelism: 1, dcn_fsdp_parallelism: 1, dcn_data_parallelism: -1
llama3_70B-pretrain.yamlici_fsdp_parallelism: 8, ici_data_parallelism: 1, dcn_fsdp_parallelism: 1, dcn_data_parallelism: -1
llama3_8B-pretrain.yamlici_fsdp_parallelism: 8, ici_data_parallelism: 1, dcn_fsdp_parallelism: 1, dcn_data_parallelism: -1
mixtral_8x7B-pretrain.yamlici_fsdp_parallelism: 1, ici_data_parallelism: 1, dcn_fsdp_parallelism: 1, dcn_data_parallelism: -1
qwen3_14B-pretrain.yamlici_fsdp_parallelism: 8, ici_data_parallelism: 1, dcn_fsdp_parallelism: 1, dcn_data_parallelism: -1
qwen3_30B_A3B-pretrain.yamlici_fsdp_parallelism: 1, ici_data_parallelism: 1, dcn_fsdp_parallelism: 1, dcn_data_parallelism: -1

The llama2_7B-pretrain.yaml example also sets dataset_type: "synthetic" and hf_access_token: ${HF_TOKEN:""} for gated Hugging Face assets when you switch to real data.

Common patterns

Testing with mock data

Set mock_data: true (Megatron/TorchTitan) or synthetic dataset settings (MaxText) to validate the configurations and infrastructure without I/O-heavy datasets.

Real training data

  • Megatron: Configure train_data_path / valid_data_path and tokenizer assets in overrides once mock_data is false.
  • For all backends, ensure host paths are mounted in container mode (--volume or container.options.volume in YAML).
  • TorchTitan/MaxText: Follow backend-specific dataset fields in the overrides and presets.

Scaling from single-node to multi-node

  • Use Slurm mode for allocation; keep the container entry if you want the same image on every node.
  • Set environment variables consistently (NNODES, NODE_RANK, MASTER_ADDR, MASTER_PORT, GPUS_PER_NODE); the Slurm entry script injects them when using primus-cli slurm.
  • Increase values in the parallelism fields (Megatron TP/PP/EP; TorchTitan parallelism; MaxText ici_* / dcn_*) to match topology.

Hugging Face token for gated models

Export HF_TOKEN on the host before launching container mode; runner/.primus.yaml lists HF_TOKEN under container.options.env so it can be forwarded into the container. MaxText configurations may reference ${HF_TOKEN:""} directly.

hipBLASLt autotuning (three stages)

Controlled with PRIMUS_HIPBLASLT_TUNING_STAGE (see examples/README.md):

StagePurpose
1Dump GEMM shapes seen during training (reduce train_iters for faster collection).
2Tune kernels from dumped shapes (offline tooling under examples/offline_tune).
3Train using tuned kernel artifacts from ./output/tune_hipblaslt/....

Example (from in-repo docs):

export PRIMUS_HIPBLASLT_TUNING=1        # master switch (required; tuning is skipped without it)
export PRIMUS_HIPBLASLT_TUNING_STAGE=1
./runner/primus-cli direct -- train pretrain \
  --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml

Supported models

The tables above in the Megatron, TorchTitan, and MaxText sections are curated MI300X examples from the Primus repository. Use examples/<backend>/configs/ in the repository as the authoritative inventory, as new presets and hardware-specific examples may be added there before this document is updated to reflect their additions.

BackendExample regionParallelism vocabulary
Megatron-LMexamples/megatron/configs/MI300X/tensor_model_parallel_size, pipeline_model_parallel_size, expert_model_parallel_size (and env-driven ${PRIMUS_TP:…} variants).
TorchTitanexamples/torchtitan/configs/MI300X/parallelism.* (e.g. tensor_parallel_degree, pipeline_parallel_degree, expert_parallel_degree, FSDP shard settings).
MaxTextexamples/maxtext/configs/MI300X/ici_fsdp_parallelism, ici_data_parallelism, dcn_fsdp_parallelism, dcn_data_parallelism.

For scripting patterns that predate primus-cli, the repository still documents examples/run_local_pretrain.sh and examples/run_slurm_pretrain.sh in examples/README.md; equivalent launches are shown above using ./runner/primus-cli.