Due numeri su una GPU
Sai già come fare c = a + b in Python. L'hai fatto mille volte. Qui ne parleremo con i tensori di PyTorch. Un tensore è semplicemente un array di numeri. Metterlo su una GPU significa che quell'array vive nella memoria della GPU invece che nella normale memoria della CPU. Quando a e b sono due tensori su una GPU, quella singola riga viene eseguita abbastanza velocemente che non ci pensi mai.
Ora riduci tutto. Supponiamo che a e b siano due singoli float, entrambi sulla GPU. Stessa riga. c = a + b. Cosa viene effettivamente eseguito sul chip?
La risposta è un kernel. In questo contesto, un kernel è un piccolo programma che la GPU esegue su un certo insieme di dati. Non è il kernel del sistema operativo che avvia il tuo portatile. Non è il kernel matematico di un libro di algebra lineare. La parola viene riutilizzata molto e non è colpa tua. Nel mondo delle GPU, un kernel significa semplicemente: una piccola funzione che la GPU riceve l'istruzione di eseguire, subito, in parallelo, sui dati che le hai passato.
Alla fine di questo articolo, sarai in grado di guardare un frammento di codice PyTorch e contare quanti kernel la GPU eseguirà. Sembra un piccolo trucco, e lo è, ma è anche il primo passo per scrollarsi di dosso la sensazione della "GPU come scatola nera", quella sensazione per cui il tuo modello è lento e non sai perché. Ogni domanda che puoi porre sulle prestazioni della GPU alla fine torna a "quali kernel sono stati eseguiti e cosa stavano facendo". Quindi è da qui che iniziamo.
Il tuo primo kernel
Ingrandiamo un po' a e b. Tensori di lunghezza 8 questa volta. Ancora una riga PyTorch: c = a + b.
Quando esegui questo, la tua CPU (la macchina che esegue effettivamente Python) dice alla GPU: ehi, esegui questo programma su questi dati. Questa istruzione si chiama lancio. Ciò che viene lanciato è un kernel: un programma, pronto per essere eseguito. I lanci in sé sono economici, microsecondi ciascuno. Ciò che circonda il lancio (i dati che vanno alla GPU, i risultati che tornano) è dove risiede il costo reale, ed è ciò che conteremo man mano.
All'interno del kernel, il lavoro vero è svolto da minuscoli lavoratori chiamati thread. Una GPU ne ha migliaia a disposizione. Per la nostra addizione di lunghezza 8, 8 thread raccolgono il lavoro: il thread 0 gestisce l'elemento 0, il thread 1 gestisce l'elemento 1, e così via fino al thread 7. Ogni thread esegue lo stesso minuscolo programma: legge un elemento di a, legge l'elemento corrispondente di b, li somma, scrive il risultato in c.
(In pratica, la GPU lancia i thread in gruppi di dimensione fissa chiamati warp, sempre 32 thread sulle schede NVIDIA, e maschera gli extra quando il tuo array non è divisibile uniformemente. Si può ignorare per ora.)
Quindi abbiamo una riga PyTorch, un lancio, un kernel, 8 thread che fanno 8 addizioni. Ora contiamo cosa ha effettivamente attraversato il chip. Per fare l'addizione, ogni thread ha bisogno del suo elemento di a e del suo elemento di b. Sono 8 letture di a e 8 letture di b. Poi ogni thread scrive il suo risultato in c. Sono 8 scritture.
Queste letture e scritture vanno alla grande memoria situata proprio accanto al chip della GPU. Sulle schede da datacenter (A100, H100) quella memoria si chiama HBM (high-bandwidth memory). Sulle schede consumer (RTX 4090) e sulle T4 in stile Colab con cui è probabile che tu provi, si chiama GDDR. In ogni caso, è una memoria veloce accanto al chip, e la chiameremo semplicemente memoria GPU. È veloce, ma non è gratuita, e ogni viaggio verso di essa ha un costo.
Un kernel = un lancio = un passaggio sui dati. Qualunque cosa il kernel faccia al suo interno, le letture e scritture ai suoi bordi (i viaggi verso la memoria GPU per recuperare gli input, il viaggio di ritorno per scrivere gli output) sono la parte che costa. Questa è l'intera struttura.
Niente di tutto ciò cambia quando i tensori diventano più grandi. Stessa riga PyTorch, stesso kernel, solo più thread. Se a e b hanno un milione di elementi ciascuno, la GPU lancia lo stesso kernel con uno sciame più grande di thread. La matematica scala, i byte spostati scalano, il modello mentale no. Una riga, un kernel.
Cosa succede tra due operazioni
1c = (a + b).relu()
Conosci Python abbastanza bene da sapere che si tratta di due operazioni, un'addizione e poi una relu, concatenate. In un interprete, sono due chiamate di funzione. Su una GPU, in modalità eager di PyTorch, sono due lanci di kernel: uno per l'addizione, uno per la relu. Fin qui, niente di sorprendente.
Ciò che è veramente interessante è ciò che succede tra i due kernel.
Quando l'addizione termina, deve mettere il suo risultato da qualche parte. Quel "da qualche parte" è la memoria GPU. L'addizione scrive un array intermedio completo (chiamiamolo tmp) in memoria. Poi un attimo dopo viene lanciato il kernel della relu, e il suo primo compito è rileggere lo stesso array tmp dalla memoria. Legge l'intero array, applica la relu a ogni elemento, scrive il risultato in c.
Contiamo il traffico di memoria per questi due kernel:
- L'addizione: legge
a, leggeb, scrivetmp. Tre trasferimenti di dimensione dell'array. - La relu: legge
tmp, scrivec. Altri due.
Cinque trasferimenti di dimensione dell'array in totale. Confrontalo con la sola addizione di lunghezza 8 della sezione precedente, che erano tre. Aggiungere .relu() alla catena non ti è costato solo il calcolo della relu. Ti è costato un intero andata e ritorno dell'array attraverso la memoria GPU, perché tmp ha dovuto essere scritto solo così il kernel successivo poteva rileggerlo.
Non è stato memorizzato nella cache nulla. tmp non è rimasto in un registro o in una cache locale veloce. È uscito verso la memoria GPU (quella lenta e lontana) ed è tornato subito. I due kernel sono estranei l'uno all'altro. Devono passarsi i dati attraverso l'unico mezzo che entrambi sanno usare: la memoria GPU.
Perché PyTorch lo fa in questo modo? Perché in modalità eager, quando scrivi a + b, PyTorch lo esegue subito. Non sa che stai per chiamare .relu() dopo. Ogni operazione viene inviata nel momento in cui la sua riga Python viene eseguita. Non c'è un piano, nessuna anticipazione. Ogni operazione sta da sola, produce un array reale e lo passa a ciò che viene dopo attraverso la memoria.
Questo è lo schema da tenere a mente. Ogni valore intermedio nel tuo codice PyTorch viene fisicamente scritto nella memoria GPU e riletto dall'operazione successiva. Ognuno di essi. Questo è ciò che "conteggio dei kernel" misura davvero. Ogni kernel in più è un altro andata e ritorno che i tuoi dati devono fare attraverso la memoria GPU.
Fusione: due operazioni, un kernel
Immagina un kernel che faccia tutto in un colpo solo: legge il suo elemento di a, legge il suo elemento di b, li somma, applica la relu al risultato (tutto all'interno del kernel, su un minuscolo spazio di lavoro privato per thread che non lascia mai il chip), e solo allora scrive il valore finale in c. L'intermedio (a + b) esiste ancora, ma solo all'interno del kernel, nello spazio di lavoro privato di ogni thread. Non viene mai scritto nella memoria GPU. tmp, come array reale, non esiste affatto.
Conta i trasferimenti ora. Letture di a: 1 per elemento. Letture di b: 1 per elemento. Scritture di c: 1 per elemento. Tre trasferimenti di dimensione dell'array. Stessa matematica della versione a due kernel, ma due andata e ritorno in meno.
Con lunghezza 8, è un errore di arrotondamento. A nessuno importa. Con lunghezza 1 milione, o 100 milioni, quegli andata e ritorno extra diventano una grossa fetta del tempo di esecuzione, e l'orologio da parete lo riflette. Il motivo per cui il traffico di memoria finisce per dominare in questo modo è l'intero argomento dell'Articolo 2 di questa serie, quindi lascerò il "perché" da parte qui. Il punto ora è solo: stessa matematica, meno viaggi, più veloce in pratica.
Questo trucco, combinare operazioni che sarebbero state kernel separati in un unico kernel in modo che l'intermedio non debba mai visitare la memoria GPU, ha un nome. Si chiama fusione. Questa è tutta la parola. Questa è tutta l'idea.
Ora la parte scomoda. Scrivere a mano quel kernel combinato sembra facile per addizione + relu. Due operazioni. Una riga di "calcolo" in mezzo. Ma il codice PyTorch reale ha dozzine di operazioni concatenate, ognuna con le proprie forme, tipi e regole di broadcasting. Scrivere un kernel fuso che gestisca tutto correttamente è un vero lavoro di ingegneria. Di solito non scriveresti a mano questi kernel elementwise di routine.
Buona notizia: PyTorch include uno strumento che fa il re-write per te, esattamente per questo tipo di casi, in automatico. Si chiama torch.compile.
Probabilmente hai già visto torch.compile usato così:
1model = torch.compile(model)
Una riga. Qualcuno su Internet ti ha detto che rende le cose più veloci. Ecco cosa fa realmente, in parole povere: invece di eseguire le tue operazioni una alla volta come fa la modalità eager, torch.compile cattura le operazioni tensoriali che la tua funzione esegue, cerca opportunità per combinarle e genera codice ottimizzato. Le chiamate successive che corrispondono alle stesse ipotesi possono riutilizzare quel codice.
La fusione che abbiamo fatto su carta sopra (addizione e relu, che condividono un kernel, tmp che non tocca mai la memoria) è esattamente il tipo di cosa che torch.compile farà automaticamente al tuo codice, purché le operazioni siano abbastanza semplici. Quando la gente dice che torch.compile "rende PyTorch più veloce", questa è una parte enorme di ciò che intende.
Per i casi che torch.compile non può fondere da solo (operazioni personalizzate che non riconosce, riduzioni insolite, layout di memoria strani), qualcuno deve ancora scrivere un kernel a mano. È a questo che servono strumenti come Triton e CUDA. Un articolo a parte.
Vedilo con i tuoi occhi
Tutto quanto sopra è stato un conteggio di kernel su carta. È ora di contarli su una GPU reale. Se hai una macchina con una GPU CUDA a portata di mano (una workstation, un notebook Colab, un'istanza cloud), puoi eseguire questo da solo in pochi minuti.
Lo strumento è torch.profiler. È integrato in PyTorch. Tutto ciò che fa è registrare ciò che la GPU ha effettivamente fatto mentre il tuo codice veniva eseguito e restituirti una tabella che puoi leggere.
Passo 1: la versione eager.
Avvolgi la riga a due operazioni in una funzione in modo da avere qualcosa da chiamare:
1import torch2from torch.profiler import profile, ProfilerActivity34def add_relu(a, b):5 return (a + b).relu()67a = torch.randn(1_000_000, device="cuda")8b = torch.randn(1_000_000, device="cuda")910with profile(activities=[ProfilerActivity.CUDA]) as prof:11 add_relu(a, b)12 torch.cuda.synchronize()1314print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
torch.cuda.synchronize() serve solo per assicurarsi che la GPU abbia finito prima di leggere i tempi. Il lavoro della GPU viene eseguito in modo asincrono e senza la sincronizzazione a volte misureresti il sovraccarico del lancio invece del lavoro effettivo del kernel.
Passo 2: leggi l'output.
Il tuo output effettivo del profiler avrà più righe di quanto potresti aspettarti. Una serie di righe di allocazione di memoria e di contabilità di PyTorch saranno mescolate. Le righe che ci interessano sono i kernel CUDA, le funzioni effettive eseguite dalla GPU. Cerca le righe con kernel nel nome. Le due righe saranno più o meno così:
1vectorized_elementwise_kernel<...BinaryFunctor_add...> 1 12us2vectorized_elementwise_kernel<...threshold_kernel...> 1 9us
I nomi esatti dei template cambiano tra le versioni di PyTorch (relu appare spesso come threshold perché è l'operazione sottostante, e add talvolta come CUDAFunctor_add). Non cercare di analizzare l'intera stringa. Conta solo le righe. Due righe. Due kernel. Uno per l'addizione, uno per la relu. Esattamente quello che abbiamo detto sarebbe successo una sezione fa.
Passo 3: la versione compilata.
Un cambiamento di riga. Avvolgi la funzione in torch.compile:
1compiled = torch.compile(add_relu)
Ora ecco un tranello che vale la pena conoscere prima di eseguirlo. La prima chiamata a una funzione compilata con torch.compile è lenta, a volte stranamente lenta, perché è in quel momento che torch.compile fa il suo lavoro: analizzare il tuo codice, capire cosa fondere, generare il kernel fuso. Se fai il profiling della primissima chiamata, stai misurando la fase di compilazione, non il kernel. Quindi lo schema è: chiamala una volta per riscaldarla, butta via il risultato, poi fai il profiling.
1compiled(a, b) # riscaldamento, da buttare2torch.cuda.synchronize()34with profile(activities=[ProfilerActivity.CUDA]) as prof:5 compiled(a, b)6 torch.cuda.synchronize()78print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
Passo 4: leggi di nuovo l'output.
Questa volta:
1triton_poi_fused_add_relu_0 1 14us
Una riga. Un kernel. Il nome ti dice anche cosa ha fatto: addizione e relu fuse. Stessa matematica di prima, un lancio invece di due.
Hai appena fatto la cosa di cui questo articolo ha parlato, in una frase: hai chiesto a PyTorch di combinare due operazioni in un kernel, hai guardato il profiler e hai confermato che il conteggio è passato da due a uno. Fusione, in natura, sulla tua macchina.
Se vuoi vederlo in modo più drammatico, provalo con alcune dimensioni di tensori diverse. Con lunghezza 100, entrambe le versioni girano così velocemente che la differenza si perde nel rumore. Con lunghezza 10 milioni o 100 milioni, la versione compilata inizia a staccarsi chiaramente, perché l'andata e ritorno che abbiamo tagliato è una parte reale del lavoro a quella scala.
Contare i kernel non è più un consiglio astratto. Hai un modo per verificarlo.
Riassumiamo
Ecco il tutto in un colpo solo.
Il tuo codice PyTorch, quando viene eseguito su una GPU, si trasforma in una sequenza di kernel. Ogni kernel è un lancio, un passaggio sui tuoi dati, un andata e ritorno attraverso la memoria GPU per recuperare gli input e scrivere gli output. Le operazioni semplici diventano un kernel. Le catene di operazioni diventano un kernel per operazione di default, con gli intermedi che fanno andata e ritorno attraverso la memoria in mezzo. torch.compile può fondere catene semplici per te in modo che quegli intermedi non tocchino mai la memoria. Meno kernel di solito significa meno traffico di memoria. E meno traffico di memoria di solito significa più veloce.
