# NeMo Automodel

[NeMo Automodel](https://github.com/NVIDIA-NeMo/Automodel) là thư viện huấn luyện PyTorch DTensor-native từ NVIDIA, dùng để fine-tuning và pretraining các mô hình diffusion ở quy mô lớn. Nó hoàn toàn native trên Hugging Face — bạn có thể huấn luyện bất kỳ mô hình nào ở định dạng Diffusers từ Hub mà không cần chuyển đổi checkpoint. Cùng một công thức YAML và script huấn luyện có thể tùy chỉnh chạy ở mọi quy mô, từ 1 GPU đến hàng trăm node, với huấn luyện phân tán [FSDP2](https://pytorch.org/docs/stable/fsdp.html), tải dữ liệu theo bucket đa độ phân giải và huấn luyện trong không gian latent đã mã hóa sẵn để tối đa hóa hiệu suất sử dụng GPU. Thư viện này sử dụng [flow matching](https://huggingface.co/papers/2210.02747) cho quá trình huấn luyện, hoàn toàn mã nguồn mở (Apache 2.0), được NVIDIA hỗ trợ và được bảo trì tích cực.

NeMo Automodel tích hợp trực tiếp với Diffusers. Nó tải các mô hình được huấn luyện trước từ Hugging Face Hub bằng cách sử dụng các lớp mô hình của Diffusers và tạo ra các kết quả đầu ra bằng [DiffusionPipeline](/docs/diffusers/v0.43.0/vi/api/pipelines/overview#diffusers.DiffusionPipeline).

Quy trình làm việc điển hình là cài đặt NeMo Automodel (pip hoặc Docker), chuẩn bị dữ liệu bằng cách mã hóa chúng thành các tệp `.meta`, cấu hình công thức YAML, khởi động huấn luyện bằng `torchrun` và chạy suy luận với checkpoint kết quả.

## Các mô hình được hỗ trợ

| Mô hình | ID Hugging Face | Nhiệm vụ | Tham số | Trường hợp sử dụng |
|-------|----------------|------|------------|----------|
| 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) | Text-to-Video | 1.3B | tạo video trên phần cứng hạn chế (chạy được trên một thẻ A100 40GB duy nhất) |
| FLUX.1-dev | [black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) | Văn bản sang Hình ảnh | 12B | tạo hình ảnh chất lượng cao |
| HunyuanVideo 1.5 | [hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v](https://huggingface.co/hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v) | Text-to-Video | 13B | tạo video chất lượng cao |

## Cài đặt

### Yêu cầu phần cứng

| Thành phần | Tối thiểu | Được khuyến nghị |
|-----------|---------|-------------|
| GPU | A100 40GB | A100 80GB / H100 |
| GPU | 4 | 8+ |
| RAM | 128 GB | 256 GB+ |
| Lưu trữ | 500 GB SSD | 2 TB NVMe |

Cài đặt NeMo Automodel bằng pip. Để xem đầy đủ các phương pháp cài đặt (bao gồm cài đặt từ mã nguồn), hãy xem [hướng dẫn cài đặt NeMo Automodel](https://docs.nvidia.com/nemo/automodel/latest/guides/installation.html).

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

Ngoài ra, hãy sử dụng container Docker có sẵn bao gồm tất cả các phụ thuộc.

```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]
> Checkpoint sẽ bị mất khi container thoát trừ khi bạn bind-mount thư mục checkpoint vào host. Ví dụ, thêm `-v /host/path/checkpoints:/workspace/checkpoints` vào lệnh `docker run`.

## Chuẩn bị dữ liệu

NeMo Automodel huấn luyện các mô hình khuếch tán trong không gian tiềm ẩn. Ảnh hoặc video thô phải được tiền xử lý thành các tệp `.meta` chứa các giá trị tiềm ẩn của VAE và các vector nhúng văn bản trước khi huấn luyện. Điều này giúp tránh việc mã hóa lại ở mỗi bước huấn luyện.

Sử dụng công cụ tiền xử lý tích hợp để mã hóa dữ liệu của bạn. Công cụ này tự động phân phối công việc trên tất cả các GPU khả dụng.

Lệnh tiền xử lý video giống nhau cho cả Wan 2.1 và HunyuanVideo, nhưng các cờ khác nhau. Wan 2.1 sử dụng `--processor wan` cùng với `--resolution_preset` và `--caption_format sidecar`, trong khi HunyuanVideo sử dụng `--processor hunyuan` cùng với `--target_frames` để đặt số khung hình và `--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
```

### Định dạng đầu ra

Preprocessing tạo ra một thư mục cache được tổ chức theo từng nhóm độ phân giải. NeMo Automodel hỗ trợ huấn luyện đa độ phân giải thông qua lấy mẫu theo nhóm. Các mẫu được nhóm lại theo độ phân giải không gian để mỗi batch chứa các mẫu có cùng kích thước, tránh lãng phí do đệm (padding).

```
/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]
> Xem hướng dẫn [Chuẩn bị bộ dữ liệu Diffusion](https://docs.nvidia.com/nemo/automodel/latest/datasets/diffusion-dataset) để biết về định dạng chú thích, yêu cầu dữ liệu đầu vào và tất cả các đối số tiền xử lý khả dụng.

## Cấu hình huấn luyện

Fine-tuning được điều khiển bởi hai thành phần:

1. Một script công thức ([finetune.py](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/diffusion/finetune/finetune.py)) là điểm vào Python chứa vòng lặp huấn luyện: nạp mô hình, xây dựng dataloader, thực hiện các bước truyền thuận/truyền ngược, tính toán tổn thất flow matching, lưu checkpoint và ghi log.
2. Một tệp cấu hình YAML xác định tất cả các cài đặt mà công thức sử dụng: mô hình nào để tinh chỉnh, dữ liệu nằm ở đâu, siêu tham số của bộ tối ưu hóa, chiến lược song song hóa, và nhiều hơn nữa. Bạn tùy chỉnh quá trình huấn luyện bằng cách chỉnh sửa tệp này thay vì sửa đổi mã, cho phép bạn mở rộng quy mô từ 1 đến hàng trăm GPU.

Bất kỳ trường YAML nào cũng có thể được ghi đè từ 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
```

Dưới đây là cấu hình có chú thích để tinh chỉnh Wan 2.1 T2V 1.3B, với mỗi phần được giải thích.

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

### Tham chiếu trường cấu hình

Bảng dưới đây liệt kê các cấu hình tối thiểu cần thiết. Xem [các ví dụ NeMo Automodel](https://github.com/NVIDIA-NeMo/Automodel/tree/main/examples/diffusion/finetune) để có các cấu hình ví dụ đầy đủ cho tất cả các mô hình.

| Mục | Bắt buộc? | Những gì cần thay đổi |
|---------|-----------|----------------|
| `model` | Có | Đặt `pretrained_model_name_or_path` thành ID mô hình Hugging Face. Đặt `mode: finetune` hoặc `mode: pretrain`. |
| `step_scheduler` | Có | `global_batch_size` là kích thước batch hiệu dụng trên tất cả các GPU. `ckpt_every_steps` kiểm soát tần suất tạo checkpoint. Gradient accumulation được tính toán tự động. |
| `data` | Có | Đặt `cache_dir` thành đường dẫn chứa các tệp `.meta` đã tiền xử lý của bạn. Thay đổi `_target_` và `model_type` cho các mô hình khác nhau. |
| `optim` | Có | `learning_rate: 5e-6` là giá trị mặc định tốt cho việc tinh chỉnh. Hãy điều chỉnh cho phù hợp với bộ dữ liệu và mô hình của bạn. |
| `lr_scheduler` | Có | Chọn `cosine`, `linear` hoặc `constant` cho `lr_decay_style`. Đặt `lr_warmup_steps` để warmup dần. |
| `flow_matching` | Có | `adapter_type` phải khớp với model (`simple` cho Wan, `flux` cho FLUX, `hunyuan` cho HunyuanVideo). Xem cấu hình cụ thể của model cho `adapter_kwargs`. |
| `fsdp` | Có | Đặt `dp_size` bằng số lượng GPU. Đối với cấu hình đa nút, đặt bằng tổng số GPU trên tất cả các nút. |
| `checkpoint` | Được khuyến nghị | Đặt `checkpoint_dir` thành một đường dẫn bền vững, đặc biệt là trong Docker. Sử dụng `restore_from` để tiếp tục từ một checkpoint trước đó. |
| `wandb` | Tùy chọn | Cấu hình để bật theo dõi thí nghiệm Weights & Biases. Đặt `mode: disabled` để tắt. |

## Khởi động huấn luyện

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

Chạy lệnh sau trên mỗi node, đặt `NODE_RANK` cho phù hợp:

```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]
> Để huấn luyện đa nút, đặt `fsdp.dp_size` trong YAML thành tổng số GPU **tổng** trên tất cả các nút (ví dụ: 16 cho 2 nút, mỗi nút có 8 GPU).

## Sinh

Sau khi huấn luyện, hãy tạo video hoặc hình ảnh từ các prompt văn bản bằng checkpoint đã tinh chỉnh.

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

Với một checkpoint đã được tinh chỉnh:

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

Với một checkpoint đã được tinh chỉnh:

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

Với một checkpoint đã được tinh chỉnh:

```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"]'
```

## Tích hợp Diffusers

NeMo Automodel được xây dựng dựa trên Diffusers và sử dụng nó làm nền tảng cho việc tải mô hình và suy luận. Nó tải mô hình trực tiếp từ Hugging Face Hub bằng cách sử dụng các lớp mô hình của Diffusers như [WanTransformer3DModel](/docs/diffusers/v0.43.0/vi/api/models/wan_transformer_3d#diffusers.WanTransformer3DModel), [FluxTransformer2DModel](/docs/diffusers/v0.43.0/vi/api/models/flux_transformer#diffusers.FluxTransformer2DModel) và [HunyuanVideoTransformer3DModel](/docs/diffusers/v0.43.0/vi/api/models/hunyuan_video_transformer_3d#diffusers.HunyuanVideoTransformer3DModel), đồng thời tạo ra kết quả thông qua các pipeline của Diffusers như [WanPipeline](/docs/diffusers/v0.43.0/vi/api/pipelines/wan#diffusers.WanPipeline) và [FluxPipeline](/docs/diffusers/v0.43.0/vi/api/pipelines/flux#diffusers.FluxPipeline).

Tích hợp này mang lại một số lợi ích cho người dùng Diffusers:

- **Không cần chuyển đổi checkpoint**: trọng số được huấn luyện trước từ Hub hoạt động ngay lập tức. Chỉ `pretrained_model_name_or_path` đến bất kỳ ID mô hình nào ở định dạng Diffusers và bắt đầu huấn luyện ngay.
- **Hỗ trợ mô hình ngay từ ngày đầu**: khi một mô hình khuếch tán mới được thêm vào Diffusers và tải lên Hub, nó có thể được tinh chỉnh với NeMo Automodel mà không cần chờ đợi một script huấn luyện riêng biệt.
- **Đầu ra tương thích với Pipeline**: các checkpoint đã fine-tune được lưu ở định dạng có thể nạp trực tiếp trở lại vào các pipeline của Diffusers để suy luận, chia sẻ trên Hub, hoặc tối ưu hóa thêm bằng các công cụ như lượng tử hóa và biên dịch.
- **Đào tạo có khả năng mở rộng cho các mô hình Diffusers**: NeMo Automodel bổ sung các khả năng đào tạo phân tán (FSDP2, đa nút, phân loại bucket đa độ phân giải) vượt xa những gì các script đào tạo Diffusers tích hợp sẵn cung cấp, trong khi vẫn giữ nguyên các giao diện mô hình và pipeline.
- **Hệ sinh thái chung**: bất kỳ mô hình, bộ thích ứng LoRA hay thành phần pipeline nào từ hệ sinh thái Diffusers đều duy trì khả năng tương thích xuyên suốt quy trình huấn luyện và suy luận.

## Đội ngũ NVIDIA

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

## Tài nguyên

- [NeMo Automodel GitHub](https://github.com/NVIDIA-NeMo/Automodel)
- [Hướng dẫn tinh chỉnh Diffusion](https://docs.nvidia.com/nemo/automodel/latest/recipes-e2e-examples/diffusion-fine-tuning)
- [Chuẩn bị bộ dữ liệu Diffusion](https://docs.nvidia.com/nemo/automodel/latest/datasets/diffusion-dataset)
- [Phạm vi hỗ trợ Mô hình khuếch tán](https://docs.nvidia.com/nemo/automodel/latest/model-coverage/diffusion.html)
- [NeMo Automodel cho Transformers (fine-tuning LLM/VLM)](https://huggingface.co/docs/transformers/en/community_integrations/nemo_automodel_finetuning)