Ridurre i tempi di addestramento di Qwen di 6 volte con AI Runtime

Contesto e obiettivi di performance

Allenare o fare fine‑tuning di un modello di fondazione da diversi miliardi di parametri richiede risorse computazionali enormi, spesso centinaia o migliaia di ore di GPU. Tuttavia, migliorare le prestazioni di un LLM è complesso: metriche di alto livello come l’utilizzo della GPU possono risultare fuorvianti.

In questo post dimostriamo come, partendo da un utilizzo medio della GPU del 98 %, sia possibile ridurre il tempo di fine‑tuning del modello Qwen 3.8 27B su AI Runtime da 7 h 44 min a 1 h 17 min, grazie a ottimizzazioni sistematiche.

Misurare ciò che conta davvero

Il nostro training di riferimento, eseguito con la libreria transformers e il dataset HuggingFace H4/ultrachat_200k, impiegava 7 h 44 min con una velocità di 9 396 token al secondo. Sebbene la GPU mostrasse il 98 % di utilizzo, tale valore non rifletteva l’effettiva efficienza del modello.

Per approfondire, abbiamo introdotto il meccanismo di torch trace nel ciclo di addestramento, limitando il numero di step per ottenere risposte rapide. I risultati di un singolo step sono riassunti nelle tabelle seguenti:

• Tabella 1 (CPU): 14,4 s con 1 714 382 chiamate – operazioni principali per tempo totale.
• Tabella 2 (GPU): 10,3 s con 269 448 operazioni – operazioni principali per tempo totale.

Dalle tabelle emergono due problemi critici: un numero eccessivo di kernel eseguiti e un padding superfluo nei batch.

Problema del padding

Ogni batch deve essere un tensore rettangolare; le righe vengono quindi riempite fino al token massimo presente nel batch, causando spreco di memoria.

Fusione di kernel per ridurre le operazioni

Nel caso di Qwen, il kernel di moltiplicazione matriciale batch viene eseguito oltre 30 000 volte in un singolo step. La fusione di kernel aggrega più piccoli kernel GPU in un unico unità, diminuendo i trasferimenti tra la memoria ad alta larghezza di banda e quella on‑die e migliorando la pianificazione dei core.

Pur potendo utilizzare torch.compile, il numero elevato di operazioni renderebbe il processo troppo lento. Pertanto, abbiamo optato per fusi manuali, soprattutto per Flash Attention, che torch.compile non identifica automaticamente.

La libreria transformers utilizza di default torch.sdpa per l’attenzione, ma molti piccoli kernel provengono da layer non standard come gated delta network. Fortunatamente, esistono due kernel ottimizzati:

• causal‑conv1d
• flash‑linear‑attention

Questi kernel sono già integrati nella implementazione Qwen di transformers. Per installarli, si consiglia di includere i pacchetti flash-linear-attention e causal-conv1d nel container Docker, così da ridurre i tempi di avvio del job.

Una volta installati, il tempo totale è sceso a 2 h 5 min (miglioramento di 3,7×). Le nuove tabelle mostrano la riduzione delle chiamate:

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

Nota: per il training sono necessari kernel sia per il forward che per il backward; alcuni kernel sono disponibili solo per il forward, limitando la loro utilità.

Gestione avanzata del padding

Il batch predefinito usa tensori 2D rettangolari di 8 192 token, il massimo gestibile senza esaurire la memoria. Tuttavia, questo comporta un significativo padding.

Se Qwen usasse solo meccanismi di attenzione standard, torch.NestedTensor sarebbe una soluzione efficace, ma flash-linear-attention supporta solo tensori 1D “packed”.

Per il fine‑tuning, è necessario garantire che i token di un esempio attendano solo i token dello stesso esempio. Abbiamo quindi introdotto position_ids corretti, che il modello passa ai kernel fusi di attenzione, applicando il giusto mask.

Con questa modifica, il tempo di training è sceso da 2 h 5 min a 1 h 17 min, pari a un ulteriore risparmio del 36 %.

Trade‑off della complessità

L’introduzione di tensori 1D e position_ids aumenta la complessità della pipeline di input, incrementando il rischio di errori che potrebbero degradare la qualità del modello. Altre tecniche, come l’ordinamento per lunghezza dei campioni, riducono il padding ma possono introdurre correlazioni indesiderate tra campioni vicini.

Riepilogo delle ottimizzazioni

• Misurazione dettagliata: utilizzare torch trace per individuare colli di bottiglia, non solo l’utilizzo globale della GPU.
• Fusione di kernel: installare e sfruttare flash‑linear‑attention e causal‑conv1d per ridurre le operazioni da 269 k a 18 k.
• Riduzione del padding: passare a tensori 1D packed con position_ids corretti, ottenendo un risparmio del 36 %.

Complessivamente, queste tre leve hanno permesso di passare da 9 396 tok/s a una velocità equivalente a 6 × l’efficienza, riducendo il tempo totale da 7 h 44 min a 1 h 17 min.

Benefici di AI Runtime

AI Runtime (AIR) incorpora molte di queste ottimizzazioni di default: ambienti con le ultime versioni di CUDA e Flash Attention, monitoraggio automatico di metriche critiche (utilizzo GPU, temperatura, salute), versioni aggiornate di torch e profiler, e un agente interno Genie Code che individua ulteriori opportunità di ottimizzazione.

Inoltre, la lettura dei dati da Databricks è ottimizzata tramite UCVolumeDataset, riducendo il tempo di I/O.

Provate AI Runtime per il vostro prossimo carico di lavoro di training. Per iniziare, consultate il AI Runtime cookbook per dettagli sull’implementazione.