Riduzione del tempo di addestramento di Qwen di 6 volte con AI Runtime

Introduzione al problema di addestramento dei LLM

L’addestramento o il fine‑tuning di un modello di linguaggio di dimensioni multimiliardi di parametri richiede un consumo di risorse considerevole, spesso misurato in centinaia o migliaia di ore di GPU. In questo contesto, migliorare le prestazioni non è solo una questione di aumentare la percentuale di utilizzo della GPU, ma di capire come la GPU viene effettivamente impiegata durante il ciclo di addestramento.

Nel post originale Databricks, gli autori hanno mostrato come, partendo da un utilizzo medio della GPU del 98 %, siano riusciti a ridurre il tempo di addestramento di Qwen 3.8 27B da 7 ore 44 minuti a 1 ora 17 minuti, ottenendo un’accelerazione complessiva di 6×.

Misurare ciò che conta: oltre la semplice GPU utilization

Il primo passo è stato capire che la metrica “GPU utilization” è ingannevole. Essa indica solo la percentuale di tempo in cui la GPU è occupata, ma non la qualità dell’utilizzo. Per analizzare più a fondo, è stato impiegato il meccanismo di tracing di torch, limitando il numero di step per ottenere risultati rapidi.

Dati di riferimento del tracing

• Tabella 1 – CPU: 14,4 s totali con 1 714 382 chiamate.
• Tabella 2 – GPU: 10,3 s totali con 269 448 operazioni.

Queste tabelle hanno evidenziato due problemi principali: un elevato numero di operazioni elementwise a bassa intensità e un eccessivo padding nei batch, che sprecava memoria e tempo di calcolo.

Fusion di kernel: ridurre le chiamate isolate

Un kernel fuso combina più kernel piccoli in un’unica unità, riducendo i trasferimenti di dati tra la memoria ad alta larghezza di banda e quella on‑chip della GPU. Per Qwen, il kernel di moltiplicazione di matrici in batch veniva eseguito più di 30 000 volte per singolo step, indicando una potenziale area di ottimizzazione.

Pur potendo affidarsi a torch.compile, il tempo di compilazione sarebbe stato proibitivo. La soluzione è stata quindi l’uso manuale di kernel fusi, in particolare Flash Attention e causal‑conv1d, già disponibili per le architetture gated‑delta network di Qwen.

Metodi di installazione dei kernel

• Creare un container Docker con tutti i kernel pre‑installati al momento della compilazione (approccio consigliato).
• Utilizzare kernelhub, più comodo ma dipendente dalla specifica versione di torch e CUDA, con il rischio di ricadere su kernel non ottimizzati.

Impatto dell’installazione di flash‑linear‑attention e causal‑conv1d

Dopo aver aggiunto i pacchetti Python flash-linear-attention e causal-conv1d, il tempo di addestramento è sceso a 2 ore 5 minuti, una velocità migliorata di 3,7× rispetto al punto di partenza.

Statistiche post‑ottimizzazione

• Tabella 3 – CPU: 3,9 s totali con 118 316 chiamate.
• Tabella 4 – GPU: 3,9 s totali con 18 060 chiamate.

Il numero totale di operazioni è diminuito di circa 15×, dimostrando l’efficacia delle fusioni di kernel.

Nota importante: è necessario disporre di kernel sia per il passaggio forward sia per quello backward; alcuni kernel sono disponibili solo per l’inferenza.

Riduzione del padding mediante tensor packing

Il batch di default utilizza tensori rettangolari 2D di 8 192 token, il massimo gestibile senza esaurire la memoria. Questo approccio genera notevoli quantità di padding, soprattutto quando le sequenze sono più corte.

Se Qwen usasse solo meccanismi di attenzione standard, il NestedTensor di PyTorch sarebbe una soluzione ideale, eliminando del tutto il padding. Tuttavia, flash-linear-attention supporta solo un approccio tradizionale di tensori 1D “packed”.

Come gestire il packing senza perdere la correttezza del fine‑tuning

• Passare i position_ids corretti per ogni esempio, così che i kernel di attenzione fusi applichino la maschera adeguata.
• Modificare il codice di esempio (disponibile nei repository AI Runtime) per calcolare le entrate necessarie al packing.

Con questa modifica, il tempo di addestramento è sceso ulteriormente a 1 ora 17 minuti.

Compromessi introdotti dal packing

Il packing rende la pipeline di input più complessa e aumenta il rischio di errori che potrebbero degradare la qualità del modello. L’autore sottolinea che, per questo caso specifico, il risparmio del 36 % di tempo giustifica il trade‑off.

Altre tecniche per ridurre il padding includono:

• Ordinare le sequenze per lunghezza prima del batching, riducendo lo spazio sprecato.
• Attenzione: l’ordinamento può introdurre correlazioni indesiderate tra campioni adiacenti, potenzialmente biasando il modello.

Conclusioni chiave

Quando si addestra un LLM su dataset con sequenze più corte rispetto al limite di token massimo (8192 in questo caso), l’efficienza del packing dei dati nei batch è fondamentale per risparmiare tempo.

Per il dataset SFTHuggingFaceH4/ultrachat_200k, le ottimizzazioni descritte hanno prodotto:

• Un miglioramento complessivo della velocità di addestramento di ≈ 6×.
• Una riduzione del consumo di risorse GPU pur mantenendo un utilizzo medio del 98 %.
• Una diminuzione drastica delle chiamate CPU/GPU, da centinaia di migliaia a poche decine di migliaia.

Il risultato dimostra che, anche in scenari dove la GPU sembra già “piena”, una profonda analisi dei kernel, l’uso di kernel fusi e la gestione intelligente del padding possono portare a guadagni di prestazioni significativi.

Passi pratici per replicare i risultati

Chi desidera provare subito le ottimizzazioni può seguire questi passaggi:

• Clonare il repository AI Runtime di Databricks.
• Installare i pacchetti flash-linear-attention e causal-conv1d nel proprio ambiente.
• Abilitare il tracing di torch per identificare i colli di bottiglia.
• Applicare il packing dei tensori 1D con i corretti position_ids.
• Eseguire il fine‑tuning su ultrachat_200k e verificare il nuovo tempo di addestramento.

Seguendo queste linee guida, è possibile ottenere miglioramenti analoghi su altri modelli LLM di grandi dimensioni.