NeMo Automodel

NeMo Automodel est une bibliothèque d’entraînement native PyTorch DTensor d’NVIDIA, conçue pour l’ajustement fin et la préformation de modèles de diffusion à grande échelle. Elle est native Hugging Face : entraînez n’importe quel modèle au format Diffusers issu du Hub, sans conversion de point de contrôle. La même recette YAML et le même script d’entraînement modifiable fonctionnent à toute échelle, d’un seul GPU à des centaines de nœuds, avec l’entraînement distribué FSDP2, le chargement de données par lots multi-résolution et l’entraînement dans l’espace latent pré-encodé pour une utilisation maximale des GPU. Elle utilise le flow matching pour l’entraînement et est entièrement open source (Apache 2.0), soutenue par NVIDIA et activement maintenue.

NeMo Automodel s'intègre directement avec Diffusers. Il charge des modèles préentraînés depuis le Hugging Face Hub en utilisant les classes de modèles Diffusers et génère des sorties avec le DiffusionPipeline.

Le flux de travail typique consiste à installer NeMo Automodel (via pip ou Docker), à préparer vos données en les encodant dans des fichiers .meta, à configurer une recette YAML, à lancer l’entraînement avec torchrun et à effectuer l’inférence avec le point de contrôle résultant.

Modèles pris en charge

Installation

Exigences matérielles

Installez NeMo Automodel avec pip. Pour l’ensemble des méthodes d’installation (y compris à partir du code source), consultez le guide d’installation de NeMo Automodel.

pip3 install nemo-automodel

Sinon, utilisez le conteneur Docker préconstruit qui inclut toutes les dépendances.

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

Les points de contrôle sont perdus lorsque le conteneur se termine, sauf si le répertoire des points de contrôle est monté en bind-mount sur l'hôte. Par exemple, ajoutez -v /host/path/checkpoints:/workspace/checkpoints à la commande docker run.

Préparation des données

NeMo Automodel entraîne des modèles de diffusion dans l'espace latent. Les images ou vidéos brutes doivent être prétraitées en fichiers .meta contenant les latents VAE et les embeddings de texte avant l'entraînement. Cela évite de réencoder à chaque étape d'entraînement.

Utilisez l’outil de prétraitement intégré pour encoder vos données. L’outil répartit automatiquement le travail sur toutes les 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

Format de sortie

La prétraitement génère un répertoire de cache organisé par résolutions. NeMo Automodel prend en charge l’entraînement multi-résolution par échantillonnage par lots. Les échantillons sont regroupés par résolution spatiale afin que chaque lot contienne des échantillons de même taille, évitant ainsi le gaspillage dû au bourrage.

/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)

Consultez le guide Préparation des jeux de données de diffusion pour les formats de légendes, les exigences relatives aux données d’entrée et tous les arguments de prétraitement disponibles.

Configuration de l'entraînement

Le fine-tuning est piloté par deux composants :

1. Un script de recette (finetune.py) est un point d’entrée Python qui contient la boucle d’entraînement : chargement du modèle, construction du dataloader, exécution des passes avant/arrière, calcul de la perte de flow matching, sauvegarde des points de contrôle et journalisation.
2. Un fichier de configuration YAML spécifie tous les paramètres utilisés par la recette : le modèle à affiner, l'emplacement des données, les hyperparamètres de l'optimiseur, la stratégie de parallélisme, et plus encore. Vous personnalisez l'entraînement en modifiant ce fichier plutôt qu'en changeant le code, ce qui vous permet de passer d'un seul GPU à des centaines.

Tout champ YAML peut également être remplacé depuis l'interface en ligne de commande :

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

Voici la configuration annotée pour l’ajustement fin de Wan 2.1 T2V 1.3B, avec chaque section expliquée.

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

Référence des champs de configuration

Le tableau ci-dessous répertorie les configurations minimales requises. Consultez les exemples NeMo Automodel pour obtenir des configurations complètes pour tous les modèles.

Lancer l'entraînement

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

Génération

Après l’entraînement, générez des vidéos ou des images à partir de prompts texte en utilisant le point de contrôle affiné.

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

Intégration de Diffusers

NeMo Automodel est construit sur la base de Diffusers et l’utilise comme socle pour le chargement des modèles et l’inférence. Il charge les modèles directement depuis le Hugging Face Hub en utilisant les classes de modèles Diffusers telles que WanTransformer3DModel, FluxTransformer2DModel et HunyuanVideoTransformer3DModel, et génère des sorties via les pipelines Diffusers comme WanPipeline et FluxPipeline.

Cette intégration offre plusieurs avantages aux utilisateurs de Diffusers :

• Aucune conversion de point de contrôle : les poids préentraînés du Hub fonctionnent immédiatement. Pointez pretrained_model_name_or_path vers n’importe quel identifiant de modèle au format Diffusers et commencez l’entraînement sans délai.
• Prise en charge des modèles dès le jour 0 : lorsqu'un nouveau modèle de diffusion est ajouté à Diffusers et téléversé sur le Hub, il peut être affiné avec NeMo Automodel sans attendre un script d'entraînement dédié.
• Sorties compatibles avec les pipelines : les points de contrôle affinés sont enregistrés dans un format qui peut être rechargé directement dans les pipelines Diffusers pour l'inférence, le partage sur le Hub, ou une optimisation supplémentaire avec des outils tels que la quantification et la compilation.
• Entraînement évolutif pour les modèles Diffusers : NeMo Automodel ajoute des capacités d'entraînement distribué (FSDP2, multi-nœuds, bucketing multi-résolution) qui dépassent ce que fournissent les scripts d'entraînement intégrés de Diffusers, tout en conservant les mêmes interfaces de modèle et de pipeline.
• Écosystème partagé : tout modèle, adaptateur LoRA ou composant de pipeline issu de l'écosystème Diffusers reste compatible tout au long du flux de travail d'entraînement et d'inférence.

Équipe NVIDIA

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

Ressources

• GitHub de NeMo Automodel
• Guide d’affinage des modèles de diffusion
• Préparation du jeu de données de diffusion
• Couverture des modèles de diffusion
• NeMo Automodel pour Transformers (ajustement fin de LLM/VLM)