Primus integrates several large-model backends (Megatron-LM, TorchTitan, JAX MaxText, …) and applies a lightweight patch layer to keep configuration flags consistent with the Primus CLI. This page captures those backend-specific switches so they live alongside the rest of the documentation (instead of the historical primus/README_patch.md file).
- Start with the Base Module Parameters table below—every backend module inherits these knobs.
- Jump to the backend-specific section for details on extra CLI/config options and links to the patched source files.
- When editing configs or CLI presets, cross-reference the Primus CLI Reference so command examples and backend parameters stay in sync.
This section lists, at a high level, the model families Primus currently targets on each backend. For more details and configuration examples, refer to the backend-specific patch notes below.
- LLaMA family: LLaMA2, LLaMA3, LLaMA3.1, LLaMA3.3, LLaMA4 (various sizes from 7B up to 405B+)
- DeepSeek family: DeepSeek-V2 (lite/base/full) and DeepSeek-V3
- MoE / Mixtral: Mixtral-8x7B / 8x22B, large MoE configs (515B, 1T, 2T, 4T) and DeepSeek-style MoE
- Qwen family: Qwen2.5 (7B/72B) and Qwen3 (8B/30B/235B variants)
- Other GPT-style models: Grok1/2, GPT-OSS 20B and generic
language_model.yaml
- LLaMA family: LLaMA3, LLaMA3.1, LLaMA3.3 (various sizes, including FP8 variants)
- DeepSeek family: DeepSeek-V3 (16B and 671B, FP8 and BF16 configs)
- Qwen family: Qwen3 small/medium models (0.6B, 1.7B, 32B)
- LLaMA family: LLaMA2 (7B/70B), LLaMA3 (8B/70B), LLaMA3.3 (70B)
- DeepSeek family: DeepSeek-V2 16B
- MoE / Mixtral: Mixtral-8x7B
- Other models: Grok1 and additional MaxText-supported transformers (see MaxText docs for the full list)
All modules inherit the options defined in primus/configs/modules/module_base.yaml:
| Argument Name | Default Value | Description |
|---|---|---|
trainable |
false |
Whether the module participates in training. |
sink_level |
null |
Global sink level for logging. Overrides file_sink_level and stderr_sink_level if set. |
file_sink_level |
DEBUG |
Logging level for file sink (log files). |
stderr_sink_level |
INFO |
Logging level for stderr/console output. |
Primus keeps a curated patch layer on top of upstream Megatron-LM so CLI presets and configs can expose additional controls. Use this section with the Base Module Parameters above for shared module parameters, and the Primus CLI Reference for CLI/config usage patterns.
ℹ️ The Version column maps to Primus internal patch tags (v0.x.y) so you know when a flag landed.
These arguments are introduced in the Megatron module logic (e.g., training loop, logging, resume logic). They are defined via patching and can be configured to control training behavior and logging utilities.
| New Argument | Default Value | Version | Description | Patched Files | Notes |
|---|---|---|---|---|---|
disable_tensorboard |
true |
v0.1.0 | Whether to disable TensorBoard. Set to false if you want to enable profiling or torch trace. |
NA | Required for timeline and performance debugging. |
disable_wandb |
true |
v0.1.0 | Whether to disable Weights & Biases logging. | NA | Useful for internal benchmarking. |
disable_compile_dependencies |
true |
v0.1.0 | Disables Megatron’s custom kernel compilation. Most ops are already covered by TE. | NA | Avoids redundant compilation steps. |
auto_continue_train |
false |
v0.1.0 | Automatically resume training from the latest checkpoint if found in the --save path. |
NA | Simplifies job restarts. |
disable_last_saving |
false |
v0.1.0 | Skip saving the final checkpoint at the last iteration. | NA | Useful for profiling or benchmarking runs. |
no_fp8_weight_transpose_cache |
false |
v0.2.0 | Disable the FP8 weight transpose cache to save memory. | megatron.core.extensions.transformer_engine.TELinear, megatron.core.extensions.transformer_engine.TELayerNormColumnParallelLinear, megatron.core.extensions.transformer_engine.TEDelayedScaling |
May affect performance but reduce memory use. |
decoder_pipeline_manual_split_list |
null |
v0.2.0 | Enable manual pipeline split in (interleaved) 1F1B pipeline parallelism. | megatron.core.transformer.transformer_block.get_num_layers_to_build, megatron.core.transformer.transformer_layer.get_transformer_layer_offset |
Deprecated. Use pipeline_model_parallel_layout instead. |
pp_warmup |
false |
v0.2.0 | Add fwd/bwd warmup to save iter1's time when pp degree is large. | NA | Can save much time for pipeline debug. |
dump_pp_data |
false |
v0.2.0 | Enable dumping pp schedule data for visualization. | megatron.core.pipeline_parallel.schedules.forward_step, megatron.core.pipeline_parallel.schedules.backward_step, megatron.core.pipeline_parallel.schedules.forward_backward_pipelining_with_interleaving, megatron.core.pipeline_parallel.schedules.forward_backward_pipelining_without_interleaving |
Useful for pipeline schedule visualization. |
disable_profiler_activity_cpu |
false |
v0.2.0 | Disable CPU activity in torch profiling. | NA | If you only want to trace CUDA kernels and get a smaller trace JSON file, you can enable this option. However, if you plan to run with TraceLen, please do not enable it. more torch profiler args: torch_profiler_record_shapes: true, torch_profiler_with_stack: true, torch_profiler_use_gzip: true |
use_rocm_mem_info |
false |
v0.2.0 | Logging ROCm memory information in Megatron-LM Trainer | NA | If use_rocm_mem_info = True, ROCm memory information will be collected with rocm-smi at every iteration. |
use_rocm_mem_info_iters |
[1,2] |
v0.2.0 | Logging ROCm memory information in Megatron-LM Trainer for some iterations | NA | If use_rocm_mem_info = False, ROCm memory information will be collected at the iterations specified in use_rocm_mem_info_iters. |
patch_zero_bubble |
false |
v0.2.0 | Using Zero-Bubble pipeline parallism | megatron.core.optimizer.ChainedOptimizer, megatron.core.pipeline_parallel.get_forward_backward_func, megatron.core.tensor_parallel.layers.LinearWithGradAccumulationAndAsyncCommunication, megatron.core.parallel_stat.default_embedding_ranks, megatron.core.parallel_stat.is_pipeline_last_stage, megatron.core.parallel_stat.is_rank_in_embedding_group, megatron.core.distributed.finalize_model_grads, megatron.core.transformer.transformer_layer.get_transformer_layer_offset |
If patch_zero_bubble = True, Zero bubble pipeline parallism will be enable to use. See more detail at ZeroBubble User Guide |
disable_mlflow |
true |
v0.3.0 | Track model development using MLflow | NA | Envs: export DATABRICKS_TOKEN=your_tokenexport DATABRICKS_HOST=your_hostexport MLFLOW_TRACKING_URI=databricksexport MLFLOW_REGISTRY_URI=databricks-ucArguments: mlflow_run_name: null,mlflow_experiment_name: null |
recompute_layer_ids |
null |
v0.4.0 | Specify the exact IDs of layers to recompute, enabling more flexible memory reduction | NA | Using recompute_layer_ids=[layer_id_0, layer_id_1,...] together with recompute_granularity=full, where layer_id ranges from 0 to num_layers - 1. |
dataloader_mp_context |
null |
v0.5.0 | Set DataLoader.multiprocessing_context to avoid SIGSEGV caused by fork()-hostile native state (RDMA MRs, HIP runtime, IPC handles). |
torch.utils.data.DataLoader.__init__ |
Accepted values: "forkserver", "spawn", "fork", or null (keep PyTorch default). Only takes effect when num_workers > 0 and no explicit multiprocessing_context is passed. |
These arguments affect the internal architecture or layer implementations. They are patched into the model construction logic and used for tuning or debugging specific variants.
| New Argument | Default Value | Version | Description | Patched Files | Notes |
|---|---|---|---|---|---|
disable_primus_topk_router |
false |
v0.1.0 | Disable PrimusTopkRouter and use TopkRouter implemented by megatron. | megatron.core.transformer.moe.router.TopKRouter |
Used to debug internal. |
moe_router_force_load_balancing |
false |
v0.1.0 | Force token redistribution in MoE to achieve load balance across experts. | megatron.core.transformer.moe.router.TopKRouter |
Use to debug MoE imbalance issues. |
use_deprecated_20241209_moe_layer |
false |
v0.1.0 | Enable legacy MoE implementation for debugging/perf comparison. | megatron.core.transformer.moe.moe_layer.MoELayer, megatron.core.transformer.moe.moe_layer.MoESubmodules, megatron.core.transformer.moe.experts.GroupedMLP, megatron.core.transformer.moe.experts.SequentialMLP, megatron.core.transformer.moe.experts.TEGroupedMLP, megatron.core.transformer.moe.router.TopKRouter |
Deprecated, used for internal testing only. |
moe_permute_fusion |
false |
v0.1.0 | Permutation and unpermutation fusion. | megatron.core.extensions.transformer_engine, megatron.core.transformer.moe.moe_utils |
Fuse permutation and unpermutation in moe layer. |
moe_use_fused_router_with_aux_score |
false |
v0.2.0 | Fused router topk and calculation of moe aux loss score. Need Primus turbo backend | megatron.core.transformer.moe.router.TopKRouter |
Used to reduce launch overhead of the small kernels in router. |
| New Argument | Default Value | Version | Description | Patched Files | Notes |
|---|---|---|---|---|---|
use_turbo_gemm |
false |
v0.8.0 | Use Primus-Turbo linear modules (PrimusTurboLinear, PrimusTurboColumnParallelLinear, PrimusTurboRowParallelLinear, PrimusTurboLayerNormColumnParallelLinear) in place of the TE linear modules. |
megatron.core.extensions.transformer_engine.TELinear, megatron.core.extensions.transformer_engine.TEColumnParallelLinear, megatron.core.extensions.transformer_engine.TERowParallelLinear, megatron.core.extensions.transformer_engine.TELayerNormColumnParallelLinear |
Accelerates dense GEMMs. Supports FP8 recipes (tensorwise, blockwise, mxfp8) and FP4 recipe (mxfp4). Replaces the deprecated use_turbo_parallel_linear. Please set enable_primus_turbo=True first. |
use_turbo_grouped_gemm |
false |
v0.8.0 | Use Primus-Turbo grouped GEMM (PrimusGroupedMLP with PrimusTurboColumnParallelGroupedLinear / PrimusTurboRowParallelGroupedLinear) for MoE experts in place of TEGroupedMLP. |
megatron.core.transformer.moe.experts.TEGroupedMLP, megatron.core.extensions.transformer_engine.TEColumnParallelGroupedLinear, megatron.core.extensions.transformer_engine.TERowParallelGroupedLinear |
Accelerates MoE grouped GEMMs. Incompatible with moe_use_legacy_grouped_gemm=True. Required by Sync-Free MoE stage 2/3. Replaces the deprecated use_turbo_grouped_mlp. Please set enable_primus_turbo=True first. |
use_turbo_permute_padding |
false |
v0.8.0 | Pad tokens of every experts to 16 or 32 multiple to reduce d2h and h2d. | megatron.core.transformer.moe.token_dispatcher.MoEFlexTokenDispatcher, megatron.core.transformer.moe.experts.TEGroupedMLP |
Only effective under FP8/FP4 with use_turbo_deepep=True. Pad multiple is 16/32 (FP8, depending on recipe) or 32 (FP4). Please set enable_primus_turbo=True first. |
use_turbo_deepep |
false |
v0.4.0 | Use Primus-turbo DeepEPTokenDispatcher. |
megatron.core.transformer.moe.token_dispatcher.MoEFlexTokenDispatcher |
Used Primus-Turbo DeepEP to accelerate MoE token dispatcher. You must both setenable_primus_turbo=True and use_turbo_deepep=True to enable this function. |
turbo_deepep_num_cu |
32 |
v0.4.0 | Set the number of CUs to use for Primus-Turbo DeepEP. | 64 or 80 for ep8, 32 for ep16-64 is best practice. | |
turbo_deepep_use_comm_stream |
false |
v0.4.0 | Primus-Turbo DeepEP will use an internal stream to dispatch/combine when enabled, default used current_stream |
Please both setenable_primus_turbo=True and use_turbo_deepep=True first |
|
turbo_sync_free_moe_stage |
0 |
v0.4.0 | Primus Sync-Free MoE has 4 stages. See RFC: Primus-Megatron SyncFree MoE for more details. | stage 2 is recommended for better performance. Please setenable_primus_turbo=True first |
TorchTitan integration uses the same Primus configuration surface (CLI flags + YAML) but exposes a few extra knobs via patches. Pair this with the Base Module Parameters above for shared module parameters.
| New Argument | Default Value | Version | Description | Patched Files | Notes |
|---|---|---|---|---|---|
primus_turbo.enable_embedding_autocast |
true |
v0.4.0 | Automatically casts nn.Embedding outputs to the AMP dtype (e.g., bf16) during training so downstream layers stay in sync. |
(Primus TorchTitan patch set) | Disable only if you manage casting manually. |
Primus integrates JAX MaxText as a backend for running LLaMA and related transformer models on AMD GPUs. At the moment, Primus does not apply any additional patch-layer arguments on top of MaxText—the MaxText configuration surface (YAML + CLI) is used as-is.
Use this section together with:
- The Base Module Parameters and Supported Models above for a high-level model overview
- The Primus CLI Reference for Primus CLI usage patterns
- The official MaxText documentation for the full set of MaxText-specific arguments
- Primus currently wires MaxText via
primus/configs/models/maxtextandprimus/configs/modules/maxtext. - Model families currently exercised in examples include:
- LLaMA2 7B/70B
- LLaMA3 8B/70B
- LLaMA3.3 70B
- DeepSeek-V2 16B
- Mixtral-8x7B
- Grok1
- There are no extra Primus-only flags for MaxText yet; as we add MaxText-specific patches (e.g., ROCm optimizations, logging helpers), they will be documented in tables here in the same style as the Megatron-LM and TorchTitan patch notes.