NeMo Automodel

NeMo Automodel adalah pustaka pelatihan native PyTorch DTensor dari NVIDIA untuk fine-tuning dan pretraining model difusi berskala besar. Pustaka ini native Hugging Face — latih model berformat Diffusers mana pun dari Hub tanpa perlu konversi checkpoint. Resep YAML dan skrip pelatihan yang dapat dimodifikasi yang sama berjalan pada skala apa pun, mulai dari 1 GPU hingga ratusan node, dengan pelatihan terdistribusi FSDP2, pemuatan data berbucket multiresolusi, dan pelatihan ruang laten pra-terkodekan untuk pemanfaatan GPU maksimal. Pustaka ini menggunakan flow matching untuk pelatihan dan sepenuhnya open source (Apache 2.0), didukung oleh NVIDIA, serta dipelihara secara aktif.

NeMo Automodel terintegrasi langsung dengan Diffusers. Ia memuat model pralatih dari Hugging Face Hub menggunakan kelas model Diffusers dan menghasilkan output dengan DiffusionPipeline.

Alur kerja yang umum adalah menginstal NeMo Automodel (pip atau Docker), menyiapkan data Anda dengan mengenkodnya ke dalam file .meta, mengonfigurasi resep YAML, meluncurkan pelatihan dengan torchrun, dan menjalankan inferensi dengan checkpoint yang dihasilkan.

Model yang didukung

Instalasi

Kebutuhan perangkat keras

Install NeMo Automodel dengan pip. Untuk seluruh metode instalasi (termasuk dari source), lihat panduan instalasi NeMo Automodel.

pip3 install nemo-automodel

Sebagai alternatif, gunakan kontainer Docker yang sudah dibuat sebelumnya yang mencakup semua dependensi.

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 akan hilang saat container keluar kecuali Anda bind-mount direktori checkpoint ke host. Misalnya, tambahkan -v /host/path/checkpoints:/workspace/checkpoints ke perintah docker run.

Persiapan data

NeMo Automodel melatih model difusi di ruang laten. Gambar atau video mentah harus dipreproses menjadi file .meta yang berisi laten VAE dan embedding teks sebelum pelatihan. Hal ini menghindari pengkodean ulang pada setiap langkah pelatihan.

Gunakan alat pra-pemrosesan bawaan untuk mengkodekan data Anda. Alat ini secara otomatis mendistribusikan pekerjaan di semua GPU yang tersedia.

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

Format output

Preprocessing menghasilkan direktori cache yang diorganisasi berdasarkan bucket resolusi. NeMo Automodel mendukung pelatihan multi-resolusi melalui bucketed sampling. Sampel dikelompokkan berdasarkan resolusi spasial sehingga setiap batch berisi sampel berukuran sama, menghindari pemborosan 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)

Lihat panduan Persiapan Dataset Difusi untuk format caption, persyaratan data input, dan semua argumen pra-pemrosesan yang tersedia.

Konfigurasi pelatihan

Fine-tuning didorong oleh dua komponen:

1. Skrip resep (finetune.py) adalah titik masuk Python yang berisi loop pelatihan: memuat model, membangun dataloader, menjalankan langkah maju/mundur, menghitung kerugian flow matching, membuat checkpoint, dan mencatat log.
2. File konfigurasi YAML menentukan semua pengaturan yang digunakan oleh resep: model mana yang akan dihaluskan, di mana data berada, hiperparameter optimizer, strategi paralelisme, dan lainnya. Anda menyesuaikan pelatihan dengan mengedit file ini, bukan dengan memodifikasi kode, sehingga memungkinkan Anda untuk berskala dari 1 hingga ratusan GPU.

Setiap field YAML juga dapat ditimpa dari 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

Di bawah ini adalah konfigurasi anotasi untuk fine-tuning Wan 2.1 T2V 1.3B, dengan setiap bagian dijelaskan.

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

Referensi bidang konfigurasi

Tabel di bawah ini mencantumkan konfigurasi minimal yang diperlukan. Lihat contoh NeMo Automodel untuk contoh konfigurasi lengkap untuk semua model.

Mulai pelatihan

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

Generasi

Setelah pelatihan, generate video atau gambar dari prompt teks menggunakan checkpoint yang telah dihaluskan.

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

Integrasi Diffusers

NeMo Automodel dibangun di atas Diffusers dan menggunakannya sebagai tulang punggung untuk pemuatan model dan inferensi. Model dimuat langsung dari Hugging Face Hub menggunakan kelas model Diffusers seperti WanTransformer3DModel, FluxTransformer2DModel, dan HunyuanVideoTransformer3DModel, serta menghasilkan output melalui pipeline Diffusers seperti WanPipeline dan FluxPipeline.

Integrasi ini memberikan beberapa manfaat bagi pengguna Diffusers:

• Tidak ada konversi checkpoint: bobot pretrained dari Hub dapat langsung digunakan. Arahkan pretrained_model_name_or_path ke ID model berformat Diffusers mana pun dan mulai pelatihan segera.
• Dukungan model sejak hari pertama: ketika model difusi baru ditambahkan ke Diffusers dan diunggah ke Hub, model tersebut dapat dihaluskan (fine-tuned) dengan NeMo Automodel tanpa perlu menunggu skrip pelatihan khusus.
• Output yang kompatibel dengan Pipeline: checkpoint hasil fine-tuning disimpan dalam format yang dapat dimuat kembali secara langsung ke dalam pipeline Diffusers untuk inferensi, berbagi di Hub, atau optimasi lebih lanjut dengan alat seperti kuantisasi dan kompilasi.
• Pelatihan yang dapat diskalakan untuk model Diffusers: NeMo Automodel menambahkan kemampuan pelatihan terdistribusi (FSDP2, multi-node, bucketing multiresolution) yang melampaui apa yang disediakan oleh skrip pelatihan bawaan Diffusers, sambil tetap mempertahankan antarmuka model dan pipeline yang sama.
• Ekosistem bersama: model, adapter LoRA, atau komponen pipeline mana pun dari ekosistem Diffusers tetap kompatibel di seluruh alur kerja pelatihan dan inferensi.

Tim NVIDIA

• Pranav Prashant Thombre, [email protected]
• Linnan Wang, [email protected]
• Alexandros Koumparoulis, [email protected]

Sumber daya

• NeMo Automodel GitHub
• Panduan Fine-Tuning Diffusion
• Persiapan Dataset Difusi
• Cakupan Model Difusi
• NeMo Automodel untuk Transformers (fine-tuning LLM/VLM)