Ring Attention

A atenção em anel calcula a autoatenção exata numa sequência demasiado longa para um único acelerador, dividindo-a em blocos distribuídos por dispositivos organizados em anel. Cada dispositivo mantém um bloco de consultas e passa blocos de chaves e valores ao dispositivo vizinho, sobrepondo essa comunicação à computação no bloco que já tem, pelo que o comprimento máximo da sequência que um modelo consegue processar aumenta proporcionalmente ao número de dispositivos. Hao Liu, Matei Zaharia e Pieter Abbeel, da UC Berkeley, apresentaram-na em outubro de 2023 no artigo Ring Attention with Blockwise Transformers for Near-Infinite Context, com base no seu transformador paralelo por blocos anterior, e relataram treino e inferência com janelas de contexto de milhões de tokens, sem aproximar a atenção.

Como funciona a atenção em anel#

Uma sequência é dividida em N blocos, um por GPU ou TPU. Cada dispositivo mantém o seu bloco de consultas durante todo o cálculo. Em cada um dos N passos, calcula a atenção entre as suas consultas e o bloco de chaves e valores que tem nesse momento, envia esse bloco ao dispositivo seguinte e recebe outro do dispositivo anterior:

Após N passos, cada bloco de consultas terá sido combinado com todos os blocos de chaves e valores. Nenhum dispositivo armazena a matriz completa de atenção nem o conjunto completo de chaves e valores, pelo que a memória por dispositivo depende do tamanho do bloco, e não do comprimento da sequência. A transferência fica oculta pela computação desde que o cálculo da atenção num bloco demore mais do que o seu envio, o que define o tamanho mínimo do bloco para determinado hardware. Os resultados parciais combinam-se com a mesma softmax acumulada que a atenção flash usa numa única GPU: cada dispositivo mantém uma soma acumulada das pontuações exponenciadas e dos valores ponderados, e divide uma pela outra no fim. A simulação de quatro dispositivos numa máquina mostra que o resultado é exato:

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))

Apresenta Matches full attention: True: cada dispositivo só manteve um bloco de pontuações de 8 × 8 e, em conjunto, calcularam a atenção sobre a matriz completa de 32 × 32. Os kernels de produção também subtraem um máximo acumulado antes de calcular a exponencial para manter os números estáveis. Com uma máscara causal, os dispositivos que contêm os primeiros blocos de consultas ignoram a maioria dos blocos de chaves e ficam inativos; a Striped Attention resolve este problema distribuindo uniformemente os tokens da sequência por cada dispositivo, aumentando até 1,45 vezes o débito de treino de ponta a ponta com 256K tokens.

Aplicações#

A atenção em anel é importante sempre que uma entrada é demasiado longa para um único dispositivo, e o vídeo é o exemplo mais claro. O mesmo grupo de Berkeley usou-a para treinar o Large World Model (LWM), uma família de modelos de visão e linguagem (VLMs) com 7 mil milhões de parâmetros, cujo contexto aumentou gradualmente de 4K para 1 milhão de tokens de texto e vídeo intercalados. O LWM codifica cada fotograma como 256 tokens, pelo que uma hora de vídeo amostrada a um fotograma por segundo corresponde a 921.600 tokens; o modelo respondeu a perguntas sobre compilações do YouTube com uma hora, uma tarefa de compreensão de vídeo de longo alcance. O treino de modelos de linguagem de grande porte (LLMs) que leem bases de código ou coleções de documentos inteiras num único prompt atinge o mesmo limite. A API context_parallel() do PyTorch substitui a atenção de produto escalar escalado por atenção em anel no seu contexto, e a equipa TorchTitan usou-a para treinar o Llama 3 8B em sequências de 1 milhão de tokens, distribuídas por 32 GPUs H100.

Atenção em anel vs. outros tipos de paralelismo#

• Paralelismo de contexto: A estratégia geral de dividir uma sequência entre dispositivos. A atenção em anel é o algoritmo usado na implementação do PyTorch, e o DeepSpeed Ulysses é uma alternativa que usa trocas todos-para-todos para fornecer a cada dispositivo todos os tokens de um subconjunto de cabeças de atenção.
• Paralelismo de tensores e paralelismo de pipeline: Dividem os pesos do modelo dentro das camadas ou entre elas, o que ajuda quando os parâmetros excedem a memória de um dispositivo.
• Atenção flash: Um kernel de dispositivo único. A atenção em anel executa um kernel por blocos deste tipo em cada dispositivo, para o bloco local.

O paralelismo de dados divide o lote em vez da sequência e é o modo usado pelo Ultralytics YOLO26 para o treino multi-GPU: passar device=[0, 1] inicia o DistributedDataParallel do PyTorch, que mantém uma cópia completa do modelo em cada GPU, atribui a cada GPU uma parte igual de cada lote de imagens e calcula a média dos gradientes após cada passo.