Backend patch notes

Backend patch notes#

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

How to read these notes#

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.

Supported models#

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.

Megatron-LM#

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

TorchTitan#

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)

JAX MaxText#

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)

Base module parameters#

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.

Megatron-LM patch notes#

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.

1. Module-level parameters#

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_token` `export DATABRICKS_HOST=your_host` `export MLFLOW_TRACKING_URI=databricks` `export MLFLOW_REGISTRY_URI=databricks-uc` Arguments: `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.

2. Model-definition parameters#

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.

3. Primus-Turbo related options#

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 set`enable_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 set`enable_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 set`enable_primus_turbo=True` first

TorchTitan patch notes#

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.

JAX MaxText patch notes#

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

MaxText-specific notes#

Primus currently wires MaxText via primus/configs/models/maxtext and primus/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.