Context Parallelism

O paralelismo de contexto é uma técnica de computação distribuída que divide uma sequência de entrada longa por vários aceleradores. Cada GPU processa apenas parte da sequência e coopera com as restantes durante a atenção. Isto reduz a memória de ativações por dispositivo e permite treinar um transformer com entradas que podem exceder a memória de uma única GPU, como documentos muito longos, vídeos extensos ou grandes coleções de recortes de imagem.

Ao contrário de simplesmente aumentar a janela de contexto de um modelo, o paralelismo de contexto não altera a quantidade de informação que a arquitetura pode aceitar teoricamente. Em vez disso, torna computacionalmente viável processar esse contexto ao distribuir a dimensão da sequência.

Como funciona o paralelismo de contexto#

Imagina uma sequência com 32 000 tokens e o paralelismo de contexto distribuído por quatro GPUs. Cada dispositivo recebe inicialmente cerca de 8 000 tokens e armazena as ativações intermédias correspondentes.

A maioria das operações, como a normalização e as camadas feed-forward, pode processar estes segmentos locais da sequência de forma independente. O desafio é o mecanismo de atenção: uma consulta local pode precisar de aceder a chaves e valores guardados em todos os outros dispositivos.

Por isso, as implementações trocam blocos de chaves e valores, ou KV, entre GPUs. Na atenção em anel, cada dispositivo calcula a atenção parcial com os seus dados locais, passa um bloco KV ao dispositivo seguinte e repete o processo até ter processado toda a sequência. O tutorial de paralelismo de contexto do PyTorch demonstra este comportamento através da atenção distribuída de produto escalar escalado, enquanto o pacote de paralelismo de contexto da NVIDIA descreve opções de comunicação com all-gather, reduce-scatter e baseada em anel.

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

Porque é importante o paralelismo de contexto#

As sequências longas criam dois grandes problemas de escalabilidade. Primeiro, as ativações guardadas consomem mais memória à medida que o comprimento da sequência aumenta. Segundo, a atenção própria padrão compara os tokens ao longo da sequência, gerando muita computação e dados temporários.

O paralelismo de contexto resolve o problema da memória ao distribuir as ativações pelos dispositivos. Também pode distribuir o trabalho da atenção, embora a comunicação introduza um novo custo. Os sistemas eficientes sobrepõem as transferências KV à computação, usando operações como as documentadas no guia das operações coletivas NCCL.

Esta técnica é mais útil quando o comprimento da sequência, e não os pesos do modelo nem o tamanho do lote, causa um erro de falta de memória. Faz parte do domínio mais amplo do treino distribuído e é normalmente combinada com outras estratégias para escalar várias dimensões em simultâneo.

Paralelismo de contexto vs. técnicas relacionadas#

• O paralelismo de tensores divide operações ou matrizes de pesos dentro de camadas individuais. O paralelismo de contexto, por sua vez, divide os tokens ao longo da dimensão da sequência.
• O paralelismo de pipeline atribui diferentes grupos de camadas do modelo a dispositivos distintos. Divide a profundidade do modelo, não o comprimento da sequência.
• O paralelismo de dados replica o modelo e atribui diferentes exemplos de treino a cada réplica. A visão geral da computação distribuída no PyTorch recomenda esta abordagem quando o modelo completo e cada amostra cabem numa GPU.
• O paralelismo de sequências divide frequentemente as ativações de operações selecionadas associadas ao paralelismo de tensores. O paralelismo de contexto aplica a divisão da sequência de forma mais abrangente às entradas e ativações da rede.

Estas abordagens são complementares. O guia da NVIDIA sobre estratégias de paralelismo mostra como o paralelismo de contexto, de tensores, de pipeline e de dados pode formar uma disposição multidimensional dos dispositivos.

Aplicações no mundo real#

• IA para documentos longos: Um modelo de linguagem jurídica ou médica pode precisar de processar um processo completo, o historial de um paciente ou um manual técnico. O paralelismo de contexto distribui os milhares de tokens do documento pelos aceleradores, reduzindo a pressão sobre a memória das ativações e preservando a atenção entre secções distantes.

• Compreensão de vídeos longos e conteúdos multimodais: Os transformers de vídeo e os grandes modelos de visão podem representar fotogramas, recortes de imagem, segmentos de áudio e texto como uma única sequência longa de tokens. A distribuição dessa sequência ajuda os modelos a analisar gravações extensas sem reduzir drasticamente o número de fotogramas nem os detalhes espaciais. A visão geral da AWS Neuron sobre paralelismo de contexto ilustra como grupos de aceleradores podem trocar fragmentos KV nestas cargas de trabalho com contexto longo.

Para arquiteturas compactas de visão computacional, como o Ultralytics YOLO26, o paralelismo de contexto costuma ser desnecessário. O paralelismo de dados padrão em várias GPUs, através do fluxo de trabalho de treino de modelos do Ultralytics, é geralmente a forma mais adequada de acelerar o treino.

Utilização prática e compensações#

O exemplo mínimo seguinte utiliza a API experimental de paralelismo de contexto do PyTorch e os controlos de atenção de produto escalar escalado. Guarda-o como cp_example.py e executa-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, pelo que cada processo recebe um fragmento da sequência enquanto a atenção é coordenada através da malha de dispositivos.

Na prática, os engenheiros devem confirmar que a redução de memória compensa a sobrecarga de comunicação, usar interligações rápidas e avaliar comprimentos de sequência representativos. Para projetos de visão padrão, a Ultralytics Platform oferece fluxos de trabalho locais e na cloud mais simples para anotação de conjuntos de dados, treino, implementação e monitorização, sem exigir configuração manual do paralelismo de contexto.