BigMac - Xiaohongshuによるオープンソースのマルチモーダル大規模モデルパイプライン並列トレーニングフレームワーク - AiBoss

ビッグマックとは何ですか？

BigMacは、XiaohongshuのDots Infraチームが開発したオープンソースのマルチモーダル大規模モデルパイプライン並列トレーニングフレームワークです。MLLMトレーニングにおける計算効率とメモリ使用量のバランスを取るという課題に対し、安全なネスト型パイプラインパラダイムを提案しています。このフレームワークは、成熟したLLMパイプラインを基盤とし、エンコーダーとジェネレーターの計算を組み込みながらも実行順序を崩しません。これにより、メモリの安定性を維持しながらトレーニング速度を1.08～1.9倍向上させています。Dotsマルチモーダルトレーニングの中核コンポーネントとして、既に本番環境に導入されています。

ビッグマックの主な機能

安全なネストされたパイプラインへの依存LLMパイプラインを安定した基盤として使用し、エンコーダー/ジェネレーターの計算は、モジュール間のバブル干渉を回避するために、依存関係が満たされる場所に埋め込まれています。
ビデオメモリのバウンドモードアクティベーションエンコーダとジェネレータはビデオメモリをアクティブ化します。 O(1)LLMのアクティベーション動作は変更されておらず、メモリ不足（OOM）を起こすことなく、より大きなバッチサイズをサポートします。
グローバルスケジューリングと実行の分離スケジューラは、すべてのランク、マイクロバッチ、およびモジュールを網羅するグローバルなオペレータテーブルを生成します。エグゼキュータは、ローカルシーケンスを解釈し、対応するバックエンドに配信します。
生産ライン向けのシームレスなインターフェースアルゴリズムエンジニアはモジュールの入出力境界を記述するだけでよく、BigMacがスケジューリング、ハンドオフ、デバイス間通信を自動的に処理するため、モデルコードのモジュール性を維持できます。
スケジュールを考慮したツールチェーンプロファイラ、シミュレータ、および視覚化ツールを内蔵。オペレーターごとのトレースエクスポートをサポート（Perfetto UIと互換性あり）。バブル、処理の遅いオペレーター、および通信待ちを迅速に特定します。

ビッグマックの技術的原理

安全なネストされたパイプラインスケジューリングへの依存BigMacはこのタイムラインを基本として利用します。スケジューラは、エンコーダの順方向、LLMの順方向、ジェネレータの順方向、およびそれらに対応する逆方向演算間のデータ依存関係を分析し、入力が準備できており、LLMの実行順序を妨げない「ギャップ」期間にのみモーダル計算を挿入します。LLMは元のスケジュールに従って継続され、勾配が準備でき次第、エンコーダとジェネレータの起動を解除できます。逆方向演算は迅速に実行され、「計算効率は必然的にGPUメモリ使用量の増加につながり、GPUメモリ使用量の増加は必然的にバブルの増加につながる」というパレートのトレードオフを打破します。
グローバルなオペレーターテーブルをバックエンドから切り離すシステムは、スケジューリング戦略をグローバルな演算子テーブルとして明示的にカプセル化します。このテーブルには、すべてのパイプラインランクにわたるすべてのLLM/エンコーダ/ジェネレータの順方向/逆方向演算子、およびモジュール境界通信が含まれます。実行レイヤーはローカル演算子シーケンスのみを解釈します。LLM演算子はメガトロンコアを再利用し、モーダル演算子は独自のデータ並列/FSDP実装を保持します。スケジュールの調整には、モデルランタイムの書き換えは必要ありません。

WeChatでフォローして「オープンソース「、参加するAIオープンソースプロジェクトに関するディスカッショングループ

BigMacの使い方

モジュール境界を定義するencoder_forward、llm_forward、generator_forwardの各関数をモジュール式に記述し、各モジュールの入力と出力を宣言してください。
BigMacインターフェースに接続しますBigMacは、PP透過インターフェースを介してモジュールを登録することで、パイプラインのステージ分割、アクティベーションの引き継ぎ、グラジエントの引き継ぎ、およびランク間通信を自動的に処理します。
並列戦略を構成するPP/VPP、マイクロバッチ数、モジュール配置戦略などのパラメータを設定します。
トレーニングを実施するスケジューラが起動してグローバルプランを生成し、エグゼキュータが各ランクで演算子シーケンスを実行します。
パフォーマンス診断内蔵のプロファイラを使用してオペレーターごとの処理時間を収集し、ランクに合わせたタイムラインをPerfetto UIにエクスポートして、バブルやボトルネックを分析できます。または、シミュレーターを使用して、トレーニングを開始する前に並列処理戦略を迅速に反復できます。

ビッグマックの主な利点

パレート最適解の打破高い演算効率と低いビデオメモリ使用量を両立させ、どちらか一方を選択する必要がなくなる。
干渉のないLLMパイプライン成熟したLLMスケジュールを維持します。エンコーダー/ジェネレーターのレイテンシ変動はパイプライン全体に伝播しません。
ビデオメモリ拡張可能エンコーダー/ジェネレーターはGPUメモリをO(1)でアクティブ化し、本番レベルの大規模バッチトレーニングをサポートします。
アルゴリズムフレンドリーモデルコードはパイプラインの実行時間を意識する必要がなく、シングルカード検証の実験はパイプライン並列処理にも自然に拡張できる。
デバッグ可能でデプロイ可能グローバルスケジュールは可視化され、検査可能であり、付属のトレース／シミュレーションツールチェーンにより、パフォーマンス最適化のコストが削減されます。

BigMacのプロジェクトアドレス

プロジェクト公式サイト：https://dots-infra.github.io/BigMac/
GitHubリポジトリ：https://github.com/Dots-Infra/BigMac/
arXiv技術論文：https://arxiv.org/pdf/2605.25451

ビッグマックの競合製品比較

寸法	BigMac	Megatron-LM
スケジューリングパラダイム	LLMを基盤とした、安全なネスト型パイプラインに依存しています。	モーダルモジュールを独立したステージとして構成した、高効率なビデオメモリ用シングルパイプライン。
計算効率	高い、モジュール間のバブルなし	このような状況では、エンコーダー／ジェネレーターの変動によってパイプラインバブルが発生する可能性がある。
ビデオメモリ使用量	安定、モード活性化 O(1)	性能は低いが、モジュール間の結合によりメモリ最適化の余地が限られている。
拡張性	理解と生成を組み合わせた二段階アプローチを用いた大規模トレーニングを支援する。	生成経路の拡張は、バブルのサイズと記憶圧を大幅に増加させる。
インターフェースの使いやすさ	PP透過性、ゼロ侵入アルゴリズムコード	パイプラインの実行時間とステージ分割を理解し、それに合わせて対応することが必要です。
ツールチェーン	スケジュールを考慮したプロファイラとシミュレーターを内蔵	従来のnsys/torchトレースの手動相関分析に依存する

BigMac アプリケーションシナリオ

マルチモーダル理解モデルの事前学習例えば、テキストや画像の理解、動画の理解のためのMLLMのエンドツーエンドの事前学習では、ViT/WhisperなどのエンコーダとLLM間のパイプラインによる連携を効率的に処理する必要がある。
マルチモーダル生成モデルのトレーニングまた、理解と生成の経路（画像生成や音声合成など）も含まれており、LLMとMMDiT/LDMなどのジェネレーター間の双方向の依存関係の調整が必要となる。
大規模バッチトレーニングメモリ制約のあるクラスタ上でより大きなグローバルバッチサイズを追求し、従来の高効率コンピューティング設計におけるアクティベーション蓄積によって引き起こされるメモリ不足を回避するため。
迅速な実験と反復アルゴリズムチームは、シングルカード環境またはDP環境でモデル構造を検証し、コードを書き直すことなく、BigMacインターフェースを介してパイプライン並列処理にシームレスに拡張できます。
トレーニングパフォーマンスの調整エンジニアリングチームは、スケジュールを考慮したプロファイラとシミュレータを使用して、パイプラインのボトルネックを特定し、並列構成とモジュール配置戦略を最適化した。