Ring Attention

La atención en anillo calcula la autoatención exacta en una secuencia demasiado larga para un acelerador, dividiéndola en bloques distribuidos entre dispositivos dispuestos en anillo. Cada dispositivo conserva un bloque de consultas y pasa bloques de claves y valores a su vecino; esa comunicación se solapa con el cálculo del bloque que ya tiene, de modo que la longitud máxima de secuencia que puede manejar un modelo aumenta proporcionalmente al número de dispositivos. Hao Liu, Matei Zaharia y Pieter Abbeel, de UC Berkeley, la presentaron en octubre de 2023 en Ring Attention with Blockwise Transformers for Near-Infinite Context, basándose en su Transformer paralelo por bloques anterior, y comunicaron entrenamiento e inferencia con ventanas de contexto de millones de tokens sin aproximar la atención.

Cómo funciona la atención en anillo#

Una secuencia se divide en N bloques, uno por GPU o TPU. Cada dispositivo conserva su bloque de consultas durante todo el cálculo. En cada uno de los N pasos, calcula la atención entre sus consultas y el bloque de claves y valores que tiene en ese momento, envía ese bloque al dispositivo siguiente y recibe otro del anterior:

Tras N pasos, cada bloque de consultas ha interactuado con todos los bloques de claves y valores. Ningún dispositivo almacena la matriz de atención completa ni el conjunto completo de claves y valores, por lo que la memoria por dispositivo depende del tamaño del bloque y no de la longitud de la secuencia. La transferencia queda oculta tras el cálculo siempre que calcular la atención sobre un bloque lleve más tiempo que enviarlo, lo que establece un tamaño mínimo de bloque para cada hardware. Los resultados parciales se combinan con la misma softmax acumulativa que utiliza la atención flash dentro de una sola GPU: cada dispositivo mantiene una suma acumulada de las puntuaciones exponenciadas y de los valores ponderados, y al final divide una por la otra. La simulación de cuatro dispositivos en una sola máquina muestra que el resultado es exacto:

import torch

torch.manual_seed(0)
n_devices, block, dim = 4, 8, 16
q, k, v = torch.randn(3, n_devices * block, dim)
q_blocks, k_blocks, v_blocks = q.chunk(n_devices), k.chunk(n_devices), v.chunk(n_devices)

outputs = []
for i in range(n_devices):  # device i keeps query block i for the whole computation
    numerator, denominator = torch.zeros(block, dim), torch.zeros(block, 1)
    for step in range(n_devices):
        j = (i - step) % n_devices  # key-value block that has arrived from the previous device
        p = torch.exp(q_blocks[i] @ k_blocks[j].T / dim**0.5)
        numerator += p @ v_blocks[j]
        denominator += p.sum(dim=-1, keepdim=True)
    outputs.append(numerator / denominator)

full = torch.softmax(q @ k.T / dim**0.5, dim=-1) @ v
print(f"Largest score block per device: {block}x{block}, full matrix: {len(q)}x{len(k)}")
print("Matches full attention:", torch.allclose(torch.cat(outputs), full, atol=1e-5))

Imprime Matches full attention: True: cada dispositivo solo llegó a tener un bloque de puntuaciones de 8 × 8 y, entre todos, generaron atención sobre la matriz completa de 32 × 32. Los kernels de producción también restan un máximo acumulado antes de calcular la exponencial para mantener la estabilidad numérica. Con una máscara causal, los dispositivos que contienen bloques de consultas iniciales omiten la mayoría de los bloques de claves y quedan inactivos; Striped Attention lo resuelve asignando a cada dispositivo tokens repartidos uniformemente por la secuencia, lo que aumenta hasta 1,45 veces el rendimiento del entrenamiento de principio a fin con 256K tokens.

Aplicaciones#

La atención en anillo es útil siempre que una entrada sea demasiado larga para un solo dispositivo, y el vídeo es el ejemplo más claro. El mismo grupo de Berkeley la utilizó para entrenar Large World Model (LWM), una familia de modelos de visión y lenguaje (VLM) de 7.000 millones de parámetros cuyo contexto creció progresivamente de 4K a 1 millón de tokens de texto y vídeo intercalados. LWM codifica cada fotograma como 256 tokens, por lo que una hora de vídeo muestreada a un fotograma por segundo equivale a 921.600 tokens, y el modelo respondió preguntas sobre recopilaciones de YouTube de una hora, una tarea de comprensión de vídeo a largo plazo. El entrenamiento de modelos de lenguaje grandes (LLM) que leen bases de código completas o colecciones de documentos en una sola indicación se topa con el mismo límite. La API context_parallel() de PyTorch sustituye la atención de producto escalar escalado por la atención en anillo dentro de su contexto, y el equipo TorchTitan la utilizó para entrenar Llama 3 8B con secuencias de 1 millón de tokens en 32 GPU H100.

Atención en anillo frente a otros tipos de paralelismo#

• Paralelismo de contexto: Estrategia general para dividir una secuencia entre dispositivos. La atención en anillo es el algoritmo que utiliza la implementación de PyTorch, y DeepSpeed Ulysses es una alternativa que emplea intercambios de todos con todos para proporcionar a cada dispositivo todos los tokens de un subconjunto de cabezales de atención.
• Paralelismo de tensores y paralelismo de canalización: Dividen los pesos del modelo dentro de las capas o entre ellas, lo que resulta útil cuando los parámetros superan la memoria de un dispositivo.
• Atención flash: Kernel para un único dispositivo. La atención en anillo ejecuta un kernel por bloques de este tipo en cada dispositivo para su bloque local.

El paralelismo de datos divide el lote en lugar de la secuencia, y es el modo que utiliza Ultralytics YOLO26 para el entrenamiento con varias GPU: pasar device=[0, 1] inicia PyTorch DistributedDataParallel, que mantiene una copia completa del modelo en cada GPU, asigna a cada GPU la misma parte de cada lote de imágenes y promedia los gradientes después de cada paso.