# Megatron Bridge backend configuration reference

Megatron Bridge integrates [Megatron-Core](https://github.com/NVIDIA/Megatron-LM) training with Hugging Face–centric workflows. In Primus, the **`megatron_bridge`** framework is used for post-training with module preset `sft_trainer.yaml`, and the repository also ships a pretraining preset at `primus/configs/modules/megatron_bridge/pretrain_trainer.yaml`.

## Recipe system

Megatron Bridge resolves training defaults through a **recipe** and **flavor**:

- `recipe` is a Python module path under `megatron.bridge.recipes` (e.g. `qwen.qwen3`).
- `flavor` is the function name inside that module (e.g. `qwen3_8b_finetune_config`) that returns a `ConfigContainer`.

At runtime, `load_recipe_config` in `primus/backends/megatron_bridge/config_utils.py`:

1. Imports `megatron.bridge.recipes.<recipe>` and calls `<flavor>(**filtered_backend_args)` to build the baseline `ConfigContainer`.
2. **Deep-merges** Primus `backend_args` (from YAML + CLI) into that dataclass via `_merge_dict_to_dataclass`, so user overrides sit on top of recipe defaults.

You normally specify `recipe`, `flavor`, `hf_path`, and `dataset` in the model YAML; training hyperparameters and parallelism go in module overrides or experiment module overrides (`modules.post_trainer.overrides` for SFT/post-training, `modules.pre_trainer.overrides` for pretraining examples).

---

## 1. Base module parameters

From `primus/configs/modules/megatron_bridge/sft_trainer.yaml` (extends `module_base.yaml`). Pretraining examples use `pretrain_trainer.yaml` instead.

| Parameter | Default | Description |
|-----------|---------|-------------|
| `trainable` | `true` | Module participates in the training graph. |
| `sink_level` | `null` | Inherited from `module_base.yaml`; structured logging sink level. |
| `file_sink_level` | `DEBUG` | File sink verbosity. |
| `stderr_sink_level` | `INFO` | Stderr sink verbosity. |

---

## 2. Training

| Parameter | Default | Description |
|-----------|---------|-------------|
| `stage` | `"sft"` | Backend stage selector. Primus dispatches post-training via `primus train posttrain` and loads the Megatron Bridge posttrain trainer when this module is used under `post_trainer`. |
| `trainable` | `true` | See Base module parameters. |

**CLI note:** The user-facing suite is **`posttrain`** (`primus train posttrain --config ...`). The YAML `stage` field selects the Megatron Bridge trainer implementation (`sft`), not the CLI suite name.

For Bridge pretraining, use the normal pretraining suite (`primus train pretrain --config ...`) with experiments that reference `modules.pre_trainer.config: pretrain_trainer.yaml`.

---

## 3. Fine-tuning method (PEFT)

Primus examples set these under `modules.post_trainer.overrides` (see `examples/megatron_bridge/configs/`).

| Parameter | Example | Description |
|-----------|---------|-------------|
| `peft` | `"none"`, `"lora"` | Parameter-efficient fine-tuning mode. |
| `peft_dim` | `16` | LoRA rank (example: `llama31_70b_lora_posttrain.yaml`). |
| `peft_alpha` | `32` | LoRA scaling alpha (same example). |
| `packed_sequence` | `false` | Pack multiple short sequences per microbatch when supported. |

Additional keys such as `pretrained_checkpoint`, `use_distributed_optimizer`, or `cross_entropy_loss_fusion` appear in larger examples and are merged into the recipe `ConfigContainer` when the dataclass exposes matching fields.

---

## 4. Parallelism

Typical overrides from Megatron Bridge examples:

| Parameter | Example | Description |
|-----------|---------|-------------|
| `tensor_model_parallel_size` | `1`, `2`, `8` | Tensor parallelism degree. |
| `pipeline_model_parallel_size` | `1` | Pipeline parallelism degree. |
| `virtual_pipeline_model_parallel_size` | `null` | Virtual pipeline stages per rank when PP > 1. |
| `context_parallel_size` | `1` | Context parallelism degree. |
| `sequence_parallel` | `false` | Sequence parallelism within TP groups. |
| `use_megatron_fsdp` | `false` | Optional Megatron FSDP path. |

---

## 5. Training hyperparameters

| Parameter | Example | Description |
|-----------|---------|-------------|
| `train_iters` | `200`, `1000` | Total training iterations. |
| `global_batch_size` | `8`, `128` | Global batch across data-parallel groups. |
| `micro_batch_size` | `1`, `2` | Per-GPU microbatch before gradient accumulation. |
| `seq_length` | `2048`, `8192` | Training sequence length. |
| `eval_interval` | `30` | Steps between evaluations. |
| `save_interval` | `50` | Steps between checkpoint saves. |

---

## 6. Learning rate

| Parameter | Example | Description |
|-----------|---------|-------------|
| `finetune_lr` | `1.0e-4`, `5.0e-6` | Peak learning rate for fine-tuning. |
| `min_lr` | `0.0` | Floor learning rate after decay. |
| `lr_warmup_iters` | `50` | Linear warmup length in iterations. |
| `lr_decay_iters` | `null` | Optional decay span; `null` defers to recipe defaults. |

---

## 7. Precision

| Parameter | Example | Description |
|-----------|---------|-------------|
| `precision_config` | `bf16_mixed`, `fp16_mixed`, `fp32` | Mixed-precision recipe for Megatron Bridge. |
| `comm_overlap_config` | `null` | Optional communication/compute overlap policy object. |
| `pipeline_dtype` | `null` | Dtype for pipeline stages when PP is enabled. |

---

## 8. Memory optimization

| Parameter | Example | Description |
|-----------|---------|-------------|
| `recompute_granularity` | `full` | Activation recomputation granularity. |
| `recompute_method` | `uniform` | How recomputation is scheduled across layers. |
| `recompute_num_layers` | `1` | Number of layers per recompute group (workload-dependent). |

---

## 9. Primus-Turbo

From `sft_trainer.yaml` (defaults shown).

| Parameter | Default | Description |
|-----------|---------|-------------|
| `enable_primus_turbo` | `true` | Master flag for Primus-Turbo optimized kernels and paths. |
| `use_turbo_attention` | `false` | Turbo attention implementation. |
| `use_turbo_parallel_linear` | `false` | Turbo parallel linear layers. |
| `use_turbo_grouped_gemm` | `false` | Turbo grouped GEMM flag for MoE paths (the preset ships this key). The former `use_turbo_grouped_mlp` alias has been removed. |
| `moe_use_fused_router_with_aux_score` | `false` | Fused MoE router with auxiliary loss handling. |
| `enable_turbo_attention_float8` | `false` | FP8 path inside Turbo attention. |
| `use_turbo_deepep` | `false` | DeepEP-style expert-parallel integration. |
| `turbo_deepep_num_cu` | `32` | Compute-unit count hint for DeepEP. |
| `turbo_deepep_use_comm_stream` | `false` | Use dedicated communication streams. |
| `turbo_sync_free_moe_stage` | `0` | Sync-free MoE scheduling stage. |
| `use_turbo_fused_act_with_probs` | `false` | Fuse activation with probability tensors where applicable. |
| `use_turbo_rms_norm` | `false` | Turbo RMSNorm path. |

**Environment:** `PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND` (default `TURBO`) is read in `primus/backends/megatron/patches/args/rocm_arg_validation.py` to select MoE dispatch/combine behavior when Turbo MoE is active.

---

## 10. Model and dataset

Model YAML files (`qwen3_8b.yaml`, `qwen3_32b.yaml`, `llama31_70b.yaml`) supply:

| Parameter | Example | Description |
|-----------|---------|-------------|
| `recipe` | `qwen.qwen3`, `llama.llama3` | Recipe module under `megatron.bridge.recipes`. |
| `flavor` | `qwen3_8b_finetune_config`, `llama31_70b_finetune_config` | Flavor function producing the baseline `ConfigContainer`. |
| `hf_path` | `Qwen/Qwen3-8B`, `meta-llama/Meta-Llama-3.1-70B` | Hugging Face model id for weights/tokenizer flows. |
| `dataset` | nested | Example: `dataset_name: "rajpurkar/squad"` for SQuAD-style fine-tuning. |

**Logging (optional overrides in examples):** `wandb_project`, `wandb_entity`, `wandb_exp_name` might be set under `overrides` for experiment tracking when Weights & Biases is configured.

---

## Argument merge mechanics

`MegatronBridgeArgBuilder` (`primus/backends/megatron_bridge/argument_builder.py`) performs a **deep merge** of CLI and YAML into a single dict/namespace before `load_recipe_config` runs. Nested dicts (for example dataset or optimizer sections) combine recursively; explicit `None` in the merged structure can clear fields depending on merge rules in `_merge_dict_to_dataclass`.

---

## Example layouts

Under `examples/megatron_bridge/configs/`, per-GPU directories (for example `MI300X/`, `MI355X/`) contain full experiment YAMLs that set `work_group`, `user_name`, `exp_name`, `workspace`, and Megatron Bridge modules. Post-training examples use `modules.post_trainer` with `config: sft_trainer.yaml`; MI300X pretraining examples use `modules.pre_trainer` with `config: pretrain_trainer.yaml`. Both patterns set `framework: megatron_bridge`, `model: <preset>.yaml`, and an `overrides` block for parallelism, LR, precision, and related options.