Context Parallelism

Il parallelismo del contesto è una tecnica di calcolo distribuito che suddivide una lunga sequenza di input tra più acceleratori. Ogni GPU elabora solo una parte della sequenza e coopera con le altre durante l'attenzione. Questo riduce la memoria delle attivazioni per dispositivo, consentendo a un Transformer di addestrarsi su input che potrebbero superare la memoria di una singola GPU, come documenti molto lunghi, video estesi o grandi raccolte di patch di immagini.

A differenza del semplice ampliamento della finestra di contesto di un modello, il parallelismo del contesto non modifica la quantità di informazioni che l'architettura può teoricamente accettare. Rende invece computazionalmente praticabile l'elaborazione di quel contesto distribuendo la dimensione della sequenza.

Come funziona il parallelismo del contesto#

Supponiamo che una sequenza contenga 32.000 token e che il parallelismo del contesto utilizzi quattro GPU. Ogni dispositivo riceve inizialmente circa 8.000 token e memorizza le relative attivazioni intermedie.

La maggior parte delle operazioni, come la normalizzazione e i layer feed-forward, può elaborare in modo indipendente questi segmenti locali della sequenza. La difficoltà è rappresentata dal meccanismo di attenzione: una query locale potrebbe dover prestare attenzione a chiavi e valori conservati su ciascuno degli altri dispositivi.

Le implementazioni scambiano quindi blocchi chiave-valore, o KV, tra le GPU. Con la ring attention, ogni dispositivo calcola l'attenzione parziale usando i propri dati locali, passa un blocco KV al dispositivo successivo e ripete il processo finché non ha elaborato l'intera sequenza. Il tutorial di PyTorch sul parallelismo del contesto illustra questo comportamento tramite l'attenzione distribuita a prodotto scalare scalato, mentre il pacchetto NVIDIA per il parallelismo del contesto descrive le opzioni di comunicazione basate su all-gather, reduce-scatter e ring.

Il risultato finale è matematicamente equivalente all'attenzione sull'intera sequenza, fatte salve le normali differenze numeriche, ma nessuna singola GPU deve conservare tutte le attivazioni della sequenza.

Perché il parallelismo del contesto è importante#

Le sequenze lunghe causano due importanti problemi di scalabilità. Primo, le attivazioni salvate occupano più memoria man mano che aumenta la lunghezza della sequenza. Secondo, l'attenzione standard tra sé confronta i token lungo la sequenza, generando un notevole carico di calcolo e dati temporanei.

Il parallelismo del contesto affronta il problema della memoria suddividendo le attivazioni tra i dispositivi. Può anche suddividere il lavoro di attenzione, sebbene la comunicazione introduca un nuovo costo. I sistemi efficienti sovrappongono i trasferimenti KV ai calcoli, usando operazioni come quelle documentate nella guida alle operazioni collettive NCCL.

Questa tecnica è più utile quando l'errore di memoria esaurita è causato dalla lunghezza della sequenza, anziché dai pesi del modello o dalla dimensione del batch. Rientra nell'ambito più ampio dell'addestramento distribuito e viene comunemente combinata con altre strategie per scalare simultaneamente più dimensioni.

Parallelismo del contesto e tecniche correlate#

• Il parallelismo tensoriale suddivide le operazioni o le matrici dei pesi all'interno dei singoli layer. Il parallelismo del contesto suddivide invece i token lungo la dimensione della sequenza.
• Il parallelismo a pipeline assegna gruppi diversi di layer del modello a dispositivi diversi. Suddivide la profondità del modello, non la lunghezza della sequenza.
• Il parallelismo dei dati replica il modello e assegna a ogni replica esempi di addestramento diversi. La panoramica sul calcolo distribuito in PyTorch lo consiglia quando il modello completo e ogni campione rientrano in una singola GPU.
• Il parallelismo della sequenza suddivide spesso le attivazioni per alcune operazioni associate al parallelismo tensoriale. Il parallelismo del contesto applica la partizione della sequenza in modo più ampio agli input e alle attivazioni della rete.

Questi approcci sono complementari. La guida NVIDIA alle strategie di parallelismo mostra come il parallelismo del contesto, tensoriale, a pipeline e dei dati possa formare una disposizione multidimensionale dei dispositivi.

Applicazioni nel mondo reale#

• IA per documenti lunghi: Un modello linguistico legale o medico potrebbe dover elaborare un intero fascicolo, una cartella clinica o un manuale tecnico. Il parallelismo del contesto distribuisce le migliaia di token del documento tra gli acceleratori, riducendo la pressione sulla memoria delle attivazioni e preservando l'attenzione tra sezioni distanti.

• Comprensione di video lunghi e contenuti multimodali: I Transformer per i video e i grandi modelli visivi possono rappresentare fotogrammi, patch di immagini, segmenti audio e testo come un'unica lunga sequenza di token. Distribuire la sequenza aiuta i modelli ad analizzare registrazioni estese senza ridurre drasticamente il numero di fotogrammi o i dettagli spaziali. La panoramica di AWS Neuron sul parallelismo del contesto illustra come gruppi di acceleratori possano scambiarsi frammenti KV per questi carichi di lavoro con contesti lunghi.

Per le architetture compatte di visione artificiale come Ultralytics YOLO26, il parallelismo del contesto in genere non è necessario. Il parallelismo standard dei dati su più GPU, tramite il flusso di lavoro di addestramento dei modelli Ultralytics, è generalmente il metodo più adatto per accelerare l'addestramento.

Uso pratico e compromessi#

Il seguente esempio essenziale usa l'API sperimentale di PyTorch per il parallelismo del contesto e i controlli dell'attenzione a prodotto scalare scalato. Salvalo come cp_example.py e avvialo su due GPU con torchrun --standalone --nproc-per-node=2 cp_example.py.

import os

import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.experimental import context_parallel
from torch.nn.attention import SDPBackend, sdpa_kernel

rank = int(os.environ["RANK"])
world_size = int(os.environ["WORLD_SIZE"])

torch.cuda.set_device(rank)
torch.cuda.manual_seed(0)
dist.init_process_group("nccl")

mesh = init_device_mesh("cuda", (world_size,))
qkv = [torch.randn(1, 4, 4096, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True) for _ in range(3)]

with sdpa_kernel(SDPBackend.FLASH_ATTENTION), context_parallel(mesh, buffers=tuple(qkv), buffer_seq_dims=(2, 2, 2)):
    output = F.scaled_dot_product_attention(*qkv, is_causal=True)

output.float().square().mean().backward()
dist.destroy_process_group()

Qui la dimensione 2 è quella della sequenza, quindi ogni processo riceve una partizione della sequenza, mentre l'attenzione si coordina tra i dispositivi della mesh.

Nella pratica, gli ingegneri dovrebbero verificare che la riduzione dell'uso di memoria compensi il sovraccarico della comunicazione, usare interconnessioni veloci e misurare prestazioni su lunghezze di sequenza rappresentative. Per i progetti di visione standard, la Ultralytics Platform offre flussi di lavoro cloud e locali più semplici per l'annotazione dei dataset, l'addestramento, la distribuzione e il monitoraggio, senza richiedere la configurazione manuale del parallelismo del contesto.