Megatron-HF Bridge Backend

Megatron-HF Bridge Backend#

AReaL currently supports two bridge backends for MegatronEngine:

  • mbridge (default)

  • megatron-bridge

Set the backend with:

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:

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.