Ring Attention

环形注意力通过将序列拆分为多个区块，并分布到按环形连接的设备上，从而在单个加速器无法容纳的超长序列上计算精确的自注意力。每台设备保留一个查询区块，并将键和值区块传给相邻设备，同时将通信与当前区块的计算重叠，因此模型可处理的最长序列会随设备数量成比例增长。UC Berkeley 的 Hao Liu、Matei Zaharia 和 Pieter Abbeel 于 2023 年 10 月在《使用分块 Transformer 实现近乎无限上下文的环形注意力》中提出该方法，该工作基于他们早期的分块并行 Transformer，并报告了在数百万个词元的上下文窗口上进行训练和推理，且无需近似注意力。

环形注意力的工作方式#

序列会拆分为 N 个区块，每个 GPU 或 TPU 负责一个区块。整个计算过程中，每台设备都会保留自己的查询区块。在 N 个步骤的每一步中，设备都会计算其查询与当前持有的键值区块之间的注意力，将该区块发送给下一台设备，并从上一台设备接收另一个区块：

经过 N 个步骤后，每个查询区块都会与每个键值区块计算一次。没有任何设备会存储完整的注意力矩阵或完整的键和值集合，因此每台设备的内存用量取决于区块大小，而非序列长度。只要计算一个区块注意力所需的时间长于发送一个区块所需的时间，传输就能被计算过程隐藏；这也为特定硬件设定了最小区块大小。部分结果会使用与单个 GPU 内的闪存注意力相同的运行中 softmax 进行合并：每台设备会持续累加指数化分数及其加权值，并在最后将两者相除。在一台机器上模拟四台设备，可以看出结果是精确的：

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 倍。

应用场景#

只要单个输入长到无法放入一台设备，环形注意力就有用；视频是最明显的例子。同一 Berkeley 研究团队使用它训练了大型世界模型（LWM），这是一个包含 70 亿参数视觉语言模型（VLM）的系列，上下文长度逐步从 4K 扩展到 100 万个交错排列的文本和视频词元。LWM 将每帧编码为 256 个词元，因此以每秒一帧采样的一小时视频共包含 921,600 个词元；该模型能够回答有关一小时 YouTube 合辑的问题，这属于长时程视频理解任务。训练能够在单个提示中读取整个代码库或文档集合的大型语言模型（LLM）也会遇到同样的限制。PyTorch 的 context_parallel() API 会在其上下文中用环形注意力替换缩放点积注意力；TorchTitan 团队则使用它在 32 个 H100 GPU 上以 100 万词元序列训练 Llama 3 8B。

环形注意力与其他并行方式#

• 上下文并行：将一个序列拆分到多台设备上的通用策略。环形注意力是 PyTorch 实现所用的算法；DeepSpeed Ulysses 则是另一种方案，它通过全对全交换，让每台设备获得某些注意力头对应的所有词元。
• 张量并行和流水线并行：在层内或跨层拆分模型权重，适用于参数量超出单台设备内存容量的情况。
• 闪存注意力：一种单设备内核。环形注意力会在每台设备上为本地区块运行这种分块内核。

数据并行拆分的是批次而非序列，Ultralytics YOLO26 多 GPU 训练时采用的就是这种模式：传入 device=[0, 1] 会启动 PyTorch DistributedDataParallel，在每个 GPU 上保留完整的模型副本，让每个 GPU 平分每个图像批次，并在每一步之后对梯度求平均。