Ring Attention

リングアテンションは、1台のアクセラレーターで処理できない長さのシーケンスに対して、デバイスをリング状に配置し、ブロックに分割して厳密なセルフアテンションを計算します。各デバイスはクエリのブロックを1つ保持し、キーと値のブロックを隣のデバイスに渡します。この通信と、手元にあるブロックの計算を重ねることで、モデルが処理できる最大シーケンス長はデバイス数に比例して伸びます。UC BerkeleyのHao Liu氏、Matei Zaharia氏、Pieter Abbeel氏は、2023年10月にRing Attention with Blockwise Transformers for Near-Infinite Contextでこの手法を発表しました。以前に発表したブロック単位の並列トランスフォーマーを発展させたもので、アテンションを近似せずに数百万トークンのコンテキストウィンドウで学習と推論を行えると報告しています。

リングアテンションの仕組み#

シーケンスをN個のブロックに分割し、各GPUまたはTPUに1つずつ割り当てます。各デバイスは、計算中ずっと自分のクエリブロックを保持します。Nステップの各ステップで、現在保持しているキー・バリューブロックとクエリとの間のアテンションを計算し、そのブロックを次のデバイスに送り、前のデバイスから別のブロックを受け取ります。

Nステップ後には、すべてのクエリブロックがすべてのキー・バリューブロックと処理されています。どのデバイスも、アテンション行列全体やキーと値の全データを保持しません。そのため、デバイスあたりのメモリ使用量はシーケンス長ではなくブロックサイズに依存します。1ブロックのアテンション計算にかかる時間が、そのブロックの転送時間より長ければ、転送は計算に隠されます。この条件が、特定のハードウェアに必要な最小ブロックサイズを決めます。部分的な結果は、1台のGPU内でflash attentionが使うものと同じ逐次的なsoftmaxで統合されます。各デバイスは、指数化したスコアの累積和と、重み付き値の累積和を保持し、最後に前者を後者で割ります。1台のマシンで4台のデバイスをシミュレーションすると、厳密な結果が得られます。

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倍に高めます。

用途#

リングアテンションは、1つの入力が単一デバイスで処理できる長さを超える場面で重要であり、動画が最も分かりやすい例です。同じBerkeleyの研究グループは、リングアテンションを使ってLarge World Model（LWM）を学習しました。LWMは、テキストと動画を交互に含むコンテキストを4Kから100万トークンまで段階的に拡張した、70億パラメーターのビジョン・ランゲージモデル（VLM）のファミリーです。LWMは各フレームを256トークンとしてエンコードするため、1秒に1フレームをサンプリングした1時間の動画は921,600トークンになります。このモデルは、長時間のYouTube動画集について質問に回答しました。これは長期的な動画理解のタスクです。コードベース全体や文書コレクション全体を1つのプロンプトで読み込む大規模言語モデル（LLM）の学習でも、同じ制約に直面します。PyTorchのcontext_parallel() APIは、コンテキスト内でscaled dot-product attentionをリングアテンションに置き換えます。TorchTitanチームはこれを使い、32基のH100 GPUにまたがる100万トークンのシーケンスでLlama 3 8Bを学習しました。

リングアテンションとその他の並列化手法#

• コンテキスト並列化： 1つのシーケンスを複数のデバイスに分割する一般的な手法です。リングアテンションは、PyTorchの実装で使われているアルゴリズムです。DeepSpeed Ulyssesは、all-to-all交換を使って、各デバイスに一部のアテンションヘッドに対応するすべてのトークンを割り当てる別の手法です。
• テンソル並列化とパイプライン並列化： レイヤー内またはレイヤー間でモデルの重みを分割します。パラメーターが1台のデバイスのメモリ容量を超える場合に役立ちます。
• Flash Attention： 単一デバイス用のカーネルです。リングアテンションでは、各デバイスがローカルブロックに対して、この種のブロック単位カーネルを実行します。

データ並列化では、シーケンスではなくバッチを分割します。Ultralytics YOLO26は、マルチGPU 学習にこの方式を使用します。device=[0, 1]を渡すと、PyTorchのDistributedDataParallelが起動します。各GPUにモデル全体のコピーを保持し、画像の各バッチを均等に分配して、各ステップ後に勾配を平均化します。