Ridurre di 6 volte il tempo di addestramento di Qwen su AI Runtime

L'addestramento o il fine-tuning di un modello di base da molti miliardi di parametri richiede una notevole potenza di calcolo, spesse volte pari a centinaia o migliaia di ore di addestramento su GPU. Tuttavia, l'ottimizzazione delle prestazioni di addestramento per un LLM è una sfida e le metriche di alto livello, come l'utilizzo della GPU, potrebbero non mostrare l'effettivo utilizzo. In questo post del blog, mostriamo come superare i livelli di astrazione e velocizzare di 6 volte il fine-tuning del più recente modello Qwen 3.8 27B su AI Runtime. Attraverso un'ottimizzazione sistematica delle prestazioni, anche con un utilizzo della GPU già al 98%, passiamo da 7 ore e 44 minuti a 1 ora e 17 minuti nel fine-tuning di un dataset fisso.

Analizzeremo nel dettaglio le tecniche alla base di questi miglioramenti, da come misurare ciò che conta davvero ai dettagli di implementazione e ai compromessi di progettazione legati a ciascuna ottimizzazione.

Se vuoi iniziare subito, passa direttamente al codice di esempio su AI Runtime e fai una prova!

Configurazione del modello

Iniziamo con un loop di addestramento standard per Qwen 3.8 27B utilizzando la libreria transformers e il dataset di addestramento SFT HuggingFaceH4/ultrachat_200k.

• Applichiamo FSDP2, la più recente tecnica di sharding del modello di torch. Ciò è necessario per far rientrare i pesi del modello e lo stato dell'ottimizzatore nelle GPU. Un approccio data parallel puramente distribuito non riuscirebbe a far rientrare i pesi del modello, a queste dimensioni, su una singola GPU.
• Abilitiamo anche l'activation checkpointing, una tecnica che scambia spazio di memoria per potenza di calcolo nel passaggio all'indietro per consentire una dimensione del batch maggiore per GPU.
• Il dataloader impacchetta batch 2D con un numero massimo di token per batch pari a 8192, eseguendo il padding fino alla dimensione massima degli esempi nel batch.

Benchmark di addestramento iniziali

L'addestramento baseline del modello sull'intero dataset impiega 7 ore e 44 minuti e, durante l'addestramento, la GPU mostra un utilizzo medio del 98%. Il risultato è un addestramento a 9.396 tok/s.

Questo utilizzo della GPU sembra elevato, ma come accennato in precedenza, questo dato è fuorviante. La metrica standard di utilizzo della GPU misura solo la percentuale di tempo in cui la GPU sta facendo qualcosa, quindi non ci dice quanto bene il modello stesse utilizzando la GPU mentre era attiva. Per capirlo dobbiamo analizzare la situazione più a fondo.

Un modo per farlo è il meccanismo di trace di torch; possiamo modificare il nostro loop di addestramento come segue (limitando anche il numero di passaggi per ottenere una risposta rapida):

Per un passaggio di esempio, questo ci fornisce i seguenti dati:

Tabella 1: CPU - 14,4 s con 1.714.382 chiamate. La tabella contiene le operazioni principali per tempo totale.

Tabella 2: GPU - 10,3 s con 269.448 operazioni. La tabella contiene le operazioni principali per tempo totale

Aree di miglioramento

Da queste tabelle notiamo alcuni problemi generali con il modello attuale:

• La CPU impiega più tempo a preparare ogni passaggio di addestramento di quanto la GPU impieghi ad eseguirlo. Questo è inatteso poiché generalmente il lato CPU dell'esecuzione su GPU è asincrono, ovvero la CPU dovrebbe procedere con un buon anticipo rispetto alla GPU. L'effetto principale è lo stallo della GPU, per cui abbiamo l'opportunità di migliorare questo aspetto.
• Ci sono troppe operazioni! Nelle tabelle sopra abbiamo centinaia di migliaia di piccole operazioni; in questo caso il problema è che la CPU non riesce a mettere in coda le operazioni della GPU abbastanza velocemente da mantenerla occupata. L'idea chiave è che dobbiamo ridurre il numero di operazioni che la CPU invia alla GPU. Esistono alcune tecniche generali per questo, e una si adatta perfettamente alla nostra situazione:
  • I CUDA graph di torch registrano i kernel cuda incontrati e li riproducono con un'unica chiamata CPU. Sebbene sia uno strumento potente, presenta diverse limitazioni ed entra in conflitto con il dinamismo di FSDP2. Quindi per ora è da escludere.
  • Torch.compile traccia il codice e utilizza Triton per fondere automaticamente i kernel ove possibile. In questo esempio, utilizzeremo kernel fusi manualmente e, nel caso degli LLM, ciò non lascia spazio a miglioramenti con torch.compile.
  • I kernel fusi manualmente utilizzano tecniche di fusione specializzate per operazioni particolari che il compilatore di torch non può applicare autonomamente. Fortunatamente nella community esistono già molti kernel predefiniti e collaudati. Ne è un esempio "Flash Attention", che in questo caso garantisce un piccolo miglioramento delle prestazioni.

Inoltre, riscontriamo un problema con il padding in eccesso nel batch: il tensore di input deve essere rettangolare e ogni riga viene completata con il padding fino al numero massimo di token per quel batch. Mostreremo un trucco per gli LLM che può essere utilizzato per minimizzarlo.

Nelle sezioni seguenti approfondiremo ciascun argomento tecnico, illustrando di cosa si tratta, le best practice consigliate e i principali compromessi da considerare durante l'applicazione.

Kernel fusi

Noterai che alcuni kernel vengono eseguiti molto frequentemente! In particolare, per Qwen, il kernel di moltiplicazione di matrici a batch viene eseguito oltre 30.000 volte in un singolo passaggio.

In questo caso esamineremo prima i kernel fusi. Un kernel fuso prende più kernel GPU più piccoli e li combina in un'unica unità. Ciò può migliorare le prestazioni evitando costosi trasferimenti tra la memoria ad alta larghezza di banda della GPU e la sua memoria on-die a larghezza di banda ancora superiore. Inoltre, può consentire una pianificazione più efficiente dei core GPU.

Con un numero così elevato di kernel aggiuntivi, anche se torch.compile potrebbe funzionare, impiegherà molto tempo per l'esecuzione. Quindi, ridurre il numero grezzo di operazioni tramite kernel fusi manualmente è un ottimo primo passo. Inoltre, fusioni note come Flash Attention non sono qualcosa che torch.compile sarebbe in grado di trovare da solo.

La libreria transformers di default utilizza SDPA di torch per i layer di attention standard. Tuttavia, molte di queste piccole operazioni che vediamo provengono dai layer non standard di tipo gated delta network. Fortunatamente, per questi layer esistono due kernel, causal-conv1d e flash-linear-attention, scritti appositamente per fondere molte di queste operazioni per layer. Come fare a scoprirlo? L'implementazione della libreria transformers di Qwen ci dice già quali kernel supporta:

• flash-linear-attention
• causal-conv1d

Esistono due metodi principali per ottenere questi kernel:

• Installare autonomamente i kernel. Ciò può spaziare da:
  • Una semplice installazione tramite pip (ad es. per flash-linear-attention) da PyPI
  • Un git clone e la compilazione del repository pertinente
  • Cercare un'altra fonte per un pacchetto precompilato nel caso in cui la compilazione richieda troppo tempo
• Utilizzare Kernel Hub di Hugging Face per scaricare i kernel precompilati.

In genere consigliamo il primo approccio, soprattutto combinato con il nostro supporto per i container Docker: crea il tuo container con tutti i kernel già installati in fase di build per ridurre il tempo necessario all'avvio del job.

Sebbene Kernel Hub sia comodo, dato che i kernel spesso dipendono in modo esplicito dalle versioni di torch e CUDA installate, potrebbe non esserci un kernel disponibile per la tua configurazione specifica; in tal caso, il modello di transformers passerà in modo trasparente a un kernel non ottimizzato con prestazioni inferiori.

Dopo aver installato i pacchetti Python flash-linear-attention e causal-conv1d, abbiamo migliorato il tempo end-to-end di un fattore di 3,7x, riducendolo a 2 ore e 5 minuti. Nelle tabelle seguenti puoi vedere i nuovi kernel (evidenziati) che prendono il posto delle funzioni element-wise delle tabelle precedenti. Anche il numero totale di operazioni eseguite è stato notevolmente ridotto, con un fattore di 15x:

Tabella 3: CPU - 3,9 s con 118.316 chiamate. La tabella contiene le operazioni principali per tempo totale.

Tabella 4: GPU - 3,9 s con 18.060 chiamate. La tabella contiene le operazioni principali per tempo totale.

Avvertenza: durante l'addestramento avrai bisogno sia di un kernel per il passaggio forward sia per quello backward, ma alcuni kernel sono solo per il forward, in quanto sono stati implementati esclusivamente per velocizzare l'inferenza.

Punti chiave

• Valuta la possibilità di utilizzare kernel altamente ottimizzati e ben consolidati già disponibili per i modelli più comuni. L'uso di questi kernel può offrire un notevole miglioramento delle prestazioni con uno sforzo minimo. In questo caso specifico, abbiamo riscontrato un miglioramento di 3,7x nella velocità di addestramento, riducendo il tempo di addestramento da 7 ore e 44 minuti a 2 ore e 5 minuti.

Batching, tokenizzazione e padding

Il nostro modello Qwen ora è piuttosto veloce, ma possiamo farlo andare ancora più veloce. Di default, la nostra configurazione di batching utilizza tensor rettangolari 2D da 8.192 token, poiché è la dimensione massima utilizzabile senza esaurire la memoria. Lo svantaggio di questo approccio è che si tende a sprecare una quantità notevole di spazio nel tensor a causa del padding.

Se Qwen utilizzasse solo meccanismi di attention standard, NestedTensor di torch sarebbe ben supportato e aiuterebbe a ridurre questo costo rimuovendo completamente il padding. Tuttavia, flash-linear-attention supporta solo un approccio a tensor impacchettati 1D più tradizionale:

Tuttavia non possiamo farlo in modo arbitrario: passare semplicemente il tensor 1D al modello può andare bene per il pre-addestramento, ma per il fine-tuning dobbiamo assicurarci che, all'interno di un esempio, i token prestino attention solo ai token dello stesso esempio. Per risolvere questo problema dobbiamo passare al modello i position_ids corretti per l'esempio. A sua volta, il modello li passerà ai kernel di attention fusi, applicando la maschera di attention corretta. Se dai un'occhiata al codice di esempio qui, puoi vedere come calcoliamo gli input necessari per far funzionare il tutto.

Con questa modifica all'apparenza semplice, il nostro tempo di addestramento passa da 2 ore e 5 minuti a 1 ora e 17 minuti.

I motivi sono due:

1. Il motivo principale è che questo impacchettamento aumenta il numero di sequenze per batch e riduce il numero totale di batch di circa il 36%. Nota che questo dipende dalla distribuzione e dalla lunghezza degli esempi nel tuo dataset.
2. Un motivo secondario è che i kernel di flash linear attention devono essere compilati just-in-time per ciascuna forma di input; prima la forma era dinamica, il che comportava alcuni minuti aggiuntivi di compilazione.

Compromesso: Questo tipo di modifica rende la pipeline di input più complessa e sussiste un elevato rischio di introdurre errori che influiscono sulla qualità del modello. Tuttavia, in questo caso, un risparmio del 36% vale la pena. Esistono altre tecniche per ridurre il padding, come mantenere il meccanismo di batching originale ordinando però tutti gli input per lunghezza in modo da minimizzare il padding; questo ordine potrebbe tuttavia introdurre correlazioni inaspettate tra campioni adiacenti, introducendo bias nel modello.

Punti chiave: Quando si addestrano LLM su un dataset con input più brevi del numero massimo di token addestrabili (ad es. 8192 in questo caso), l'impacchettamento efficiente dei dati nei batch può far risparmiare un tempo significativo. Per questo dataset abbiamo riscontrato un miglioramento di 1,6x nella velocità di addestramento, riducendo il tempo di addestramento da 2 ore e 5 minuti a 1 ora e 17 minuti.

Conclusioni

In questo post del blog abbiamo mostrato come moltiplicare per 6 l'efficienza dell'addestramento. I punti chiave sono stati: (1) misurare innanzitutto ciò che conta, non solo l'utilizzo della GPU, ma anche dove il modello trascorre il tempo durante l'addestramento, (2) ottimizzare e fondere i kernel utilizzando librerie note come Flash Linear Attention e (3) personalizzare il modo in cui i batch di dati vengono forniti al modello.

Per farlo su vasta scala, serve una piattaforma che renda predefinito un percorso ad alte prestazioni. In AIR, molti di questi miglioramenti delle prestazioni sono integrati. Includiamo le ultime versioni di CUDA e flash attention nei nostri ambienti di base, tracciamo automaticamente per te le metriche fondamentali come l'utilizzo della GPU, la temperatura e lo stato di salute, e garantiamo che le versioni di torch (e i profiler) siano aggiornate. Anche il nostro agente interno, Genie Code, è ottimizzato per individuare opportunità come queste per ottimizzare ulteriormente il tuo addestramento. Inoltre, la lettura dei dati da Databricks è ottimizzata tramite il nostro UCVolumeDataset.

Prova AIR per il tuo prossimo carico di lavoro di addestramento! Vuoi provarlo? Consulta il cookbook di AI Runtime, con i dettagli su come eseguire questo modello in autonomia.

(Questo post sul blog è stato tradotto utilizzando strumenti basati sull'intelligenza artificiale) Post originale