Ehi GPU, cosa succede con la mia matrice?

What's happening to my matrix, GPU?

Una guida delicata per capire come le GPU effettuano la moltiplicazione di matrici

Foto di Thomas Foster su Unsplash

La moltiplicazione di matrici; il Santo Graal delle reti neurali profonde e dei giganti moderni della comprensione del linguaggio. Come MLE o scienziati dei dati, le nostre dita sono troppo veloci a digitare tf.matmul o torch.matmul e non guardiamo mai indietro. Ma non ditemi che non avete mai avuto la milliseconda infatuazione di sapere cosa potrebbe succedere a quella matrice quando entra nella GPU! Se lo avete fatto, siete nel posto giusto. Venite con me in un viaggio attraverso le affascinanti complessità all’interno di una GPU.

Vi spiegherò come queste potenze di calcolo triturano i numeri. Imparerete tre cose impressionanti e poco conosciute che le GPU fanno, quando si trovano faccia a faccia con le matrici. Alla fine di questo post del blog, avrete una buona comprensione di come funziona la moltiplicazione di matrici all’interno delle GPU.

GEMM: Un vero tesoro 💎 per una GPU

GEMM o moltiplicazione generalizzata di matrici è il kernel che viene eseguito quando le GPU effettuano la moltiplicazione di matrici.

C = a (A.B) + b C

Qui, a e b sono scalari, A è una matrice MxK, B è una matrice KxN, e quindi C è una matrice MxN. Facile come bere un bicchier d’acqua! Potreste chiedervi perché esiste quell’addizione finale. Si scopre che questo è un modello piuttosto comune per le reti neurali (ad es. aggiunta di bias, applicazione di ReLU, aggiunta di connessioni residue).

Trucco n. 1: Il prodotto esterno è fuori da questo mondo 👽

Se vi viene chiesto di scrivere un algoritmo di moltiplicazione di matrici dai primi principi, ecco cosa farete (a meno che non siate dotati di una GPU al posto del cervello – non risparmierebbe soldi per un MLE!).

for (int i = 0; i < M; ++i)    for (int j = 0; j < N; ++j)        for (int k = 0; k < K; ++k)             C[i][j] += A[i][k] * B[k][j];

Ecco una visualizzazione animata che mostra cosa fa questo.

Moltiplicazione di due matrici basata sul prodotto interno (ricreato dall'autore - fonte d'ispirazione: https://www.adityaagrawal.net/blog/architecture/matrix_multiplication)

Ma sapevate che le GPU disprezzano questa implementazione 🤔? Per capire perché è così, è necessario comprendere l’architettura della memoria della GPU,

Per tutti i confronti e le specifiche, userò le specifiche della GPU Nvidia A100.

Una GPU ha tre livelli di memoria principali,

  • Memoria globale o HBM (quello a cui solitamente ci si riferisce come memoria GPU e quello che si vede quando si esegue nvidia-smi)
  • Memoria condivisa (una memoria locale dedicata a un singolo multiprocessore di streaming [o SM] e condivisa tra i thread in esecuzione in quel SM)
  • Registri (allocati individualmente ai thread per svolgere il loro carico di lavoro)

Ecco com’è fatto,

La tipica gerarchia di memoria di una GPU (cache L0/L1/L2 ignorate per semplicità)

La prima cosa da notare è che la memoria condivisa (da ora in poi indicata come SRAM) è molto più piccola dell’HBM, figuriamoci dei registri. Quindi la tua matrice non ci entrerà (nella maggior parte dei casi). Se torniamo alla nostra animazione, per una singola riga di A tutte le colonne di B devono essere recuperate e il processo deve essere ripetuto per tutte le righe in A. Ciò significa che la GPU deve fare molte letture per calcolare l’output. L’HBM (~1,5 TB/s) è diverse grandezze più lento della SRAM (~19 TB/s).

Per metterlo in numeri, diciamo che vuoi moltiplicare una matrice 10x20 e una matrice 20x30, devi leggere le colonne di B 10x30=300 volte. C’è un modo migliore per farlo?

Scopriamo che un semplice trucco può andare molto lontano qui! Semplicemente inverti l’ordine dei cicli, in modo che k diventi il ciclo esterno. E sei fatto! 😮

for (int k = 0; k < K; ++k)     for (int i = 0; i < M; ++i)        for (int j = 0; j < N; ++j)            C[i][j] += A[i][k] * B[k][j];

Non abbiamo toccato il calcolo effettivo, solo l’ordine dei cicli, quindi dovremmo ottenere lo stesso risultato di prima. Ecco come appare ora la moltiplicazione delle matrici!

Moltiplicazione di due matrici basata sul prodotto esterno (ricreato dall'autore - fonte di ispirazione: https://www.adityaagrawal.net/blog/architecture/matrix_multiplication )

Vedi, portiamo solo una colonna di A e una riga di B alla volta e non guardiamo mai indietro. Ciò richiede molto meno letture rispetto all’implementazione originale. L’unica differenza è che prima stavamo calcolando il prodotto interno tra due vettori, ora stiamo calcolando il prodotto esterno.

La differenza tra il prodotto interno e il prodotto esterno mostrata in verde per due vettori (blu e giallo).

Ma ancora, abbiamo bisogno dell’intera C nella SRAM, che potrebbe essere troppo grande per entrare nella SRAM. Cosa fa CUDA in questo caso? Questo ci porta al secondo trucco.

Trucco n. 2: Dividi e conquista (e accumula)

Niente paura! Non ti sovraccaricherò con alcuna matematica complessa o algoritmi Leetcode. La cosa principale da tenere a mente è che una matrice è una disposizione 2D di singoli blocchi. L’animazione seguente fa giustizia a ciò che sto cercando di spiegare.

Puoi iterare ogni blocco in A e B e ancora calcolare l'esatto risultato del blocco corrispondente di C

Il risultato del blocco verde 💚 è la striscia azzurra chiaro di A 💙 e la striscia gialla chiaro di B 💛. Andando oltre, per calcolare l’output, puoi portare un blocco di quella striscia di A e un blocco dalla striscia di B alla volta, calcolare l’output e accumulare il risultato nella casella verde.

Ciò ci dà un framework flessibile in cui possiamo caricare un blocco (o una tessera) di A e B di dimensioni arbitrarie e comunque calcolare la risposta finale. Non dobbiamo fermarci qui, possiamo continuare a dividere il problema in problemi ancora più piccoli. vale a dire, la matrice è suddivisa in tessere, le tessere in frammenti e i frammenti in valori individuali.

Usando l'approccio della suddivisione in tessere, il problema può essere scomposto in modo ricorsivo

E questo si presta bene all’architettura di esecuzione del processo in una GPU. Ci sono tre livelli di esecuzione del kernel in una GPU. Per semplicità, diremo che un SM esegue un singolo blocco di thread (anche se in pratica li esegue contemporaneamente, per ridurre qualcosa noto come effetto di coda).

  • Thread
  • Warps (una raccolta di 32 thread)
  • Blocchi di thread (una raccolta di diversi warps)

Il numero esatto di thread in un blocco di thread dipende da un’architettura specifica. Ad esempio, un’A100 ha le seguenti specifiche.

  • Massimo di 2048 thread per SM
  • Massimo di 1024 thread per blocco
  • Massimo di 32 blocchi di thread per SM

Tornando alla suddivisione in tessere, è stato scoperto che (euristicamente) una tessera di matrice di dimensioni 256x128 per blocco di thread dà un’efficienza ragionevole per la maggior parte dei problemi. Pertanto, è una dimensione di tessera comune utilizzata da CUDA.

Potresti aver sentito parlare di una pratica comune di mantenere la dimensione della batch, la dimensione della dimensione nascosta come potenze di due. Ecco da dove viene! Quando le dimensioni della tua matrice sono potenze di due, saranno completamente divisibili in un insieme di tessere senza resto. Se non lo sono, il tuo codice sarà meno efficiente.

I calcoli della GPU sono più efficienti quando le dimensioni della matrice sono potenze di 2

Cosa succede quando non è una potenza di 2?

Ciò che succede è un effetto noto come quantizzazione delle tessere. In altre parole, se hai una dimensione di riga di tessera di 128 ma la tua matrice ha 257 elementi in una riga, avrai bisogno non di due, ma di tre tessere in una riga (ovvero 256+1). Questo è illustrato di seguito.

Solo perché avevamo un elemento in più nelle righe, dobbiamo dedicare due interi blocchi di thread

Il problema di questo è che il blocco di thread esegue la stessa quantità di calcoli indipendentemente dai dati utili che vi risiedono. Quindi, stai perdendo l’opportunità di fare calcoli utili dalla tua GPU, portando a inefficienze.

Un effetto simile è noto come quantizzazione dell’onda, dove la matrice è sovradimensionata e gli SM non possono adattarla tutti in una volta. Quindi la GPU deve fare il calcolo in 2 “onde”. Tuttavia, questo è meno preoccupante per le GPU moderne in quanto sfruttano la concorrenza per ridurre la quantizzazione dell’onda.

La quantizzazione delle tessere si verifica quando un blocco di thread deve spillare dati parzialmente, la quantizzazione dell’onda si verifica quando gli SM devono spillare dati.

Trucco #3: Uno è meglio di due

Il trucco finale è la fusione del kernel. Molto spesso, è più veloce fare tutti i calcoli in un unico kernel che avere due kernel chiamati uno dopo l’altro. Perché? Perché un kernel deve scrivere i dati su HBM e l’altro deve leggerli di nuovo. Abbiamo già parlato di quanto sia lento questo. Un approccio migliore è semplicemente combinare le due operazioni in una sola. Alcuni esempi sono,

  • matmul + bias + relu
  • conv + bias + relu
  • batchnorm + relu

Quindi, come si vede qui (sono sicuro che Pytorch ha un glossario simile), ci sono molti kernel fusi offerti attraverso TensorFlow che combinano operazioni comode in un singolo kernel. Nel codice, significa qualcosa del genere,

for (int m = 0; m < M; m += Mtile)     for (int n = 0; n < N; n += Ntile)        for (int k = 0; k < K; ++k)            float tmp = 0            for (int i = 0; i < Mtile; ++i)                for (int j = 0; j < Ntile; ++j)                     int row = m + i;                    int col = n + j;                    tmp += A[row][k] * B[k][col];                    // Do other things                    C[row][col] = tmp

In altre parole, teniamo cara la nostra variabile tmp fino a quando abbiamo completato tutti i nostri calcoli. Solo allora scriveremo il risultato su C.

Conclusione

Ecco fatto. Spero che questa sia stata un’escursione piacevole tra i dettagli di una GPU. Se sei interessato alla versione audio-visiva, qui c’è il link al mio video su YouTube.

Per riassumere, abbiamo discusso tre cose che rendono le GPUs molto veloci nella moltiplicazione tra matrici.

  • Le GPUs abbandonano l’implementazione più amichevole del prodotto interno di matmul e adottano l’implementazione del prodotto esterno più efficiente in lettura di matmul.
  • Le GPUs dividono le matrici in blocchi più piccoli (e i blocchi in frammenti) e distribuiscono il carico di calcolo tra blocchi di thread, warp e thread.
  • Le GPUs utilizzano la fusione di kernel per portare funzionalità comunemente co-occorrenti insieme, migliorando l’efficienza della GPU.

Se ti è piaciuta questa storia, sentiti libero di iscriverti a Nisoo e riceverai notifiche di nuovi contenuti da me, oltre a sbloccare l’accesso completo a migliaia di storie di qualità di altri autori.

Come membro di Nisoo, una parte della tua quota di iscrizione va agli scrittori che leggi, e avrai accesso completo a ogni storia…

thushv89.medium.com

Salvo diversa indicazione, tutte le immagini sono dell’autore

Riferimenti:

  • https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html
  • https://developer.nvidia.com/blog/cutlass-linear-algebra-cuda/
  • https://developer.nvidia.com/blog/nvidia-ampere-architecture-in-depth/