Entraîner FLUX ou Wan à grande échelle sans convertir de checkpoint

NVIDIA NeMo Automodel permet d'entrainer les modeles de Hugging Face sans convertir de checkpoint. Une simplification majeure pour les pipelines de production.

NeMo Automodel, la bibliothèque d'entraînement open source de NVIDIA, se branche désormais sur les modèles au format Diffusers hébergés sur Hugging Face. Fruit d'un travail commun entre les deux entreprises, l'intégration vise l'entraînement distribué des modèles de diffusion sans conversion de checkpoint ni réécriture de modèle, sous licence Apache 2.0.

Le principe consiste à pointer vers un identifiant de modèle du Hub pour lancer un entraînement, les poids affinés repartant ensuite directement dans une pipeline Diffusers pour l'inférence. Le parallélisme relève de la configuration plutôt que du code : FSDP2, tensor, expert, context et pipeline parallel se déclarent dans un fichier YAML. S'y ajoutent l'orchestration multi-nœuds via SLURM, un dataloading à résolutions multiples par buckets, et un entraînement en espace latent qui consomme des latents VAE et des embeddings texte pré-calculés au lieu d'encoder les images à chaque étape. La bibliothèque ne prend pour l'instant en charge que les modèles à flow matching.

Six familles disposent de recettes prêtes à l'emploi : Wan 2.1 en 1,3 et 14 milliards de paramètres, Wan 2.2 A14B, FLUX.1-dev et FLUX.2-dev, HunyuanVideo 1.5 et Qwen-Image. Chacune accepte un fine-tuning complet ou une approche LoRA, selon que l'on privilégie la qualité maximale sur cluster ou l'efficacité sur un seul nœud.

La démonstration porte sur FLUX.1-dev spécialisé en deux cents étapes sur soixante-dix-huit cartes du tarot Rider-Waite tombées dans le domaine public. Les prompts contenant le token déclencheur adoptent la palette crème, rouge et noir, les contours d'encre appuyés et la composition allégorique des cartes, tandis que les mêmes prompts sans ce token restent photographiques, signe que l'apprentissage ne recouvre pas le modèle de base. Une interface Python est annoncée en complément des fichiers YAML.