# NeMo Automodel

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

| Modèle | ID Hugging Face | Tâche | Paramètres | Cas d'usage |
|-------|----------------|------|------------|----------|
| 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) | Texte vers vidéo | 1.3B | génération de vidéos sur du matériel limité (compatible avec une seule A100 de 40 Go) |
| FLUX.1-dev | [black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) | Texte vers image | 12B | génération d'images de haute qualité |
| HunyuanVideo 1.5 | [hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v](https://huggingface.co/hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v) | Texte vers vidéo | 13B | génération de vidéos de haute qualité |

## Installation

### Exigences matérielles

| Composant | Minimum | Recommandé |
|-----------|---------|-------------|
| GPU | A100 40GB | A100 80 Go / H100 |
| GPU | 4 | 8+ |
| RAM | 128 Go | 256 Go+ |
| Stockage | SSD de 500 Go | 2 To NVMe |

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

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

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

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

La commande de prétraitement vidéo est identique pour Wan 2.1 et HunyuanVideo, mais les drapeaux diffèrent. Wan 2.1 utilise `--processor wan` avec `--resolution_preset` et `--caption_format sidecar`, tandis que HunyuanVideo utilise `--processor hunyuan` avec `--target_frames` pour définir le nombre d’images et `--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
```

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

> [!TIP]
> Consultez le guide [Préparation des jeux de données de diffusion](https://docs.nvidia.com/nemo/automodel/latest/datasets/diffusion-dataset) 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](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/diffusion/finetune/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 :

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

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

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

### Référence des champs de configuration

Le tableau ci-dessous répertorie les configurations minimales requises. Consultez les [exemples NeMo Automodel](https://github.com/NVIDIA-NeMo/Automodel/tree/main/examples/diffusion/finetune) pour obtenir des configurations complètes pour tous les modèles.

| Section | Requis ? | À modifier |
|---------|-----------|----------------|
| `model` | Oui | Définissez `pretrained_model_name_or_path` sur l’identifiant du modèle Hugging Face. Définissez `mode: finetune` ou `mode: pretrain`. |
| `step_scheduler` | Oui | `global_batch_size` est la taille de lot effective sur l’ensemble des GPU. `ckpt_every_steps` contrôle la fréquence des points de contrôle. L’accumulation des gradients est calculée automatiquement. |
| `data` | Oui | Définissez `cache_dir` sur le chemin contenant vos fichiers `.meta` prétraités. Modifiez `_target_` et `model_type` pour d'autres modèles. |
| `optim` | Oui | `learning_rate: 5e-6` est une valeur par défaut appropriée pour le fine-tuning. Ajustez-la en fonction de votre jeu de données et de votre modèle. |
| `lr_scheduler` | Oui | Choisissez `cosine`, `linear` ou `constant` pour `lr_decay_style`. Définissez `lr_warmup_steps` pour un réchauffage progressif. |
| `flow_matching` | Oui | `adapter_type` doit correspondre au modèle (`simple` pour Wan, `flux` pour FLUX, `hunyuan` pour HunyuanVideo). Consultez les configurations spécifiques au modèle pour `adapter_kwargs`. |
| `fsdp` | Oui | Définissez `dp_size` sur le nombre de GPU. Pour un déploiement multi-nœuds, définissez-le sur le nombre total de GPU sur tous les nœuds. |
| `checkpoint` | Recommandé | Définissez `checkpoint_dir` sur un chemin persistant, en particulier dans Docker. Utilisez `restore_from` pour reprendre depuis un point de contrôle précédent. |
| `wandb` | Facultatif | Configurez pour activer le suivi des expériences Weights & Biases. Définissez `mode: disabled` pour le désactiver. |

## Lancer l'entraînement

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

Exécutez ce qui suit sur chaque nœud, en définissant `NODE_RANK` en conséquence :

```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]
> Pour l’entraînement multi-nœuds, définissez `fsdp.dp_size` dans le fichier YAML sur le nombre **total** de GPU sur l’ensemble des nœuds (par exemple, 16 pour 2 nœuds avec 8 GPU chacun).

## 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é.

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

Avec un point de contrôle ajusté finement :

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

Avec un point de contrôle ajusté finement :

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

Avec un point de contrôle ajusté finement :

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

## 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](/docs/diffusers/v0.43.0/fr/api/models/wan_transformer_3d#diffusers.WanTransformer3DModel), [FluxTransformer2DModel](/docs/diffusers/v0.43.0/fr/api/models/flux_transformer#diffusers.FluxTransformer2DModel) et [HunyuanVideoTransformer3DModel](/docs/diffusers/v0.43.0/fr/api/models/hunyuan_video_transformer_3d#diffusers.HunyuanVideoTransformer3DModel), et génère des sorties via les pipelines Diffusers comme [WanPipeline](/docs/diffusers/v0.43.0/fr/api/pipelines/wan#diffusers.WanPipeline) et [FluxPipeline](/docs/diffusers/v0.43.0/fr/api/pipelines/flux#diffusers.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, pthombre@nvidia.com
- Linnan Wang, linnanw@nvidia.com
- Alexandros Koumparoulis, akoumparouli@nvidia.com

## Ressources

- [GitHub de NeMo Automodel](https://github.com/NVIDIA-NeMo/Automodel)
- [Guide d’affinage des modèles de diffusion](https://docs.nvidia.com/nemo/automodel/latest/recipes-e2e-examples/diffusion-fine-tuning)
- [Préparation du jeu de données de diffusion](https://docs.nvidia.com/nemo/automodel/latest/datasets/diffusion-dataset)
- [Couverture des modèles de diffusion](https://docs.nvidia.com/nemo/automodel/latest/model-coverage/diffusion.html)
- [NeMo Automodel pour Transformers (ajustement fin de LLM/VLM)](https://huggingface.co/docs/transformers/en/community_integrations/nemo_automodel_finetuning)