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ネイティブライブラリです。その目的は、大規模モデルのための2つの主要なエコシステム間のブリッジとして機能することです。

  1. 🤗 Hugging Face モデル形式(オープンソースLLM、視覚言語モデルなどの事実上の標準)
  2. Megatron‑Core(NVIDIAのMegatron‑LMプロジェクトから派生した、高性能なテンソルおよびパイプライン並列トレーニングエンジン)。

このブリッジを使用すると、チェックポイントの双方向変換、変換の正確性の検証、そしてMegatron‑Coreの強力な並列処理(テンソル、パイプライン、混合精度サポート)を活用したモデルのトレーニングや微調整が可能になります。また、多くの最先端モデル(Llama 3.2、GLM‑5、Qwen‑3、Nemotron‑3、Gemma‑4 VLなど)向けのすぐに使えるトレーニング「レシピ」も同梱されています。


なぜ重要なのか

  • 速度とスケール – Megatron‑Coreは、CUDA‑graph実行、FP8/MXFP4量子化、コンテキスト並列などの高度な機能を備え、数十台のGPUで兆単位のパラメータを持つモデルを実行できます。このブリッジにより、あらゆるHugging Faceモデルをその世界に持ち込むことができます。
  • ワンストップワークフロー – 公開されているHFチェックポイントの変換 → Megatronチェックポイント → トレーニング/微調整 → HFへのエクスポートまで、すべて同じPython APIで行えます。
  • 幅広いモデル対応 – リポジトリは、数十種類の最新LLM、マルチモーダルモデル、拡散モデルのDay-0サポートを追跡しており、頻繁に更新されています(ニュースセクションを参照)。
  • コミュニティ対応 – Dockerコンテナ、CI、コードカバレッジバッジ、およびAIコーディングアシスタントが新しいモデルを追加するのを支援するための「skills」ディレクトリが含まれています。

主な機能 (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/ ディレクトリ。

一般的なワークフロー (クイックスタート)

  1. モデルを選択 – 例: Hugging Faceの meta-llama/Llama-3.2-1B。
  2. from megatron.bridge import AutoBridge
    bridge = AutoBridge.from_hf_pretrained("meta-llama/Llama-3.2-1B", trust_remote_code=True)
    
  3. Megatronプロバイダーを作成し、並列サイズを設定してモデルを実体化します。
  4. トレーニング – 提供されたレシピ (pretrain(cfg, forward_step)) を実行するか、独自のトレーニングループを作成します。
  5. エクスポート – 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やマルチモーダルモデルの大規模トレーニング、微調整、チェックポイント変換を可能にする、真にアクティブにメンテナンスされているオープンソースプロジェクトです。

関連

  • Dispatch
  • プロジェクト
  • Dispatch
  • プロジェクト
  • プロジェクト