Ulysses与Ring-Attention序列并行技术详解

在大模型训练的工程实践中，一个日益棘手的问题正不断挑战硬件极限：如何高效处理超长序列输入？随着Qwen3、Llama4等模型支持32k甚至131k上下文，多模态场景中一张高分辨率图像也能轻易生成上万token，传统注意力机制的$O(n^2)$复杂度让显存消耗呈爆炸式增长。一块H100的80GB显存，在处理4096长度序列时可能仅够跑一个微小批次。

正是在这种背景下，Ulysses和Ring-Attention作为两种先进的序列并行策略脱颖而出。它们不再依赖昂贵的单卡扩容，而是通过分布式协作“化整为零”，将原本集中式的全局注意力拆解到多个设备上协同完成。这种思路不仅缓解了显存压力，更打开了通向超长上下文建模的大门。

以魔搭社区的ms-swift框架为例，其显存优化模块原生集成了这两种技术，用户只需一行配置即可启用。但这背后的实现远非“开箱即用”那么简单——理解其原理，才能在实际部署中做出合理权衡。

我们不妨从最直观的问题开始：为什么标准自注意力会成为瓶颈？

假设输入序列长度为 $N=8192$，隐藏维度 $D=4096$，使用BF16精度。仅Key和Value缓存就需要：
$$
2 \times N \times D \times 2\,\text{bytes} = 2 \times 8192 \times 4096 \times 2 \approx 128\,\text{GB}
$$
这还只是激活值，不包括模型参数和梯度。显然，任何单卡都无法承受。而数据并行对此无能为力，因为它复制的是整个计算图；张量并行主要切分权重维度，对序列维度毫无帮助。

这就是序列并行的用武之地。

先看Ulysses——它像是一个“聚合再计算”的模式。设想你有8张GPU，每张负责1/8的输入token。每个设备独立计算本地Query（Q），以及本地的Key（K）和Value（V）。但关键在于下一步：所有设备通过All-Gather操作交换各自的K和V，最终每台设备都拥有了完整的全局K和V矩阵。

有了全局K/V，每个设备就可以用自己的局部Q去参与完整注意力计算。比如第0号GPU虽然只持有前1024个token的Q，但它能计算出这些Q对全部8192个KV的关注权重。最后，将输出结果按序列维度再次划分，并通过Reduce-Scatter分发回各设备，保持分布一致性。

整个过程等价于单机全量计算，但显存占用从 $O(N^2)$ 降到了接近 $O(N^2/P)$，其中 $P$ 是设备数。通信方面，All-Gather会带来一次总量为 $O(Nd)$ 的跨设备数据交换（$d$为头维度），在NVLink或InfiniBand这类高速互联下，延迟可控。

下面是其核心逻辑的简化实现：

import torch import torch.distributed as dist def ulysses_attention_forward(q, k, v, rank, world_size): # All-Gather K and V to form global context global_k = [torch.zeros_like(k) for _ in range(world_size)] global_v = [torch.zeros_like(v) for _ in range(world_size)] dist.all_gather(global_k, k) dist.all_gather(global_v, v) global_k = torch.cat(global_k, dim=2) # [B, H, L, D] global_v = torch.cat(global_v, dim=2) # Local Q attends to global K/V attn_weights = torch.matmul(q, global_k.transpose(-1, -2)) / (q.size(-1) ** 0.5) attn_weights = torch.softmax(attn_weights, dim=-1) attn_output = torch.matmul(attn_weights, global_v) return attn_output

这段代码虽未实现真正的Reduce-Scatter，但已体现精髓。在 ms-swift 中，只需设置sequence_parallel=ulysses即可自动启用，框架会接管通信调度与内存管理。

然而，Ulysses 并非没有代价。All-Gather要求所有设备同时持有全局K/V副本，这意味着每张卡仍需存储完整的 $(N \times d)$ 级别中间状态，尤其当注意力头较多时，显存峰值依然可观。此外，随着设备数量增加，All-Gather的广播式通信容易引发网络拥塞，限制了其在百卡以上集群中的扩展性。

于是，Google 提出的Ring-Attention给出了另一种解法：不追求一次性获取全局信息，而是通过环形流水线逐步累积。

想象P个GPU排成一个环。初始时，每个设备拥有自己的K和V。然后进入 $P-1$ 轮通信循环：每轮中，每个设备将其当前持有的K/V发送给下一个节点（顺时针），同时从前一个节点接收新的K/V块。收到后，立即用本地Q与新到的K/V计算部分注意力贡献，并累加到输出缓冲区中。

经过 $P-1$ 轮后，每个设备都“见过”所有分片，完成了等效的全局注意力。由于每次只处理一对分片，无需缓存完整K/V，显存占用极为平稳。

更重要的是，环形通信天然负载均衡——没有中心节点，也没有广播风暴，非常适合大规模集群部署。当然，这也带来了更高的实现复杂度：必须精确控制每轮的发送/接收顺序，避免死锁；且总通信轮数随设备数线性增长，延迟高于Ulysses。

数值稳定性是另一个挑战。由于输出是多轮softmax结果的累加，直接相加会导致溢出。因此，Ring-Attention通常采用对数空间累加（log-space accumulation）来保证精度：

def ring_attention_forward(q, k, v, rank, world_size): B, H, S, D = q.shape device = q.device output_acc = torch.zeros_like(q) logsumexp_acc = torch.zeros(B, H, S, device=device) current_k, current_v = k.clone(), v.clone() for step in range(world_size): attn_scores = torch.matmul(q, current_k.transpose(-1, -2)) / (D ** 0.5) # Numerically stable accumulation using log-sum-exp max_prev = logsumexp_acc.max(dim=-1, keepdim=True)[0] max_curr = attn_scores.max(dim=-1, keepdim=True)[0] max_both = torch.maximum(max_prev, max_curr) exp_diff_acc = torch.exp(logsumexp_acc - max_both) exp_diff_curr = torch.exp(attn_scores - max_both.unsqueeze(-1)) logsumexp_acc = max_both + torch.log(exp_diff_acc + exp_diff_curr.sum(dim=-1, keepdim=True)) partial_out = torch.matmul(torch.softmax(attn_scores - max_both.unsqueeze(-1), dim=-1), current_v) output_acc = (output_acc * exp_diff_acc + partial_out * exp_diff_curr.sum(dim=-1, keepdim=True)) \ / (exp_diff_acc + exp_diff_curr.sum(dim=-1, keepdim=True)) if step < world_size - 1: next_rank = (rank + 1) % world_size prev_rank = (rank - 1 + world_size) % world_size dist.send(tensor=current_k, dst=next_rank) dist.send(tensor=current_v, dst=next_rank) current_k = torch.zeros_like(k) current_v = torch.zeros_like(v) dist.recv(tensor=current_k, src=prev_rank) dist.recv(tensor=current_v, src=prev_rank) return output_acc

可以看到，该实现通过维护logsumexp_acc实现跨轮次的概率归一化，确保最终输出与原始attention数学等价。

在实际系统如 ms-swift 中，这类细节已被封装。开发者可通过简洁的YAML配置切换策略：

parallel: sequence_parallel: "ring" # or "ulysses" tensor_parallel_size: 4 pipeline_parallel_size: 2 data_parallel_size: 8

但选择哪种方式，仍需结合具体场景判断。

如果你的集群规模较小（<32卡）、网络带宽充足（如NVLink全连接），且希望快速落地，Ulysses 是更稳妥的选择。它的逻辑清晰，调试简单，在Megatron-LM和DeepSpeed中均有成熟实现。

而当你面对百卡级训练任务，或受限于RDMA网络的吞吐能力，Ring-Attention 的通信效率优势就会显现。它对显存波动更友好，适合长时间稳定运行的大规模作业。PaLM和Chinchilla等超大规模模型已验证其可行性。

值得注意的是，这两种技术并非孤立存在。在真实训练流水线中，它们常与其他并行策略组合使用：

与TP结合：在head维度进一步切分注意力头，形成“双层并行”；
与PP协同：在层间做流水线划分，提升设备利用率；
叠加DP：在批次维度复制模型，增强数据吞吐。

例如，在训练Qwen3-VL这样的多模态模型时，视觉编码器输出的4096个patch token进入LLM的cross-attention层。此时若启用ulysses或ring，系统会自动将序列切分到8卡集群，每卡仅处理512个token的Q，却能感知全局视觉上下文。反向传播时，梯度也沿相同路径同步，确保数学一致性。

这种架构解决了三大痛点：