NeMo Automodel

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, 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 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.

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ợ

Cài đặt

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

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.

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.

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

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.

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

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

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)

Xem hướng dẫn Chuẩn bị bộ dữ liệu Diffusion để 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) 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:

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.

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 để có các cấu hình ví dụ đầy đủ cho tất cả các mô hình.

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

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

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

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.

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

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

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

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

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

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, FluxTransformer2DModel và HunyuanVideoTransformer3DModel, đồng thời tạo ra kết quả thông qua các pipeline của Diffusers như WanPipeline và 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, [email protected]
• Linnan Wang, [email protected]
• Alexandros Koumparoulis, [email protected]

Tài nguyên

• NeMo Automodel GitHub
• Hướng dẫn tinh chỉnh Diffusion
• Chuẩn bị bộ dữ liệu Diffusion
• Phạm vi hỗ trợ Mô hình khuếch tán
• NeMo Automodel cho Transformers (fine-tuning LLM/VLM)