マルチモーダル大規模言語モデルは、次世代AIの基盤となっており、そのトレーニングシステムは長い間、古い問題に苦しんでいます。つまり、計算が早いものはメモリを十分に確保できないし、メモリを節約するものでは速度が落ちてしまいます。小紅書（Xiaohongshu）のdots infraチームはこの二難を「マルチモーダルパイプライントレーニングのパレート前線」と呼んでおり、現在彼らはこの前線に割り込みをかけました。7月22日、チームは「BigMac」という新しいパイプライン並列トレーニングの枠組みを正式にオープンソース化しました。これは元のマルチモーダル環境に特化したものです。オープンソースのアドレスはGitHubのDots-Infra配下に掲載されています。

マルチモーダルモデルは常に整ったtransformerではありません。典型的なMLLMは3つの部分で構成されます：画像や音声をembeddingに変換するモダリティエンコーダー、それをもとに推論を行うLLM本体、そしてLLMの出力を画像や音声などのターゲットモダリティに戻す生成器です。これらの3つの部分は形が大きく異なり、同じパイプラインに入れようとすると問題が起きます。

image.png

業界ではこれまで通常、2つの道がありました。1つは計算効率が高い方法で、エンコーダーと生成器をLLMパイプラインから外して別に実行します。この方法の利点は、モダリティモジュールの処理時間がLLMパイプラインで空泡を生まないことです。ただしデメリットとして、アクティブメモリがmicrobatch数に伴って膨らむため、生産規模では特に高コストになります。もう1つはメモリ効率が高く、すべてのモジュールを一つのパイプラインに詰め込んで先頭と末尾のフェーズを行います。これによりアクティブライフサイクルが短くなりメモリ使用量も減少しますが、エンコーダーまたは生成器が遅ければ全体のLLMパイプラインが待たされるため、末尾の空泡が発生します。規模が大きくなると、どちらの設計でもボトルネックが露呈します。

BigMacの出発点は単純ながらも重要なものです。スケジューリングの主幹は、すでに高度に最適化されたLLMパイプラインであるべきです。大規模LLMトレーニングでは1F1Bやinterleaved1F1Bといった成熟したスケジュールがすでに使われており、これらは生産レベルのトレーニングスタックと深く結びついています。BigMacはそれらを置き換えるのではなく、LLMスケジュールをベースタイムラインとして利用し、エンコーダーと生成器の計算を入力が準備できているかつLLMの実行順序を乱さない場所に挿入します。チームはこの設計を「準依存安全なネストパイプライン」と呼びます。

この設計によって得られる2つの性質があります。第一に、エンコーダーと生成器の処理時間の波動がLLMパイプラインを通じて伝播しなくなり、LLMは自分自身の本来のペースで進み続けます。第二に、モダリティアクティブメモリはアルゴリズム面でO(1)に抑えられます。エンコーダーはすべてのmicrobatchのアクティブをパイプライン終了まで保持する必要がなくなり、生成器も長く続くアクティブの尾を引きずる必要がありません。言い換えれば、BigMacはメモリと空泡の間で交換をしているのではなく、スケジュール構造を書き換えて、これらの2つの目標が互いに排除されなくなったのです。

image.png

スケジュールを設計できるところから実際にトレーニングできるところまでの間に、システム統合、インターフェース定義、性能調査という一連の工学的なハードルがあります。BigMacはこの関門を突破するために作られました。ここでは3つの要素を提供しています。まず、全体のスケジュールを明示的にします。実行時にSchedulerはすべてのpipeline rank、microbatch、およびモジュールタイプをカバーするグローバルoperator表を生成し、Executorはこれを各rankのローカルシーケンスに分割し、Megatron CoreなどのLLMバックエンド、モダリティランタイム、通信バックエンドに配布します。これによりスケジュール戦略が可視化され、検証可能となり、バックエンドにも優しくなります。

第二に、アルゴリズムエンジニアにとって無感覚なパイプライン並列インターフェースです。エンジニアはそれぞれのモジュールが何を生成し、何を消費するかだけを記述すれば、後はstageの分割、activationとgradientの受け渡し、跨デバイス通信などはすべてBigMacが下部で処理します。これにより、単一GPUで検証されたマルチモーダル実験が、より自然にパイプライン並列に拡張できます。第三に、スケジュール構造を理解するためのprofiler、simulator、ビジュアルツールチェーンです。一度のトレーニングイテレーションをoperatorレベルに分解し、各rankがどの時間に何を計算しているのか、どこで空転しているのか、どの依存関係が後続を妨げているのかを確認します。simulatorは大きな費用をかけてトレーニングを開始する前に、さまざまなPP設定とmicrobatchの組み合わせを試し、空泡とスループットへの影響を予測することができます。

効果は2つの代表的なワークロードにおいて検証されました。MLLM-UnderstandingではQwen3-30B-A3Bを本体にし、1.3BのViTエンコーダーを組み合わせました。Optimusという計算効率の高いベースラインに対して、BigMacは1.08から1.1倍の加速を達成しました。一方、メモリ効率の高いベースラインであるMegatron-DistTrainに対しては1.6から1.9倍の加速を達成しました。特に注目すべきはメモリの面です。per-GPUバッチサイズが増えるにつれて、BigMacのピークメモリは安定していますが、Optimusはエンコーダーのアクティブを保持する必要があるため、メモリが急激に増加し、さらに大きなバッチサイズではオーバーフロー（OOM）に陥ります。MLLM-Generationはさらに複雑で、20BのMMDiT生成器を追加しています。違いは顕著です。OptimusはすべてのテストバッチサイズでOOMし、BigMacは生成器のbackwardを適切に実行し、生成器側のアクティブを迅速に解放することで回避しました。Megatron-DistTrainとの比較では1.5から1.9倍の加速を達成し、メモリは依然として安定しています。これはネストパイプラインが最も価値のある状況です。システムはエンコーダーと生成器の依存関係を同時に扱う必要がありますが、大量のアクティブを保持したり、深刻な空泡を許容することはできません。

現在、BigMacは小紅書の生産環境で、マルチモーダルモデルトレーニングの主要なコンポーネントの一つとして動作しています。関連する論文「BigMac: Breaking the Pareto Frontier of Compute and Memory in Multimodal LLM Training」はarXivに掲載されており、チームはインタラクティブなPPProfilerのtrace例も公開しています。マルチモーダルトレーニングにおけるメモリと速度の両方の制約に苦しんでいる研究者やエンジニアにとって、この生産環境での検証付きオープンソースは、再び車輪を再発明する時間を大幅に節約するかもしれません。