This project is being deprecated and no new models will be supported. Please contribute your suggestions/ideas to github.com/nvidia-nemo/megatron-bridge. Thank you.
MBridge provides a seamless bridge between Hugging Face models and Megatron-Core's optimized implementation for efficient distributed training and inference. It also offers necessary tools and processes for integrating Reinforcement Learning (RL) with Megatron.
MBridge is a prototype project, the idea has been adopted as Megatron-Bridge. For more advanced features such as training loop, mixed precision(FP8, BF16, FP4 etc.), PEFT, please refer to Megatron-Bridge.
MBridge allows you to convert popular Hugging Face models to Megatron-Core format, enabling you to leverage advanced parallelism strategies for large-scale training and inference. The library supports various model architectures and simplifies the process of transitioning between these frameworks. For Reinforcement Learning workflows, MBridge provides interfaces and tools needed to connect RL algorithms with Megatron-optimized models.
pip install mbridge
from megatron.core import parallel_state as mpu
from mbridge import AutoBridge
# Initialize distributed environment
mpu.initialize_model_parallel(
tensor_model_parallel_size=tp,
pipeline_model_parallel_size=pp,
virtual_pipeline_model_parallel_size=vpp,
context_parallel_size=cp,
expert_model_parallel_size=ep,
)
# Load a model from Hugging Face
HF_MODEL_PATH = "/path/to/Qwen/Qwen2.5-7B-Instruct"
# or llama model
HF_MODEL_PATH = "/path/to/llama/llama3-8b-instruct"
bridge = AutoBridge.from_pretrained(HF_MODEL_PATH)
# Get a Megatron-Core model and load weights from Hugging Face
model = bridge.get_model(weight_path=HF_MODEL_PATH)
# Export weights back to Hugging Face format for inference engine
for key, weight in bridge.export_weights(model):
# Process or save the exported weights
print(f"Exported: {key}")
# save model with HF format
bridge.save_weights(model, "path/to/save/model", memory_efficient=False) # set memory_efficient=True if the model is vary large
Currently supported models:
The example directory contains scripts demonstrating common use cases:
0.load_model_and_generate_single_gpu.py: Loading a model and generating text on a single GPU1.load_model_and_export_single_gpu.py: Loading a model and exporting weights on a single GPU2.load_model_and_export_multiple_gpus.py: Loading a model and exporting weights using multiple GPUs with TP/PP/CP/VPP parallelismMBridge provides a set of post model creation callbacks to customize the model after it is created.
make_value_model: Add a value model to the modelfreeze_moe_router: Freeze the router of the modelfrom mbridge.utils.post_creation_callbacks import make_value_model, freeze_moe_router
bridge = AutoBridge.from_pretrained(HF_MODEL_PATH)
model = bridge.get_model(weight_path=HF_MODEL_PATH, post_model_creation_callbacks=[make_value_model, freeze_moe_router])
Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
(top 30 of 31)
Python
95.6%
JavaScript
2.2%
This project is being deprecated and no new models will be supported. Please contribute your suggestions/ideas to github.com/nvidia-nemo/megatron-bridge. Thank you.
MBridge provides a seamless bridge between Hugging Face models and Megatron-Core's optimized implementation for efficient distributed training and inference. It also offers necessary tools and processes for integrating Reinforcement Learning (RL) with Megatron.
MBridge is a prototype project, the idea has been adopted as Megatron-Bridge. For more advanced features such as training loop, mixed precision(FP8, BF16, FP4 etc.), PEFT, please refer to Megatron-Bridge.
MBridge allows you to convert popular Hugging Face models to Megatron-Core format, enabling you to leverage advanced parallelism strategies for large-scale training and inference. The library supports various model architectures and simplifies the process of transitioning between these frameworks. For Reinforcement Learning workflows, MBridge provides interfaces and tools needed to connect RL algorithms with Megatron-optimized models.
pip install mbridge
from megatron.core import parallel_state as mpu
from mbridge import AutoBridge
# Initialize distributed environment
mpu.initialize_model_parallel(
tensor_model_parallel_size=tp,
pipeline_model_parallel_size=pp,
virtual_pipeline_model_parallel_size=vpp,
context_parallel_size=cp,
expert_model_parallel_size=ep,
)
# Load a model from Hugging Face
HF_MODEL_PATH = "/path/to/Qwen/Qwen2.5-7B-Instruct"
# or llama model
HF_MODEL_PATH = "/path/to/llama/llama3-8b-instruct"
bridge = AutoBridge.from_pretrained(HF_MODEL_PATH)
# Get a Megatron-Core model and load weights from Hugging Face
model = bridge.get_model(weight_path=HF_MODEL_PATH)
# Export weights back to Hugging Face format for inference engine
for key, weight in bridge.export_weights(model):
# Process or save the exported weights
print(f"Exported: {key}")
# save model with HF format
bridge.save_weights(model, "path/to/save/model", memory_efficient=False) # set memory_efficient=True if the model is vary large
Currently supported models:
The example directory contains scripts demonstrating common use cases:
0.load_model_and_generate_single_gpu.py: Loading a model and generating text on a single GPU1.load_model_and_export_single_gpu.py: Loading a model and exporting weights on a single GPU2.load_model_and_export_multiple_gpus.py: Loading a model and exporting weights using multiple GPUs with TP/PP/CP/VPP parallelismMBridge provides a set of post model creation callbacks to customize the model after it is created.
make_value_model: Add a value model to the modelfreeze_moe_router: Freeze the router of the modelfrom mbridge.utils.post_creation_callbacks import make_value_model, freeze_moe_router
bridge = AutoBridge.from_pretrained(HF_MODEL_PATH)
model = bridge.get_model(weight_path=HF_MODEL_PATH, post_model_creation_callbacks=[make_value_model, freeze_moe_router])
Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
(top 30 of 31)
Python
95.6%
JavaScript
2.2%