# 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/v0.43.0/zh/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
```

**混元视频：**

```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
```

### 输出格式

预处理会生成一个按分辨率桶组织的缓存目录。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 数据集准备](https://docs.nvidia.com/nemo/automodel/latest/datasets/diffusion-dataset) 指南，了解字幕格式、输入数据要求以及所有可用的预处理参数。

## 训练配置

微调由两个组件驱动：

1. 配方脚本（[finetune.py](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/diffusion/finetune/finetune.py)）是一个 Python 入口点，其中包含训练循环：加载模型、构建数据加载器、执行前向/反向传播、计算流匹配损失、保存检查点以及记录日志。
2. YAML 配置文件指定了该配方（recipe）所使用的全部设置：要微调的模型、数据所在位置、优化器超参数、并行策略等。您只需编辑此文件即可自定义训练，而无需修改代码，从而可以将规模从 1 块 GPU 扩展到数百块 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/v0.43.0/zh/api/models/wan_transformer_3d#diffusers.WanTransformer3DModel)、[FluxTransformer2DModel](/docs/diffusers/v0.43.0/zh/api/models/flux_transformer#diffusers.FluxTransformer2DModel) 和 [HunyuanVideoTransformer3DModel](/docs/diffusers/v0.43.0/zh/api/models/hunyuan_video_transformer_3d#diffusers.HunyuanVideoTransformer3DModel)）直接从 Hugging Face Hub 加载模型，并通过 [WanPipeline](/docs/diffusers/v0.43.0/zh/api/pipelines/wan#diffusers.WanPipeline) 和 [FluxPipeline](/docs/diffusers/v0.43.0/zh/api/pipelines/flux#diffusers.FluxPipeline) 等 Diffusers 流水线生成输出。

此集成为 Diffusers 用户提供了多项优势：

- **无需转换检查点**：来自 Hub 的预训练权重可直接使用。将 `pretrained_model_name_or_path` 指向任意 Diffusers 格式的模型 ID，即可立即开始训练。
- **Day-0 模型支持**：当新的扩散模型被添加到 Diffusers 并上传到 Hub 时，无需等待专用的训练脚本，即可使用 NeMo Automodel 对其进行微调。
- **与 Pipeline 兼容的输出**：微调后的检查点会以一种格式保存，该格式可以直接加载回 Diffusers 的 pipeline 中用于推理、在 Hub 上分享，或使用量化工具和编译工具进行进一步优化。
- **Diffusers 模型的可扩展训练**：NeMo Automodel 增加了分布式训练功能（FSDP2、多节点、多分辨率分桶），其能力超出了 Diffusers 内置训练脚本所提供的范围，同时保持了相同的模型和流水线接口。
- **共享生态系统**：来自 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)
- [扩散模型微调指南](https://docs.nvidia.com/nemo/automodel/latest/recipes-e2e-examples/diffusion-fine-tuning)
- [扩散数据集准备](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)