Les modèles de diffusion ont propulsé certaines des versions open source les plus excitantes des deux dernières années — pensez à FLUX.1-dev pour le texte-à-image, ou Wan 2.1 et HunyuanVideo pour le texte-à-vidéo. La bibliothèque 🤗 Diffusers est devenue le foyer de référence pour ces modèles, offrant aux chercheurs et aux développeurs une interface unique et cohérente pour l'inférence, l'adaptation et la composition de pipelines.

L'entraînement et le fine-tuning de ces modèles sont également en plein essor, ce qui nécessite des outils capables de gérer le sharding économe en mémoire, la mise en cache des latents, le bucketing multi-résolution et des configurations qui s'adaptent facilement d'un seul GPU à des centaines.

C'est là qu'intervient la bibliothèque open source NVIDIA NeMo Automodel. Aujourd'hui, nous mettons en lumière la collaboration entre NVIDIA et Hugging Face qui apporte un entraînement distribué de diffusion de qualité production à tout modèle au format Diffusers sur le Hub Hugging Face — sans conversion de checkpoint, sans réécriture de modèle pour les nouveaux modèles. L'intégration est documentée dans le guide d'entraînement Diffusers et est entièrement open source sous licence Apache 2.0.

Qu'est-ce que NeMo Automodel ?

NeMo Automodel est une bibliothèque d'entraînement open source native PyTorch DTensor, faisant partie du framework NVIDIA NeMo. Elle est construite autour de deux principes de conception qui comptent pour l'écosystème Diffusers :

  • Natif Hugging Face. Pointez pretrained_model_name_or_path vers n'importe quel ID de modèle Diffusers sur le Hub et commencez l'entraînement. NeMo Automodel utilise les classes de modèles Diffusers (par exemple, WanTransformer3DModel) pour le chargement et les pipelines Diffusers (WanPipeline) pour la génération. Les checkpoints font l'aller-retour proprement dans l'écosystème Diffusers.

  • Un programme, n'importe quelle échelle. Les recettes et scripts d'entraînement peuvent être facilement modifiés pour s'adapter à n'importe quelle échelle d'entraînement. Le parallélisme est un choix de configuration, pas une réécriture de code — passez de FSDP2 au parallélisme tensoriel, parallélisme d'experts, parallélisme de contexte et parallélisme de pipeline en déclarant des configurations, sans réécrire les modèles.

AutoModel ne prend actuellement en charge que les modèles à flux matching (flow-matching). Sous le capot, il utilise le flux matching comme objectif d'entraînement, avec un entraînement dans l'espace latent (via les sorties VAE pré-encodées) et un chargement de données par buckets multi-résolution pour accélérer le débit.

Modèles de diffusion pris en charge

L'intégration NeMo Automodel est livrée avec des recettes de fine-tuning prêtes à l'emploi pour les modèles de diffusion ouverts ci-dessous. La liste reflète les recettes actuellement dans examples/diffusion/finetune.

ModèleID Hugging FaceTâcheParamètresRecette LoRAWan 2.1 T2V 1.3B / 14BWan-AI/Wan2.1-T2V-1.3B-Diffusers / Wan-AI/Wan2.1-T2V-14B-DiffusersTexte-à-Vidéo1,3B (tient sur un seul A100 40 Go) / 14BOuiWan 2.2 T2V A14BWan-AI/Wan2.2-T2V-A14B-DiffusersTexte-à-Vidéo27B total (MoE), 14B actifs par étapeNonFLUX.1-devblack-forest-labs/FLUX.1-devTexte-à-Image12BOuiFLUX.2-devblack-forest-labs/FLUX.2-devTexte-à-Image32BOuiHunyuanVideo 1.5hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2vTexte-à-Vidéo13BOuiQwen-ImageQwen/Qwen-ImageTexte-à-Image20B (MMDiT)Oui

Ce que cette collaboration permet

Pour les utilisateurs de Diffusers, les gains pratiques se déclinent en quelques capacités concrètes.

Pas de conversion de checkpoint. Les poids pré-entraînés du Hub fonctionnent directement. Il n'y a pas de « format d'entraînement » séparé vers lequel convertir, puis reconvertir. Votre checkpoint fine-tuné se charge directement dans un DiffusionPipeline pour l'inférence, ou retourne sur le Hub pour le partage. Les outils en aval — quantification, compilation, adaptateurs LoRA, échantillonneurs personnalisés — continuent tous de fonctionner.

Chemin rapide vers la prise en charge de nouveaux modèles. Lorsqu'un nouveau modèle de diffusion arrive dans Diffusers, l'activer dans NeMo Automodel nécessite un petit ajout de code contenu — un gestionnaire de prétraitement des données et un adaptateur de modèle — plutôt qu'un script d'entraînement personnalisé complet. Le reste de la pile de recettes (FSDP2, chargement par buckets, checkpointing, génération) est conservé inchangé, et le même flux de travail piloté par YAML s'applique.

Fine-tuning complet et efficace en paramètres. Le fine-tuning complet et le PEFT de style LoRA sont tous deux pris en charge, vous pouvez donc choisir entre une qualité maximale (FT complet sur un grand cluster) ou une efficacité maximale (LoRA sur un seul nœud). La même structure de recette gère les deux.

Entraînement scalable qui va au-delà de ce que les scripts intégrés offrent. NeMo Automodel ajoute des schémas de sharding tels que FSDP2, parallélisme tensoriel, de contexte et de pipeline, l'orchestration multi-nœuds (SLURM aujourd'hui, Kubernetes à venir) et le bucketing multi-résolution. Ces capacités rendent possible l'entraînement de modèles plus grands comme FLUX.1-dev (12B) et HunyuanVideo (13B).

Un aperçu du flux de travail de fine-tuning

Dans cette section, nous parcourons le flux de travail typique pour le fine-tuning de l'un des modèles pris en charge. La méthode recommandée pour installer Automodel est le conteneur Docker NeMo Automodel (nvcr.io/nvidia/nemo-automodel:26.06), qui est livré avec PyTorch, TransformerEngine et d'autres dépendances compilées CUDA pré-construites. Alternativement, installez avec pip3 install nemo-automodel ou à partir des sources (pip3 install git+https://github.com/NVIDIA-NeMo/Automodel.git) ; consultez le guide d'installation pour toutes les options.

Ce guide parcourt un fine-tuning complet du transformateur de FLUX.1-dev sur le jeu de données de tarot Rider–Waite à 78 cartes, puis génère à partir du checkpoint résultant. Il réutilise les configurations YAML archivées et applique les paramètres spécifiques à l'exécution comme des surcharges en ligne de commande, donc aucun nouveau fichier de configuration n'est nécessaire.

1. Pré-encoder le jeu de données

La recette de diffusion consomme des latents VAE mis en cache et des embeddings de texte au lieu d'encoder les images sources à chaque étape d'entraînement. Diffusez les 78 images Rider–Waite directement depuis Hugging Face et répartissez le prétraitement sur tous les GPU visibles :

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

Les légendes contiennent déjà le token déclencheur trtcrd. Avec ce budget de pixels et le rapport d'aspect portrait du jeu de données, le prétraitement assigne les échantillons au bucket 384×640 utilisé par l'exécution de démonstration.

Pour l'entraînement d'images, le prétraitement produit des fichiers cache .pt et des métadonnées shardées :

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

2. Lancer l'entraînement avec le YAML FLUX existant

Utilisez examples/diffusion/finetune/flux_t2i_flow.yaml directement. Le YAML sélectionne déjà FLUX.1-dev, le fine-tuning complet du transformateur, l'adaptateur de flux matching FLUX, une taille de batch effective de 32 et FSDP2 à huit voies.

Fournissez les chemins et paramètres spécifiques au tarot comme surcharges en ligne de commande :

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

L'exécution produit des checkpoints aux étapes 50, 100, 150 et 200. Le checkpoint final est étiqueté epoch_66_step_199 ; l'étiquette est basée sur zéro même si elle représente la 200e étape d'optimiseur terminée.

3. Générer à partir du checkpoint fine-tuné

Utilisez le YAML de génération FLUX existant et pointez model.checkpoint vers le checkpoint d'entraînement complet :

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

Incluez trtcrd pour invoquer le style tarot appris. Pour une comparaison de contrôle, gardez la graine et la scène fixes mais omettez le déclencheur :

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

Résultats

À l'étape 200, les invites d'astronaute déclenchées conservent leur contenu demandé tout en acquérant une palette vintage crème, rouge et noire, des contours d'encre épais, des champs de couleur plats, des tons de papier vieilli et une composition de carte allégorique. L'astronaute non déclenché reste photographique, démontrant que l'effet appris est substantiellement associé à trtcrd plutôt que de remplacer le modèle de base globalement.

Invite (graine)BaselineFine-tuné (étape 200)Astronaute s'occupant d'un jardin de roses sur Mars (graine 2026)

4. Performances

Toutes les mesures ont été collectées sur un nœud avec 8× GPU NVIDIA H100 80 Go. Les résultats sont des moyennes ± écart-type d'échantillon sur trois fenêtres de 10 étapes en régime permanent.

Texte-à-image — 512×512

ModèleEntraînementParallélismeGBS / LBSTemps par étapeImages/sImages/s/GPUAllocation max/GPUFLUX.1-devCompletFSDP232 / 40,902 ± 0,039 s35,51 ± 1,554,44 ± 0,1963,88 GioFLUX.1-devLoRA r64DDP48 / 60,894 ± 0,008 s53,73 ± 0,486,72 ± 0,0667,43 GioQwen-ImageCompletFSDP240 / 50,974 ± 0,075 s41,21 ± 3,065,15 ± 0,3853,55 GioQwen-ImageLoRA r64DDP24 / 30,515 ± 0,006 s46,63 ± 0,545,83 ± 0,0766,33 Gio

Texte-à-vidéo — 512×512×49 images

Chaque échantillon est un clip vidéo de 49 images.

ModèleEntraînementGBS / LBSActivation checkpointingTemps par étapeClips/sClips/s/GPUAllocation max/GPUWan 2.1 1.3BComplet8 / 1Désactivé0,942 ± 0,038 s8,50 ± 0,351,06 ± 0,046,09 GioWan 2.1 14BComplet8 / 1Activé3,798 ± 0,017 s2,107 ± 0,0060,263 ± 0,00633,35 GioWan 2.1 14BLoRA r6416 / 2Activé7,585 ± 0,014 s2,110 ± 0,0000,263 ± 0,00024,07 GioWan 2.2 A14B, bruit élevéComplet8 / 1Activé4,628 ± 0,031 s1,730 ± 0,0100,217 ± 0,00623,57 GioHunyuanVideo 1.5Complet8 / 1Activé5,926 ± 0,046 s1,350 ± 0,0100,170 ± 0,00015,90 GioHunyuanVideo 1.5LoRA r648 / 1Activé5,575 ± 0,006 s1,433 ± 0,0060,180 ± 0,00010,58 Gio

Détails des mesures

  • Matériel : 8× H100 80 Go HBM3, entièrement connectés NVLink.

  • Jeu de données d'images : lambda/naruto-blip-captions, 256 échantillons mis en cache.

  • Jeu de données vidéo : svjack/Lelouch_Vi_Britannia_FramePack_First_Last_Frame_Video_Captioned, 112 échantillons mis en cache.

  • Les lots complets ont été appliqués avec drop_last=true.

  • Les écritures de checkpoints ont été désactivées.

  • Le temps par étape inclut le chargement des données, la passe avant, la passe arrière, le clipping de gradient, l'optimiseur et le travail du scheduler.

  • La mémoire est le pic d'utilisation de l'allocateur CUDA PyTorch, pas la mémoire totale du périphérique NVML.

Autres exemples de fine-tuning/LoRA

Les résultats du fine-tuning et du LoRA montrent la puissance de NeMo Automodel pour la spécialisation de domaine. Par exemple, le fine-tuning du modèle Wan 2.1 sur un jeu de données vidéo Ghibli a adapté avec succès le style de sortie, démontré par un changement notable dans l'apparence d'une fleur par rapport à la baseline.

Baseline :

Fine-tuné sur les vidéos de Ghibli :

Nous avons également observé l'impact distinct de l'utilisation de LoRA, où l'application de l'adaptateur à Wan 2.1 a fait adopter à la vidéo un style Ghibli caractéristique, particulièrement visible dans la mise en évidence des yeux des personnages.

Sans LoRA :

Avec LoRA :

Ces exemples, y compris ceux pour FLUX.2, confirment que les utilisateurs peuvent atteindre à la fois une qualité maximale via le fine-tuning complet et une efficacité maximale via le PEFT de style LoRA, adaptant la sortie à des domaines stylistiques spécifiques.

Essayez-le dès aujourd'hui

Apprenez-en plus sur l'intégration et trouvez d'autres exemples de fine-tuning dans la documentation de NeMo Automodel

À venir : API de recettes Pythoniques

YAML est un bon choix pour une configuration reproductible, en particulier pour les équipes qui veulent des fichiers qu'elles peuvent archiver, réviser et réutiliser, mais de nombreuses équipes ont également besoin d'une interface programmatique.

Dans une prochaine version de NeMo Automodel, nous prévoyons de proposer les recettes de diffusion via une API Pythonique entièrement typée également. Les utilisateurs pourront composer les mêmes pièces de modèle, données, optimiseur, PEFT/LoRA, parallélisme, checkpointing et génération directement depuis Python.

Le chemin Pythonique vise à rendre les recettes plus faciles à utiliser à partir du code d'entraînement existant, des notebooks et des workflows d'expérimentation, et à offrir une interface Pythonique de première classe aux côtés du chemin de démarrage rapide YAML.