Questo tutorial spiega come implementare xFormers, un ambiente versatile per costruire modelli Transformer efficienti in termini di memoria su GPU. Si inizia validando l'attenzione memory-efficient rispetto alle implementazioni standard e si confrontano velocità e occupazione di memoria a diverse lunghezze di sequenza. Ci si concentra quindi su tecniche avanzate come le maschere causali, le sequenze variabili, l'attenzione con query raggruppate (GQA), i bias posizionali ALiBi e i modelli con SwiGLU. In conclusione, viene assemblato un modello GPT-style per dimostrare l'applicazione pratica di queste tecnologie.

Impostare xFormers e Validare l'Attenzione Memory-Efficient

Per utilizzare xFormers efficacemente, iniziamo installandolo, importando le librerie necessarie e verificando che la GPU sia disponibile. Utilizziamo anche xformers.info per verificare i kernel di attenzione supportati. Definiamo funzioni di supporto per misurare il tempo di esecuzione CUDA e il consumo di memoria.

L'implementazione di xFormers, grazie a memoryefficientattention, è verificata testando le differenze tra essa e un'implementazione standard. Si osserva che l'output prodotto è esatto (a parte arrotondamenti in fp16), grazie a una matrice di punteggio non memorizzata.

```python

def cuda_time(fn, iters=20, warmup=5):

for _ in range(warmup): fn()

torch.cuda.synchronize()

s, e = (torch.cuda.Event(enabletiming=True) for in range(2))

s.record()

for _ in range(iters): fn()

e.record(); torch.cuda.synchronize()

return s.elapsed_time(e) / iters

def peakmemmb(fn):

torch.cuda.emptycache(); torch.cuda.resetpeakmemorystats()

fn(); torch.cuda.synchronize()

return torch.cuda.maxmemoryallocated() / 1e6

```

Benchmark Memoria e Velocità dell'Attenzione Classica

Confrontiamo l'implementazione naive dell'attenzione con xFormers, utilizzando sequenze di lunghezza crescente. Misuriamo il tempo di esecuzione e la memoria massima occupata per dimostrare l'efficienza di xFormers.

Seqlen Naive MB xformers MB Naive ms xf ms
51280017017.5512.78
1024320032037.5515.67
20481280064078.3419.12
4096512001280160.0124.24

Dall'analisi emerge che la memoria richiesta da xFormers cresce in modo lineare, a differenza della crescita quadratica dell'attenzione naive.

Applicare Attenzione Causale con Maschere Inferiori

L'attenzione causale prevede l'uso di una maschera inferiore, che si genera implicitamente senza il bisogno di memorizzare una matrice booleana quadratica. La versione xFormers è più efficiente in termini di tempo e memoria.

```python

B, M, H, K = 2, 256, 8, 64

q, k, v = (torch.randn(B, M, H, K, device=device, dtype=torch.float16) for _ in range(3))

outcausal = xops.memoryefficientattention(q, k, v, attnbias=ab.LowerTriangularMask())

```

Utilizzare Sequenze di Lunghezza Variabile

Bloccando le sequenze variabili in batch singoli senza padding, si riduce inutili sprechi di memoria. xFormers usa BlockDiagonalMask per raggruppare e processare sequenze di lunghezze diverse senza penalizzazioni di input.

    • La tecnica permette di evitare sprechi di padding.
    • Viene utilizzato BlockDiagonalCausalMask per l'attenzione causale su sequenze di lunghezza variabile.
    • Si recupera facilmente il segmento originale tramite la divisione dei risultati.

Dimotrazione di Gruppo Query Attention (GQA)

Nel GQA, diversi head di query condividono head di chiave-valore limitati, per ridurre il consumo di memoria. Questo permette l'uso di cache chiave-valore piccole, utili per l'inferenza su modelli Llama/Mistral.

```python

B, M, K = 2, 256, 64

nqheads, nkvheads = 8, 2

qg = torch.randn(B, M, G, Hq, K, device=device, dtype=torch.float16)

kg = torch.randn(B, M, G, 1, K, device=device, dtype=torch.float16)

vg = torch.randn(B, M, G, 1, K, device=device, dtype=torch.float16)

outgqa = xops.memoryefficient_attention(qg, kg, vg)

print("GQA output shape :", tuple(out_gqa.shape), "= [B, M, G, Hq, K]")

```

Utilizzare SwiGLU e ALiBi

Integrare strati feed-forward di tipo SwiGLU offre miglioramenti in accuratezza e efficienza rispetto agli strati standard. Inoltre, ALiBi introduce bias posizionali avanzati, consentendo attenzione non dipendente da posizione e modelli scalabili.

    • L'ALiBi si applica facilmente ai kernel di attenzione di xFormers.
    • I modelli con ALiBi sono più resistenti a dati mal configurati.
    • L'uso di ALiBi è particolarmente rilevante in contesti di modelli linguistici generativi.

Pipeline Automatica di Addestramento con Mixed-Precision

I calcoli in mixed-precision migliorano la velocità e la gestione della memoria rispetto agli utilizzi a precisione singola o doppia. Questo permette di addestrare modelli più grandi con risorse ridotte.

    • Supporta automaticamente la selezione di dati float16 e float32.
    • Consente l'uso di GPU più moderne e ottimizzate per mixed-precision.
    • Mantiene alta precisione senza penalizzazioni in tempo.

Applicazioni Pratiche: Modelli GPT-Style

Combinando le tecniche menzionate, si può costruire un modello del tipo GPT, che include xFormers attention, SwiGLU e training con addestramento misto. Questo tipo di modello si adatta bene a compiti complessi di generazione linguistica.

L'integrare tecnologie xFormers offre modelli più scalabili, velocissimi e risparmiatori di memoria