Model support matrix#
This document summarizes which model families Primus targets per backend and lists representative checked-in model presets and example experiment YAML under the repository. It distinguishes curated examples from theoretical support (a preset or upstream stack may exist without a matching examples/ entry). Use the filesystem under primus/configs/models/ and examples/*/configs/ as the authoritative live inventory.
For how to add presets, see Adding model configurations. Backend parameter references: Megatron, TorchTitan, MaxText, Megatron Bridge.
Overview: Supported model families (high level)#
The following aligns with the backend overview and the configs present in this tree.
Backend |
Model families (documentation / stack scope) |
|---|---|
Megatron-LM |
LLaMA2 / LLaMA3 / LLaMA3.1 / LLaMA3.3 / LLaMA4 (sizes from small to 405B+), DeepSeek-V2 (including lite), DeepSeek-V3, and DeepSeek-V4 (flash / pro—from v26.7 also on |
TorchTitan |
LLaMA3 family (including 3.1), LLaMA4 examples, DeepSeek-V3 examples, and Qwen3 examples including 0.6B, 1.7B, 4B, 8B, 14B, and 32B variants where present. Additional presets exist under |
MaxText (JAX) |
LLaMA2 / LLaMA3 / LLaMA3.3, DeepSeek-V2 16B, Mixtral-8x7B, Grok1, Qwen3 14B / 30B-A3B (per presets and examples). Broader coverage may exist in upstream MaxText; see MaxText. |
Megatron Bridge |
Qwen3 pretraining and post-training examples, plus post-training examples for Zebra LLaMA and Mamba where present. LLaMA 3.1 70B Bridge examples appear under MI355X. |
Diffusion |
Flux.1 ( |
World Mirror |
HunyuanWorld-Mirror geometry training and Gaussian-head post-training. The model checkout is the unmodified submodule |
AutoModel |
Flux.1 ( |
HummingbirdXT |
Registered backend with a post-training trainer and one checked-in example; user-facing support level still needs maintainer confirmation. |
Interpretation: “Supported” in upstream code can exceed what this repository ships as YAML. Rows below reference representative files that exist under primus/configs/models/ and examples/; they should not be treated as a complete generated inventory.
Megatron model configs#
Model presets live in primus/configs/models/megatron/. Example experiments that reference those presets appear under examples/megatron/configs/MI300X/, MI325X/, and MI355X/.
For TorchTitan, the MI300X, MI325X, and MI355X example directories carry the same model set (25 configs each). For Megatron, MI300X has 82 example configs and MI325X 70: MI325X has everything MI300X has except the pure GDN/KDA and the GDN/KDA/300M Zebra LLaMA hybrid recipes, and qwen3_5_35B_A3B (BF16 and FP8). MI355X has 124, adding models such as glm5, glm5_3_flash, gpt_oss_120B, kimi_k2, lfm2_8B_A1B, minimax_m2.5, and minimax_m3. gfx1250 has seven single-GPU LoRA-SFT configs: the BF16 recipes, plus FP8 hybrid proxies for Llama-2-70B, Qwen2.5-72B, and Qwen3-235B-A22B added in v26.9. Each row’s SKU list in the following table reflects exactly which SKUs ship a curated example (see, for example, qwen3_5_35B_A3B, which is MI300X/MI355X only).
Model name (file) |
Preset path |
Role |
Example experiment dirs |
Precision in examples |
|---|---|---|---|---|
|
|
Dense model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Base fragment ( |
— |
— |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
MoE model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Base fragment |
— |
— |
|
|
MoE model preset |
MI355X (pretrain BF16/FP8, plus a 4-layer BF16 SFT) |
BF16, FP8 |
|
|
MoE model preset |
No curated example in this repo |
— |
|
|
Base fragment |
— |
— |
|
|
Model preset |
No curated example in this repo |
— |
|
|
Model preset |
MI355X |
BF16, FP8 |
|
|
Hybrid MoE preset (GLM-5.3-Flash, 45 layers, KDA + DSA with mHC) |
MI355X |
BF16 |
|
|
First 8 layers of GLM-5.3-Flash, for smoke tests |
MI355X ( |
BF16 |
|
|
Base fragment |
— |
— |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Model preset |
MI355X |
BF16, FP8 |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Base fragment |
— |
— |
|
|
Base fragment |
— |
— |
|
|
MoE model preset |
MI355X |
BF16, FP8 |
|
|
Generic Megatron LM defaults |
Used via |
— |
|
|
MoE model preset |
MI355X |
BF16, FP8 |
|
|
Base fragment |
— |
— |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
4-layer single-GPU proxy |
gfx1250 (LoRA SFT) |
BF16, FP8 |
|
|
Base fragment |
— |
— |
|
|
Base fragment |
— |
— |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Base fragment |
— |
— |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Hugging Face-tokenizer variant for native SFT |
gfx1250 (LoRA SFT) |
BF16 |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
MoE model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
MoE model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Base fragment |
— |
— |
|
|
Model preset |
MI300X, MI325X, MI355X |
Set in experiment overrides |
|
|
Base fragment |
— |
— |
|
|
MoE model preset |
MI355X |
BF16, FP8 |
|
|
MoE model preset |
MI355X |
BF16, FP8 |
|
|
MoE model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
MoE model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Base fragment |
— |
— |
|
|
Large MoE template |
No curated example in this repo |
— |
|
|
Large MoE template |
No curated example in this repo |
— |
|
|
Large MoE template |
No curated example in this repo |
— |
|
|
Large MoE template |
No curated example in this repo |
— |
|
|
MoE proxy / test template |
No curated example in this repo |
— |
|
|
Primus Megatron root defaults |
Used via |
— |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 (MXFP4 on MI355X) |
|
|
4-layer single-GPU proxy |
gfx1250 (LoRA SFT) |
BF16, FP8 |
|
|
Base fragment |
— |
— |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
MoE model preset |
MI300X, MI325X, MI355X |
BF16, FP8 (MXFP4 on MI355X) |
|
|
MoE model preset |
MI300X, MI355X |
BF16, FP8 |
|
|
MoE model preset |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
4-layer single-GPU MoE proxy |
gfx1250 (LoRA SFT) |
BF16, FP8 |
|
|
MoE model preset (ZAYA1-8B, CCA + MoE) |
MI300X, MI325X, MI355X |
BF16, FP8 |
|
|
Base fragment |
— |
— |
|
|
Model preset |
MI300X, MI325X, MI355X |
Set in experiment overrides |
|
|
Model preset |
MI300X, MI325X, MI355X |
Set in experiment overrides |
|
|
Model preset |
MI300X, MI325X, MI355X |
Set in experiment overrides |
|
|
Model preset |
MI300X |
Set in experiment overrides |
|
|
Model preset |
MI300X, MI355X |
Set in experiment overrides |
|
|
Model preset |
MI300X, MI355X |
Set in experiment overrides |
|
|
Model preset |
MI300X |
Set in experiment overrides |
|
|
Pure KDA preset |
MI300X, MI355X |
|
|
|
Pure KDA preset |
MI300X |
|
|
|
Pure GDN preset |
MI300X, MI355X |
|
|
|
Pure GDN preset |
MI300X |
|
ZAYA1-8B: zaya1.yaml follows the released Zyphra/ZAYA1-base checkpoint (80 decoder stages alternating CCA attention and MoE). The MI300X, MI325X, and MI355X BF16 and FP8 examples are 30-iteration mock-data runs. On each GPU the micro-batch size is the largest that completed on one 8-GPU node with rocm/primus:v26.8, when the recipes were sized. The zaya1_8B-FP8-pretrain.yaml recipes set fp8: hybrid with tensorwise scaling, so FP8 is accurate from a cold start: CCA projections and expert SwiGLU GEMMs run as Transformer Engine FP8 linears, and the router, convolutions, norms, attention, embeddings, and tied logits stay in BF16. They also set NVTE_USE_CAST_TRANSPOSE_TRITON=0, because the Triton cast kernel autotunes for every new expert token count. The BF16 port landed on main in #1199.
GLM-5.3-Flash: glm5_3_flash.yaml (model_type: glm5_next) is a 45-layer hybrid: KDA linear attention on 34 layers, DSA sparse attention on 11, manifold-constrained hyper-connections (mHC) as the residual, and a 288-expert MoE. examples/megatron/configs/MI355X/glm5_3_flash-BF16-pretrain.yaml and examples/megatron/models/glm5-3-flash/run_glm5_3_flash.sh run it, and can load the released Hugging Face checkpoint for any TP / EP / PP (GLM5_HF_LOAD). Its forward-only log-probabilities match SGLang on the released checkpoint, and the 8-layer glm5_3_flash_8L.yaml trains from scratch. Training the full 45-layer model needs at least four MI355X nodes and has not been run. The DSA indexer is frozen (no gradient, no indexer distillation loss), and no Primus Turbo paths are enabled (enable_primus_turbo: false) (#1230). Packed sequences (--qkv-format thd) run under context parallelism, and the loader also accepts text-only Hugging Face exports (#1265, #1254).
MiniMax M3: minimax_m3.yaml covers the text tower only. The swigluoai activation and Gemma-style RMSNorm map onto existing Megatron options, and MiniMax Sparse Attention (MSA) runs on FlyDSL kernels on MI355X by default (msa_backend: flydsl), with a plain-PyTorch reference backend (msa_backend: eager) for other GPUs. Tensor, context and sequence parallelism for MSA, MTP, and the vision tower are not supported yet. The preset records the unmapped config.json values as comments.
Parallelism: Tensor, pipeline, and expert parallel sizes are not fixed in model presets; they are set in experiment overrides (for example tensor_model_parallel_size, pipeline_model_parallel_size, expert_model_parallel_size). MoE presets such as qwen3_235B_A22B.yaml typically require non-default expert parallelism in real runs—see the matching experiment YAML.
TorchTitan model configs#
Presets: primus/configs/models/torchtitan/. Examples: examples/torchtitan/configs/MI300X/, MI325X/, and MI355X/.
Model name (file) |
Preset path |
Example experiment dirs |
Precision in examples |
|---|---|---|---|
|
|
MI300X, MI325X, MI355X |
BF16 |
|
|
MI300X, MI325X, MI355X |
FP8 |
|
|
MI300X, MI325X, MI355X |
BF16 |
|
|
MI300X, MI325X, MI355X |
FP8 |
|
|
MI300X, MI325X, MI355X |
(see experiment) |
|
|
Preset only; stock examples use |
— |
|
|
No example in this repo |
— |
|
|
No example in this repo |
— |
|
|
No example in this repo |
— |
|
|
No example in this repo |
— |
|
|
MI300X, MI325X, MI355X |
BF16 |
|
|
MI300X, MI325X, MI355X |
FP8 |
|
|
MI300X, MI325X, MI355X |
BF16 |
|
|
MI300X, MI325X, MI355X |
FP8 |
|
|
MI300X, MI325X, MI355X |
BF16 |
|
|
MI300X, MI325X, MI355X |
FP8 |
|
|
No example in this repo |
— |
|
|
No example in this repo |
— |
|
|
No example in this repo |
— |
|
|
MoE; MI300X, MI325X, MI355X |
BF16 |
|
|
MoE; MI300X, MI325X, MI355X |
FP8 |
|
|
MoE; MI300X, MI325X, MI355X |
BF16 |
|
|
MoE; MI300X, MI325X, MI355X |
FP8 |
|
|
MI300X, MI325X, MI355X |
(see experiment) |
|
|
MI300X, MI325X, MI355X |
(see experiment) |
|
|
MI300X, MI325X, MI355X |
(see experiment) |
|
|
MI300X, MI325X, MI355X |
(see experiment) |
|
|
MI300X, MI325X, MI355X |
(see experiment) |
|
|
MI300X, MI325X, MI355X |
(see experiment) |
Parallelism: Controlled by TorchTitan launch configuration and Primus module overrides (see TorchTitan patch notes and TorchTitan parameters); not embedded in the small job / model preset alone.
MaxText model configs#
Presets: primus/configs/models/maxtext/. Examples: examples/maxtext/configs/MI300X/ and examples/maxtext/configs/MI355X/.
Model name (file) |
Preset path |
Example experiment dirs |
|---|---|---|
|
|
MI300X, MI355X |
|
|
MI300X |
|
|
MI300X, MI355X |
|
|
MI300X, MI355X |
|
|
MI300X, MI355X |
|
|
MI300X, MI355X |
|
|
MI355X |
|
|
Convergence test only ( |
|
|
MI300X, MI355X |
|
|
MI300X, MI355X |
|
|
MI300X, MI355X |
|
|
MI300X, MI355X |
|
|
Extended by other presets (not a standalone run) |
Parallelism: JAX / MaxText sharding is configured in experiment overrides (for example ici_fsdp_parallelism, ici_data_parallelism, dcn_* in sample experiments). See MaxText parameters.
Megatron Bridge model configs#
Presets: primus/configs/models/megatron_bridge/. Examples: examples/megatron_bridge/configs/MI300X/ and examples/megatron_bridge/configs/MI355X/.
Model name (file) |
Preset path |
Recipe / flavor (from preset) |
Example experiment dirs |
|---|---|---|---|
|
|
|
MI300X pretrain, MI355X posttrain |
|
|
|
MI300X, MI355X |
|
|
|
MI355X |
|
|
Zebra LLaMA presets |
MI300X posttrain |
|
|
Mamba preset |
MI300X posttrain |
Example filenames include *_pretrain.yaml, *_sft_posttrain.yaml, and *_lora_posttrain.yaml; precision such as bf16_mixed is set in experiment overrides.
Hardware compatibility (example directories)#
Curated example layouts under examples/ use GPU SKU subdirectories. As of this document:
GPU SKU |
|
|
|
|
|---|---|---|---|---|
MI300X |
Yes |
Yes |
Yes |
Yes |
MI355X |
Yes |
Yes |
Yes |
Yes |
MI325X |
Yes |
Yes |
Yes |
No |
gfx1250 |
Yes (LoRA-SFT proxies) |
No |
No |
No |
Megatron, TorchTitan and MaxText ship MI325X example directories in addition to MI300X and MI355X examples. MaxText includes MI355X-only entries such as llama3.1_405B-pretrain.yaml. Megatron’s gfx1250 examples are single-GPU LoRA-SFT proxies for bring-up. Megatron Bridge MI300X examples include Qwen3 8B and 32B pretraining plus Qwen3 32B, Zebra LLaMA, and Mamba post-training examples; LLaMA 3.1 70B Bridge examples appear under MI355X.
Absence of a SKU directory for a given backend does not imply the backend cannot run there; it means this tree does not currently provide a checked-in example path to copy from.
Model architecture reference (Megatron presets)#
Values below come from primus/configs/models/megatron/ presets (merged through extends). Vocabulary size is usually defined by the tokenizer / Hugging Face config, not duplicated in every YAML; context is max_position_embeddings where set in the chain. Use this table as a quick reference for common sizes—not an exhaustive spec of every parameter.
Model family |
Example preset |
Hidden size |
Layers |
Attention heads |
KV heads (GQA) |
Max position (context) |
|---|---|---|---|---|---|---|
LLaMA 2 7B |
|
4096 |
32 |
32 |
32 (no GQA) |
From |
LLaMA 3 8B |
|
4096 |
32 |
32 |
8 |
8192 ( |
LLaMA 3 70B |
|
8192 |
80 |
64 |
8 |
8192 |
LLaMA 3.1 405B |
|
16384 |
126 |
128 |
8 |
8192 |
Qwen3 8B |
|
4096 |
36 |
32 |
8 |
131072 ( |
Mixtral 8x7B |
|
4096 |
32 |
32 |
— |
4096 |
DeepSeek-V3 (MoE) |
|
7168 |
61 |
128 (MLA) |
— |
See preset / HF |
Mamba 370M |
|
(Mamba stack) |
— |
— |
— |
— |
For MoE and hybrid architectures (LLaMA 4, Qwen3-MoE, large moe_*.yaml templates), refer to the full YAML and upstream model cards; headline dimensions alone do not capture expert layout or MLA.