# 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](../02-user-guide/cli-reference.md) 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`](https://github.com/AMD-AGI/Primus/blob/main/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. | ### Backend index - [Megatron-LM patch notes](#megatron-lm-patch-notes) - [TorchTitan patch notes](#torchtitan-patch-notes) - [JAX MaxText patch notes](#jax-maxtext-patch-notes) --- ## 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](#base-module-parameters) above for shared module parameters, and the [Primus CLI Reference](../02-user-guide/cli-reference.md) 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](../../primus/backends/megatron/core/pipeline_parallel/zerobubble/README.md) | | `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](https://github.com/AMD-AGI/Primus/issues/203) 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](#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](#base-module-parameters) and [Supported Models](#supported-models) above for a high-level model overview - The [Primus CLI Reference](../02-user-guide/cli-reference.md) 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.