Training FLUX or Wan at scale without checkpoint conversion

NVIDIA NeMo Automodel allows training Hugging Face models without converting checkpoints. A major simplification for production pipelines.

NeMo Automodel, NVIDIA’s open-source training library, now plugs into Diffusers-format models hosted on Hugging Face. The result of a joint effort between the two companies, the integration aims at the distributed training of diffusion models without checkpoint conversion or model rewriting, under the Apache 2.0 license.

The principle consists of pointing to a model ID on the Hub to launch training, with the fine-tuned weights then going directly back into a Diffusers pipeline for inference. Parallelism is a matter of configuration rather than code: FSDP2, tensor, expert, context, and pipeline parallel are declared in a YAML file. Added to this are multi-node orchestration via SLURM, multi-resolution dataloading using buckets, and latent-space training that consumes pre-computed VAE latents and text embeddings instead of encoding images at every step. The library currently only supports flow-matching models.

Six families have ready-to-use recipes: Wan 2.1 in 1.3B and 14B parameters, Wan 2.2 A14B, FLUX.1-dev and FLUX.2-dev, HunyuanVideo 1.5, and Qwen-Image. Each accepts full fine-tuning or a LoRA approach, depending on whether the priority is maximum quality on a cluster or efficiency on a single node.

The demonstration focuses on FLUX.1-dev specialized in two hundred steps on seventy-eight Rider-Waite tarot cards that have fallen into the public domain. Prompts containing the trigger token adopt the cream, red, and black palette, the heavy ink outlines, and the allegorical composition of the cards, while the same prompts without this token remain photographic, indicating that the training does not overwrite the base model. A Python interface has been announced to complement the YAML files.