Context Parallelism

Il context parallelism è una tecnica di calcolo distribuito che suddivide una lunga sequenza di input su più acceleratori. Ciascuna GPU elabora solo una parte della sequenza cooperando con le altre durante l'attenzione. Questo riduce la memoria di attivazione per dispositivo, consentendo a un transformer di eseguire l'addestramento su input che potrebbero superare la memoria di una sola GPU, come documenti molto lunghi, video estesi o grandi raccolte di patch di immagini.

A differenza del semplice ampliamento del context window di un modello, il context parallelism non modifica la quantità di informazioni che l'architettura può teoricamente accettare. Piuttosto, rende l'elaborazione di tale contesto computazionalmente pratica distribuendo la dimensione della sequenza.

Come funziona il Context Parallelism#

Supponiamo che una sequenza contenga 32.000 token e che il context parallelism utilizzi quattro GPU. Ciascun dispositivo riceve inizialmente circa 8.000 token e memorizza le attivazioni intermedie corrispondenti.

La maggior parte delle operazioni, come la normalizzazione e i layer feed-forward, può elaborare questi blocchi di sequenza locali in modo indipendente. La sfida è l'attention mechanism: una query locale potrebbe dover prestare attenzione a chiavi e valori detenuti da ogni altro dispositivo.

Le implementazioni scambiano pertanto blocchi chiave-valore, o KV, tra le GPU. Nel ring attention, ciascun dispositivo calcola l'attenzione parziale utilizzando i propri dati locali, passa un blocco KV al dispositivo successivo e ripete l'operazione fino ad aver elaborato l'intera sequenza. Il PyTorch context parallel tutorial dimostra questo comportamento tramite l'attention a prodotto scalare scalato distribuito, mentre il NVIDIA context parallel package descrive le opzioni di comunicazione all-gather, reduce-scatter e basate su anello.

Il risultato finale è matematicamente equivalente all'attenzione a sequenza completa, fatte salve normali differenze numeriche, ma nessuna singola GPU deve trattenere ogni attivazione della sequenza.

Perché il Context Parallelism è importante#

Le sequenze lunghe creano due problemi di ridimensionamento principali. Primo, le attivazioni salvate consumano più memoria man mano che la lunghezza della sequenza aumenta. Secondo, la self-attention standard confronta i token lungo la sequenza, producendo calcoli consistenti e dati temporanei.

Il context parallelism risolve il problema della memoria suddividendo le attivazioni tra i dispositivi. Può anche dividere il lavoro di attenzione, sebbene la comunicazione introduca un nuovo costo. I sistemi efficienti sovrappongono i trasferimenti KV al calcolo utilizzando operazioni come quelle documentate nel NCCL collective operations guide.

La tecnica è particolarmente preziosa quando la lunghezza della sequenza, anziché i pesi del modello o la dimensione del batch, provoca un errore di esaurimento della memoria. Appartiene al campo più ampio del distributed training e viene comunemente combinata con altre strategie per scalare più dimensioni contemporaneamente.

Context Parallelism vs. Tecniche correlate#

• Tensor parallelism divide le operazioni o le matrici di pesi all'interno dei singoli layer. Il context parallelism divide invece i token lungo la dimensione della sequenza.
• Pipeline parallelism assegna diversi gruppi di layer del modello a differenti dispositivi. Partiziona la profondità del modello anziché la lunghezza della sequenza.
• Il data parallelism replica il modello e assegna a ciascuna replica esempi di addestramento diversi. Il PyTorch distributed overview lo consiglia quando il modello completo e ciascun campione rientrano in una singola GPU.
• Lo sequence parallelism spesso frammenta (shards) le attivazioni per operazioni selezionate associate al tensor parallelism. Il context parallelism applica il partizionamento della sequenza in modo più ampio sugli input e sulle attivazioni della rete.

Questi approcci sono complementari. Il NVIDIA parallelism strategies guide mostra come context, tensor, pipeline e data parallelism possano formare una disposizione di dispositivi multidimensionale.

Applicazioni nel mondo reale#

• IA per documenti lunghi: Un modello linguistico legale o medico potrebbe dover elaborare un intero fascicolo di un caso, la cronologia di un paziente o un manuale tecnico. Il context parallelism distribuisce le migliaia di token del documento tra gli acceleratori, riducendo la pressione sulla memoria di attivazione pur preservando l'attenzione tra sezioni distanti.

• Comprensione di video lunghi e multimodale: I transformer video e i large vision models possono rappresentare fotogrammi, patch di immagini, segmenti audio e testo come un'unica lunga sequenza di token. La distribuzione di tale sequenza aiuta i modelli ad analizzare registrazioni estese senza ridurre aggressivamente il numero di fotogrammi o il dettaglio spaziale. L'AWS Neuron context parallelism overview illustra come i gruppi di acceleratori possano scambiare frammenti KV per questi carichi di lavoro a contesto lungo.

Per le architetture di computer vision compatte come Ultralytics YOLO26, il context parallelism è solitamente non necessario. Il data parallelism multi-GPU standard tramite il Ultralytics model training workflow rappresenta in genere il modo più appropriato per accelerare l'addestramento.

Uso pratico e compromessi#

Il seguente esempio minimale utilizza l'API sperimentale di context-parallel di PyTorch e i scaled dot-product attention controls. 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 è la dimensione della sequenza, quindi ogni processo riceve un frammento di sequenza mentre l'attenzione si coordina attraverso la mesh di dispositivi.

In pratica, gli ingegneri dovrebbero verificare che la memoria ridotta superi l'overhead di comunicazione, utilizzare interconnessioni veloci ed effettuare il benchmark di lunghezze di sequenza rappresentative. Per i progetti di vision standard, Ultralytics Platform offre flussi di lavoro cloud e locali più semplici per l'annotazione di dataset, l'addestramento, il deployment e il monitoraggio senza richiedere una configurazione context-parallel manuale.