Megatron-Bridge#

RLinf supports Megatron-Bridge through the Megatron-LM training backend. This integration lets users start Megatron-LM training directly from HuggingFace-format checkpoints, use model architectures supported by Megatron-Bridge, and keep RLinf’s training loop, data pipeline, logging, and checkpoint workflow unchanged.

Use Megatron-Bridge when:

  • the actor-side model is large and FSDP or FSDP2 becomes a performance bottleneck;

  • the model architecture is not yet supported by RLinf’s native Megatron-LM integration.

Megatron-Bridge resources:

Environment Setup#

MBridge currently uses RLinf’s agentic environment. Install the base environment first:

bash requirements/install.sh agentic
source .venv/bin/activate

Install the extra Python packages required by the MBridge path:

uv pip install transformers==4.57.1 bitsandbytes

The reasoning image does not include the megatron.bridge package. Clone Megatron-Bridge and the matching Megatron-LM revision, then add both source trees to PYTHONPATH:

export MBRIDGE_ROOT=/path/to/Megatron-Bridge-0.3.0
export MEGATRON_LM_ROOT=/path/to/Megatron-LM-b0cc2706ddc60d2aefd5fff346445b5c013036a8

mkdir -p "$(dirname "${MBRIDGE_ROOT}")" "$(dirname "${MEGATRON_LM_ROOT}")"
git clone --branch v0.3.0 https://github.com/NVIDIA-NeMo/Megatron-Bridge.git "${MBRIDGE_ROOT}"
git clone https://github.com/NVIDIA/Megatron-LM.git "${MEGATRON_LM_ROOT}"
git -C "${MEGATRON_LM_ROOT}" checkout b0cc2706ddc60d2aefd5fff346445b5c013036a8

export PYTHONPATH="${MBRIDGE_ROOT}/src:${MEGATRON_LM_ROOT}:${PYTHONPATH}"
export CUDA_DEVICE_MAX_CONNECTIONS=1
python -c "from megatron.bridge import AutoBridge; print('Megatron-Bridge OK')"

If your cluster image already mounts these repositories, keep the same PYTHONPATH exports and skip the two git clone commands.

Download the model and dataset used by the reasoning example:

# For faster downloads in mainland China, you can set:
# export HF_ENDPOINT=https://hf-mirror.com
hf download deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B \
    --local-dir /path/to/model/DeepSeek-R1-Distill-Qwen-1.5B

mkdir -p /dataset/boba
hf download inclusionAI/AReaL-boba-Data AReaL-boba-106k.jsonl \
    --repo-type dataset \
    --local-dir /dataset/boba

Download the model and dataset used by the VLM SFT example:

# For faster downloads in mainland China, you can set:
# export HF_ENDPOINT=https://hf-mirror.com
hf download Qwen/Qwen2.5-VL-3B-Instruct \
    --local-dir /path/to/Qwen2.5-VL-3B-Instruct

hf download keplerccc/Robo2VLM-1 \
    --repo-type dataset \
    --local-dir /path/to/Robo2VLM-1

Warning

The Robo2VLM download places training and evaluation files in the same directory, for example train-00000-of-00262.parquet and test-0000X-of-00003.parquet. Move them into separate directories before training. Otherwise, RLinf will read the whole dataset as training data.

The example SFT config expects /path/to/Robo2VLM-1/data and /path/to/Robo2VLM-1/eval_data. If you store the dataset elsewhere, update data.train_data_paths and data.eval_data_paths accordingly.

Overview#

After MBridge is enabled, RLinf imports and builds the Megatron-LM model through Megatron-Bridge instead of relying on the traditional Megatron checkpoint conversion workflow.

The key configuration is different for reasoning/RL and SFT tasks.

For reasoning tasks:

actor:
  training_backend: megatron
  megatron:
    mbridge: True
    use_hf_ckpt: True
    ckpt_convertor:
      hf_model_path: /path/to/model/DeepSeek-R1-Distill-Qwen-1.5B

When actor.megatron.mbridge is True and use_hf_ckpt is True, RLinf reads the model path from actor.megatron.ckpt_convertor.hf_model_path and lets MBridge build the Megatron model provider.

For SFT tasks:

actor:
  training_backend: megatron
  model:
    model_path: /path/to/Qwen2.5-VL-3B-Instruct
    megatron_checkpoint: null
  megatron:
    use_hf_ckpt: True
    mbridge: True

When actor.megatron.mbridge is True, RLinf reads the model path from actor.model.model_path and lets MBridge build the Megatron model provider.

Quick Start#

  1. Export the MBridge paths before launching training:

export PYTHONPATH=/path/to/Megatron-Bridge-0.3.0/src:$PYTHONPATH
export PYTHONPATH=/path/to/Megatron-LM-b0cc2706ddc60d2aefd5fff346445b5c013036a8:$PYTHONPATH
export CUDA_DEVICE_MAX_CONNECTIONS=1

  1. Prepare the HuggingFace model and data directories:

# Reasoning tasks need:
/path/to/model/DeepSeek-R1-Distill-Qwen-1.5B
/dataset/boba/AReaL-boba-106k.jsonl
# SFT tasks need:
/path/to/Qwen2.5-VL-3B-Instruct
/path/to/Robo2VLM-1/data
/path/to/Robo2VLM-1/eval_data

  1. Update the model, tokenizer, and dataset paths in the config.

Path Differences#

MBridge reads HuggingFace checkpoint paths from different config entries for different training tasks:

  • Reasoning / RL tasks usually read the HuggingFace model path from actor.megatron.ckpt_convertor.hf_model_path;

  • SFT tasks usually read the HuggingFace model path from actor.model.model_path;

  • the tokenizer path is still specified by actor.tokenizer.tokenizer_model. We recommend keeping it consistent with the model directory.

Therefore, do not only copy mbridge: True when migrating configs. Also check whether the model path is configured in the entry used by the current task type.

Reasoning task example:

actor:
  tokenizer:
    tokenizer_model: "/path/to/model/DeepSeek-R1-Distill-Qwen-1.5B"
  training_backend: megatron
  megatron:
    mbridge: True
    use_hf_ckpt: True
    ckpt_convertor:
      hf_model_path: /path/to/model/DeepSeek-R1-Distill-Qwen-1.5B

data:
  train_data_paths: ["/dataset/boba/AReaL-boba-106k.jsonl"]
  val_data_paths: ["/dataset/boba/AReaL-boba-106k.jsonl"]

SFT example:

actor:
  model:
    model_type: "qwen2.5_vl"
    model_path: "/path/to/Qwen2.5-VL-3B-Instruct"
    megatron_checkpoint: null

  tokenizer:
    tokenizer_model: "/path/to/Qwen2.5-VL-3B-Instruct"

  megatron:
    use_hf_ckpt: True
    mbridge: True

data:
  train_data_paths: "/path/to/Robo2VLM-1/data"
  eval_data_paths: "/path/to/Robo2VLM-1/eval_data"

  1. Launch the corresponding training script.

Start reasoning training from the repository root. Use the Python entrypoint so the MBridge override is applied explicitly:

python examples/reasoning/main_grpo.py \
    --config-path "$(pwd)/examples/reasoning/config/math" \
    --config-name qwen2.5-1.5b-grpo-megatron \
    +actor.megatron.mbridge=True

Start VLM SFT training from the repository root:

bash examples/sft/run_vlm_sft.sh qwen2_5_vl_megatron_sft_vlm

Checkpoint Loading#

When Megatron-Bridge is used in RLinf, RLinf saves both checkpoint formats:

  • HuggingFace checkpoint;

  • Megatron checkpoint.

The checkpoint directory is organized as follows:

/path/to/logs/qwen2.5-1.5b-grpo-megatron/checkpoints/
β”œβ”€β”€ global_step_10/
β”‚   └── actor/
β”‚       β”œβ”€β”€ hf_model/
β”‚       β”‚   β”œβ”€β”€ model.safetensors
β”‚       β”‚   └── tokenizer.json
β”‚       β”œβ”€β”€ iter_0000010/
β”‚       β”‚   β”œβ”€β”€ mp_rank_00/
β”‚       β”‚   β”‚   β”œβ”€β”€ distrib_optim.pt
β”‚       β”‚   β”‚   └── model_optim_rng.pt
β”‚       β”‚   └── mp_rank_01/
β”‚       β”‚       β”œβ”€β”€ distrib_optim.pt
β”‚       β”‚       └── model_optim_rng.pt
β”‚       └── latest_checkpointed_iteration.txt
└── global_step_20/
    └── ...

The hf_model directory stores HuggingFace-format model weights and tokenizer files. The iter_XXXXXXX directory stores Megatron model weights and optimizer states. latest_checkpointed_iteration.txt records the latest checkpointed iteration. In this example, global_step_10/ and global_step_20/ are two different checkpoints for step 10 and step 20.

For resume training, you can load only the Megatron checkpoint. The HuggingFace-format checkpoint is not required.

runner:
  resume_dir: /path/to/logs/qwen2.5-1.5b-grpo-megatron/checkpoints/global_step_10

Practical Notes#

  • Keep actor.model.megatron_checkpoint: null when use_hf_ckpt: True.

  • Set actor.megatron.use_hf_ckpt: False only when loading a prepared Megatron checkpoint.

  • For Qwen3-VL models, keep actor.model.apply_rope_fusion: False.

  • For Qwen2.5 models, qkv_bias is forced on for model compatibility.

  • For Qwen3 models, qk_layernorm is forced on for model compatibility.

  • Make sure the tokenizer path matches the HuggingFace model directory.

Troubleshooting#

model.megatron_checkpoint is required if use_hf_ckpt is False

use_hf_ckpt is disabled, but no Megatron checkpoint path was provided. Set actor.megatron.use_hf_ckpt: True or provide runner.resume_dir.

model.megatron_checkpoint should be None if use_hf_ckpt is True

HuggingFace loading and Megatron checkpoint loading are both enabled. Set actor.model.megatron_checkpoint: null.

Qwen3-VL fails with a deepstack_visual_indexes assertion

The model’s visual deepstack configuration does not match the current pipeline split. First try pipeline_model_parallel_size: 1. If pipeline parallelism is required, make sure the first language pipeline stage has enough layers to contain all deepstack_visual_indexes. If you are using a reduced-layer checkpoint, also verify that the visual deepstack configuration matches the number of language model layers.