NeMo Automodel

NeMo Automodel es una biblioteca de entrenamiento nativa de PyTorch DTensor de NVIDIA para el ajuste fino y el preentrenamiento de modelos de difusión a gran escala. Es nativa de Hugging Face: entrena cualquier modelo en formato Diffusers del Hub sin necesidad de convertir puntos de control. La misma receta YAML y el mismo script de entrenamiento modificable se ejecutan a cualquier escala, desde 1 GPU hasta cientos de nodos, con entrenamiento distribuido FSDP2, carga de datos por cubos multiresolución y entrenamiento en el espacio latente precodificado para una máxima utilización de la GPU. Utiliza flow matching para el entrenamiento y es completamente de código abierto (Apache 2.0), con soporte de NVIDIA y mantenimiento activo.

NeMo Automodel se integra directamente con Diffusers. Carga modelos preentrenados desde Hugging Face Hub utilizando las clases de modelos de Diffusers y genera salidas con DiffusionPipeline.

El flujo de trabajo típico consiste en instalar NeMo Automodel (pip o Docker), preparar los datos codificándolos en archivos .meta, configurar una receta YAML, iniciar el entrenamiento con torchrun y ejecutar la inferencia con el punto de control resultante.

Modelos compatibles

Instalación

Requisitos de hardware

Instale NeMo Automodel con pip. Para el conjunto completo de métodos de instalación (incluido desde el código fuente), consulte la guía de instalación de NeMo Automodel.

pip3 install nemo-automodel

Como alternativa, use el contenedor de Docker precompilado que incluye todas las dependencias.

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

Los puntos de control se pierden al salir del contenedor, a menos que se monte el directorio de puntos de control en el host mediante un bind mount. Por ejemplo, añade -v /host/path/checkpoints:/workspace/checkpoints al comando docker run.

Preparación de datos

NeMo Automodel entrena modelos de difusión en el espacio latente. Las imágenes o videos en bruto deben preprocesarse en archivos .meta que contienen latentes de VAE y embeddings de texto antes del entrenamiento. Esto evita re-codificar en cada paso de entrenamiento.

Utilice la herramienta de preprocesamiento integrada para codificar sus datos. La herramienta distribuye automáticamente el trabajo entre todas las GPU disponibles.

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 salida

El preprocesamiento genera un directorio de caché organizado por cubo de resolución. NeMo Automodel admite el entrenamiento multi-resolución mediante muestreo por cubos. Las muestras se agrupan por resolución espacial para que cada lote contenga muestras del mismo tamaño, evitando el desperdicio por relleno.

/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 la guía de Preparación de conjuntos de datos de difusión para conocer los formatos de descripción, los requisitos de los datos de entrada y todos los argumentos de preprocesamiento disponibles.

Configuración de entrenamiento

El ajuste fino se basa en dos componentes:

1. Un script de receta (finetune.py) es un punto de entrada en Python que contiene el bucle de entrenamiento: carga del modelo, construcción del dataloader, ejecución de las pasadas hacia adelante y hacia atrás, cálculo de la pérdida de flow matching, creación de checkpoints y registro de logs.
2. Un archivo de configuración YAML especifica todas las configuraciones que utiliza la receta: qué modelo se debe ajustar, dónde se encuentra los datos, los hiperparámetros del optimizador, la estrategia de paralelismo y más. Personaliza el entrenamiento editando este archivo en lugar de modificar el código, lo que te permite escalar desde 1 hasta cientos de GPUs.

Cualquier campo YAML también se puede sobrescribir desde la 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

A continuación se muestra la configuración anotada para el ajuste fino de Wan 2.1 T2V 1.3B, con cada sección 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

Referencia de campos de configuración

La tabla siguiente enumera las configuraciones mínimas requeridas. Consulte los ejemplos de NeMo Automodel para ver configuraciones de ejemplo completas para todos los modelos.

Iniciar el entrenamiento

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

Generación

Después del entrenamiento, genera videos o imágenes a partir de prompts de texto utilizando el punto de control 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"]'

Integración de Diffusers

NeMo Automodel se construye sobre Diffusers y lo utiliza como base para la carga de modelos y la inferencia. Carga los modelos directamente desde Hugging Face Hub mediante clases de modelos de Diffusers como WanTransformer3DModel, FluxTransformer2DModel y HunyuanVideoTransformer3DModel, y genera salidas a través de pipelines de Diffusers como WanPipeline y FluxPipeline.

Esta integración ofrece varios beneficios para los usuarios de Diffusers:

• Sin conversión de puntos de control: los pesos preentrenados del Hub funcionan directamente. Apunta pretrained_model_name_or_path a cualquier ID de modelo en formato Diffusers y comienza a entrenar de inmediato.
• Soporte de modelos desde el día 0: cuando se añade un nuevo modelo de difusión a Diffusers y se sube al Hub, se puede ajustar con NeMo Automodel sin esperar a un guion de entrenamiento dedicado.
• Salidas compatibles con Pipeline: los puntos de control ajustados se guardan en un formato que puede cargarse directamente en las pipelines de Diffusers para inferencia, para compartir en el Hub o para una optimización adicional con herramientas como la cuantización y la compilación.
• Entrenamiento escalable para modelos de Diffusers: NeMo Automodel añade capacidades de entrenamiento distribuido (FSDP2, multinodo, bucketing multiresolución) que van más allá de lo que ofrecen los scripts de entrenamiento integrados de Diffusers, manteniendo las mismas interfaces de modelo y pipeline.
• Ecosistema compartido: cualquier modelo, adaptador LoRA o componente de pipeline del ecosistema de Diffusers permanece compatible durante todo el flujo de trabajo de entrenamiento e inferencia.

Equipo de NVIDIA

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

Recursos

• NeMo Automodel GitHub
• Guía de ajuste fino de difusión
• Preparación del conjunto de datos de difusión
• Cobertura de modelos de difusión
• NeMo Automodel para Transformers (ajuste fino de LLM/VLM)