Ring Attention

La ring attention calcola la self-attention esatta su una sequenza troppo lunga per un singolo acceleratore, suddividendola in blocchi distribuiti tra dispositivi disposti ad anello. Ogni dispositivo mantiene un blocco di query e passa al dispositivo adiacente i blocchi di key e value, sovrapponendo questa comunicazione al calcolo sul blocco già in suo possesso; così, la lunghezza massima della sequenza gestibile da un modello cresce in proporzione al numero di dispositivi. Hao Liu, Matei Zaharia e Pieter Abbeel dell’Università della California a Berkeley l’hanno presentata nell’ottobre 2023 nell’articolo Ring Attention with Blockwise Transformers for Near-Infinite Context, sviluppando il loro precedente transformer parallelo a blocchi; hanno riferito di aver eseguito training e inferenza con finestre di contesto di milioni di token senza approssimare l’attention.

Come funziona la ring attention#

Una sequenza viene suddivisa in N blocchi, uno per ogni GPU o TPU. Ogni dispositivo conserva il proprio blocco di query per tutta l’elaborazione. In ognuno degli N passaggi calcola l’attention tra le proprie query e il blocco key-value che detiene in quel momento, invia quel blocco al dispositivo successivo e ne riceve un altro dal dispositivo precedente:

Dopo N passaggi, ogni blocco di query ha interagito con ogni blocco di key e value. Nessun dispositivo memorizza mai l’intera matrice di attention né l’insieme completo di key e value, quindi la memoria per dispositivo dipende dalla dimensione dei blocchi, non dalla lunghezza della sequenza. Il trasferimento viene nascosto dal calcolo, purché il calcolo dell’attention su un blocco duri più dell’invio del blocco; ciò determina la dimensione minima del blocco per l’hardware specifico. I risultati parziali vengono combinati con la stessa softmax progressiva usata dalla flash attention all’interno di una singola GPU: ogni dispositivo mantiene una somma progressiva dei punteggi esponenziati e dei valori ponderati, quindi alla fine divide la prima per la seconda. La simulazione di quattro dispositivi su una sola macchina dimostra che il risultato è esatto:

import torch

torch.manual_seed(0)
n_devices, block, dim = 4, 8, 16
q, k, v = torch.randn(3, n_devices * block, dim)
q_blocks, k_blocks, v_blocks = q.chunk(n_devices), k.chunk(n_devices), v.chunk(n_devices)

outputs = []
for i in range(n_devices):  # device i keeps query block i for the whole computation
    numerator, denominator = torch.zeros(block, dim), torch.zeros(block, 1)
    for step in range(n_devices):
        j = (i - step) % n_devices  # key-value block that has arrived from the previous device
        p = torch.exp(q_blocks[i] @ k_blocks[j].T / dim**0.5)
        numerator += p @ v_blocks[j]
        denominator += p.sum(dim=-1, keepdim=True)
    outputs.append(numerator / denominator)

full = torch.softmax(q @ k.T / dim**0.5, dim=-1) @ v
print(f"Largest score block per device: {block}x{block}, full matrix: {len(q)}x{len(k)}")
print("Matches full attention:", torch.allclose(torch.cat(outputs), full, atol=1e-5))

Stampa Matches full attention: True: ogni dispositivo ha mantenuto solo un blocco di punteggi 8 × 8 e, insieme, hanno calcolato l’attention sull’intera matrice 32 × 32. I kernel di produzione sottraggono anche un massimo progressivo prima di esponenziare, per mantenere stabili i valori. Con una maschera causale, i dispositivi che contengono i primi blocchi di query saltano la maggior parte dei blocchi di key e restano inattivi; Striped Attention risolve il problema assegnando a ogni dispositivo token distribuiti uniformemente nella sequenza, aumentando fino a 1,45 volte il throughput di training end-to-end con 256K token.

Applicazioni#

La ring attention è utile ogni volta che un input è troppo lungo per un singolo dispositivo, e il video ne è l’esempio più evidente. Lo stesso gruppo di Berkeley l’ha usata per addestrare il Large World Model (LWM), una famiglia di modelli visione-linguaggio (VLMs) da 7 miliardi di parametri, il cui contesto è cresciuto gradualmente da 4K a 1 milione di token di testo e video alternati. LWM codifica ogni fotogramma in 256 token, perciò un’ora di video campionata a un fotogramma al secondo equivale a 921.600 token; il modello ha risposto a domande su raccolte di video YouTube della durata di un’ora, un task di comprensione video a lungo raggio. Anche l’addestramento di modelli linguistici di grandi dimensioni (LLMs) che leggono intere basi di codice o raccolte di documenti in un solo prompt incontra lo stesso limite. L’API context_parallel() di PyTorch sostituisce la scaled dot-product attention con la ring attention nel proprio contesto e il team TorchTitan l’ha usata per addestrare Llama 3 8B su sequenze da 1 milione di token, distribuite su 32 GPU H100.

Ring attention a confronto con altre forme di parallelismo#

• Parallelismo del contesto: strategia generale che suddivide una sequenza tra più dispositivi. La ring attention è l’algoritmo alla base dell’implementazione PyTorch; DeepSpeed Ulysses è un’alternativa che usa scambi all-to-all per assegnare a ogni dispositivo tutti i token relativi a un sottoinsieme di teste di attention.
• Parallelismo dei tensori e parallelismo della pipeline: suddividono i pesi del modello all’interno dei livelli o tra i livelli, utile quando i parametri superano la memoria di un singolo dispositivo.
• Flash attention: un kernel per un singolo dispositivo. La ring attention esegue su ogni dispositivo un kernel a blocchi di questo tipo per il blocco locale.

Il parallelismo dei dati suddivide il batch invece della sequenza ed è la modalità usata da Ultralytics YOLO26 per il training multi-GPU: passare device=[0, 1] avvia DistributedDataParallel di PyTorch, che mantiene una copia completa del modello su ogni GPU, assegna a ogni GPU una parte uguale di ogni batch di immagini e calcola la media dei gradienti dopo ogni passaggio.