# NeMo Automodel

[NeMo Automodel](https://github.com/NVIDIA-NeMo/Automodel) は、NVIDIA による PyTorch DTensor ネイティブの学習ライブラリで、大規模な拡散モデルのファインチューニングと事前学習に使用されます。Hugging Face ネイティブであり、チェックポイントの変換なしに、Hub 上の Diffusers 形式のモデルをそのまま学習できます。同じ YAML レシピとハック可能な学習スクリプトが、1 GPU から数百ノードまでのあらゆる規模で動作し、[FSDP2](https://pytorch.org/docs/stable/fsdp.html) による分散学習、マルチ解像度バケット化されたデータローディング、最大限の GPU 利用率を実現する事前エンコードされた潜在空間での学習に対応しています。学習には [flow matching](https://huggingface.co/papers/2210.02747) を使用し、完全にオープンソース（Apache 2.0）で、NVIDIA がサポートし、積極的にメンテナンスされています。

NeMo Automodel は Diffusers と直接統合されています。Diffusers のモデルクラスを使用して Hugging Face Hub から事前学習済みモデルを読み込み、[DiffusionPipeline](/docs/diffusers/main/ja/api/pipelines/overview#diffusers.DiffusionPipeline) を用いて出力を生成します。

典型的なワークフローは、NeMo Automodel をインストール（pip または Docker）し、データを `.meta` ファイルにエンコードして準備し、YAML レシピを設定し、`torchrun` でトレーニングを起動し、生成されたチェックポイントで推論を実行することです。

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

| モデル | Hugging Face ID | タスク | パラメータ | ユースケース |
|-------|----------------|------|------------|----------|
| Wan 2.1 T2V 1.3B | [Wan-AI/Wan2.1-T2V-1.3B-Diffusers](https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B-Diffusers) | テキストから動画へ | 1.3B | 限られたハードウェアでの動画生成（単一の40GB A100に収まる） |
| FLUX.1-dev | [black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) | テキストから画像へ | 12B | 高品質な画像生成 |
| HunyuanVideo 1.5 | [hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v](https://huggingface.co/hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v) | テキストから動画へ | 13B | 高品質な動画生成 |

## インストール

### ハードウェア要件

| コンポーネント | 最小 | 推奨 |
|-----------|---------|-------------|
| GPU | A100 40GB | A100 80GB / H100 |
| GPU | 4 | 8+ |
| RAM | 128 GB | 256 GB以上 |
| ストレージ | 500 GB SSD | 2 TB NVMe |

pip を使用して NeMo Automodel をインストールします。ソースからのインストールを含む、すべてのインストール方法については、[NeMo Automodel インストールガイド](https://docs.nvidia.com/nemo/automodel/latest/guides/installation.html)を参照してください。

```bash
pip3 install nemo-automodel
```

代わりに、すべての依存関係が含まれている事前構築済みのDockerコンテナを使用できます。

```bash
docker pull nvcr.io/nvidia/nemo-automodel:26.02.00
docker run --gpus all -it --rm --shm-size=8g nvcr.io/nvidia/nemo-automodel:26.02.00
```

> [!WARNING]
> コンテナが終了すると、チェックポイントディレクトリをホストにバインドマウントしていない場合、チェックポイントは失われます。例えば、`docker run` コマンドに `-v /host/path/checkpoints:/workspace/checkpoints` を追加してください。

## データ準備

NeMo Automodel は、潜在空間で拡散モデルを学習します。学習を開始する前に、生の画像や動画を、VAE 潜在変数とテキスト埋め込みを含む `.meta` ファイルに前処理する必要があります。これにより、各学習ステップでの再エンコードを回避できます。

組み込みの前処理ツールを使用してデータをエンコードします。このツールは、利用可能なすべてのGPUに自動的にワークを分散します。

動画の前処理コマンドは Wan 2.1 と HunyuanVideo で同じですが、フラグが異なります。Wan 2.1 は `--processor wan` を `--resolution_preset` と `--caption_format sidecar` と組み合わせて使用し、HunyuanVideo は `--processor hunyuan` を `--target_frames` と組み合わせてフレーム数を設定し、`--caption_format meta_json` を使用します。

**Wan 2.1:**

```bash
python -m tools.diffusion.preprocessing_multiprocess video \
    --video_dir /data/videos \
    --output_dir /cache \
    --processor wan \
    --resolution_preset 512p \
    --caption_format sidecar
```

**HunyuanVideo:**

```bash
python -m tools.diffusion.preprocessing_multiprocess video \
    --video_dir /data/videos \
    --output_dir /cache \
    --processor hunyuan \
    --target_frames 121 \
    --caption_format meta_json
```

```bash
python -m tools.diffusion.preprocessing_multiprocess image \
    --image_dir /data/images \
    --output_dir /cache \
    --processor flux \
    --resolution_preset 512p
```

### 出力形式

Preprocessing は、解像度バケットごとに整理されたキャッシュディレクトリを生成します。NeMo Automodel は、バケット化されたサンプリングによるマルチ解像度トレーニングをサポートしています。サンプルは空間解像度ごとにグループ化され、各バッチが同じサイズのサンプルのみを含むようにすることで、パディングによる無駄を回避します。

```
/cache/
├── 512x512/                      # Resolution bucket
│   ├── <hash1>.meta              # VAE latents + text embeddings
│   ├── <hash2>.meta
│   └── ...
├── 832x480/                      # Another resolution bucket
│   └── ...
├── metadata.json                 # Global config (processor, model, total items)
└── metadata_shard_0000.json      # Per-sample metadata (paths, resolutions, captions)
```

> [!TIP]
> キャプションの形式、入力データの要件、および利用可能なすべての前処理引数については、[Diffusion Dataset Preparation](https://docs.nvidia.com/nemo/automodel/latest/datasets/diffusion-dataset)ガイドを参照してください。

## トレーニング設定

ファインチューニングは、以下の2つのコンポーネントによって駆動されます。

1. レシピスクリプト（[finetune.py](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/diffusion/finetune/finetune.py)）は、学習ループを含む Python のエントリポイントです。モデルの読み込み、dataloader の構築、フォワード/バックワードパスの実行、flow matching ロスの計算、チェックポイントの保存、およびログ出力を行います。
2. YAML設定ファイルは、レシピが使用するすべての設定を指定します。ファインチューニングするモデル、データの保存場所、オプティマイザーのハイパーパラメータ、並列化戦略などです。コードを変更するのではなく、このファイルを編集してトレーニングをカスタマイズすることで、1台から数百台のGPUまでスケールできます。

YAMLのフィールドは、CLIから上書きすることもできます:

```bash
torchrun --nproc-per-node=8 examples/diffusion/finetune/finetune.py \
  -c examples/diffusion/finetune/wan2_1_t2v_flow.yaml \
  --optim.learning_rate 1e-5 \
  --step_scheduler.num_epochs 50
```

以下は、Wan 2.1 T2V 1.3B のファインチューニング用のアノテーション付き設定です。各セクションについて説明します。

```yaml
seed: 42

# ── Experiment tracking (optional) ──────────────────────────────────────────
# Weights & Biases integration for logging metrics, losses, and learning rates.
# Set mode: "disabled" to turn off.
wandb:
  project: wan-t2v-flow-matching
  mode: online
  name: wan2_1_t2v_fm

# ── Model ───────────────────────────────────────────────────────────────────
# pretrained_model_name_or_path: any Hugging Face model ID or local path.
# mode: "finetune" loads pretrained weights; "pretrain" trains from scratch.
model:
  pretrained_model_name_or_path: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
  mode: finetune

# ── Training schedule ───────────────────────────────────────────────────────
# global_batch_size: effective batch across all GPUs.
# Gradient accumulation is computed automatically: global / (local × num_gpus).
step_scheduler:
  global_batch_size: 8
  local_batch_size: 1
  ckpt_every_steps: 1000     # Save a checkpoint every N steps
  num_epochs: 100
  log_every: 2               # Log metrics every N steps

# ── Data ────────────────────────────────────────────────────────────────────
# _target_: the dataloader factory function.
#   Use build_video_multiresolution_dataloader for video models (Wan, HunyuanVideo).
#   Use build_text_to_image_multiresolution_dataloader for image models (FLUX).
# model_type: "wan" or "hunyuan" (selects the correct latent format).
# base_resolution: target resolution for multiresolution bucketing.
data:
  dataloader:
    _target_: nemo_automodel.components.datasets.diffusion.build_video_multiresolution_dataloader
    cache_dir: PATH_TO_YOUR_DATA
    model_type: wan
    base_resolution: [512, 512]
    dynamic_batch_size: false  # When true, adjusts batch per bucket to maintain constant memory
    shuffle: true
    drop_last: false
    num_workers: 0

# ── Optimizer ───────────────────────────────────────────────────────────────
# learning_rate: 5e-6 is a good starting point for fine-tuning.
# Adjust weight_decay and betas for your dataset.
optim:
  learning_rate: 5e-6
  optimizer:
    weight_decay: 0.01
    betas: [0.9, 0.999]

# ── Learning rate scheduler ─────────────────────────────────────────────────
# Supports cosine, linear, and constant schedules.
lr_scheduler:
  lr_decay_style: cosine
  lr_warmup_steps: 0
  min_lr: 1e-6

# ── Flow matching ───────────────────────────────────────────────────────────
# adapter_type: model-specific adapter — must match the model:
#   "simple" for Wan 2.1, "flux" for FLUX.1-dev, "hunyuan" for HunyuanVideo.
# timestep_sampling: "uniform" for Wan, "logit_normal" for FLUX and HunyuanVideo.
# flow_shift: shifts the flow schedule (model-dependent).
# i2v_prob: probability of image-to-video conditioning during training (video models).
flow_matching:
  adapter_type: "simple"
  adapter_kwargs: {}
  timestep_sampling: "uniform"
  logit_mean: 0.0
  logit_std: 1.0
  flow_shift: 3.0
  num_train_timesteps: 1000
  i2v_prob: 0.3
  use_loss_weighting: true

# ── FSDP2 distributed training ──────────────────────────────────────────────
# dp_size: number of GPUs for data parallelism (typically = total GPUs on node).
# tp_size, cp_size, pp_size: tensor, context, and pipeline parallelism.
#   For most fine-tuning, dp_size is all you need; leave others at 1.
fsdp:
  tp_size: 1
  cp_size: 1
  pp_size: 1
  dp_replicate_size: 1
  dp_size: 8

# ── Checkpointing ──────────────────────────────────────────────────────────
# checkpoint_dir: where to save checkpoints (use a persistent path with Docker).
# restore_from: path to resume training from a previous checkpoint.
checkpoint:
  enabled: true
  checkpoint_dir: PATH_TO_YOUR_CKPT_DIR
  model_save_format: torch_save
  save_consolidated: false
  restore_from: null
```

### 設定フィールドの参照

下表には、最小限の必須設定をリストしています。すべてのモデルの完全な設定例については、[NeMo Automodel の例](https://github.com/NVIDIA-NeMo/Automodel/tree/main/examples/diffusion/finetune)を参照してください。

| セクション | 必須? | 変更すべき点 |
|---------|-----------|----------------|
| `model` | はい | `pretrained_model_name_or_path` に Hugging Face のモデル ID を設定します。`mode: finetune` または `mode: pretrain` を設定します。 |
| `step_scheduler` | はい | `global_batch_size` は、すべての GPU にわたる実効バッチサイズです。`ckpt_every_steps` はチェックポイントの頻度を制御します。勾配の蓄積は自動的に計算されます。 |
| `data` | はい | `cache_dir` に、前処理済みの`.meta`ファイルを含むパスを設定します。モデルが異なる場合は、`_target_`と`model_type`を変更してください。 |
| `optim` | はい | `learning_rate: 5e-6` はファインチューニングの適切なデフォルト値です。データセットやモデルに応じて調整してください。 |
| `lr_scheduler` | はい | `lr_decay_style` には `cosine`、`linear`、または `constant` を指定します。段階的なウォームアップには `lr_warmup_steps` を設定します。 |
| `flow_matching` | はい | `adapter_type` はモデルと一致する必要があります（Wan は `simple`、FLUX は `flux`、HunyuanVideo は `hunyuan`）。`adapter_kwargs` については、モデル固有の設定を参照してください。 |
| `fsdp` | はい | `dp_size` を GPU の数に設定します。マルチノードの場合、すべてのノードにまたがる GPU の総数に設定します。 |
| `checkpoint` | 推奨 | `checkpoint_dir` は永続的なパスに設定してください。Docker を使用する場合、特に重要です。`restore_from` を使用すると、以前のチェックポイントから再開できます。 |
| `wandb` | オプション | Weights & Biases の実験トラッキングを有効にするよう設定します。`mode: disabled` を設定すると無効になります。 |

## トレーニングを開始

```bash
torchrun --nproc-per-node=8 \
  examples/diffusion/finetune/finetune.py \
  -c examples/diffusion/finetune/wan2_1_t2v_flow.yaml
```

各ノードで、`NODE_RANK`を適切に設定して、以下を実行してください。

```bash
export MASTER_ADDR=node0.hostname
export MASTER_PORT=29500
export NODE_RANK=0  # 0 on master, 1 on second node, etc.

torchrun \
  --nnodes=2 \
  --nproc-per-node=8 \
  --node_rank=${NODE_RANK} \
  --rdzv_backend=c10d \
  --rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} \
  examples/diffusion/finetune/finetune.py \
  -c examples/diffusion/finetune/wan2_1_t2v_flow_multinode.yaml
```

> [!NOTE]
> マルチノード学習では、YAML の `fsdp.dp_size` に、すべてのノードにまたがる GPU の**総数**を設定します（例: ノード 2 台で各 8 GPU の場合は 16）。

## 生成

学習後、ファインチューニングされたチェックポイントを使用して、テキストプロンプトから動画や画像を生成します。

```bash
python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_wan.yaml
```

ファインチューニング済みのチェックポイントを使用する場合:

```bash
python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_wan.yaml \
  --model.checkpoint ./checkpoints/step_1000 \
  --inference.prompts '["A dog running on a beach"]'
```

```bash
python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_flux.yaml
```

ファインチューニング済みのチェックポイントを使用する場合:

```bash
python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_flux.yaml \
  --model.checkpoint ./checkpoints/step_1000 \
  --inference.prompts '["A dog running on a beach"]'
```

```bash
python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_hunyuan.yaml
```

ファインチューニング済みのチェックポイントを使用する場合:

```bash
python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_hunyuan.yaml \
  --model.checkpoint ./checkpoints/step_1000 \
  --inference.prompts '["A dog running on a beach"]'
```

## Diffusersとの統合

NeMo Automodel は Diffusers を基盤として構築されており、モデルの読み込みと推論のバックボーンとして Diffusers を利用しています。モデルは [WanTransformer3DModel](/docs/diffusers/main/ja/api/models/wan_transformer_3d#diffusers.WanTransformer3DModel)、[FluxTransformer2DModel](/docs/diffusers/main/ja/api/models/flux_transformer#diffusers.FluxTransformer2DModel)、[HunyuanVideoTransformer3DModel](/docs/diffusers/main/ja/api/models/hunyuan_video_transformer_3d#diffusers.HunyuanVideoTransformer3DModel) などの Diffusers のモデルクラスを使用して Hugging Face Hub から直接読み込まれ、[WanPipeline](/docs/diffusers/main/ja/api/pipelines/wan#diffusers.WanPipeline) や [FluxPipeline](/docs/diffusers/main/ja/api/pipelines/flux#diffusers.FluxPipeline) などの Diffusers パイプラインを通じて出力が生成されます。

この統合により、Diffusersのユーザーには以下の利点があります。

- **チェックポイント変換不要**：Hubからの事前学習済み重みはそのまま利用できます。`pretrained_model_name_or_path`にDiffusers形式のモデルIDを指定して、すぐにトレーニングを開始できます。
- **Day-0 モデルサポート**: 新しい拡散モデルが Diffusers に追加され、Hub にアップロードされると、専用の学習スクリプトを待つことなく、NeMo Automodel を使用してファインチューニングを行うことができます。
- **Pipeline 互換の出力**:ファインチューニングされたチェックポイントは、Diffusers のパイプラインに直接読み込んで推論を実行したり、Hub で共有したり、量子化やコンパイルなどのツールでさらに最適化したりできる形式で保存されます。
- **Diffusersモデルのスケーラブルな学習**: NeMo Automodelは、組み込みのDiffusers学習スクリプトが提供する機能を超えた分散学習機能（FSDP2、マルチノード、マルチリゾリューションバケット化）を追加します。その際、モデルおよびパイプラインのインターフェースはそのまま維持されます。
- **共有エコシステム**: Diffusersエコシステム由来のモデル、LoRAアダプター、パイプラインコンポーネントはすべて、トレーニングおよび推論ワークフロー全体を通じて互換性を保ちます。

## NVIDIA チーム

- Pranav Prashant Thombre, pthombre@nvidia.com
- Linnan Wang, linnanw@nvidia.com
- Alexandros Koumparoulis, akoumparouli@nvidia.com

## リソース

- [NeMo Automodel GitHub](https://github.com/NVIDIA-NeMo/Automodel)
- [Diffusionファインチューニングガイド](https://docs.nvidia.com/nemo/automodel/latest/recipes-e2e-examples/diffusion-fine-tuning)
- [Diffusionデータセットの準備](https://docs.nvidia.com/nemo/automodel/latest/datasets/diffusion-dataset)
- [拡散モデルのカバレッジ](https://docs.nvidia.com/nemo/automodel/latest/model-coverage/diffusion.html)
- [Transformers 向け NeMo Automodel（LLM/VLM のファインチューニング）](https://huggingface.co/docs/transformers/en/community_integrations/nemo_automodel_finetuning)