Ring Attention

Ring attention은 링 형태로 연결된 장치에 시퀀스를 블록 단위로 분할해, 단일 가속기에서 처리하기에는 너무 긴 시퀀스에 대해 정확한 셀프 어텐션을 계산합니다. 각 장치는 쿼리 블록 하나를 보유하고 키와 값 블록을 이웃 장치에 전달합니다. 이때 이미 보유한 블록의 연산과 통신을 겹쳐 수행하므로, 모델이 처리할 수 있는 최대 시퀀스 길이는 장치 수에 비례해 늘어납니다. UC Berkeley의 Hao Liu, Matei Zaharia, Pieter Abbeel은 2023년 10월 Ring Attention with Blockwise Transformers for Near-Infinite Context에서 이 방법을 소개했습니다. 이들은 이전의 블록 단위 병렬 트랜스포머 연구를 바탕으로, 어텐션을 근사하지 않고 수백만 개의 토큰을 포함하는 컨텍스트 윈도우로 학습 및 추론을 수행했다고 보고했습니다.

Ring Attention의 작동 방식#

시퀀스를 N개 블록으로 분할하고 GPU 또는 TPU마다 하나씩 할당합니다. 각 장치는 전체 계산 과정에서 자신의 쿼리 블록을 보유합니다. N단계에 걸쳐 각 장치는 쿼리와 현재 보유한 키-값 블록 간의 어텐션을 계산하고, 해당 블록을 다음 장치로 보내며 이전 장치에서 다른 블록을 받습니다.

N단계를 거치면 각 쿼리 블록이 모든 키 및 값 블록과 연산을 마칩니다. 어떤 장치도 전체 어텐션 행렬이나 키와 값 전체를 저장하지 않으므로, 장치별 메모리 사용량은 시퀀스 길이가 아니라 블록 크기에 따라 달라집니다. 한 블록의 어텐션 계산이 블록 전송보다 오래 걸리면 전송이 계산에 가려지며, 이는 주어진 하드웨어에서 필요한 최소 블록 크기를 결정합니다. 부분 결과는 단일 GPU 안에서 Flash Attention이 사용하는 것과 같은 누적 소프트맥스로 결합합니다. 각 장치는 지수화된 점수의 누적합과 가중치가 적용된 값의 누적합을 유지하고, 마지막에 두 값을 나눕니다. 단일 머신에서 장치 네 개를 시뮬레이션하면 결과가 정확함을 확인할 수 있습니다.

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은 각 장치에 시퀀스 전체에 고르게 분산된 토큰을 할당해 이 문제를 해결하며, 토큰 256K에서 종단 간 학습 처리량을 최대 1.45배 높입니다.

응용 분야#

Ring attention은 단일 장치로 처리할 수 없을 만큼 입력이 긴 모든 경우에 중요하며, 비디오는 가장 대표적인 사례입니다. Berkeley의 같은 연구 그룹은 이를 사용해 70억 개 파라미터를 가진 비전-언어 모델(VLM) 계열인 Large World Model(LWM)을 학습했습니다. LWM의 컨텍스트 길이는 텍스트와 비디오가 교차하는 시퀀스 기준으로 4K에서 100만 토큰까지 단계적으로 늘어났습니다. LWM은 각 프레임을 토큰 256개로 인코딩하므로 초당 1프레임을 샘플링한 1시간 분량의 비디오는 토큰 921,600개에 해당합니다. 이 모델은 장시간 비디오 이해 작업인 1시간 길이 YouTube 편집 영상에 관한 질문에 답했습니다. 전체 코드베이스나 문서 모음을 하나의 프롬프트에서 읽는 대규모 언어 모델(LLM)을 학습할 때도 같은 제약에 부딪힙니다. PyTorch의 context_parallel() API는 컨텍스트 내부의 scaled dot-product attention을 ring attention으로 바꾸며, TorchTitan 팀은 이를 사용해 H100 GPU 32개에서 토큰 100만 개 길이의 시퀀스로 Llama 3 8B를 학습했습니다.

Ring Attention과 다른 병렬화 기법 비교#

• 컨텍스트 병렬화: 하나의 시퀀스를 여러 장치에 분할하는 일반적인 전략입니다. Ring attention은 PyTorch 구현의 기반 알고리즘이며, DeepSpeed Ulysses는 올-투-올 교환을 사용해 각 장치에 일부 어텐션 헤드의 모든 토큰을 할당하는 대안입니다.
• 텐서 병렬화 및 파이프라인 병렬화: 모델 가중치를 계층 내부 또는 계층 간에 분할합니다. 파라미터가 단일 장치의 메모리 용량을 초과할 때 유용합니다.
• Flash Attention: 단일 장치에서 실행되는 커널입니다. Ring attention은 각 장치의 로컬 블록에서 이러한 유형의 블록 단위 커널을 실행합니다.

데이터 병렬화는 시퀀스가 아닌 배치를 분할하며, Ultralytics YOLO26은 멀티 GPU 학습에 이 방식을 사용합니다. device=[0, 1]을 전달하면 PyTorch DistributedDataParallel이 실행됩니다. 이 방식은 각 GPU에 전체 모델 복사본을 유지하고, 모든 이미지 배치를 각 GPU에 동일하게 분배하며, 각 단계가 끝난 후 그래디언트를 평균합니다.