Context Parallelism

El paralelismo de contexto es una técnica de computación distribuida que divide una secuencia de entrada larga entre varios aceleradores. Cada GPU procesa solo una parte de la secuencia y colabora con las demás durante la atención. Esto reduce la memoria de activaciones por dispositivo y permite entrenar un Transformer con entradas que podrían superar la memoria de una sola GPU, como documentos muy largos, vídeos extensos o grandes colecciones de parches de imagen.

A diferencia de limitarse a ampliar la ventana de contexto de un modelo, el paralelismo de contexto no cambia la cantidad de información que la arquitectura puede aceptar en teoría. En su lugar, distribuye la dimensión de secuencia para que procesar ese contexto resulte viable desde el punto de vista computacional.

Cómo funciona el paralelismo de contexto#

Supón que una secuencia contiene 32 000 tokens y el paralelismo de contexto utiliza cuatro GPU. 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 locales de la secuencia de forma independiente. El reto es el mecanismo de atención: una consulta local puede necesitar atender a claves y valores almacenados en todos los demás dispositivos.

Por tanto, las implementaciones intercambian bloques de clave-valor, o KV, entre las GPU. En la atención en anillo, cada dispositivo calcula una atención parcial con sus datos locales, pasa un bloque KV al siguiente dispositivo y repite el proceso hasta que ha procesado toda la secuencia. El tutorial de PyTorch sobre paralelismo de contexto muestra este comportamiento mediante atención distribuida de producto escalar escalado, mientras que el paquete de NVIDIA para paralelismo de contexto describe opciones de comunicación basadas en all-gather, reduce-scatter y anillos.

El resultado final es matemáticamente equivalente a la atención sobre toda la secuencia, salvo por las diferencias numéricas habituales, pero ninguna GPU tiene que conservar todas las activaciones de la secuencia.

Por qué importa el paralelismo de contexto#

Las secuencias largas plantean dos grandes problemas de escalabilidad. En primer lugar, las activaciones almacenadas consumen más memoria a medida que aumenta la longitud de la secuencia. En segundo lugar, la autoatención estándar compara los tokens de toda la secuencia, lo que genera un volumen considerable de cómputo y datos temporales.

El paralelismo de contexto aborda el problema de memoria al dividir las activaciones entre dispositivos. También puede repartir el trabajo de atención, aunque la comunicación añade un coste. Los sistemas eficientes solapan las transferencias KV con el cómputo mediante operaciones como las documentadas en la guía de operaciones colectivas de NCCL.

Esta técnica es más útil cuando la longitud de la secuencia, y no los pesos del modelo o el tamaño del lote, provoca un error de falta de memoria. Forma parte del ámbito más amplio del entrenamiento distribuido y suele combinarse con otras estrategias para escalar varias dimensiones a la vez.

Paralelismo de contexto frente a técnicas relacionadas#

• El paralelismo tensorial divide las operaciones o las matrices de pesos dentro de cada capa. En cambio, el paralelismo de contexto divide los tokens a lo largo de la dimensión de secuencia.
• El paralelismo de canalización asigna distintos grupos de capas del modelo a distintos dispositivos. Divide la profundidad del modelo, no la longitud de la secuencia.
• El paralelismo de datos replica el modelo y asigna distintos ejemplos de entrenamiento a cada réplica. La introducción al entrenamiento distribuido de PyTorch lo recomienda cuando el modelo completo y cada muestra caben en una GPU.
• El paralelismo de secuencia suele dividir en fragmentos las activaciones de determinadas operaciones asociadas al paralelismo tensorial. El paralelismo de contexto aplica la partición de secuencia de forma más general a las entradas y activaciones de la red.

Estos enfoques son complementarios. La guía de NVIDIA sobre estrategias de paralelismo muestra cómo el paralelismo de contexto, tensorial, de canalización y de datos puede formar una distribución multidimensional entre dispositivos.

Aplicaciones en el mundo real#

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

• Comprensión de vídeos largos y contenido multimodal: Los Transformer de vídeo y los grandes modelos de visión pueden representar fotogramas, parches de imagen, segmentos de audio y texto como una única secuencia larga de tokens. Distribuir esa secuencia ayuda a los modelos a analizar grabaciones extensas sin reducir drásticamente el número de fotogramas ni el detalle espacial. La introducción de AWS Neuron al paralelismo de contexto muestra cómo los grupos de aceleradores pueden intercambiar fragmentos KV para estas cargas de trabajo con contextos largos.

Para arquitecturas compactas de visión artificial como Ultralytics YOLO26, el paralelismo de contexto suele ser innecesario. El paralelismo de datos estándar entre varias GPU mediante el flujo de trabajo de entrenamiento de modelos de Ultralytics suele ser una forma más adecuada de acelerar el entrenamiento.

Uso práctico y compromisos#

El siguiente ejemplo mínimo utiliza la API experimental de paralelismo de contexto de PyTorch y los controles de atención de producto escalar escalado. Guárdalo como cp_example.py y ejecútalo en dos 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()

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

En la práctica, los ingenieros deben confirmar que el ahorro de memoria compensa la sobrecarga de comunicación, utilizar interconexiones rápidas y comparar longitudes de secuencia representativas. Para los 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, el entrenamiento, la implementación y la supervisión, sin necesidad de configurar manualmente el paralelismo de contexto.