NVIDIA-NeMo/Megatron-Bridge

Training library for Megatron-based models with bidirectional Hugging Face conversion capability

NVIDIA‑NeMo / Megatron‑Bridge – 简介

NeMo Megatron Bridge 是一个内置于 NVIDIA NeMo 框架中的 PyTorch 原生库。其目的是作为两个大型模型生态系统之间的桥梁：

🤗 Hugging Face 模型格式（开源 LLM、视觉语言模型等的行业标准）
Megatron‑Core（源自 NVIDIA Megatron‑LM 项目的高性能张量与流水线并行训练引擎）。

此桥接工具让您可以双向转换检查点、验证转换正确性，并利用 Megatron‑Core 的强大并行能力（支持张量、流水线及混合精度）进行模型训练或微调。它同时提供许多最先进模型的现成训练“配方”（如 Llama 3.2、GLM‑5、Qwen‑3、Nemotron‑3、Gemma‑4 VL 等）。

为何它很重要

速度与规模 – Megatron‑Core 可在数十个 GPU 上运行万亿参数模型，并具备 CUDA‑graph 执行、FP8/MXFP4 量化及上下文并行等高级功能。此桥接工具让您能将任何 Hugging Face 模型带入该领域。
一站式工作流程 – 转换公开的 HF 检查点 → Megatron 检查点 → 训练/微调 → 导出回 HF，全程使用相同的 Python API。
广泛的模型覆盖 – 仓库追踪数十种近期 LLM、多模态模型与扩散模型的首日支持，并频繁更新（请参阅新闻区块）。
社区就绪 – 包含 Docker 容器、CI、代码覆盖率徽章，以及供 AI 编程辅助工具使用的“技能”目录，协助您新增模型。

核心功能 (如 README 所述)

类别	桥接工具的功能
转换	Hugging Face 与 Megatron 格式之间的双向检查点转换；权重一致性验证。
训练	基于 Megatron‑Core 的 PyTorch 原生训练循环；支持预训练、监督式微调 (SFT) 及 LoRA/PEFT 微调。
并行处理	张量模型并行、流水线模型并行及上下文并行执行；适用于单 GPU 或多 GPU 集群。
精度	混合精度训练/推理：BF16、FP8、FP4、MXFP4 等。
配方	适用于多种模型的现成配置文件 (例如 `llama32_1b_pretrain_1gpu_h100_bf16_config`)。
多模态	支持视觉语言、音频及扩散模型 (例如 Gemma‑4 VL、Nemotron‑Diffusion)。
PEFT	LoRA 及其他参数高效微调工作流程已内置于配方中。
性能工具	性能配方、确定性执行、CUDA‑graph 优化及与 MLPerf 风格基准测试的集成。
容器支持	官方 NeMo Framework Docker 镜像提供“最佳体验”。
扩展性	Megatron‑Core 子模块（可在 `main` 与 `dev` 之间切换）；用于 AI 助手指导的 `skills/` 目录。

典型工作流程 (快速入门风格)

选择模型 – 例如来自 Hugging Face 的 meta-llama/Llama-3.2-1B。

from megatron.bridge import AutoBridge
bridge = AutoBridge.from_hf_pretrained("meta-llama/Llama-3.2-1B", trust_remote_code=True)

创建 Megatron 提供者，设置并行大小，并实例化模型。
训练 – 执行提供的配方 (pretrain(cfg, forward_step)) 或编写您自己的训练循环。
导出 – bridge.save_hf_pretrained(...) 可写入完整的 HF 文件夹，或使用 export_hf_weights 流式传输权重。

所有这些都可以在 NeMo 容器内使用 torch.distributed.run 启动。

谁应该使用它？

研究人员：需要大规模实验最新 LLM，并希望利用 NVIDIA 的并行技术而无需重写模型者。
工程师：建立必须在多 GPU 集群上运行的自定义微调流水线 (SFT, LoRA) 者。
ML‑ops 团队：希望为大规模训练和检查点转换建立可重现、容器化环境者。

安装与入门

首选：拉取官方 NeMo Framework Docker 镜像 (nvcr.io/nvidia/nemo:<TAG>) 并依照 README 说明运行容器。
开发：克隆仓库，初始化 Megatron‑Core 子模块，并使用 uv sync (或 pip install -e .) 安装 Python 包。
切换 Megatron‑Core 版本 – scripts/switch_mcore.sh 让您在稳定的 main 提交与实验性的 dev 分支之间切换。

目前模型覆盖范围 (精选)

LLMs: Llama 3.2, GLM‑5.2, Qwen 3.5/3.6, DeepSeek V4, Falcon H1, Gemma 4, Nemotron‑3 Super/Ultra/Nano Omni, K‑EXAONE‑2, MiniMax‑M2.5, Kimi‑K2.5 VL 等。
多模态: Gemma 4 VL, Nemotron‑Diffusion, Qwen‑3 ASR, Bailing MoE V2。
通过「Nemotron‑Labs Diffusion」系列支持扩散与生成式视觉模型。

授权与社区

本仓库依据 NVIDIA Software License 开源（请参阅 LICENSE）。
CI 徽章显示测试覆盖率与 CI 状态；欢迎通过 CONTRIBUTING.md 指南贡献代码。

总结： NeMo Megatron Bridge 是一个真正且积极维护的开源项目，利用 NVIDIA 的 Megatron‑Core 引擎，为现代 LLM 与多模态模型实现大规模训练、微调与检查点转换。