[Paper] 在多GPU上使用张量并行扩展状态空间模型

Source: arXiv - 2602.21144v1

概述

选择性状态空间模型（SSMs）已成为处理大型语言模型（LLMs）中超长上下文的强大替代方案，相较于传统的 Transformers。论文 “Scaling State‑Space Models on Multiple GPUs with Tensor Parallelism” 针对一个实际瓶颈：当单个 GPU 的显存或带宽不足时，如何高效运行基于 SSM 的 LLM。通过改编张量并行（TP）——一种已在扩展 Transformers 中广受欢迎的技术——作者展示了如何将繁重的 SSM 计算分布到多块 GPU 上，同时保持关键的递归路径高速且通信开销低。

关键贡献

• TP‑aware SSM 设计 将大型投影矩阵进行切分，但保持每个 GPU 上的每个 token 循环状态更新本地化。
• 状态缓存机制 通过在预填充和后续解码阶段复用 SSM 隐状态，实现 “prefill‑to‑decode” (TTFT) 加速。
• 量化 All‑Reduce 用于 TP 聚合步骤，将同步带宽削减约 18 %，且不牺牲数值稳定性。
• 全面评估 在三类真实世界基于 SSM 的 LLM 系列（Mamba、Falcon‑Mamba、Zamba）上，使用 NVIDIA A6000 与 A100 集群，展示从 1 到 4 GPU 扩展时 1.6–4.0× 吞吐提升。
• 开源实现（或详细伪代码），可最小改动地嵌入现有推理栈。

方法论

1. 张量划分 – 作者将 SSM mixer 的打包权重张量（包含长程投影和局部混合核）视为一个“巨大的矩阵”。他们沿特征维度在 GPU 之间拆分该矩阵，类似于 Transformer 的经典张量并行（TP）。
2. 局部递归 – 与 Transformer 的自注意力不同，SSM 对每个 token 顺序更新隐藏状态。该设计确保每个 GPU 保持自己的隐藏状态切片，因此递归更新在关键路径上 不需要 跨 GPU 通信。
3. 预填充与解码期间的状态缓存 – 在初始的“预填充”（处理长提示）阶段，隐藏状态被缓存。当模型切换到逐 token 解码时，使用缓存的状态，避免对已处理 token 重新计算昂贵的投影。
4. 量化 All‑Reduce – 每生成一个 token 后，需要将所有 GPU 的部分结果求和（All‑Reduce）以形成最终的隐藏表示。作者在归约前将张量量化为 8 位，归约后再反量化，从而显著降低跨互连传输的数据量。
5. 基准套件 – 他们在三类模型上进行推理测试，变化上下文长度（2 K–64 K token）和批量大小，测量原始吞吐量（tokens/s）以及端到端请求延迟。

结果与发现

• Long‑context advantage（长上下文优势）: 随着上下文长度的增加，速度提升会更明显，因为 SSM 的每 token 成本占主导，且 TP 切分降低了每块 GPU 的显存压力。
• TTFT benefit（首词到首词时间收益）: 在预填充和解码阶段缓存状态，可将预填充延迟平均降低约 30 %。
• Communication efficiency（通信效率）: 量化 All‑Reduce 减少了 PCIe/NVLink 流量，使得该方法即使在互连带宽一般的集群上也能可行。

实际意义

• 部署者现在可以在 2‑4 GPU 节点上运行基于 SSM 的 LLM，拥有 32 K+ 上下文，而不会触及显存上限，这打开了文档级问答、代码库搜索或长文本生成等使用场景。
• 成本效益的扩展方式： 与其购买单块巨型 GPU（例如 H100 80 GB），团队可以将更实惠的 A6000/A100 卡拼接在一起，仍然实现近线性加速。
• 框架集成： 该设计能够干净地映射到现有的 TP 库（如 Megatron‑LM、DeepSpeed），意味着开发者只需进行少量配置更改即可加入 SSM 支持，而无需完整重写代码。
• 实时应用的更低延迟： TTFT 缓存降低了从提示摄取切换到逐 token 生成时的“热身”惩罚，这对聊天式助手至关重要。
• 量化通信 为其他内存占用大的模型（例如检索增强的 Transformer）提供了模板，在这些模型中 All‑Reduce 往往成为瓶颈。

限制与未来工作

• 硬件依赖性: 在具备高速 NVLink 的 GPU 上观察到最大的收益；在较慢的互连上，量化的 All‑Reduce 仍可能成为瓶颈。
• 模型特定调优: 划分策略假设 SSM mixer 具有特定形状；对于形状不规则或稀疏化的 mixer，可能需要自定义切片逻辑。
• 精度权衡: 虽然 8 位量化在归约步骤中有效，但作者指出在某些下游任务上会出现小幅 (<0.2 BLEU) 的下降——未来工作可以探索混合精度或自适应量化。
• 超出推理的扩展: 本文聚焦于推理；将相同的张量并行方案应用于训练（尤其是梯度累积）仍是一个未解的挑战。

底线: 通过将张量并行与巧妙的状态缓存及量化通信相结合，这项工作使大规模、长上下文的 SSM 推理在当今多 GPU 集群上变得实用——这一进展有望加速下一波基于 LLM 的应用。

作者

• Anurag Dutt
• Nimit Shah
• Hazem Masarani
• Anshul Gandhi

论文信息

• arXiv ID: 2602.21144v1
• Categories: cs.DC, cs.LG
• Published: 2026年2月24日
• PDF: 下载 PDF