Context Parallelism

El paralelismo de contexto es una técnica de computación distribuida que divide una secuencia de entrada larga entre múltiples aceleradores. Cada GPU procesa solo una parte de la secuencia mientras coopera con las demás durante la atención. Esto reduce la memoria de activaciones por dispositivo, lo que permite que un transformer se entrene con entradas que pueden superar la memoria de una sola GPU, como documentos muy largos, vídeos extendidos o grandes colecciones de fragmentos de imagen.

A diferencia de simplemente ampliar la context window de un modelo, el paralelismo de contexto no cambia la cantidad de información que la arquitectura puede aceptar teóricamente. En su lugar, hace que procesar dicho contexto sea viable a nivel computacional al distribuir la dimensión de la secuencia.

Cómo funciona el paralelismo de contexto#

Supón que una secuencia contiene 32.000 tokens y el paralelismo de contexto utiliza cuatro GPUs. Cada dispositivo recibe inicialmente unos 8.000 tokens y almacena las activaciones intermedias correspondientes.

La mayoría de las operaciones, como la normalización y las capas feed-forward, pueden procesar estos fragmentos de secuencia locales de manera independiente. El desafío es el attention mechanism: una consulta local puede necesitar atender a las claves y valores que posee cualquier otro dispositivo.

Por lo tanto, las implementaciones intercambian bloques de clave-valor (KV) entre las GPUs. En el ring attention, cada dispositivo calcula la atención parcial utilizando sus datos locales, pasa un bloque KV al siguiente dispositivo y repite el proceso hasta haber procesado la secuencia completa. El PyTorch context parallel tutorial demuestra este comportamiento mediante la atención de producto escalar escalado distribuido, mientras que el NVIDIA context parallel package describe las opciones de comunicación all-gather, reduce-scatter y basadas en anillo.

El resultado final es matemáticamente equivalente a la atención de secuencia completa, sujeto a las diferencias numéricas normales, pero ninguna GPU individual tiene que retener cada activación de la secuencia.

Por qué importa el paralelismo de contexto#

Las secuencias largas generan dos problemas principales de escala. Primero, las activaciones guardadas consumen más memoria a medida que aumenta la longitud de la secuencia. Segundo, la autoatención estándar compara tokens a lo largo de la secuencia, lo que produce un cálculo sustancial y datos temporales.

El paralelismo de contexto aborda el problema de la memoria dividiendo las activaciones entre los dispositivos. También puede dividir el trabajo de atención, aunque la comunicación introduce un nuevo coste. Los sistemas eficientes solapan las transferencias de KV con el cálculo utilizando operaciones como las documentadas en la NCCL collective operations guide.

La técnica es más valiosa cuando la longitud de la secuencia, en lugar de los pesos del modelo o el tamaño del lote, provoca un error de falta de memoria (out-of-memory). Pertenece al campo más amplio del distributed training y se suele combinar con otras estrategias para escalar múltiples dimensiones de forma simultánea.

Paralelismo de contexto frente a técnicas relacionadas#

• Tensor parallelism divide las operaciones o matrices de pesos dentro de capas individuales. En su lugar, el paralelismo de contexto divide los tokens a lo largo de la dimensión de la secuencia.
• Pipeline parallelism asigna diferentes grupos de capas del modelo a distintos dispositivos. Parciona la profundidad del modelo en lugar de la longitud de la secuencia.
• Data parallelism replica el modelo y proporciona a cada réplica diferentes ejemplos de entrenamiento. El PyTorch distributed overview lo recomienda cuando el modelo completo y cada muestra caben en una sola GPU.
• Sequence parallelism a menudo fragmenta (shards) las activaciones para operaciones seleccionadas asociadas con el paralelismo de tensores. El paralelismo de contexto aplica la partición de secuencias de manera más amplia en las entradas y activaciones de la red.

Estos enfoques son complementarios. La NVIDIA parallelism strategies guide muestra cómo el paralelismo de contexto, de tensores, de canalización (pipeline) y de datos pueden formar una disposición de dispositivos multidimensional.

Aplicaciones en el mundo real#

• IA para documentos largos: Un modelo de lenguaje jurídico o médico puede necesitar procesar un expediente completo, un historial de paciente o un manual técnico. El paralelismo de contexto distribuye los miles de tokens del documento entre los aceleradores, reduciendo la presión de la memoria de activaciones al tiempo que preserva la atención entre secciones distantes.

• Comprensión de vídeos largos y multimodal: Los transformers de vídeo y los large vision models pueden representar fotogramas, fragmentos de imagen, segmentos de audio y texto como una única secuencia larga de tokens. Distribuir dicha secuencia ayuda a los modelos a analizar grabaciones extensas sin reducir drásticamente el recuento de fotogramas ni el detalle espacial. El AWS Neuron context parallelism overview ilustra cómo los grupos de aceleradores pueden intercambiar fragmentos (shards) de KV para estas cargas de trabajo de contexto largo.

Para arquitecturas de visión artificial compactas como Ultralytics YOLO26, el paralelismo de contexto suele ser innecesario. El paralelismo de datos multi-GPU estándar a través del Ultralytics model training workflow es, por lo general, la forma más adecuada de acelerar el entrenamiento.

Uso práctico y contrapartidas#

El siguiente ejemplo mínimo utiliza la API experimental de paralelismo de contexto de PyTorch y los scaled dot-product attention controls. Guárdalo como cp_example.py y lánzalo en dos GPUs 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()

Aquí, la dimensión 2 es la dimensión de la secuencia, por lo que cada proceso recibe un fragmento de secuencia mientras la atención se coordina a través de la malla de dispositivos.

En la práctica, los ingenieros deben confirmar que la reducción de memoria supera la sobrecarga de comunicación, utilizar interconexiones rápidas y realizar pruebas de rendimiento (benchmark) con longitudes de secuencia representativas. Para proyectos de visión estándar, Ultralytics Platform ofrece flujos de trabajo locales y en la nube más sencillos para la anotación de conjuntos de datos, entrenamiento, despliegue y supervisión sin necesidad de una configuración manual de paralelismo de contexto.