Context Parallelism

O paralelismo de contexto é uma técnica de computação distribuída que divide uma sequência de entrada longa em vários aceleradores. Cada GPU processa apenas parte da sequência enquanto coopera com as demais durante a atenção. Isso reduz a memória de ativação por dispositivo, permitindo que um transformer seja treinado com entradas que possam exceder a memória de uma única GPU, como documentos muito longos, vídeos estendidos ou grandes coleções de patches de imagem.

Ao contrário de simplesmente ampliar a context window de um modelo, o paralelismo de contexto não altera a quantidade de informação que a arquitetura pode aceitar teoricamente. Em vez disso, ele torna o processamento desse contexto computacionalmente prático ao distribuir a dimensão da sequência.

Como Funciona o Paralelismo de Contexto#

Suponha que uma sequência contenha 32.000 tokens e que o paralelismo de contexto utilize quatro GPUs. Cada dispositivo recebe inicialmente cerca de 8.000 tokens e armazena as ativações intermediárias correspondentes.

A maioria das operações, como normalização e camadas feed-forward, pode processar esses pedaços de sequência locais de forma independente. O desafio é o attention mechanism: uma consulta local pode precisar prestar atenção em chaves e valores mantidos por todos os outros dispositivos.

Portanto, as implementações trocam blocos de chave-valor (KV) entre as GPUs. No ring attention, cada dispositivo calcula a atenção parcial usando seus dados locais, passa um bloco KV para o próximo dispositivo e repete até processar a sequência completa. O PyTorch context parallel tutorial demonstra esse comportamento por meio de atenção de produto escalar em escala distribuída, enquanto o NVIDIA context parallel package descreve as opções de comunicação all-gather, reduce-scatter e baseadas em anel.

O resultado final é matematicamente equivalente à atenção de sequência completa, sujeito a diferenças numéricas normais, mas nenhuma GPU individual precisa reter todas as ativações da sequência.

Por que o Paralelismo de Contexto é Importante#

Sequências longas criam dois grandes problemas de escala. Primeiro, as ativações salvas consomem mais memória à medida que o comprimento da sequência aumenta. Segundo, a auto-atenção padrão compara tokens ao longo da sequência, gerando computação substancial e dados temporários.

O paralelismo de contexto aborda o problema de memória dividindo as ativações entre os dispositivos. Ele também pode dividir o trabalho de atenção, embora a comunicação introduza um novo custo. Sistemas eficientes sobrepõem as transferências de KV com a computação usando operações como aquelas documentadas no NCCL collective operations guide.

A técnica é mais valiosa quando o comprimento da sequência, em vez dos pesos do modelo ou do tamanho do lote, causa um erro de falta de memória (out-of-memory). Ela pertence ao campo mais amplo do distributed training e é comumente combinada com outras estratégias para escalar múltiplas dimensões simultaneamente.

Paralelismo de Contexto vs. Técnicas Relacionadas#

• Tensor parallelism divide operações ou matrizes de peso dentro de camadas individuais. O paralelismo de contexto, em vez disso, divide os tokens ao longo da dimensão da sequência.
• Pipeline parallelism atribui diferentes grupos de camadas de modelo a diferentes dispositivos. Ele particiona a profundidade do modelo em vez do comprimento da sequência.
• Data parallelism replica o modelo e fornece a cada réplica diferentes exemplos de treinamento. O PyTorch distributed overview o recomenda quando o modelo completo e cada amostra cabem em uma única GPU.
• Sequence parallelism frequentemente fragmenta (shards) ativações para operações selecionadas associadas ao paralelismo de tensores. O paralelismo de contexto aplica o particionamento de sequência de forma mais ampla em entradas e ativações de rede.

Essas abordagens são complementares. O NVIDIA parallelism strategies guide mostra como o paralelismo de contexto, de tensores, de pipeline e de dados podem formar um layout de dispositivo multidimensional.

Aplicações no Mundo Real#

• IA para documentos longos: Um modelo de linguagem jurídico ou médico pode precisar processar um arquivo de caso inteiro, histórico do paciente ou manual técnico. O paralelismo de contexto distribui os milhares de tokens de documento entre os aceleradores, reduzindo a pressão de memória de ativação enquanto preserva a atenção entre seções distantes.

• Entendimento de vídeos longos e multimodalidade: Transformers de vídeo e large vision models podem representar quadros, patches de imagem, segmentos de áudio e texto como uma única sequência longa de tokens. Distribuir essa sequência ajuda os modelos a analisar gravações estendidas sem reduzir agressivamente a contagem de quadros ou o detalhe espacial. O AWS Neuron context parallelism overview ilustra como grupos de aceleradores podem trocar fragmentos KV para essas cargas de trabalho de contexto longo.

Para arquiteturas compactas de visão computacional, como Ultralytics YOLO26, o paralelismo de contexto geralmente é desnecessário. O paralelismo de dados multi-GPU padrão por meio do Ultralytics model training workflow é geralmente a maneira mais apropriada de acelerar o treinamento.

Uso Prático e Compensações#

O exemplo mínimo a seguir usa a API experimental de paralelismo de contexto do PyTorch e scaled dot-product attention controls. Salve-o como cp_example.py e execute-o em duas GPUs com 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()

Aqui, a dimensão 2 é a dimensão da sequência, portanto, cada processo recebe um fragmento de sequência enquanto a atenção se coordena através da malha de dispositivos.

Na prática, os engenheiros devem confirmar que a memória reduzida supera a sobrecarga de comunicação, usar interconexões rápidas e realizar benchmarks de comprimentos de sequência representativos. Para projetos de visão padrão, a Ultralytics Platform fornece fluxos de trabalho locais e em nuvem mais simples para anotação de conjuntos de dados, treinamento, implantação e monitoramento, sem exigir configuração manual de paralelismo de contexto.