# Megatron-HF Bridge Backend

AReaL currently supports two bridge backends for `MegatronEngine`:

- `mbridge` (default)
- `megatron-bridge`

Set the backend with:

```yaml
actor:
  megatron:
    bridge_type: mbridge
```

- Use `bridge_type=megatron-bridge` to enable the new path.
- `mbridge` is the default choice if this argument is not present

## Why this feature exists

- `mbridge` is being deprecated and does not provide PEFT/LoRA support.
- `megatron-bridge` supports more/ newer model architectures.
- `megatron-bridge` provides built-in PEFT/LoRA implementations.

## Recommendation

- For new GPU training workflows, prefer `megatron-bridge`.
- Keep `mbridge` for backward compatibility and environments that still depend on it.
- Prefer `mbridge` when using disk-based weight broadcast as it has optimized HF
  load/save path.
- If you use XCCL for weight broadcast, load/save time is less important.

## Current limitation

- Tree-attention training in `MegatronEngine` currently supports only `mbridge`.
- The `megatron-bridge` backend is not supported in the tree-attention path yet.
- `megatron-bridge` does support faster/optimized HF model load/save implementations.

## Train only the MTP head

To adapt a native MTP head to an already fine-tuned model, use the existing SFT trainer
with:

```yaml
actor:
  megatron:
    bridge_type: megatron-bridge
    enable_mtp: true
    enable_mtp_training: true
    mtp_only: true
    mtp_loss_scaling_factor: 0.1
```

`mtp_only` freezes all non-MTP parameters before distributed wrapping and optimizer
construction, including shared embeddings and the output projection. The ordinary SFT
loss remains in the graph to trigger Megatron-Core's auxiliary MTP backward; it does not
update the frozen backbone. Do not wrap the whole forward in `no_grad()`. The loss
coefficient still scales MTP gradients; it must be finite and positive.

The input checkpoint must contain both the desired fine-tuned backbone and a native MTP
head, with a matching model configuration. A checkpoint exported with MTP disabled needs
its MTP weights and configuration restored before loading. Preserve the fine-tuned
backbone, embeddings, output head, and tokenizer when restoring MTP.

For acceptance-rate recovery, use representative assistant continuations generated by
the frozen target model, with the usual SFT response mask. Validate that MTP parameters
update while every non-MTP parameter stays unchanged, then measure draft acceptance and
throughput on held-out requests. SFT loss alone does not establish speculative-decoding
speedup.

This mode requires one MTP prediction layer and the `megatron-bridge` backend. Pipeline
stages before the MTP stage remain fully frozen but still participate in the pipeline
schedule and global optimizer statistics. LoRA, critic models, FSDP wrappers, and MoE
router expert-bias updates are unsupported. Chunked LM-head loss remains unsupported
with MTP training. Full backbone forward computation is still required. Setting
`mtp_only: false` retains the default joint-training behavior.

Native MTP training requires Megatron-Core >=0.19.0 and Megatron-Bridge >=0.6.0. With
this stack, Megatron-Core derives MTP targets from the layout-aligned input IDs, while
AReaL aligns the response mask with each packed sequence's CP partition. Model-owned
Qwen THD forwards delegate the final CP slicing to Megatron-Bridge after multimodal
embedding fusion. For dense Qwen models, EP is 1.