[Paper] Untied Ulysses：内存高效上下文并行 via Headwise Chunking

发布: 3天前 (2026年2月25日 GMT+8 02:54)

6 分钟阅读

原文: arXiv

Source: arXiv - 2602.21196v1

概述

本文介绍了 UPipe，一种新的上下文并行策略，它将注意力计算按 每个头 切分，而不是按整个层切分。通过这种方式，它大幅削减了自注意力所需的激活内存，使开发者能够在不牺牲吞吐量的前提下，在更长的序列上训练大规模 Transformer。

Headwise Chunking：对注意力矩阵在单个头级别进行细粒度划分，显著降低内存使用。
Memory Savings：在 320 亿参数模型上，中间张量内存最高可降低 87.5 %。
Scalable Throughput：保持的训练速度可与现有的上下文并行方法（如 Ring Attention 和 DeepSpeed Ulysses）相媲美。
Record‑setting Context Length：展示在单个 8‑GPU H100 节点上训练 Llama‑3‑8B，使用 5 million‑token 的上下文长度，比之前的技术提升了 >25 %。
Simplicity：该技术实现仅需极少的代码修改，且不依赖特殊硬件特性。

传统的上下文并行将长序列划分到多个 GPU 上，但每个 GPU 仍需保存其切片对应的完整注意力矩阵，这会迅速耗尽显存。UPipe 改变了划分的粒度：

按头划分（Headwise Partitioning）： 将每个注意力头的查询‑键‑值（QKV）张量划分为小块（例如，1 k‑token 块）。
局部计算（Local Computation）： GPU 只计算分配给它们的块的注意力得分，然后立即丢弃中间结果。
流式归约（Streaming Reduction）： 将部分结果在 GPU 之间以环形（ring‑like）通信模式进行求和，重建完整的注意力输出，而无需在任何单个设备上实际生成完整矩阵。
与反向传播重叠（Overlap with Back‑propagation）： 将块化的前向传播与梯度计算流水线化，使 GPU 持续忙碌，保持整体吞吐量。

该方法基于已有的 “Ring Attention” 通信模式，额外加入轻量级调度器来协调头级块划分，只需对 Transformer 核心做少量修改。

模型 / 设置	最大上下文（tokens）	内存降低	训练吞吐量
32B Transformer (Ring Attention)	~1.2 M	baseline	1.0×
32B Transformer (UPipe)	5 M	≈ 87 %	0.96×
Llama‑3‑8B (8 × H100)	5 M	≈ 80 %	comparable to DeepSpeed Ulysses

通信开销： 虽然不大，但额外的 all‑reduce 步骤在高延迟互连（例如多区域集群）上会变得明显。
块大小调优： 最佳块粒度取决于模型规模和硬件；自动调优留给用户自行完成。
非 Transformer 架构： 该方法针对自注意力设计；将其扩展到卷积或 mixture‑of‑expert 层仍未探索。
未来方向： 作者建议将 headwise chunking 与激活卸载或流水线并行相结合，以将上下文长度推至 10 M 以上的 token，并进一步降低通信开销。