Xiaohongshuが「BigMac」を発表：マルチモーダル学習におけるメモリ容量と処理速度のトレードオフを打破

マルチモーダル大規模言語モデルは、次世代AIの基盤として台頭してきていますが、その学習インフラは長らく、根深いトレードオフに阻まれてきました。すなわち、速度を最適化したシステムはメモリの制約に悩まされがちであり、一方でメモリ効率を優先したシステムは処理速度が遅くなりがちです。 XiaohongshuのDots Infraチームは、このボトルネックをマルチモーダルパイプライン学習におけるパレート・フロンティアと位置づけ、ついにその壁を打ち破りました。 7月22日、同チームは、ネイティブマルチモーダルシナリオ向けに特別に設計された革新的なパイプライン並列学習パラダイム「BigMac」をオープンソース化しました。このプロジェクトは現在、GitHubの「Dots-Infra」リポジトリで公開されています。

標準的なトランスフォーマーとは異なり、マルチモーダルモデルは本質的に複雑です。典型的なマルチモーダル大規模言語モデル（MLLM）は、画像や音声を埋め込みベクトルに変換するモーダルエンコーダー、推論用のLLMバックボーン、そしてLLMの出力を画像や音声などのターゲットモダリティにマッピングするジェネレータという、3つの異なるコンポーネントで構成されています。 これらのコンポーネント間の構造的な違いは大きく、それらを単一のトレーニングパイプラインに統合する際に大きな課題となります。

現在の業界のソリューションは、一般的に2つのカテゴリーに分類されます。1つ目のアプローチは、エンコーダーとジェネレーターをLLMパイプラインから切り離し、別々に実行することで、計算効率を優先します。これにより、モーダルモジュールの処理時間の変動によってLLMパイプラインに「バブル」が生じるのを防げますが、アクティベーションメモリがマイクロバッチ数に応じて増加するため、大規模化に伴いコストが高くなります。 2つ目のアプローチは、すべてのモジュールを同一のパイプライン内に保持することでメモリ効率に重点を置いており、これによりアクティベーションのライフサイクルが短縮され、メモリ使用量が削減されます。しかし、エンコーダーやジェネレーターのいずれかが遅くなると、LLMパイプライン全体が停滞し、テールバブルが発生してしまいます。モデルの規模が大きくなるにつれ、どちらの設計にも重大なボトルネックが現れます。

BigMacは、シンプルでありながら極めて重要な原則によってこの課題に対処しています。それは、中核となるLLMパイプラインを最優先とするというものです。大規模LLMのトレーニングでは、すでに1F1Bやインターリーブ型1F1Bといった成熟したスケジューリング戦略が採用されており、これらは本番環境レベルのトレーニングスタックに深く統合されています。 BigMacは、これらの確立された手法を置き換えるのではなく、LLMのスケジューリングを基盤となるタイムラインとして利用し、入力の準備が整ったタイミングで、LLMの実行順序を乱すことなく、エンコーダとジェネレータの計算を戦略的に挿入します。チームはこのアーキテクチャを「準依存性セーフなネストされたパイプライン」と呼んでいます。

この設計には2つの重要な利点がある。第一に、エンコーダーとジェネレーターの実行時間の変動がLLMパイプライン全体に波及しなくなるため、LLMは最適なペースを維持できる。第二に、モーダル活性化に必要なメモリ要件がアルゴリズムレベルでO(1)に削減される。 エンコーダーは、パイプラインが終了するまですべてのマイクロバッチの活性化状態を保持する必要がなくなり、ジェネレーターは長い活性化状態のテールを延長することを回避できる。本質的に、BigMacはパフォーマンス向上のためにメモリを犠牲にするのではなく、スケジューリングロジックを再構築することで、これら2つの目標を相互に排他的なものではなく、両立可能なものにする。

スケジューリング設計と実際の実行との間のギャップを埋めるには、システム統合、インターフェース定義、パフォーマンスのデバッグなど、大きな技術的課題が伴います。BigMacは、3つの中核的な機能を提供することで、これらの具体的な課題に対処するように設計されています。 第一に、グローバルなスケジュールを可視化します。スケジューラは、実行時にすべてのパイプラインランク、マイクロバッチ、およびモジュールタイプを網羅する包括的なオペレータテーブルを生成します。 その後、エグゼキュータがこれらを各ランクごとのローカルシーケンスに分散し、Megatron Core のような LLM バックエンド、モーダルランタイム、通信レイヤーとシームレスに統合します。この可視性により、チェックやバックエンドの最適化が容易になります。

第二に、BigMacはアルゴリズムエンジニアには見えない形で機能するパイプライン並列インターフェースを提供します。ユーザーは各モジュールが何を生成し、何を消費するかを定義するだけで済み、ステージの分割、活性化および勾配の引き渡し、デバイス間通信はシステムが内部で処理します。これにより、単一のGPU上で検証されたマルチモーダル実験を、パイプライン並列処理へと自然に拡張することが可能になります。 第三に、BigMacはプロファイリング、シミュレーション、可視化ツールのスイートを提供します。これらのツールはトレーニングの反復処理をオペレーターレベルまで分解し、各タイムステップで各ランクが何を行っているかを正確に明らかにするとともに、アイドル期間を特定し、依存関係のボトルネックを強調表示します。 また、このシミュレーターを使用することで、チームは本格的なトレーニングに着手する前に、さまざまなPP構成やマイクロバッチの組み合わせをテストし、それらがバブルやスループットに与える影響を推定することができます。

BigMacの有効性は、2つの代表的なワークロードにおいて検証されています。MLLM-Understandingタスクにおいて、バックボーンとしてQwen3-30B-A3Bを使用し、1.3BのViTエンコーダを採用した場合、 BigMacは、計算効率に優れたベースラインであるOptimusと比較して1.08倍から1.1倍の高速化を達成し、メモリ効率に優れたベースラインであるMegatron-DistTrainと比較して1.6倍から1.9倍の高速化を実現しました。 重要な点として、GPU あたりのバッチサイズが増加してもメモリ使用量は安定しているのに対し、Optimus ではメモリ使用量がピーク時に急増し、最終的に大きなバッチサイズでメモリ不足（OOM）エラーが発生する。 MLLM-Generationタスクはより複雑で、20BのMMDiTジェネレータを伴うが、ここでは性能差がさらに拡大する。OptimusはOOMによりテストされたすべてのバッチサイズで失敗する一方、BigMacはジェネレータを逆方向に効率的に実行し、活性化値を迅速に解放することでこれを回避している。 Megatron-DistTrainと比較しても、BigMacは安定したメモリ使用量で1.5倍から1.9倍の高速化を実現しています。これは、ネストされたパイプラインの主な価値、すなわち、過度な活性化データの保持や著しいアイドル時間を生じさせることなく、エンコーダとジェネレータの両方の依存関係に対応できることを浮き彫りにしています。

BigMacは現在、Dotsのマルチモーダルモデル学習インフラの中核コンポーネントとなっており、Xiaohongshuの本番環境でも稼働しています。 付随する論文「BigMac: Breaking the Pareto Frontier of Compute and Memory in Multimodal LLM Training」は、インタラクティブなPP Profilerトレース例とともにarXivで公開されています。 マルチモーダル学習におけるメモリと速度のトレードオフに悩む研究者やエンジニアにとって、この実運用で実証済みのオープンソースプロジェクトは、時間と労力を大幅に節約する実用的な解決策を提供します。