NVIDIA NeMo と Diffusers で動画・画像モデルを大規模微調整

記事一覧に戻る

*NVIDIA と Hugging Face の共同投稿です。統合作業への貢献と共著者としての役割を果たした、Hugging Face の Sayak Paul 氏に特別なお礼を申し上げます。

ここ数年で最も注目すべきオープンソースリリースの多くは、拡散モデルによって支えられています。テキストから画像を生成する FLUX.1-dev や、テキストから動画を生成する Wan 2.1、HunyuanVideo などがその代表例です。これらのモデルは、🤗 Diffusers ライブラリを事実上のホームとしており、研究者や開発者が推論、適応、パイプラインの構築を行うための統一されたインターフェースを提供しています。

また、拡散モデルのトレーニングやファインチューニングも急増しており、メモリ効率の良いシャード処理、潜在空間のキャッシュ、マルチ解像度バケッティング、そして 1 つの GPU から数百台規模まで柔軟に拡張可能な設定を可能にするユーティリティが求められています。

これらの技術的ニーズに応えるため、私たちはオープンソースライブラリ「NVIDIA NeMo Automodel」を提供しています。本稿では、NVIDIA と Hugging Face の連携についてご紹介します。この協力により、Hugging Face Hub 上の Diffusers フォーマットモデルに対して、チェックポイントの変換や新モデルでのモデル書き換えを一切行わずに、本番環境で使える分散型拡散学習が可能になります。

詳細は Diffusers 学習ガイド に記載されており、Apache 2.0 ライセンスの下で完全にオープンソース化されています。

目次

• NeMo Automodel とは？

• 対応する拡散モデル

• この連携がもたらす革新

• ファインチューニングのワークフロー
• 1. データセットの前処理（エンコード）
• 2. 既存の FLUX YAML を用いた学習の実行
• 3. ファインチューニング済みチェックポイントからの生成
• 4. パフォーマンス

• その他のファインチューン/LoRA の事例

• 今日から試す

• 今後の予定：Pythonic レシピ API

• リソース

NeMo Automodel とは？

NeMo Automodel は、NVIDIA NeMo フレームワークの一部であるオープンソースの PyTorch DTensor ネイティブなトレーニングライブラリです。Diffusers エコシステムにとって重要な 2 つの設計原則に基づいて構築されています。

• Hugging Face ネイティブ対応。pretrained_model_name_or_path に Hugging Face Hub 上の任意の Diffusers モデル ID を指定するだけで、すぐにトレーニングを開始できます。NeMo Automodel は、モデルの読み込みに Diffusers のモデルクラス（例：WanTransformer3DModel）を、生成には Diffusers パイプライン（WanPipeline）を使用します。チェックポイントは、Diffusers エコシステム内へ円滑に往復して利用可能です。

一つのプログラムで、あらゆる規模に対応。レシピやトレーニングスクリプトは、任意のスケールでの学習に合わせて簡単に修正可能です。並列処理は設定次第で切り替えられるものであり、コードを書き直す必要はありません。FSDP2、テンソル並列、エキスパート並列、コンテキスト並列、パイプライン並列を、モデルの書き換えではなく設定宣言だけで切り替えることができます。

AutoModel は現在、フローマッチングモデルのみをサポートしています。内部では学習目標としてフローマッチングを採用し、潜在空間でのトレーニング（事前符号化された VAE の出力を利用）とマルチ解像度のバケット型データローダーを組み合わせることで、スループットを加速させています。

サポートされている拡散モデル

NeMo Automodel には、以下のオープンソース拡散モデル向けのすぐに使えるファインチューニングレシピが統合されています。このリストは、examples/diffusion/finetune ディレクトリに現在含まれているレシピを反映したものです。

この連携で実現すること

Diffusers の利用者にとって、この連携による実用的なメリットは主に以下の 3 つの機能に集約されます。

チェックポイントの変換が不要。Hugging Face Hub にある事前学習済みモデルの重みは、そのまま即座に利用可能です。変換用の「トレーニングフォーマット」を経由して一度変換し、また戻すといった手間もありません。微調整後のチェックポイントは、推論用に DiffusionPipeline に直接読み込めるため、共有用として Hub へアップロードすることも容易です。量子化やコンパイル、LoRA アダプター、カスタムサンプラーなどの下流ツールも、すべてそのまま動作し続けます。

新モデルへの対応が迅速。Diffusers に新しい拡散モデルが登場しても、NeMo Automodel で利用可能にするには、フルスクラプトのカスタムトレーニングスクリプトを書く必要はありません。データ前処理ハンドラーとモデルアダプターを追加するだけの、小さくまとまったコード変更で済みます。残りのレシピ構成（FSDP2、バケット化されたデータローダー、チェックポイント保存、生成処理）はすべてそのまま引き継がれ、YAML ドライブ型のワークフローも同じように適用可能です。

フル微調整とパラメータ効率型微調整の両方に対応。フル微調整と LoRA 形式の PEFT（Parameter-Efficient Fine-Tuning）の両方がサポートされているため、最大品質を追求する大規模クラスターでのフル FT か、単一ノードで実行する最大効率の LoRA か、用途に応じて選択できます。どちらの場合も、同じレシピ構造で処理が可能です。

ビルトインスクリプトの限界を超える、スケーラブルなトレーニング。 NeMo Automodel は、FSDP2 やテンソル、コンテキスト、パイプライン並列処理といったシャーディングスキームに加え、マルチノードオーケストレーション（現在は SLURM 対応、Kubernetes も近日追加予定）、そしてマルチ解像度バケットリングをサポートします。これらの機能により、FLUX.1-dev (12B) や HunyuanVideo (13B) のような大規模モデルのトレーニングが可能になります。

ファインチューニングワークフローの概要

ここでは、対応するすべてのモデルに対するファインチューニングの一般的なワークフローを解説します。Automodel をインストールする推奨方法は、NeMo Automodel Docker コンテナ（nvcr.io/nvidia/nemo-automodel:26.06）を利用することです。このコンテナには PyTorch、TransformerEngine、およびその他の CUDA 対応依存関係が事前にビルド済みで含まれています。あるいは、pip3 install nemo-automodel またはソースコードからのインストール（pip3 install git+https://github.com/NVIDIA-NeMo/Automodel.git）も可能です。詳細なオプションについては インストールガイド をご確認ください。

本ガイドでは、78 枚の ライダー・ウェイト塔カードデータセット を用いた FLUX.1-dev のフルトランスフォーマーファインチューニングと、その結果得られたチェックポイントからの生成プロセスを順を追って説明します。既存の YAML 設定ファイルを再利用し、実行固有の設定はコマンドライン引数で上書きする形をとるため、新規に設定ファイルを作成する必要はありません。

1. データセットの前エンコード

拡散モデルの学習レシピでは、各トレーニングステップでソース画像をエンコードするのではなく、キャッシュされた VAE ラテンツとテキスト埋め込みを利用します。78 枚の Rider–Waite カード画像を Hugging Face から直接ストリーミングし、利用可能なすべての GPU に前処理を分散させます。

uv run --locked --no-default-groups \
  --extra diffusion \
  --extra diffusion-media \
python -m tools.diffusion.preprocessing_multiprocess image \
  --dataset_name multimodalart/1920-raider-waite-tarot-public-domain \
  --dataset_media_column image \
  --dataset_caption_column caption \
  --dataset_streaming \
  --max_images 78 \
  --output_dir /cache/flux_tarot \
  --processor flux \
  --model_name black-forest-labs/FLUX.1-dev \
  --max_pixels 245760

キャプションにはすでに trtcrd トリガートークンが含まれています。このピクセル予算とデータセットの縦長アスペクト比に基づき、前処理ではサンプルをショーケース実行で使用する 384×640 のバケットに割り当てます。

画像学習用の前処理では、.pt キャッシュファイルとシャード化されたメタデータを生成します。

/cache/flux_tarot/
├── 384x640/
│   ├── <hash1>.pt
│   └── ...
├── metadata_shard_0000.json
├── metadata.json
└── _hf_dataset/
    └── images/

2. 既存の FLUX YAML を使用してトレーニングを開始する

examples/diffusion/finetune/flux_t2i_flow.yaml をそのまま使用します。この YAML ファイルには、FLUX.1-dev の選択、フルトランスフォーマーの微調整、FLUX flow-matching アダプターの適用、有効バッチサイズ 32、および 8 分割 FSDP2 が既に設定されています。

タロットデータに特化したパスや設定は、コマンドライン引数で上書きして指定します：

uv run --locked --no-default-groups --extra diffusion \
  torchrun --nproc-per-node=8 \
  examples/diffusion/finetune/finetune.py \
  -c examples/diffusion/finetune/flux_t2i_flow.yaml \
  --model.transformer_engine_fp8 false \
  --data.dataloader.cache_dir /cache/flux_tarot \
  --data.dataloader.base_resolution '[384,640]' \
  --lr_scheduler.lr_decay_style constant \
  --lr_scheduler.lr_warmup_steps 20 \
  --step_scheduler.max_steps 200 \
  --step_scheduler.ckpt_every_steps 50 \
  --checkpoint.checkpoint_dir /tmp/flux_tarot/checkpoints/full \
  --checkpoint.save_consolidated true \
  --seed 2026

実行すると、ステップ 50、100、150、200 でチェックポイントが生成されます。最終的なチェックポイントは epoch_66_step_199 というラベルが付きます。このラベルはゼロベースで付与されていますが、実際には完了した 200 番目のオプティマイザステップを指しています。

3. 微調整済みチェックポイントからの生成

既存の FLUX 生成用 YAML ファイルを利用し、model.checkpoint にトレーニング完了後のチェックポイントを指定します：

uv run --locked --no-default-groups --extra diffusion \
  python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_flux.yaml \
  --model.checkpoint /tmp/flux_tarot/checkpoints/full/epoch_66_step_199 \
  --inference.height 640 \
  --inference.width 384 \
  --inference.prompts '["a trtcrd of an astronaut tending a rose garden on Mars, \"the gardener\""]' \
  --output.output_dir /tmp/flux_tarot/generations/full/step_200 \
  --seed 2026

学習したタロットスタイルを呼び出すには trtcrd を含めます。制御比較として、シードとシーンは固定したままトリガーを省略します。

uv run --locked --no-default-groups --extra diffusion \
  python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_flux.yaml \
  --model.checkpoint /tmp/flux_tarot/checkpoints/full/epoch_66_step_199 \
  --inference.height 640 \
  --inference.width 384 \
  --inference.prompts '["an astronaut tending a rose garden on Mars, \"the gardener\""]' \
  --output.output_dir /tmp/flux_tarot/generations/control \
  --seed 2026

結果

ステップ 200 では、トリガーされた宇宙飛行士のプロンプトは要求された内容を維持しつつ、クリーム色・赤・黒のヴィンテージ調パレット、太いインクによる輪郭線、フラットな色彩領域、経年劣化を感じさせる紙の色合い、寓話的なカード構成といった特徴を獲得します。一方、トリガーされていない宇宙飛行士は依然として写真のような描写のままです。これは、学習された効果がモデル全体を置き換えるものではなく、trtcrd という特定の要素に強く関連していることを示しています。

4. パフォーマンス

すべての測定は、8 基の NVIDIA H100 (80GB) GPU を搭載したノード 1 台で行いました。結果は、3 つの定常状態における 10 ステップ区間の平均値と標準偏差を示しています。

テキストから画像へ — 512×512(秒数との整合性を確認)

テキストから動画生成 — 512×512×49 フレーム

各サンプルは、49フレームの動画クリップです。

測定の詳細

• ハードウェア：8 枚の H100（80GB HBM3）を NVLink でフル接続。
• イメージデータセット：`lambda/naruto-blip-captions` からキャッシュしたサンプル 256 件。
• ビデオデータセット：`svjack/Lelouch_Vi_Britannia_FramePack_First_Last_Frame_Video_Captioned` からキャッシュしたサンプル 112 件。
• バッチ処理は drop_last=true で強制実行。
• チェックポイントの書き出しは無効化。
• ステップ時間は、データ読み込み、順伝播、逆伝播、勾配クリッピング、オプティマイザ、スケジューラの処理時間をすべて含む。
• メモリ使用量は、PyTorch の CUDA アロケータがピーク時に消費した値であり、NVML で計測されるデバイス全体のメモリではない。

その他のファインチューニング・LoRA 事例

ファインチューニングと LoRA の結果は、NeMo Automodel がドメイン特化にどれほど強力かを如実に示しています。例えば、Wan 2.1 モデルをジブリの動画データセットでファインチューニングすると、出力スタイルが適応され、ベースラインと比較して花の外観に明確な変化が見られました。

*ベースライン：*

*ジブリの動画でファインチューニング：*

また、LoRA の使用による顕著な影響も確認できました。Wan 2.1 にアダプターを適用すると、動画全体が特徴的なジブリスタイルへと変化し、特にキャラクターの目のハイライト部分にその違いがはっきりと現れています。

*LoRA なし：*

*LoRA あり：*

これらの例、FLUX.2 のケースも含まれますが、フルファインチューニングで最高品質を達成し、LoRA スタイルの PEFT で最大効率を実現できることが確認できました。これにより、特定のスタイルドメインに合わせた出力のカスタマイズが可能になります。

今日から試す

統合の詳細や、さらに多くのファインチューニング例については、NeMo Automodel ドキュメントをご覧ください。

今後の予定：Pythonic なレシピ API

YAML は再現可能な設定に強く適しており、特にチェックインやレビュー、再利用を想定するチームには理想的です。しかし、多くのチームではプログラムによるインターフェースも必要としています。

今後の NeMo Automodel リリースでは、拡散モデルのレシピを完全に型付けされた Pythonic な API として提供することを計画しています。ユーザーは、Python から直接、同じモデル、データ、オプティマイザー、PEFT/LoRA、並列処理、チェックポイント、生成コンポーネントを組み立てることができるようになります。

この Pythonic なアプローチは、既存のトレーニングコードやノートブック、実験ワークフローからレシピをより使いやすくすることを目的としています。また、YAML によるクイックスタートパスと並行して、本格的な Pythonic なインターフェースも提供します。

リソース

• NeMo Automodel GitHub

• Diffusion Fine-Tuning Guide

• Diffusion Dataset Preparation

• Transformer 向け NeMo Automodel（LLM/VLM のファインチューニング）