スクリプトの動物園なしで重いマルチモーダルモデルを訓練する方法

Qwen2.5-VLのようなマルチモーダルモデルのファインチューニングや、拡散ベースの動画生成器の事前訓練を試みたことがあるなら、あの絶望感を覚えているでしょう。あるリポジトリでは、PyTorchとFlashAttentionのバージョン衝突で訓練がクラッシュします。別のリポジトリでは、著者が500 GBの画像を直接RAMに読み込むカスタムデータパイプラインを書いています。3つ目のリポジトリでは、FSDPが2番目の最適化ステップで固まります。

新しいアーキテクチャごとに独自の訓練スクリプト、メモリ節約のハック、並列処理の回避策がついてきます。EvolvingLMMs-Labのチームは、この散らかった経験を1つの場所に集めることを決意しました。それがlmms-engineリポジトリの始まりです。

これはスケーラブルなマルチモーダルモデル訓練用のモジュラーエンジンです。分散訓練、シーケンスパッキング、低レベルカーネル最適化の重い仕事を処理し、設定とデータだけをあなたに残します。

中身と対象者

このフレームワークは、重いマルチモーダルネットワークに取り組む研究者やMLエンジニアのために主に開発されました。通常のテキストLLMはAxolotlやLLaMA-Factoryのようなツールで快適ファインチューニングできますが、画像、オーディオ、ビデオは別の話になります。

マルチモーダル入力は巨大なコンテキスト長を生み出します。高解像度フレーム1枚やビデオ数秒で、数万のビジュアルトークンに簡単に膨れ上がります。この時点で、標準のData ParallelはすぐにGPUメモリの限界に達します。

このエンジンは、最新の並列処理技術とGPUレベルの最適化を組み合わせることでこれを解決します。サポートされているモデルのリストは印象的です：

• Vision-Languageモデル：Qwen2.5-VL、Qwen3-VL、Qwen3-VL MoE、LLaVA-OneVision
• オーディオとビデオサポート付きのマルチモーダルMoE：Qwen2.5-Omni、Qwen3-Omni MoE、Aero
• 生成・拡散アーキテクチャ：WanVideo（1.3Bおよび14Bパラメータ）、SiT（Scalable Interpolant Transformers）、dLLM
• 汎用画像理解・生成システム：BAGEL

内部の4つのエンジニアリングソリューション

プロジェクトの作成者は明らかにミニマリズムのアイデアに触発されました：コードは不必要な抽象的な魔法なしでコンパクトに書かれ、純粋なPyTorchで最新のライブラリの統合書かれています。

1. FSDP2とUlysses Sequence Parallel

重みの分散には、DTensorに基づくFully Sharded Data Parallel v2（FSDP2）の新しい実装に依存しています。古いFSDPとは異なり、2番目のバージョンは他の種類の並列処理とはるかにクリーンに構成されます。

10,000以上のビジュアルトークンがコンテキストに入ると、GPUメモリは即座に不足します。そこがUlysses Sequence Parallel（USP）の出番です。単一ノード内の複数のGPU間でトークンシーケンスを分割します。設定では、1行で設定します：

trainer_args:
  sp_ulysses_degree: 2

2. 不必要なパディングなしのシーケンスパッキング

マルチモーダル訓練における古典的な頭痛の種は、1つのバッチ内の画像とテキストのサイズ変動です。短いサンプルをパディングトークンで埋めると、GPUは最大半分もの時間を無意味なゼロ計算に費やします。

著者はFlashAttentionとの組み合わせでfirst-fit bin packingを実装しました。データはアイドル時間なしで密な長いシーケンスにパッキングされます。著者のベンチマークによると、Qwen2.5-VLのファインチューンでは、計算効率指標（MFU）が印象的な35-40%に上昇し、従来の20-25%から改善されます。

dataset_config:
  packing: true
  packing_strategy: first_fit
  packing_length: 32000

trainer_args:
  use_rmpad: true
  use_liger_kernel: true

3. MuonオプティマイザとLigerのTritonカーネル

おなじみのAdamWの代わりに、プロジェクトはMuonを提供します。このオプティマイザはTritonカーネルを通じて2D重み行列にNewton-Schulz直交化を適用します。AdamWより速く収束し、より少ないメモリを必要とします。

同時に、エンジンはLinkedInのLiger Kernelライブラリの融合カーネルで標準モデルレイヤーをオンザフライで交換できます。CrossEntropy、RMSNorm、RoPE、SwiGLUが単一操作に融合され、精度損失なしでピークVRAM消費を約30%削減します。

4. テラバイトデータセットのストリーミングロード

エポック開始前にビデオや画像の巨大な配列をメモリに読み込むことは不可能です。lmms-engineのデータパイプラインはの上に構築されています。データはArrow、JSONL、またはParquet形式からストリームで読み込まれるため、テラバイトサイズのファイルのインデックス作成を待つことなく即座に訓練が開始されます。

起動と拡張の仕組み

プロジェクトのインストールは適切にパッケージマネージャーを通じて行われますが、著者は事前にインストールされたCUDA、FlashAttention、依存関係を含む готовую Dockerイメージも提供します。

git clone https://github.com/EvolvingLMMs-Lab/lmms-engine.git
cd lmms-engine
uv pip install -e ".[all]"
uv pip install flash-attn --no-build-isolation
uv pip install liger-kernel

訓練は標準ので開始されます：

torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
  --master_addr=127.0.0.1 --master_port=12355 \
  -m lmms_engine.launch.cli config_yaml=examples/qwen3_vl/example_config.yaml

独自の特定のデータ形式やカスタムフィーチャプロセッサを追加する必要がある場合、トレーナー内部を書き換える必要はありません。コードはデコレータによる登録を伴うファクトリパターンを使用しています：

from lmms_engine.datasets import register_dataset, BaseDataset

@register_dataset("my_custom_dataset")
class MyCustomDataset(BaseDataset):
    def __init__(self, config):
        super().__init__(config)

    def __getitem__(self, idx):
        # Ваша логика чтения картинки или видео
        return item

プロジェクトが最も輝く場所

このエンジンは特定の重いシナリオのために開発されました：

1. 長いドキュメント、書籍のスキャン、インターフェーススクリーンショットでのVision-Languageモデル（Qwen2.5-VL、Qwen3-VL）のファインチューニング。
2. GPU間でのエキスパート分散を伴うスパースMixture-of-Expertsアーキテクチャの訓練（Expert Parallelism）。
3. WanVideoやSiTモデルに基づくビデオ生成の拡散実験。
4. 軽量再帰型および線形アテンションモデル（FLA / DGN）の事前訓練。

試す価値はあるか

テキストLlamaでLoRAをすばやく微調整するタスクであれば、lmms-engineはオーバースペックに感じるかもしれません。しかし、マルチモーダルネットワークの訓練中にメモリ上限に達した場合、長いビデオシーケンスでFSDP2を手動で結びつけるのに疲れている場合、または既存のクラスタから最大 FLOPSを絞り出したい場合は、このエンジンが何週間ものコーディングを節約します。

開始する最も簡単な方法は、フォルダ内の готовуюスクリプトを使用することです。これには最新のほとんどのアーキテクチャ用のbattle-tested設定が含まれています。