💬 NLP

Folding Tensor and Sequence Parallelism for Memory-Efficient Transformer Training & Inference

本文提出张量与序列并行（TSP），这是一种新颖的执行策略，它将权重分片和令牌分片折叠到单一设备轴上，以同时降低参数和激活内存开销，为长上下文和内存受限的 Transformer 模型的训练与推理提供了一种硬件高效的替代方案。

原作者： Vasu Shyam, Anna Golubeva, Quentin Anthony

发布于 2026-04-30

📖 1 分钟阅读☕ 轻松阅读

原作者： Vasu Shyam, Anna Golubeva, Quentin Anthony

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象你正和一群朋友试图拼一幅巨大的拼图，但你们只有一张非常小的桌子（计算机内存）可以工作。拼图太大，以至于没有任何一个人能一次性握住所有碎片。

本文介绍了一种让多台计算机（GPU）协同工作以训练巨型 AI 模型（本质上就是这些巨大的拼图）的新方法。作者将他们的这一新策略称为TSP（张量并行与序列并行）。

以下是使用简单类比进行的分解说明：

问题：两种旧的分工方式

为了解决这个拼图难题，团队通常采用两种旧方法之一，但两者都有缺陷：

“权重拆分”法（张量并行）：
想象拼图碎片是游戏的“规则”（模型的权重）。在这种方法中，你把规则书切成两半。A 持有前半部分规则，B 持有后半部分规则。
- 优点： 你节省了桌面上的空间，因为你没有把整本规则书存储两份。
- 缺点： 如果拼图包含一个很长的故事（长序列单词），每个人仍然必须握有整个故事才能进行游戏。如果故事巨大，你的手（内存）就会塞满，导致崩溃。
“故事拆分”法（序列并行）：
想象拼图碎片是故事本身。在这种方法中，A 持有故事的前半部分，B 持有故事的后半部分。
- 优点： 你节省了桌面上的空间，因为你没有一次性握住整个故事。
- 缺点： 每个人都必须记住整本规则书。如果规则书巨大，你的大脑（内存）就会塞满，导致崩溃。

旧的混合方案： 通常，团队会尝试通过两组不同的朋友来同时做这两件事。一组拆分规则，另一组拆分故事。但这效率低下，因为它用光了所有朋友来拆分工作，导致没有人剩下可以协助其他任务（如数据并行）。

解决方案：“折叠”法（TSP）

作者说：“为什么要用两组分开的人？让我们把工作折叠到单一轴线上。”

在TSP中，团队中的每个人同时做两件事：

他们持有规则书的一部分（权重）。
他们持有故事的一部分（序列）。

类比：
想象你正在参加一个晚宴。

旧方式： 你有一张桌子，一个人传递菜单（权重），而每个人阅读整本书。另一张桌子上，人们传递书籍（故事），而每个人背诵整份菜单。
TSP 方式： 桌上的每个人都拿到一小份菜单和一小段故事。

他们如何实现它（魔法技巧）

既然每个人都只有一小份菜单和一小段故事，他们就必须频繁交流以完成拼图。本文描述了两种聪明的方法，使他们在不被压垮的情况下做到这一点：

针对“故事”部分（注意力机制）：
想象团队需要知道整个故事才能理解特定的句子。与其让所有人同时大喊整个故事，不如让他们轮流进行。一个人向所有人广播他手中的菜单部分。然后，每个人计算他们那部分故事，并快速交换他们的故事碎片（键和值）以重建完整的上下文。这就像接力赛，他们在奔跑时传递接力棒（数据）。
针对“规则”部分（MLP）：
想象团队需要将不同的规则应用到他们的故事碎片上。与其停下来大喊规则，不如让他们把规则书页围成一圈（环形）传递。A 用第 1 页做数学运算，然后将第 1 页传给 B，同时 B 将第 2 页传给 C。在页面移动的同时，每个人都在忙着做数学运算。这保持了“交通”的流动，同时完成了“工作”。

为什么这更好？

本文声称 TSP 是一种“硬件感知”的解决方案，意味着它是专门为现代计算机芯片相互通信的方式而设计的。

内存节省： 因为每个人都持有部分规则和部分故事，每台计算机所需的内存显著下降。这使得团队能够处理更长的故事（更长的上下文），而不会耗尽内存。
速度： 尽管他们来回传递的数据更多（听起来更慢），但他们以与思考重叠的方式进行传递。“传递”发生在他们“思考”的同时，因此总时间不会增加太多。
适应房间容量： 在计算机集群中，最快的连接通常位于同一台机器上的芯片之间（就像坐在同一张桌子旁的人）。较慢的连接位于不同机器之间（不同房间里的人）。
- 旧方法经常迫使团队分散到不同的房间，从而拖慢速度。
- TSP 允许整个“拆分”团队容纳在单台机器（一张桌子）上，使他们保持在快车道上。

结果

作者在由 1,024 块强大 GPU（MI300X）组成的巨大集群上测试了这种方法。

内存： 在每次测试中，TSP 使用的内存最少，特别是在故事非常长的时候。
速度： TSP 与旧方法一样快，甚至更快。
可扩展性： 随着团队中计算机数量的增加，TSP 继续表现良好，而旧方法开始因内存限制而挣扎。

简而言之： TSP 是一种更聪明的组织计算机团队的方式。它不再将“规则”和“故事”拆分成独立的组，而是将它们结合起来，使每台计算机都只持有少量两者。这节省了空间，允许更长的故事，并使团队在相同的快速网络上高效工作。

以下是论文《折叠张量与序列并行以实现内存高效的 Transformer 训练与推理》的详细技术总结。

1. 问题陈述

大型语言模型（LLM）的训练和推理面临两个主要的内存瓶颈：

参数内存：存储模型权重、梯度和优化器状态所需的内存。
激活内存：存储中间激活值所需的内存，在注意力机制中随序列长度呈二次方增长（ $S^2$ ）。

现有的并行策略要么单独解决这些问题，要么存在权衡：

张量并行（TP）：将模型权重分片到各个设备上，减少了参数内存，但激活内存完全复制。由于激活内存高且通信量大，它在长序列长度下效率低下。
序列并行（SP）：将输入序列分片到各个设备上，减少了激活内存，但模型权重完全复制。当参数内存成为瓶颈时，它效果不佳。
传统 TP+SP：通过将两者分配给设备网格中的正交轴（例如 $T \times \Sigma$ 网格）来结合两者。虽然有效，但这需要大量设备用于模型并行（ $T \times \Sigma$ ），从而减少了可用于数据并行的设备数量。此外，它经常迫使并行组跨越较慢的节点间链路，未能充分利用高带宽的节点内互连（如 NVLink 或 AMD Infinity Fabric）。

核心挑战：如何在单个设备轴上同时减少参数和激活内存，以最大化内存效率并释放设备用于数据并行，同时不产生过高的通信成本。

2. 方法论：张量与序列并行（TSP）

作者提出了 TSP，这是一种将张量并行和序列并行“折叠”到单个设备轴上的策略。TSP 不是将 TP 和 SP 分配给正交维度，而是让每个秩（GPU）同时持有模型权重的一部分和序列的一部分。

关键设计组件：

统一分片：每个设备持有 $1/D$ 的模型参数和 $1/D$ 的序列标记。这在单个轴上将参数比例内存和激活内存均减少了 $D$ 倍。
运行时调度：为了处理权重和激活以不同方式分片的事实，TSP 针对不同层采用特定的通信调度：
- 注意力层：使用广播调度。
  - 秩遍历权重分片（打包的 $W_Q, W_K, W_V, W_O$ ）。
  - 在每次迭代中，一个秩将其权重分片广播给所有其他秩。
  - 所有秩在其序列分片上计算局部投影。
  - 执行 All-Gather 操作以在序列维度上重建完整的 Key/Value 上下文，从而执行 FlashAttention。
  - 采用之字形分区方案来平衡因果注意力的工作负载（确保后续标记不会承担不成比例更多的工作）。
- 门控 MLP 层：使用环形通信调度。
  - 权重分片（ $W_1, W_2, W_3$ ）在环上点对点循环。
  - 每个秩在本地累积部分输出。
  - 这消除了对 All-Reduce 操作（TP 中的标准操作）的需求，因为序列保持本地化，仅权重移动。
硬件感知：该设计针对 AMD MI300X（Infinity Fabric）等拓扑进行了优化，在这些拓扑中，集体带宽随参与 GPU 的数量而扩展。TSP 确保通信原语（广播和环形）利用整个组，或与计算重叠以最大化效率。

3. 主要贡献

新颖的并行策略：引入 TSP，将两个正交的并行轴合并为一个，允许同时减少参数和激活内存。
理论分析：建立了一个严格的数学模型，将 TSP 与 DP、TP、SP 和 TP+SP 进行比较，涉及：
- 内存：TSP 是唯一一种无需 $D^2$ 个设备即可将权重比例内存和激活内存均减少 $1/D$ 的方案。
- 通信：虽然 TSP 引入了权重移动通信（TP+SP 避免了这一点），但作者证明，对于长序列或中等批量大小，激活比例通信占主导地位，使得 TSP 具有竞争力。
- FLOPs：TSP 保持了与 TP 和 SP 相同的每设备计算效率。
实现与重叠：详细实现了 TSP 注意力和 MLP 块，并明确了通信/计算重叠策略（在 GEMM 计算之后流水线化权重传输）。
拓扑优化：证明 TSP 允许整个模型副本适应单个高带宽节点（节点内），避免了 2D TP+SP 网格通常所需的跨越较慢节点间链路所带来的性能惩罚。

4. 结果

作者在多达 1,024 个 AMD MI300X GPU（128 个节点）的集群上评估了 TSP，使用了一个 70 亿参数（7B）的密集 Transformer 模型。

内存效率：
- TSP 在所有测试的序列长度（16K 到 128K 标记）中始终实现了最低的峰值内存。
- 在短上下文下，TSP 与 TP 表现相当（主要由参数内存主导）。
- 在长上下文下，TSP 显著优于 TP 和 TP+SP（主要由激活内存主导），防止了其他策略失败时出现的内存溢出（OOM）错误。
吞吐量：
- 尽管由于权重移动导致理论通信量更高，但 TSP 保持了与 TP+SP 具有竞争力的吞吐量。
- 额外的通信通过流水线化成功隐藏在计算之后。
- 随着并行度（ $D$ ）和序列长度的增加，吞吐量优势逐渐扩大。
可扩展性：
- 由于内存占用减少，TSP 允许更大的微批量，进一步提高了吞吐量。
- 它使得模型副本能够完全适应单个节点（8 个 GPU），从而在更大的集群中释放秩用于数据并行（DP）。

5. 意义

硬件感知扩展：TSP 为现代 HPC 集群提供了一种实用解决方案，在这些集群中，节点内带宽很高，但节点间带宽是瓶颈。它最大限度地利用了昂贵的节点内互连。
长上下文可行性：随着模型向更长的上下文窗口（100k+ 标记）发展，激活内存成为主要瓶颈。TSP 独特地有能力处理这一问题，而无需 2D 并行所需的巨大设备开销。
灵活性：TSP 可以与其他并行方案（流水线并行、专家并行、数据并行）正交组合，为系统设计师提供了一个灵活的“轴”，以针对特定的硬件拓扑进行优化。
成本效益：通过减少模型并行所需的设备数量，TSP 允许更多设备用于数据并行，从而可能减少给定训练任务所需的集群总规模，或增加有效批量大小。

总之，TSP 代表了分布式训练策略的重大演进，它以可管理的通信量增加为代价，换取了内存效率和硬件利用率的显著提升，使其成为训练下一代长上下文 LLM 的关键工具。