NeMo Automodel

NeMo Automodel é uma biblioteca de treinamento nativa de PyTorch DTensor da NVIDIA para ajuste fino e pré-treino de modelos de difusão em escala. É nativa do Hugging Face — treine qualquer modelo no formato Diffusers do Hub sem conversão de checkpoints. A mesma receita YAML e o script de treinamento hackeável rodam em qualquer escala, de 1 GPU a centenas de nós, com treinamento distribuído FSDP2, carregamento de dados em buckets multirresolução e treinamento em espaço latente pré-codificado para máxima utilização da GPU. Usa flow matching para o treinamento e é totalmente open source (Apache 2.0), com suporte da NVIDIA e manutenção ativa.

O NeMo Automodel se integra diretamente ao Diffusers. Ele carrega modelos pré-treinados do Hugging Face Hub usando as classes de modelo do Diffusers e gera saídas com o DiffusionPipeline.

O fluxo de trabalho típico consiste em instalar o NeMo Automodel (pip ou Docker), preparar seus dados codificando-os em arquivos .meta, configurar uma receita YAML, iniciar o treinamento com torchrun e executar a inferência com o checkpoint resultante.

Modelos suportados

Instalação

Requisitos de hardware

Instale o NeMo Automodel com pip. Para o conjunto completo de métodos de instalação (incluindo a partir do código-fonte), consulte o guia de instalação do NeMo Automodel.

pip3 install nemo-automodel

Como alternativa, use o contêiner Docker pré-compilado, que inclui todas as dependências.

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

Os checkpoints são perdidos quando o container é encerrado, a menos que você monte o diretório de checkpoints no host. Por exemplo, adicione -v /host/path/checkpoints:/workspace/checkpoints ao comando docker run.

Preparação de dados

NeMo Automodel treina modelos de difusão no espaço latente. Imagens ou vídeos brutos devem ser pré-processados em arquivos .meta contendo latentes de VAE e incorporações de texto antes do treinamento. Isso evita a re-codificação em cada etapa de treinamento.

Use a ferramenta de pré-processamento integrada para codificar seus dados. A ferramenta distribui automaticamente o trabalho entre todas as GPUs disponíveis.

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

Formato de saída

O pré-processamento gera um diretório de cache organizado por bucket de resolução. O NeMo Automodel suporta treinamento multi-resolução por meio de amostragem em buckets. As amostras são agrupadas por resolução espacial para que cada lote contenha amostras do mesmo tamanho, evitando desperdício por 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)

Consulte o guia Preparação de Dataset para Difusão para formatos de legenda, requisitos de dados de entrada e todos os argumentos de pré-processamento disponíveis.

Configuração de treinamento

O fine-tuning é impulsionado por dois componentes:

1. Um script de receita (finetune.py) é um ponto de entrada em Python que contém o loop de treinamento: carregamento do modelo, construção do dataloader, execução das passadas de forward/backward, cálculo da perda de flow matching, checkpointing e registro de logs.
2. Um arquivo de configuração YAML especifica todas as configurações usadas pela receita: qual modelo fazer o fine-tuning, onde estão os dados, hiperparâmetros do otimizador, estratégia de paralelismo e mais. Você personaliza o treinamento editando esse arquivo, em vez de modificar o código, permitindo escalar de 1 para centenas de GPUs.

Qualquer campo YAML também pode ser sobrescrito pela 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

Abaixo está a configuração anotada para o ajuste fino do Wan 2.1 T2V 1.3B, com cada seção explicada.

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

Referência de campos de configuração

A tabela abaixo lista as configurações mínimas necessárias. Consulte os exemplos do NeMo Automodel para obter configurações de exemplo completas para todos os modelos.

Iniciar o treinamento

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

Geração

Após o treinamento, gere vídeos ou imagens a partir de prompts de texto usando o checkpoint ajustado.

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

Integração com Diffusers

O NeMo Automodel é construído sobre o Diffusers e o utiliza como base para o carregamento de modelos e inferência. Ele carrega modelos diretamente do Hugging Face Hub usando classes de modelos do Diffusers, como WanTransformer3DModel, FluxTransformer2DModel e HunyuanVideoTransformer3DModel, e gera saídas por meio de pipelines do Diffusers, como WanPipeline e FluxPipeline.

Esta integração oferece vários benefícios para os usuários do Diffusers:

• Na conversão de checkpoint: os pesos pré-treinados do Hub funcionam imediatamente. Aponte pretrained_model_name_or_path para qualquer ID de modelo no formato Diffusers e comece a treinar na hora.
• Suporte a modelos no dia zero: quando um novo modelo de difusão é adicionado ao Diffusers e enviado ao Hub, ele pode ser ajustado com o NeMo Automodel sem esperar por um script de treinamento dedicado.
• Saídas compatíveis com Pipeline: os checkpoints de fine-tuning são salvos em um formato que pode ser carregado diretamente de volta nos pipelines do Diffusers para inferência, compartilhamento no Hub ou otimização adicional com ferramentas como quantização e compilação.
• Escalonável treinamento para modelos Diffusers: NeMo Automodel adiciona capacidades de treinamento distribuído (FSDP2, multi-nó, bucketing multirresolução) que vão além do que os scripts de treinamento integrados do Diffusers oferecem, mantendo as mesmas interfaces de modelo e pipeline.
• Ecossistema compartilhado: qualquer modelo, adaptador LoRA ou componente de pipeline do ecossistema Diffusers permanece compatível em todo o fluxo de trabalho de treinamento e inferência.

Equipe NVIDIA

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

Recursos

• NeMo Automodel GitHub
• Guia de Ajuste Fino de Difusão
• Preparação do Conjunto de Dados de Difusão
• Cobertura de Modelos de Difusão
• NeMo Automodel para Transformers (fine-tuning de LLM/VLM)