Skip to content

pretraining.md

Latest commit

 

History

History
360 lines (269 loc) · 20.8 KB

pretraining.md

File metadata and controls

360 lines (269 loc) · 20.8 KB
Raw

Pretraining workflows

Primus is a YAML-driven training stack for AMD GPUs. You select a backend (Megatron-LM, TorchTitan, JAX MaxText, Megatron Bridge, SpecForge), 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 backends supported by Primus and their typical uses.

Backend Framework Typical use
Megatron-LM framework: megatron Large-scale transformer pretraining with Megatron-style parallelism (TP/PP/EP).
TorchTitan framework: torchtitan PyTorch-native scaled training (FSDP / tensor / pipeline / expert parallelism per config).
MaxText (JAX) framework: maxtext JAX/MaxText single- and multi-node runs; parallelism via MaxText ici_* / dcn_* settings.
MaxDiffusion (JAX) framework: maxdiffusion JAX/MaxDiffusion diffusion pretraining (WAN 2.1, FLUX.1-dev). Source is vendored as the third_party/maxdiffusion submodule; deps/patches installed by examples/maxdiffusion/setup_maxdiffusion_env.sh.
Megatron Bridge framework: megatron_bridge Bridge-oriented workflows (configure like other backends; see parameter reference).
SpecForge framework: specforge Offline hidden states capture and draft training on a ROCm SGLang Primus image. See SpecForge on Primus.

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.6 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).

Config TP PP EP
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.yaml 1 4 8
grok1-FP8-pretrain.yaml 1 4 8
grok2-BF16-pretrain.yaml 1 4 8
grok2-FP8-pretrain.yaml 1 4 8
llama2_13B-BF16-pretrain.yaml 1 1 1
llama2_13B-FP8-pretrain.yaml 1 1 1
llama2_70B-BF16-pretrain.yaml 1 1 1
llama2_70B-FP8-pretrain.yaml 1 1 1
llama2_7B-BF16-pretrain.yaml 1 1 1
llama2_7B-FP8-pretrain.yaml 1 1 1
llama3.1_405B-BF16-pretrain.yaml 8 8 1
llama3.1_405B-FP8-pretrain.yaml 8 8 1
llama3.1_70B-BF16-pretrain.yaml 1 1 1
llama3.1_70B-FP8-pretrain.yaml 1 1 1
llama3.1_8B-BF16-pretrain.yaml 1 1 1
llama3.1_8B-FP8-pretrain.yaml 1 1 1
llama3.2_1B-BF16-pretrain.yaml 1 1 1
llama3.2_1B-FP8-pretrain.yaml 1 1 1
llama3.2_3B-BF16-pretrain.yaml 1 1 1
llama3.2_3B-FP8-pretrain.yaml 1 1 1
llama3.3_70B-BF16-pretrain.yaml 1 1 1
llama3.3_70B-FP8-pretrain.yaml 1 1 1
llama3_70B-BF16-pretrain.yaml 1 1 1
llama3_70B-FP8-pretrain.yaml 1 1 1
llama3_8B-BF16-pretrain.yaml 1 1 1
llama3_8B-FP8-pretrain.yaml 1 1 1
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.yaml 1 1 1
mixtral_8x22B_v0.1-BF16-pretrain.yaml 1 4 8
mixtral_8x22B_v0.1-FP8-pretrain.yaml 1 4 8
mixtral_8x7B_v0.1-BF16-pretrain.yaml 1 1 8
mixtral_8x7B_v0.1-FP8-pretrain.yaml 1 1 8
qwen2.5_14B-BF16-pretrain.yaml 1 1 1
qwen2.5_14B-FP8-pretrain.yaml 1 1 1
qwen2.5_32B-BF16-pretrain.yaml 1 1 1
qwen2.5_32B-FP8-pretrain.yaml 1 1 1
qwen2.5_3B-BF16-pretrain.yaml 1 1 1
qwen2.5_3B-FP8-pretrain.yaml 1 1 1
qwen2.5_72B-BF16-pretrain.yaml 1 1 1
qwen2.5_72B-FP8-pretrain.yaml 1 1 1
qwen2.5_7B-BF16-pretrain.yaml 1 1 1
qwen2.5_7B-FP8-pretrain.yaml 1 1 1
qwen3_14B-BF16-pretrain.yaml 1 1 1
qwen3_14B-FP8-pretrain.yaml 1 1 1
qwen3_235B_A22B-BF16-pretrain.yaml 1 1 8
qwen3_235B_A22B-FP8-pretrain.yaml 1 1 8
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.yaml 1 1 1
qwen3_32B-FP8-pretrain.yaml 1 1 1
qwen3_4B-BF16-pretrain.yaml 1 1 1
qwen3_4B-FP8-pretrain.yaml 1 1 1
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.yaml 1 1 1
qwen3_8B-FP8-pretrain.yaml 1 1 1
hylo_llama_mamba_1B_BF16-pretrain.yaml 1 1 1
hylo_llama_mamba_3B_BF16-pretrain.yaml 1 1 1
hylo_llama_mamba_8B_BF16-pretrain.yaml 1 1 1

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

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

Section Role
work_group, user_name, exp_name, workspace Run identity and output root (supports ${VAR:default} substitution).
modules.pre_trainer.framework megatron selects Megatron-LM integration.
config: pre_trainer.yaml Module preset under primus/configs/modules/megatron/.
model: llama2_7B.yaml Model preset under primus/configs/models/megatron/ (extends llama2_base.yaml → …).
overrides Run-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

Section Role
framework: torchtitan Selects the TorchTitan integration.
config: pre_trainer.yaml Module preset under primus/configs/modules/torchtitan/.
model: llama3.1_8B.yaml Model preset under primus/configs/models/torchtitan/.
overrides.training, lr_scheduler, activation_checkpoint, primus_turbo Run-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-bf16-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/

File Key parallelism (ici_* intra-node, dcn_* inter-node)
deepseek_v2_16B-bf16-pretrain.yaml ici_fsdp_parallelism: 1, ici_data_parallelism: 1, dcn_fsdp_parallelism: 1, dcn_data_parallelism: -1
grok1-nanoo_fp8-pretrain.yaml ici_fsdp_parallelism: 1, ici_data_parallelism: 1, dcn_fsdp_parallelism: 1, dcn_data_parallelism: -1
llama2_70B-bf16-pretrain.yaml ici_fsdp_parallelism: 8, ici_data_parallelism: 1, dcn_fsdp_parallelism: 1, dcn_data_parallelism: -1
llama2_7B-bf16-pretrain.yaml ici_fsdp_parallelism: 8, ici_data_parallelism: 1, dcn_fsdp_parallelism: 1, dcn_data_parallelism: -1
llama3.3_70B-bf16-pretrain.yaml ici_fsdp_parallelism: 8, ici_data_parallelism: 1, dcn_fsdp_parallelism: 1, dcn_data_parallelism: -1
llama3_70B-bf16-pretrain.yaml ici_fsdp_parallelism: 8, ici_data_parallelism: 1, dcn_fsdp_parallelism: 1, dcn_data_parallelism: -1
llama3_8B-bf16-pretrain.yaml ici_fsdp_parallelism: 8, ici_data_parallelism: 1, dcn_fsdp_parallelism: 1, dcn_data_parallelism: -1
mixtral_8x7B-bf16-pretrain.yaml ici_fsdp_parallelism: 1, ici_data_parallelism: 1, dcn_fsdp_parallelism: 1, dcn_data_parallelism: -1
qwen3_14B-bf16-pretrain.yaml ici_fsdp_parallelism: 8, ici_data_parallelism: 1, dcn_fsdp_parallelism: 1, dcn_data_parallelism: -1
qwen3_30B_A3B-bf16-pretrain.yaml ici_fsdp_parallelism: 1, ici_data_parallelism: 1, dcn_fsdp_parallelism: 1, dcn_data_parallelism: -1

The llama2_7B-bf16-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.

fp8 MoE (v26.6): fp8 Mixture-of-Experts configs must set pure_nnx_decoder: false in their overrides; otherwise they crash at step 1 under the v26.6 pure-NNX decoder default. See MaxText parameters → Precision and quantization. Dense fp8 and bf16 configs are unaffected.

MaxDiffusion (JAX) pretraining

The MaxDiffusion backend runs JAX diffusion pretraining (WAN 2.1, FLUX.1-dev). Environment setup depends on your image:

Image has maxdiffusion installed? What happens
Yes (e.g. MAD primus_maxdiffusion image, unified docker) setup_maxdiffusion_env.sh detects it and is a no-op. Set PRIMUS_SKIP_PIP=1 to skip calling it entirely.
No (e.g. bare rocm/jax-training:maxtext-* image) The script installs everything from the Primus checkout: torch (ROCm wheels), deps, editable submodule, and patches. Requires third_party/maxdiffusion submodule to be initialized.

The relevant pieces:

  • Source is vendored as the third_party/maxdiffusion submodule.
  • Dependencies live in requirements-maxdiffusion.txt (kept separate from requirements-jax.txt so the MaxDiffusion pins never affect MaxText runs).
  • Install + patches are applied by examples/maxdiffusion/setup_maxdiffusion_env.sh (idempotent): torch/torchvision (ROCm wheels), the requirements above, an editable install of the vendored submodule, and four source patches (Flax-T5 clip rename, TensorFlow-preload-before-TransformerEngine, Shardy-on, and the TransformerEngine empty context-parallel-axis fix).

Prerequisites

Initialize the vendored submodule (a plain clone will not populate it):

git submodule update --init third_party/maxdiffusion

Run on a JAX base image (for example rocm/jax-training) or a bare-metal JAX environment, and export HF_TOKEN for gated Hugging Face assets.

Quick start (run from a bare Primus checkout)

Use primus-cli direct with BACKEND=MaxDiffusion. When PRIMUS_SKIP_PIP is unset, the prepare hooks run setup_maxdiffusion_env.sh for you (installs the stack + applies the patches), set NVTE_FRAMEWORK=jax and MAXDIFFUSION_PATH, then launch:

BACKEND=MaxDiffusion ./primus-cli direct -- train pretrain \
  --config examples/maxdiffusion/configs/MI355X/wan2.1_1.3b-pretrain.yaml

To run the environment setup once by itself (e.g. to warm an image or a shared venv), invoke the script directly, then launch with PRIMUS_SKIP_PIP=1:

bash examples/maxdiffusion/setup_maxdiffusion_env.sh
PRIMUS_SKIP_PIP=1 BACKEND=MaxDiffusion ./primus-cli direct -- train pretrain \
  --config examples/maxdiffusion/configs/MI355X/flux_dev-pretrain.yaml

Quick start (container mode)

primus-cli bootstraps the same environment: the train/pretrain/maxdiffusion prepare hooks run setup_maxdiffusion_env.sh before training and select the plain-python launcher (JAX drives every GPU from one process, so torchrun is never used).

./primus-cli container -- train pretrain \
  --config examples/maxdiffusion/configs/MI300X/wan2.1_1.3b-pretrain.yaml --max_train_steps 10

Container launches start from a clean image each time, so the setup runs on every launch. Wheels are cached under $DATA_PATH/pip_cache inside the mounted checkout, so only the first run pays for downloads. Set PRIMUS_SKIP_PIP=1 to skip the step entirely on images that already ship the stack.

Step counts use --max_train_steps (the MaxDiffusion field name). --steps belongs to MaxText and is silently ignored here.

Example configurations under examples/maxdiffusion/configs/MI355X/

File Model Status on MI355X (gfx950)
flux_dev-pretrain.yaml FLUX.1-dev ✅ validated
wan2.1_1.3b-pretrain.yaml WAN 2.1 1.3B ✅ validated
wan2.1_14b-pretrain.yaml WAN 2.1 14B ✅ validated (requires RCCL_WARP_SPEED_AUTO=0, set in config)

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):

Stage Purpose
1 Dump GEMM shapes seen during training (reduce train_iters for faster collection).
2 Tune kernels from dumped shapes (offline tooling under examples/offline_tune).
3 Train 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.

Backend Example region Parallelism vocabulary
Megatron-LM examples/megatron/configs/MI300X/ tensor_model_parallel_size, pipeline_model_parallel_size, expert_model_parallel_size (and env-driven ${PRIMUS_TP:…} variants).
TorchTitan examples/torchtitan/configs/MI300X/ parallelism.* (e.g. tensor_parallel_degree, pipeline_parallel_degree, expert_parallel_degree, FSDP shard settings).
MaxText examples/maxtext/configs/MI300X/ ici_fsdp_parallelism, ici_data_parallelism, dcn_fsdp_parallelism, dcn_data_parallelism.

./runner/primus-cli is the only entry point. The packaged launchers under examples/customer_package/ and examples/moe_package/ reach it through the shared helper runner/helpers/launch/slurm_pretrain.sh, which translates their EXP / NNODES / DATA_PATH environment contract into a primus-cli slurm invocation; call the CLI directly as shown above.

Related documentation