Escalonamento de Ajuste Fino de Modelos de Imagem e Vídeo com NVIDIA NeMo Automodel e Hugging Face Diffusers

Os modelos de difusão impulsionaram alguns dos lançamentos open-source mais empolgantes dos últimos dois anos — pense no FLUX.1-dev para texto-para-imagem, ou no Wan 2.1 e HunyuanVideo para texto-para-vídeo. A biblioteca 🤗 Diffusers se tornou o lar padrão para esses modelos, oferecendo a pesquisadores e desenvolvedores uma interface única e consistente para inferência, adaptação e composição de pipelines.

O treinamento e o fine-tuning desses modelos também estão em alta, e isso exige ferramentas que lidem com sharding eficiente em memória, cache de latentes, bucketing multirresolução e configurações que escalam suavemente de uma GPU a centenas.

Apresentamos a biblioteca open-source NVIDIA NeMo Automodel. Hoje destacamos a colaboração entre a NVIDIA e a Hugging Face que traz treinamento distribuído de difusão de nível de produção para qualquer modelo no formato Diffusers no Hugging Face Hub — sem conversão de checkpoints, sem reescrita de modelos para novos modelos. A integração está documentada no guia de treinamento do Diffusers e é totalmente open-source sob a licença Apache 2.0.

O que é o NeMo Automodel?

O NeMo Automodel é uma biblioteca de treinamento open-source nativa em PyTorch DTensor, parte do framework NVIDIA NeMo. Ela é construída em torno de dois princípios de design que importam para o ecossistema Diffusers:

• Nativo do Hugging Face. Aponte pretrained_model_name_or_path para qualquer ID de modelo Diffusers no Hub e comece o treinamento. O NeMo Automodel usa classes de modelo Diffusers (ex.: WanTransformer3DModel) para carregamento e pipelines Diffusers (WanPipeline) para geração. Os checkpoints retornam perfeitamente ao ecossistema Diffusers.

• Um programa, qualquer escala. As receitas e scripts de treinamento podem ser facilmente modificados para se adequar ao treinamento em qualquer escala. O paralelismo é uma escolha de configuração, não uma reescrita de código — alterne entre FSDP2, paralelismo de tensor, paralelismo de especialistas, paralelismo de contexto e paralelismo de pipeline declarando configurações, sem reescrever modelos.

Atualmente, o AutoModel suporta apenas modelos de flow matching. Internamente, ele usa flow matching como objetivo de treinamento, com treinamento no espaço latente (via saídas de VAE pré-codificadas) e carregamento de dados com bucketing multirresolução para acelerar a taxa de transferência.

Modelos de difusão suportados

A integração do NeMo Automodel vem com receitas de fine-tuning prontas para uso para os modelos de difusão abertos abaixo. A lista reflete as receitas atualmente em examples/diffusion/finetune.

ModeloID no Hugging FaceTarefaParâmetrosReceita LoRAWan 2.1 T2V 1.3B / 14BWan-AI/Wan2.1-T2V-1.3B-Diffusers / Wan-AI/Wan2.1-T2V-14B-DiffusersTexto-para-Vídeo1.3B (cabe em uma única A100 de 40GB) / 14BSimWan 2.2 T2V A14BWan-AI/Wan2.2-T2V-A14B-DiffusersTexto-para-Vídeo27B total (MoE), 14B ativos por passoNãoFLUX.1-devblack-forest-labs/FLUX.1-devTexto-para-Imagem12BSimFLUX.2-devblack-forest-labs/FLUX.2-devTexto-para-Imagem32BSimHunyuanVideo 1.5hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2vTexto-para-Vídeo13BSimQwen-ImageQwen/Qwen-ImageTexto-para-Imagem20B (MMDiT)Sim

O que esta colaboração desbloqueia

Para os usuários do Diffusers, os ganhos práticos se dividem em algumas capacidades concretas.

Sem conversão de checkpoints. Pesos pré-treinados do Hub funcionam imediatamente. Não há um "formato de treinamento" separado para converter e depois converter de volta. Seu checkpoint fine-tuned carrega diretamente em um DiffusionPipeline para inferência, ou de volta ao Hub para compartilhamento. Ferramentas downstream — quantização, compilação, adaptadores LoRA, amostradores personalizados — continuam funcionando.

Caminho rápido para suporte a novos modelos. Quando um novo modelo de difusão chega ao Diffusers, habilitá-lo no NeMo Automodel requer uma adição de código pequena e contida — um manipulador de pré-processamento de dados e um adaptador de modelo — em vez de um script de treinamento personalizado completo. O restante da pilha de receitas (FSDP2, carregamento com bucketing, checkpointing, geração) permanece inalterado, e o mesmo fluxo de trabalho orientado por YAML se aplica.

Fine-tuning completo e eficiente em parâmetros. Tanto o fine-tuning completo quanto o PEFT estilo LoRA são suportados, para que você possa escolher entre qualidade máxima (FT completo em um cluster grande) ou eficiência máxima (LoRA em um único nó). A mesma estrutura de receita lida com ambos.

Treinamento escalável que vai além do que os scripts embutidos oferecem. O NeMo Automodel adiciona esquemas de sharding como FSDP2, paralelismo de tensor, contexto e pipeline, orquestração multi-nó (SLURM hoje, Kubernetes em breve) e bucketing multirresolução. Essas capacidades tornam possível treinar modelos maiores como FLUX.1-dev (12B) e HunyuanVideo (13B).

Uma olhada no fluxo de trabalho de fine-tuning

Nesta seção, percorremos o fluxo de trabalho típico para fine-tuning de qualquer um dos modelos suportados. A maneira recomendada de instalar o Automodel é o contêiner Docker do NeMo Automodel (nvcr.io/nvidia/nemo-automodel:26.06), que já vem com PyTorch, TransformerEngine e outras dependências compiladas com CUDA pré-construídas. Alternativamente, instale com pip3 install nemo-automodel ou a partir do código-fonte (pip3 install git+https://github.com/NVIDIA-NeMo/Automodel.git); veja o guia de instalação para todas as opções.

Este guia percorre um fine-tuning completo do transformer do FLUX.1-dev no conjunto de dados de tarô Rider–Waite com 78 cartas, e depois gera a partir do checkpoint resultante. Ele reutiliza os arquivos YAML de configuração já incluídos e aplica configurações específicas da execução como sobrescritas de linha de comando, sem exigir novos arquivos de configuração.

1. Pré-codificar o conjunto de dados

A receita de difusão consome latentes VAE em cache e embeddings de texto em vez de codificar imagens de origem a cada passo de treinamento. Transmita as 78 imagens Rider–Waite diretamente do Hugging Face e distribua o pré-processamento por todas as GPUs visíveis:

uv run --locked --no-default-groups \
  --extra diffusion \
  --extra diffusion-media \
python -m tools.diffusion.preprocessing_multiprocess image \
  --dataset_name multimodalart/1920-raider-waite-tarot-public-domain \
  --dataset_media_column image \
  --dataset_caption_column caption \
  --dataset_streaming \
  --max_images 78 \
  --output_dir /cache/flux_tarot \
  --processor flux \
  --model_name black-forest-labs/FLUX.1-dev \
  --max_pixels 245760

As legendas já contêm o token de gatilho trtcrd. Com este orçamento de pixels e a proporção de retrato do conjunto de dados, o pré-processamento atribui as amostras ao bucket de 384×640 usado na execução de demonstração.

Para treinamento de imagem, o pré-processamento produz arquivos de cache .pt e metadados fragmentados:

/cache/flux_tarot/
├── 384x640/
│   ├── <hash1>.pt
│   └── ...
├── metadata_shard_0000.json
├── metadata.json
└── _hf_dataset/
    └── images/

2. Iniciar o treinamento com o YAML FLUX existente

Use examples/diffusion/finetune/flux_t2i_flow.yaml diretamente. O YAML já seleciona FLUX.1-dev, fine-tuning completo do transformer, o adaptador de flow matching FLUX, um tamanho de lote efetivo de 32 e FSDP2 com oito vias.

Forneça os caminhos e configurações específicos do tarô como sobrescritas de linha de comando:

uv run --locked --no-default-groups --extra diffusion \
  torchrun --nproc-per-node=8 \
  examples/diffusion/finetune/finetune.py \
  -c examples/diffusion/finetune/flux_t2i_flow.yaml \
  --model.transformer_engine_fp8 false \
  --data.dataloader.cache_dir /cache/flux_tarot \
  --data.dataloader.base_resolution '[384,640]' \
  --lr_scheduler.lr_decay_style constant \
  --lr_scheduler.lr_warmup_steps 20 \
  --step_scheduler.max_steps 200 \
  --step_scheduler.ckpt_every_steps 50 \
  --checkpoint.checkpoint_dir /tmp/flux_tarot/checkpoints/full \
  --checkpoint.save_consolidated true \
  --seed 2026

A execução produz checkpoints nos passos 50, 100, 150 e 200. O checkpoint final é rotulado como epoch_66_step_199; o rótulo é baseado em zero, embora represente o 200º passo do otimizador concluído.

3. Gerar a partir do checkpoint fine-tuned

Use o YAML de geração FLUX existente e aponte model.checkpoint para o checkpoint de treinamento completo:

uv run --locked --no-default-groups --extra diffusion \
  python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_flux.yaml \
  --model.checkpoint /tmp/flux_tarot/checkpoints/full/epoch_66_step_199 \
  --inference.height 640 \
  --inference.width 384 \
  --inference.prompts '["a trtcrd of an astronaut tending a rose garden on Mars, \"the gardener\""]' \
  --output.output_dir /tmp/flux_tarot/generations/full/step_200 \
  --seed 2026

Inclua trtcrd para invocar o estilo de tarô aprendido. Para uma comparação de controle, mantenha a semente e a cena fixas, mas omita o gatilho:

uv run --locked --no-default-groups --extra diffusion \
  python examples/diffusion/generate/generate.py \
  -c examples/diffusion/generate/configs/generate_flux.yaml \
  --model.checkpoint /tmp/flux_tarot/checkpoints/full/epoch_66_step_199 \
  --inference.height 640 \
  --inference.width 384 \
  --inference.prompts '["an astronaut tending a rose garden on Mars, \"the gardener\""]' \
  --output.output_dir /tmp/flux_tarot/generations/control \
  --seed 2026

Resultados

No passo 200, os prompts de astronauta com gatilho mantêm o conteúdo solicitado enquanto adquirem uma paleta vintage creme, vermelha e preta, contornos de tinta pesados, campos de cor plana, tons de papel envelhecido e composição alegórica de cartas. O astronauta sem gatilho permanece fotográfico, demonstrando que o efeito aprendido está substancialmente associado a trtcrd, em vez de substituir o modelo base globalmente.

Prompt (semente)Linha de baseFine-tuned (passo 200)Astronauta cuidando de um roseiral em Marte (semente 2026)

4. Desempenho

Todas as medições foram coletadas em um nó com 8 GPUs NVIDIA H100 80GB. Os resultados são médias ± desvio padrão da amostra em três janelas de estado estacionário de 10 passos.

Texto-para-imagem — 512×512

ModeloTreinamentoParalelismoGBS / LBSTempo por passoImagens/sImagens/s/GPUPico alocado/GPUFLUX.1-devCompletoFSDP232 / 40,902 ± 0,039 s35,51 ± 1,554,44 ± 0,1963,88 GiBFLUX.1-devLoRA r64DDP48 / 60,894 ± 0,008 s53,73 ± 0,486,72 ± 0,0667,43 GiBQwen-ImageCompletoFSDP240 / 50,974 ± 0,075 s41,21 ± 3,065,15 ± 0,3853,55 GiBQwen-ImageLoRA r64DDP24 / 30,515 ± 0,006 s46,63 ± 0,545,83 ± 0,0766,33 GiB

Texto-para-vídeo — 512×512×49 quadros

Cada amostra é um clipe de vídeo de 49 quadros.

ModeloTreinamentoGBS / LBSActivation checkpointingTempo por passoClipes/sClipes/s/GPUPico alocado/GPUWan 2.1 1.3BCompleto8 / 1Desligado0,942 ± 0,038 s8,50 ± 0,351,06 ± 0,046,09 GiBWan 2.1 14BCompleto8 / 1Ligado3,798 ± 0,017 s2,107 ± 0,0060,263 ± 0,00633,35 GiBWan 2.1 14BLoRA r6416 / 2Ligado7,585 ± 0,014 s2,110 ± 0,0000,263 ± 0,00024,07 GiBWan 2.2 A14B, alto ruídoCompleto8 / 1Ligado4,628 ± 0,031 s1,730 ± 0,0100,217 ± 0,00623,57 GiBHunyuanVideo 1.5Completo8 / 1Ligado5,926 ± 0,046 s1,350 ± 0,0100,170 ± 0,00015,90 GiBHunyuanVideo 1.5LoRA r648 / 1Ligado5,575 ± 0,006 s1,433 ± 0,0060,180 ± 0,00010,58 GiB

Detalhes da medição

• Hardware: 8× H100 80GB HBM3, totalmente conectados via NVLink.

• Conjunto de dados de imagem: lambda/naruto-blip-captions, 256 amostras em cache.

• Conjunto de dados de vídeo: svjack/Lelouch_Vi_Britannia_FramePack_First_Last_Frame_Video_Captioned, 112 amostras em cache.

• Lotes completos foram aplicados com drop_last=true.

• Escritas de checkpoint foram desabilitadas.

• O tempo por passo inclui carregamento de dados, forward, backward, recorte de gradiente, otimizador e trabalho do scheduler.

• A memória é o pico de uso do alocador CUDA do PyTorch, não a memória total do dispositivo NVML.

Outros exemplos de Fine-tuned/LoRA

Os resultados do fine-tuning e LoRA mostram o poder do NeMo Automodel para especialização de domínio. Por exemplo, o fine-tuning do modelo Wan 2.1 em um conjunto de dados de vídeo Ghibli adaptou com sucesso o estilo de saída, demonstrado por uma mudança perceptível na aparência de uma flor em comparação com a linha de base.

Linha de base:

Fine-tuned em vídeos do Ghibli:

Também observamos o impacto distinto do uso de LoRA, onde a aplicação do adaptador ao Wan 2.1 fez o vídeo adotar um estilo Ghibli característico, particularmente visível no destaque dos olhos dos personagens.

Sem LoRA:

Com LoRA:

Esses exemplos, incluindo os do FLUX.2, confirmam que os usuários podem alcançar tanto a máxima qualidade via fine-tuning completo quanto a máxima eficiência via PEFT estilo LoRA, adaptando a saída a domínios estilísticos específicos.

Experimente hoje

Saiba mais sobre a integração e encontre mais exemplos de fine-tuning na documentação do NeMo Automodel

Próximos passos: APIs de receitas em Python

YAML é uma ótima opção para configuração reproduzível, especialmente para equipes que desejam arquivos que possam ser commitados, revisados e reutilizados, mas muitas equipes também precisam de uma interface programática.

Em um próximo lançamento do NeMo Automodel, planejamos disponibilizar as receitas de difusão também através de uma API Python totalmente tipada. Os usuários poderão compor os mesmos componentes de modelo, dados, otimizador, PEFT/LoRA, paralelismo, checkpointing e geração diretamente do Python.

O caminho Python tem como objetivo tornar as receitas mais fáceis de usar a partir de código de treinamento existente, notebooks e fluxos de trabalho experimentais, e oferecer uma interface Python de primeira classe junto com o caminho rápido YAML.