NVIDIA Cosmos Predict 2.5 を LoRA/DoRA でファインチューニングしてロボット動画を生成する

記事のサマリー（TL;DR）

• NVIDIA Cosmos Predict 2.5（2B パラメータ）を LoRA rank=32 でファインチューニングすると、約 50M の訓練可能パラメータで済み、単一 H100 でも動作可能
• 92本のロボット操作動画・100エポックの学習（8×H100 で約2.5時間）で、時間的安定性・物理的妥当性・指示追従性の3指標すべてが改善
• LoRA と DoRA は最終精度が同等。メモリ制約があれば LoRA r=8、低ランクで学習不安定なら DoRA r=32 が推奨

生成AI・ロボティクス領域の国内開発者が注目すべき実装ポイント

Cosmos Predict 2.5 のような大規模動画生成モデルを実際のユースケースに適応させる手法として、LoRA/DoRA による PEFT（パラメータ効率ファインチューニング）は実用的な選択肢です。本記事が示すように、2B パラメータのモデルでもフルファインチューニングは高コスト・破滅的忘却リスクが伴うため、アダプタ注入方式が標準的なアプローチになりつつあります。国内でも製造業・物流領域でロボット学習用の合成データ生成ニーズが高まっており、リアルトラジェクトリの収集コスト削減手段として同様のパイプラインが参照されています。また、diffusers + accelerate + peft という OSS スタックで実装されているため、Hugging Face エコシステムに慣れたチームであれば既存インフラをほぼそのまま流用できます。

詳細

動機（Motivation）

NVIDIA Cosmos Predict 2.5 は、テキスト・画像・動画クリップを条件として物理的に妥当な動画を生成できる大規模世界モデルです。ロボット操作や特定カメラ視点といった特定ドメインに適応させるには、ターゲットドメインでのファインチューニングが必要です。

ロボットポリシーの学習には実演データが必要ですが、実機トラジェクトリの収集は時間とコストがかかります。ファインチューニング済み動画世界モデルで合成トラジェクトリを生成することは、スケーラブルな代替手段となります。

ただし、2B パラメータモデルのフルファインチューニングはコストが高く、汎用知識の破滅的忘却リスクも伴います。LoRA と DoRA は凍結済みベースモデルに小さな訓練可能アダプタモジュールを注入し、メモリ要件を削減しながらアダプタファイルを小さく可搬性の高い状態に保ちます。単一 GPU でのファインチューニングを実用的にし、推論時に異なるドメイン向けアダプタを柔軟に交換できます。

本ガイドでは、diffusers と accelerate ライブラリを使用した LoRA/DoRA によるパラメータ効率ファインチューニングを解説し、シングル・マルチ GPU 訓練の両方をサポートします。

要件（Requirements）

• Python 3.10+
• PyTorch 2.5+（CUDA 対応）
• diffusers（transformers と peft を自動で取得）、accelerate
• オプション：学習監視のため wandb
• 最低 80 GB GPU 1基（シングル GPU 学習用）；より速い反復には 8× H100 推奨

pip install -U "diffusers[torch]" transformers accelerate peft wandb

データの準備（Preparing Data）

diffusers インストール後、examples/cosmos に移動してサンプルコードを確認します。GR00T Dreams のポストトレーニングレシピと同じデータセットを使用します。

• 学習データセット：ピックアンドプレースタスクを説明するテキストプロンプト付きのロボット操作動画 92本
• テストデータセット：（プロンプト、画像）ペア 50組。モデルはテキストプロンプトと初期フレーム画像に基づいて動画を生成します。

bash download_and_preprocess_datasets.sh

学習データセットのフォルダ構造：

gr1_dataset/train
├── metas/
│   └── *.txt
├── videos/
│   └── *.mp4
└── metadata.csv

評価データセットは .txt と .png のペアがフラットなディレクトリに格納されます。

学習（Training）

VideoDataset

VideoDataset は args.train_data_dir（例：gr1_dataset/train）から各サンプルを（キャプション、動画）ペアとして読み込みます。args.num_frames より長い動画の場合、各エポックでランダムな連続ウィンドウをサンプリングし、時間的オーグメンテーションを実現します。diffusers.video_processor の VideoProcessor が生フレームを (channels, frames, height, width) 形状のテンソルにリサイズ・正規化します。

アダプタの初期化（Initialize Adapter）

Cosmos Predict 2.5 は3つのサブモジュールで構成されます。

• 動画をレイテントに符号化する VAE
• テキストプロンプトをプロンプト埋め込みに符号化するテキストエンコーダ
• レイテント空間で拡散を行う DiT（Diffusion Transformer）

学習中、VAE・テキストエンコーダ・DiT の全重みは凍結されます。LoRA アダプタは DiT のアテンション射影（to_q、to_k、to_v、to_out.0）とフィードフォワード層（ff.net.0.proj、ff.net.2）に注入されます。use_dora=True を指定すると DoRA に切り替わります。DoRA は各重みを大きさ（magnitude）と方向（direction）に分解してから低ランク更新を適用します。

from diffusers import Cosmos2_5_PredictBasePipeline
from peft import LoraConfig

pipe = Cosmos2_5_PredictBasePipeline.from_pretrained(
    "nvidia/Cosmos-Predict2.5-2B",
    revision="diffusers/base/post-trained",
    torch_dtype=torch.bfloat16,
)
lora_config = LoraConfig(
    r=args.lora_rank,
    lora_alpha=args.lora_alpha,
    target_modules=['to_q', 'to_k', 'to_v', 'to_out.0', 'ff.net.0.proj', 'ff.net.2'],
    use_dora=args.use_dora,
)
dit.add_adapter(lora_config)
cast_training_params(dit, dtype=torch.float32)

損失関数（Loss）

Cosmos Predict 2.5 は**整流フロー（rectified flow）**を使用します。タイムステップ t において、ノイズレベル σt でのノイジー補間 xt = σt·noise + (1−σt)·clean を構築し、目標速度 noise − clean を MSE 損失で学習します。動画の最初の2フレームは条件付けに使用されるため、そのレイテントにはノイズが追加されません。

オプティマイザとスケジューラ（Optimizer and Scheduler）

torch.optim.AdamW をオプティマイザとして使用し、diffusers.optimization の get_linear_schedule_with_warmup をスケジューラとして使用します。スケジューラは scheduler_warm_up_steps 回にわたって線形にウォームアップし、scheduler_f_max × learning_rate でピークに達した後、線形に減衰します。

チェックポイント（Checkpointing）

args.checkpointing_epochs エポックごとに LoRA 重みを diffusers フォーマットで保存します。accelerator.save_state() が pytorch_lora_weights.safetensors ファイルを書き出します。これが推論時にパイプラインに渡すアダプタファイルです。

学習コマンド（Training Command）

export MODEL_NAME="nvidia/Cosmos-Predict2.5-2B"
export DATA_DIR="gr1_dataset/train"
export OUT_DIR=YOUR_OUTPUT_DIR
lora_rank=32
accelerate launch --mixed_precision="bf16" train_cosmos_predict25_lora.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --revision diffusers/base/post-trained \
  --train_data_dir=$DATA_DIR \
  --train_batch_size=1 \
  --num_train_epochs=500 \
  --checkpointing_epochs=100 \
  --seed=0 \
  --output_dir=$OUT_DIR \
  --report_to=wandb \
  --height 432 --width 768 \
  --allow_tf32 --gradient_checkpointing \
  --lora_rank $lora_rank --lora_alpha $lora_rank

• lora_rank=32 の場合、約 50M の訓練可能パラメータが生成されます
• lora_alpha = lora_rank に設定することで、LoRA 更新のスケール係数が 1.0 になります
• DoRA を使用するには --use_dora を追加します
• 100 エポック学習で既に良好な結果が得られます（シングル H100 で 17 時間、8× H100 で 2.5 時間）

LoRA による推論（Running Inference with Your LoRA）

パイプラインと LoRA/DoRA 重みの読み込み

from diffusers import Cosmos2_5_PredictBasePipeline

pipe = Cosmos2_5_PredictBasePipeline.from_pretrained(
    "nvidia/Cosmos-Predict2.5-2B",
    revision="diffusers/base/post-trained",
    device_map="cuda",
    torch_dtype=torch.bfloat16,
)
pipe.load_lora_weights("/path/to/lora/checkpoint")
pipe.fuse_lora(lora_scale=1.0)

fuse_lora はアダプタ重みをベースモデルにマージし、LoRA/DoRA 分解による推論オーバーヘッドを排除します。

初期レイテントノイズの生成

再現性のため、arch_invariant_rand 関数が NumPy 経由で初期レイテントノイズを生成し、GPU アーキテクチャに依存しないノイズを保証します。

評価指標（Evaluation Metrics）

サンプソン誤差（Sampson Error）

サンプソン誤差は、マッチしたキーポイントから対応するエピポーラ線までの距離を測定する幾何学的誤差指標です。

• 時間的サンプソン誤差（Temporal Sampson Error）：単一カメラビュー内の連続フレーム間で計算し、時間的安定性を測定
• クロスビューサンプソン誤差（Cross-view Sampson Error）：異なるカメラビューの同時フレーム間で計算し、マルチビュー幾何学的整合性を測定

値が低いほど良好です。

LLM-as-a-Judge

評価には Cosmos Reason2 を LLM ジャッジとして使用し、各例を 1〜5 でスコアリングします。2つのルーブリックを設計しています。

• 物理的妥当性（video_physics.yaml）：プロンプトなしで動画が物理的常識に従っているか評価
• 指示追従性（video_IF.yaml）：プロンプトと動画の両方を入力として、記述されたタスクが正しく完了されているか評価

結果（Results）

定性分析（Qualitative Analysis）

ファインチューニング前のベースモデルには以下の問題があります。

• ロボットハンドが学習分布外のため、後のフレームで人間の手を幻覚生成
• プロンプトで指定された正しい手（左手/右手）を確実に使用しない
• 生成動画に顕著なジッターが発生

LoRA と DoRA によるファインチューニングはこれら3つの問題をすべて解消します。

定量分析（Quantitative Analysis）

LoRA と DoRA をランク 8 と 32 の4つの設定でファインチューニングを実施。各テスト例について異なるシードで5本の動画を生成し、平均スコアを報告。

• サンプソン誤差：ファインチューニング後に時間的・クロスビュー誤差ともに低下。時間的安定性とマルチビュー幾何学的整合性が向上
• 物理的妥当性スコア：ファインチューニング済みモデルはベースモデルより物理的常識に従った動画を生成
• 指示追従スコア：ファインチューニング済みモデルは正しい手の使用や指定オブジェクトとのインタラクションを含むタスクをより確実に完了

主要な知見：

• 100 エポック（8× H100 で約 2.5 時間）で3指標すべてが大幅に改善
• LoRA と DoRA は同等の性能に収束。DoRA の大きさ・方向分解は極低ランクでは有効だが、必須ではない
• ランク 32 はランク 8 より指示追従性が向上するが、幾何学的整合性や物理的妥当性は改善しない。幾何学・物理的事前知識は凍結された世界モデルの重みに大きく捉えられており、LoRA アダプタはドメイン内のロボット外観とタスク構造へ分布をシフトさせるだけで済み、これはランク 8 で達成可能

DoRA vs LoRA の使い分け：

• メモリが非常に限られているか、アダプタファイルサイズが重要な場合は LoRA r=8 から開始
• 予算に余裕があり、低ランクの LoRA で学習不安定が観察される場合は DoRA r=32 が妥当な代替案