Ring Attention

Ring attention вычисляет точное самовнимание для последовательности, слишком длинной для одного ускорителя, разбивая её на блоки, распределённые между устройствами, соединёнными в кольцо. Каждое устройство хранит один блок запросов и передаёт блоки ключей и значений соседнему устройству, совмещая обмен данными с вычислениями над уже имеющимся блоком. Благодаря этому максимальная длина последовательности, которую может обработать модель, увеличивается пропорционально числу устройств. Хао Лю, Матей Захария и Питер Абель из UC Berkeley представили этот метод в октябре 2023 года в статье Ring Attention with Blockwise Transformers for Near-Infinite Context, развивая свой более ранний подход к блочному параллелизму трансформеров. Авторы сообщили об обучении и инференсе с контекстными окнами в миллионы токенов без приближённых вычислений внимания.

Как работает ring attention#

Последовательность разбивается на N блоков — по одному на каждый GPU или TPU. Каждое устройство хранит свой блок запросов на протяжении всех вычислений. На каждом из N шагов оно вычисляет внимание между своими запросами и находящимся у него блоком ключей и значений, отправляет этот блок следующему устройству и получает другой блок от предыдущего:

После N шагов каждый блок запросов обработан со всеми блоками ключей и значений. Ни одно устройство не хранит целиком матрицу внимания или полный набор ключей и значений, поэтому объём памяти на устройство зависит от размера блока, а не от длины последовательности. Передача данных скрыта за вычислениями, если вычисление внимания над одним блоком занимает больше времени, чем его отправка; это определяет минимальный размер блока для заданного оборудования. Частичные результаты объединяются с помощью той же текущей функции softmax, которую использует flash attention на одном GPU: каждое устройство хранит текущую сумму экспоненцированных оценок и взвешенных значений, а в конце делит одну на другую. Симуляция работы четырёх устройств на одной машине показывает, что результат точен:

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

Выводится Matches full attention: True: каждое устройство хранило только блок оценок размером 8 × 8, а вместе они вычислили внимание для полной матрицы 32 × 32. В производственных ядрах перед экспоненцированием также вычитается текущий максимум, чтобы обеспечить численную стабильность. При причинно-следственной маске устройства с ранними блоками запросов пропускают большинство блоков ключей и простаивают. Striped Attention решает эту проблему, распределяя между устройствами токены, равномерно размещённые по последовательности. Это повышает общую пропускную способность обучения до 1,45 раза при длине последовательности 256K токенов.

Применение#

Ring attention полезен везде, где входные данные слишком длинны для одного устройства; самый наглядный пример — видео. Та же группа из Berkeley использовала его для обучения Large World Model (LWM) — семейства визуально-языковых моделей (VLMs) с 7 миллиардами параметров, контекст которых постепенно увеличивался с 4K до 1 миллиона чередующихся текстовых и видеотокенов. LWM кодирует каждый кадр в виде 256 токенов, поэтому час видео, взятого с частотой один кадр в секунду, содержит 921 600 токенов. Модель отвечала на вопросы о часовых подборках YouTube — это задача понимания видео на длинной дистанции. Обучение больших языковых моделей (LLMs), которые читают целые кодовые базы или коллекции документов в одном запросе, сталкивается с тем же ограничением. API context_parallel() в PyTorch заменяет scaled dot-product attention на ring attention в соответствующем контексте. Команда TorchTitan использовала его для обучения Llama 3 8B на последовательностях длиной 1 миллион токенов с 32 GPU H100.

Ring attention и другие виды параллелизма#

• Контекстный параллелизм: общий подход, при котором одна последовательность распределяется между устройствами. Алгоритм, лежащий в основе реализации PyTorch, — ring attention. Альтернативный метод DeepSpeed Ulysses использует обмены «каждый со всеми», чтобы передать каждому устройству все токены для подмножества голов внимания.
• Тензорный параллелизм и конвейерный параллелизм: веса модели распределяются внутри слоёв или между слоями. Это помогает, когда параметры не помещаются в память одного устройства.
• Flash attention: ядро для одного устройства. Ring attention запускает на каждом устройстве такое блочное ядро для обработки локального блока.

Параллелизм по данным распределяет между устройствами пакеты, а не последовательности. Именно такой режим использует Ultralytics YOLO26 для много-GPU обучения: передача device=[0, 1] запускает DistributedDataParallel в PyTorch. При этом на каждом GPU хранится полная копия модели, каждому GPU выделяется равная доля каждого пакета изображений, а после каждого шага градиенты усредняются.