Accelerare il Fintuning dei Transformers con NVIDIA NeMo AutoModel

La libreria open-source HuggingFace Transformers ha rafforzato l’e cosistema AI con nuovi supporti per modelli di tipo Mixture-of-Experts (MoE) nella release v5, diventando oggi la base per modelli avanzati. NVIDIA NeMo AutoModel, parte integrante di NVIDIA NeMo, è una libreria open-source per creare modelli generativi su larga scala e consente finetuning rapido di modelli MoE, migliorando l'addestramento con un throughput fino a 3.7 volte maggiore e una riduzione fino al 32% dell’occupazione GPU rispetto al Transformers v5 native, mantenendosi compatibile con l'API da_pretrained().

Una sinergia chiave per il training MoE

Gli modelli MoE richiedono infrastrutture più avanzate per gestire routing tokens, fusione di moltiplicazioni esperte, sharding pesi su GPU e sovrapposizione di calcolo e comunicazione. La release Transformers v5 ha introdotto il supporto per questi modelli con backend esperti, caricamento dinamico pesi e piani paralleli tensori, integrando PyTorch's DeviceMesh direttamente in from_pretrained(). NeMo AutoModel si costruisce su questa base, aggiungendo Expert Parallelism, DeepEP dispatch fuso, e TransformerEngine kernels, sfruttando l’API dinamica di Transformers per estendere il supporto ad una gamma crescente di modelli.

Funzionamento e vantaggi chiave

NeMo AutoModel implementa DeepEP, che sovrappone communication con computo esperto e permette scalabilità di modelli in GPU multipla senza bisogno di modifiche alle API. Inoltre, grazie all’importazione diretta tramite from_pretrained(), NeMo AutoModel evita cambiamenti di codice per gli sviluppatori, rendendo il passaggio da un modello HuggingFace a NeMo AutoModel semplice e immediato. Le ottimizzazioni native, incluse TransformerEngine attention e custom kernel, sono disponibili per architetture popolari come Qwen3, NVIDIA Nemotron, GPT-OSS, e DeepSeek V3, o in alternativa il sistema passa al codice HuggingFace standard mantenendo comunque ottime prestazioni.

Performance in diversi scenari

NeMo AutoModel ha mostrato risultati promettenti in diversi scenari di training: training di modelli su larga scala come NVIDIA Nemotron 3 Ultra 550B A55B su 16 nodi, e di modelli come Qwen3-30B-A3B e Nemotron 3 Nano 30B A3B su singoli nodi. Nel caso del Nemotron 3 Ultra 550B, Transformer v5 non riesce a gestire l’addestramento a causa di limiti di memoria. NeMo AutoModel risolve questo problema grazie all’Expert Parallelism che distribuisce le esperte su GPU diverse, rendendo quindi il training pieno possibile, nonostante le dimensioni molto grandi del modello.

Vantaggi specifici sul benchmark

NeMo AutoModel ha offerto miglioramenti misurabili in velocità e memoria su modelli 30B. Con il DeepEP, il processo di training su modelli MoE ha visto benefici da: Fusione di comunicazione e calcolo esperto, con AllGather/ReduceScatter ottimizzati integrati in una serie di nuclei GPU. Utilizzo di nuclei TE per attenzione fusa, RMSNorm e matrici lineari, che accelerano gli operazioni fondamentali.

Il ruolo del gate routing

Le prestazioni delle simulazioni con NeMo AutoModel tengono in considerazione un "balanced routing gate"—dove i token vengono distribuiti uniformemente alle esperte—che riflette il punto ideale per l’addestramento. Questo scenario elimina gli effetti di token non uniformi che possono causare errori o rallentamenti in modelli addestrati standard. Per confronto, HF Transformers v4 e v5 usano router nativi con simulazioni di token dummy, il che introduce instabilità nel training parallelizzato di esperti.

Risorse hardware ottimizzate

L’Expert Parallelism di NeMo AutoModel rende possibile distribuire i pesi esperti su GPU parallele, riducendo la pressione di memoria del 29-32%, per esempio Qwen3 e Nemotron Nano. L’EP=8 permette di adattare più sequenze e batchsize più grandi. DeepEP sovrappone l’elaborazione con la comunicazione esperta usando nuclei GPU fusionati (non operazioni individuali per dispatch). TransformerEngine introduce operazioni core accelerate per layer di attenzione e matrici, offrendo un miglioramento su PyTorch e Flash Attention.

Backend esperti e configurazioni chiave

La release Transformers v5 introduce il parametro expimplementation, che include backend esperto con groupedmm. Questo riduce le iterazioni esperte a un singolo prodotto matrice fuso. NeMo AutoModel va oltre, usando grouped GEMM e DeepEP con Tensor Core per modelli customizzati. L’architettura di backend esperto si configura con BackendConfig:

• grouped_mm: ottimizzazione per training con una matrice esperta.
• GroupedGemmParallel: suddivisione di pesi esperte per GPU.
• RouterParallel: routing su token con all_reduce.
• TransformerEngine: attenzione e matrici accelerate.

Usare NeMo AutoModel: configurazione e scalabilità

Per configurare NeMo AutoModel in un ambiente di training multiplexed, gli sviluppatori aggiungono solo una riga di importazione: from nemoautomodel.auto_model import AutoModelForCausalLM. Il resto del codice rimane identico. Per abilitare il training distribuito, si aggiunge la configurazione mesh:

from huggingfacehub import frompretrained

Conclusioni

NVIDIA NeMo AutoModel semplifica l’ottimizzazione del training per modelli MoE, integrando con Transformers v5 e offrendo API compatibili, DeepEP, TransformerEngine e EP. Queste tecnologie abilitano modelli più grandi, training più rapido e gestione efficiente rispetto all’implementazione standard. L’integrazione di modelli popolari, ottimizzazione per nodi multi-GPU, e il supporto per training full e per addestratori di modelli singoli dimostrano che NVIDIA ha ottenuto un sistema versatile e performante per lo sviluppo di AI generativa su larga scala.