Una collaborazione congiunta tra NVIDIA e Hugging Face presenta la possibilità di implementare addestramento distribuito di modelli diffusi di alta qualità. L'integrazione presenta utilità avanzate come shard efficienti in termini di memoria, caching latente, bucketing a risoluzione multipla e configurazioni scalabili, adatti sia per un singolo GPU che per molteplici GPU.

Introduzione al framework NVIDIA NeMo Automodel

NVIDIA NeMo Automodel è una libreria open source, parte integrante del framework NeMo, costruita attorno a due principi chiave per l'ecosistema Diffusers:

    • Supporto per addestramento di modelli mediante corrispondenze di flusso: Al momento, AutoModel supporta modelli di corrispondenza di flusso. Funziona su spazio latente (attraverso pre-encoded VAE output) e caricamento dati a risoluzione multipla per accelerare il throughput.
    • Integrazione diffusers: L'integrazione con 🤗 Diffusers semplifica la gestione modelli esistenti. Ogni modello che diventa disponibile all'interno di Diffusers può essere rapidamente abilitato tramite NeMo Automodel, richiedendo solo aggiunte di codice contenute.

Supporto per modelli di diffusione

Tipi di addestramento

Per Diffusers, l'integrazione introduce un'ampia gamma di funzionalità operative:

    • Nessuna conversione di checkpoint: I pesi preprocessati caricano direttamente da Hugging Face senza bisogno di conversioni aggiuntive.
    • Percorso rapido per nuovi modelli: L'implementazione di nuovi modelli richiede solo un aggiornamento mirato in termini di pre-processing e di adattamento.
    • Adattamento completo e efficiente: Supporta sia l'adattamento completo che lo stile PEFT con LoRA, massimizzando la flessibilità.
    • Adeguata addestrabilità: Conclude con sharding, parallelismi, orchestrazione multi-nodo e bucketing a risoluzione multipla.

Pipeline di addestramento dettagliata

Installazione

Raccomandiamo l'installazione via container Docker (nvcr.io/nvidia/nemo-automodel:26.06), completo di PyTorch, TransformerEngine e dipendenze CUDA. Oltre a questo, la libreria è installabile via pip:

    • Utilizza pip3 install nemo-automodel
    • O via sorgente pip3 install git+https://github.com/NVIDIA-NeMo/Automodel.git

Fine-tune FLUX.1-dev

Questo esempio illustra l'adattamento di FLUX.1-dev usando il dataset Rider–Waite tattico. Vengono utilizzati YAML come parametri configurativi, senza necessità di modificare file .yml esistenti.

I dati vengono preprocessati in file .pt e distribuiti automaticamente su GPU visibili. Le configurazioni YAML include:

    • Seleziona FLUX.1-dev
    • Adattamento completo
    • Batch effectivo di 32
    • 8 FSDP2 paralleli

Parametri specifici

I checkpoint vengono prodotti a step 50, 100, 150, 200; il checkpoint finale è epoch66step_199. Per l'uso generativo, utilizza i file YAML esistenti, indicando model.checkpoint come modello fine-tuned finale.

Un confronto visivo mostra che i prompt generati dopo 200 step mantengono i contenuti richiesti, acquisendo però colori vintage, linee grafiche chiare, e composizioni stilizzate. I contenuti generati senza trigger restano fotografici.

Esempio applicato a Wan 2.1

L'adattamento di Wan 2.1 su un dataset Ghibli ha prodotto risultati stilisticamente coerenti con l'approccio artistico Ghibli. Ad esempio:

    • Fiori mostrano una notevole differenza di stile rispetto al modello base.
    • Applicazione di LoRA su Wan 2.1 introduce colorazione e luminosità tipiche del Ghibli, in particolare nella resa occhi dei personaggi.

I risultati confermano che NeMo Automodel offre addestramento ad alta qualità e ad alta efficienza, ideale per personalizzare output in modo mirato.

Estensioni a FLUX.2

Per FLUX.2, adattamenti simili hanno generato risultati coerenti con le aspettative. L'integrazione supporta sia addestramenti completi che PEFT, adattabili a specifici domini.

API Python

E' in programma un'implementazione futura di una API Python a tipi completi, permettendo il controllo diretto da Python su ogni aspetto dell'addestramento (modello, dati, ottimizzazione, shard, generazione) per adattamenti flessibili e automatizzazione.

La libreria mantiene un livello elevato di flessibilità e adattabilità, fornendo alle squadre la possibilità di personalizzare i modelli con una configurazione riproducibile e congiungibile, unico per il ciclo di sviluppo e test su varie piattaforme.