# NeMo Automodel

[NeMo Automodel](https://github.com/NVIDIA-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](https://pytorch.org/docs/stable/fsdp.html), 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](https://huggingface.co/papers/2210.02747) 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](/docs/diffusers/main/pt/api/pipelines/overview#diffusers.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

| Model | ID do Hugging Face | Tarefa | Parâmetros | Caso de uso |
|-------|----------------|------|------------|----------|
| 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) | Texto para Vídeo | 1.3B | geração de vídeo em hardware limitado (cabe em uma única A100 de 40GB) |
| FLUX.1-dev | [black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) | Texto para Imagem | 12B | geração de imagens de alta qualidade |
| HunyuanVideo 1.5 | [hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v](https://huggingface.co/hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v) | Texto para Vídeo | 13B | geração de vídeo de alta qualidade |

## Instalação

### Requisitos de hardware

| Componente | Mínimo | Recomendado |
|-----------|---------|-------------|
| GPU | A100 40GB | A100 80GB / H100 |
| GPUs | 4 | 8+ |
| RAM | 128 GB | 256 GB+ |
| Armazenamento | SSD de 500 GB | 2 TB NVMe |

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](https://docs.nvidia.com/nemo/automodel/latest/guides/installation.html).

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

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

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

O comando de pré-processamento de vídeo é o mesmo para o Wan 2.1 e o HunyuanVideo, mas as flags diferem. O Wan 2.1 usa `--processor wan` com `--resolution_preset` e `--caption_format sidecar`, enquanto o HunyuanVideo usa `--processor hunyuan` com `--target_frames` para definir a contagem de quadros e `--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
```

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

> [!TIP]
> Consulte o guia [Preparação de Dataset para Difusão](https://docs.nvidia.com/nemo/automodel/latest/datasets/diffusion-dataset) 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](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/diffusion/finetune/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:

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

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

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

### Referência de campos de configuração

A tabela abaixo lista as configurações mínimas necessárias. Consulte os [exemplos do NeMo Automodel](https://github.com/NVIDIA-NeMo/Automodel/tree/main/examples/diffusion/finetune) para obter configurações de exemplo completas para todos os modelos.

| Seção | Obrigatório? | O que alterar |
|---------|-----------|----------------|
| `model` | Sim | Defina `pretrained_model_name_or_path` como o ID do modelo no Hugging Face. Defina `mode: finetune` ou `mode: pretrain`. |
| `step_scheduler` | Sim | `global_batch_size` é o tamanho efetivo do lote em todas as GPUs. `ckpt_every_steps` controla a frequência dos checkpoints. A acumulação de gradientes é calculada automaticamente. |
| `data` | Sim | Defina `cache_dir` como o caminho que contém seus arquivos `.meta` pré-processados. Altere `_target_` e `model_type` para diferentes modelos. |
| `optim` | Sim | `learning_rate: 5e-6` é um valor padrão adequado para ajuste fino. Ajuste conforme seu conjunto de dados e modelo. |
| `lr_scheduler` | Sim | Escolha `cosine`, `linear` ou `constant` para `lr_decay_style`. Defina `lr_warmup_steps` para um aquecimento gradual. |
| `flow_matching` | Sim | `adapter_type` deve corresponder ao modelo (`simple` para Wan, `flux` para FLUX, `hunyuan` para HunyuanVideo). Consulte as configurações específicas do modelo para `adapter_kwargs`. |
| `fsdp` | Sim | Defina `dp_size` para o número de GPUs. Para multi-nó, defina para o total de GPUs em todos os nós. |
| `checkpoint` | Recomendado | Defina `checkpoint_dir` para um caminho persistente, especialmente em Docker. Use `restore_from` para retomar de um checkpoint anterior. |
| `wandb` | Opcional | Configure para habilitar o rastreamento de experimentos do Weights & Biases. Defina `mode: disabled` para desativar. |

## Iniciar o treinamento

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

Execute o seguinte em cada nó, definindo `NODE_RANK` de acordo:

```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]
> Para treinamento em múltiplos nós, defina `fsdp.dp_size` no YAML com o número **total** de GPUs em todos os nós (por exemplo, 16 para 2 nós com 8 GPUs cada).

## Geração

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

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

Com um checkpoint ajustado finamente:

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

Com um checkpoint ajustado finamente:

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

Com um checkpoint ajustado finamente:

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

## 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](/docs/diffusers/main/pt/api/models/wan_transformer_3d#diffusers.WanTransformer3DModel), [FluxTransformer2DModel](/docs/diffusers/main/pt/api/models/flux_transformer#diffusers.FluxTransformer2DModel) e [HunyuanVideoTransformer3DModel](/docs/diffusers/main/pt/api/models/hunyuan_video_transformer_3d#diffusers.HunyuanVideoTransformer3DModel), e gera saídas por meio de pipelines do Diffusers, como [WanPipeline](/docs/diffusers/main/pt/api/pipelines/wan#diffusers.WanPipeline) e [FluxPipeline](/docs/diffusers/main/pt/api/pipelines/flux#diffusers.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](https://github.com/NVIDIA-NeMo/Automodel)
- [Guia de Ajuste Fino de Difusão](https://docs.nvidia.com/nemo/automodel/latest/recipes-e2e-examples/diffusion-fine-tuning)
- [Preparação do Conjunto de Dados de Difusão](https://docs.nvidia.com/nemo/automodel/latest/datasets/diffusion-dataset)
- [Cobertura de Modelos de Difusão](https://docs.nvidia.com/nemo/automodel/latest/model-coverage/diffusion.html)
- [NeMo Automodel para Transformers (fine-tuning de LLM/VLM)](https://huggingface.co/docs/transformers/en/community_integrations/nemo_automodel_finetuning)