veScale-FSDP: Flexible and High-Performance FSDP at Scale

本文提出了 veScale-FSDP，一种通过引入灵活的 RaggedShard 分片格式和结构感知规划算法，有效解决现有 FSDP 系统难以适配块状量化及非元素级优化器等前沿需求，从而在数万张 GPU 规模下实现更高吞吐量和更低内存占用的全新分布式训练系统。

要点

这篇文章介绍了一个名为 veScale-FSDP 的新技术，它是用来训练超级巨大的人工智能模型（比如像人一样聪明的聊天机器人）的“超级引擎”。

为了让你更容易理解，我们可以把训练一个大模型想象成在一个巨大的厨房里，由成千上万个厨师（GPU 显卡）合作做一道极其复杂的“超级大餐”。

1. 以前的困境：切菜切得乱七八糟

背景：
现在的 AI 模型太大了，一台电脑（GPU）根本存不下，必须把模型“切”成几千份，分给几千台电脑一起算。以前的方法（叫 FSDP 或 ZeRO）就像是一个死板的切菜工。

• 死板的切法（旧系统）： 不管这块肉（模型参数）是什么形状，切菜工都只按“一刀切”的方式，把肉切成完全一样大小的薄片。
  • 问题 A（切坏了）： 有些特殊的烹饪方法（比如“块状量化”或“矩阵优化器”），要求必须整块整块地处理，不能切碎。旧系统把整块肉切碎了，厨师们就得把碎肉拼回去，或者重新切，非常浪费时间。
  • 问题 B（切多了）： 为了凑齐人数，切菜工经常会在肉块之间塞入很多“空气”（填充数据/Padding），导致切出来的肉片里有很多没用的空隙，既占地方（内存浪费）又浪费传送时间（通信效率低）。
  • 结果： 几千个厨师忙得团团转，但真正在做菜的时间很少，大部分时间都在搬运碎肉和塞空气。

2. 新的解决方案：veScale-FSDP（智能切菜工 + 传送带）

ByteDance（字节跳动）的工程师们发明了这个新系统，它做了三件聪明的事：

第一招：灵活的“不规则切法” (RaggedShard)

• 比喻： 以前的切菜工只切正方形，现在 veScale-FSDP 请来了一个**“随形切菜大师”**。
• 解释： 它不再强迫把模型切成一样大的块。如果模型里有一块是 32x32 的方块，它就按 32x32 切；如果有一块是长条，它就按长条切。
• 好处： 那些特殊的“烹饪技巧”（如 8-bit 量化、Muon 优化器）不需要把肉切碎再拼起来，直接整块处理。这让厨师们能直接上手干活，不用做多余的拼接工作。

第二招：智能的“排队算法” (Planning Algorithm)

• 比喻： 即使切好了，怎么把这些肉块放到传送带上（通信缓冲区）也有讲究。如果乱塞，传送带会卡住，或者为了填满传送带塞很多空气。
• 解释： veScale-FSDP 有一个**“超级调度员”**。它会在训练开始前，花几秒钟算一下：怎么把不同大小的肉块（参数）排列组合，才能让传送带塞得最满，且不需要塞太多“空气”（Padding）。
• 好处： 这是一个很难的数学题（NP 难问题），但他们用了一个聪明的“捷径”算法，算得又快又好。结果就是，几千个厨师之间的传菜速度极快，几乎没有等待时间。

第三招：零损耗的“传送带” (Distributed Buffer / DBuffer)

• 比喻： 以前的传送带，厨师拿菜要先把菜从盘子里倒出来，再倒进传送带，吃完再倒回盘子（数据拷贝）。
• 解释： veScale-FSDP 发明了一种**“悬浮传送带”**。它让数据直接在内存里“指路”，不需要搬运。
• 好处： 厨师拿菜是“零拷贝”的，瞬间完成。这不仅省了力气（计算资源），还省了盘子（显存空间）。

3. 实际效果：快了多少？省了多少？

经过测试，这个新系统带来了巨大的提升：

• 速度更快： 训练速度比以前的系统快了 5% 到 66%。这意味着以前要跑一个月的模型，现在可能只需要两周。
• 更省空间： 占用的显存（厨房空间）减少了 16% 到 30%。这意味着用同样的几千台电脑，可以训练更大的模型，或者在同样的模型下，不需要买那么多昂贵的显卡。
• 规模更大： 它可以轻松扩展到 几万台 GPU 一起工作，而不会像以前那样因为“切菜”太乱而崩溃。

4. 为什么这很重要？

• 对科学家： 以前想尝试新的模型结构或新的优化方法，往往因为系统不支持“不规则切法”而被迫放弃，或者需要花几个月去修改底层代码。现在，veScale-FSDP 像是一个万能插座，插上就能用，让科学家能专注于设计更好的模型，而不是纠结于怎么切分数据。
• 对大众： 这意味着未来我们能更快、更便宜地训练出更聪明、更强大的 AI，而且这些 AI 能更好地处理复杂的任务。

总结一句话：
veScale-FSDP 就像给 AI 训练厨房换了一套智能、灵活且零损耗的自动化流水线，让几千个厨师能完美配合，不再因为切菜和传菜而浪费时间，从而让超级 AI 的诞生变得更快、更省钱。

技术摘要

veScale-FSDP 技术总结

1. 研究背景与核心问题 (Problem)

随着大语言模型（LLM）向万亿参数规模发展，全量分片数据并行（FSDP/ZeRO） 成为训练的基础技术。然而，现有的 FSDP 系统（如 DeepSpeed ZeRO, PyTorch FSDP1/FSDP2, Megatron-FSDP）在面对现代结构感知（Structure-Aware） 训练方法时存在严重瓶颈：

1. 与新型优化器和量化方法的冲突：
   • 非元素级优化器：如 Shampoo 和 Muon，需要在完整的 2D 矩阵上进行计算，而现有 FSDP 通常按元素或行进行均匀分片，导致无法直接支持。
   • 块级量化（Block-wise Quantization）：如 DeepSeek-V3 使用的 8-bit Adam，要求张量按特定块（Block）切分。现有系统的固定分片边界往往与量化块边界不对齐，导致需要复杂的边界检查、填充（Padding）或侵入式修改模型代码。
2. 性能与内存效率低下：
   • 通信开销：现有实现（特别是 PyTorch FSDP2 的每参数分片设计）在 AllGather 和 ReduceScatter 过程中涉及大量的交错内存拷贝（Interleaved Copy），显著降低了吞吐量。
   • 内存碎片与浪费：为了对齐通信缓冲区或满足分片均匀性，系统往往需要大量填充（Padding），导致显存占用过高。在万卡集群规模下，内存限制成为主要瓶颈，导致作业频繁 OOM（显存溢出）或资源浪费。
   • 扩展性不足：现有系统在扩展到数万个 GPU 时，通信效率下降，难以维持线性扩展。

2. 方法论 (Methodology)

veScale-FSDP 通过重新设计 FSDP 系统，提出了三个核心组件来解决上述问题：

2.1 RaggedShard：灵活的分片格式

为了解决分片粒度僵化的问题，veScale-FSDP 引入了 RaggedShard 作为 DTensor 的一种新放置（Placement）格式。

• 任意粒度：支持自定义的块大小（Block Size），可以是行、列或任意 2D 块。
• 非均匀分布：允许不同设备持有不同数量的块（Ragged），完美适配块级量化和非均匀分片需求。
• 兼容性：RaggedShard 可以与现有的 TP（张量并行）、EP（专家并行）等策略无缝组合，无需侵入式修改模型代码。

2.2 结构感知规划算法 (Structure-Aware Planning)

为了在支持灵活分片的同时保证通信效率，veScale-FSDP 将张量分组和布局问题建模为一个 NP-hard 优化问题。

• 目标：最小化通信缓冲区大小（减少 Padding），同时满足三个约束：
  1. 非分片块（Non-Sharded Block）：量化块必须完整位于某个设备的分片内，避免跨设备同步。
  2. 连续内存（Contiguous Memory）：避免通信缓冲区内的张量碎片化。
  3. 负载均衡（Balanced Load）：确保各设备间的通信负载平衡。
• 求解：由于直接求解 ILP（整数线性规划）在大规模下不可行，作者设计了一种多项式时间的启发式动态规划（DP）算法。该算法利用 Transformer 模型参数的规律性（如线性权重占主导、层间块大小一致），通过排序和二分搜索快速找到近优解。

2.3 分布式缓冲区 (Distributed Buffer, DBuffer)

为了消除内存拷贝开销，veScale-FSDP 实现了 DBuffer 原语。

• 零拷贝（Zero-Copy）：RaggedShard 的张量直接映射到全局 DBuffer 的切片上，通信前后无需数据拷贝。
• 算子融合：支持组级（Group-level）操作，将多个张量的预处理（如加、乘、清零）融合，减少 CUDA 内核启动开销。
• 原地通信：结合 DBuffer 实现高效的 AllGather 和 ReduceScatter，减少显存碎片。

3. 关键贡献 (Key Contributions)

1. 原生支持结构感知训练：veScale-FSDP 是首个能原生支持非元素级优化器（如 Muon）和块级量化（如 8-bit Adam）的 FSDP 系统，无需修改模型或优化器代码。
2. RaggedShard 抽象：提出了一种通用的分片格式，打破了传统 FSDP 必须均匀分片的限制，实现了分片粒度与计算需求的完美对齐。
3. 高性能规划与通信：通过 NP-hard 问题的启发式求解和 DBuffer 设计，显著减少了填充开销和内存拷贝，实现了极高的通信效率。
4. 大规模扩展性：成功在数万个 GPU（最高 10K+）和2.4 万亿参数的模型上进行了验证，证明了其工业级可用性。

4. 实验结果 (Results)

在 1024 张 GPU 上的基准测试（包括 LLaMA-3-70B, GPT-OSS-120B, 内部 MoE 模型）显示：

• 吞吐量提升：相比现有 FSDP 系统（DeepSpeed, FSDP1, FSDP2, Megatron-FSDP），veScale-FSDP 实现了 5% ~ 66% 的吞吐量提升。
  • 在 MoE 模型上提升尤为显著（11-66%），主要得益于消除了填充开销和优化的通信重叠。
• 显存占用降低：峰值显存占用降低了 16% ~ 30%。
  • 这主要归功于 DBuffer 的确定性内存管理和减少的填充（Padding）。
• 扩展性：
  • 弱扩展（Weak Scaling）：在 1K 到 8K GPU 上，吞吐量随 GPU 数量增加呈现近线性增长。
  • 强扩展（Strong Scaling）：在固定全局 Batch Size 下，1K 到 10K GPU 依然保持高效。
  • 模型扩展：支持在 1K GPU 上训练 2.4T 参数模型，且 MFU（模型浮点运算利用率）随模型增大而微升。
• 组件消融实验：
  • 禁用规划算法导致吞吐量下降 34.6%（因量化块对齐失败）。
  • 禁用 DBuffer 导致吞吐量下降 7.2%（因拷贝开销）。
  • 禁用 RaggedShard 则无法运行 8-bit Adam（需侵入式修改）。

5. 意义与影响 (Significance)

• 工业级落地：veScale-FSDP 已在 ByteDance 的生产环境中经过实战检验，支持了超过 10K GPU 的超大规模训练任务，证明了其在真实场景下的稳定性和高效性。
• 解耦模型与系统：通过 RaggedShard 和 DTensor 的集成，veScale-FSDP 将模型定义与系统优化解耦。研究人员可以专注于模型架构创新（如稀疏 MoE、新型优化器），而无需关心底层分片细节。
• 开源与生态贡献：RaggedShard 代码已开源，且该设计思路已被纳入 PyTorch 官方路线图，有望成为未来大模型训练的标准基础设施。
• 解决“最后一公里”问题：解决了从理论上的高效算法（如 Muon、块量化）到实际大规模训练系统落地之间的鸿沟，为下一代万亿参数模型的训练提供了关键支撑。

总结：veScale-FSDP 通过引入灵活的分片格式（RaggedShard）、智能的布局规划算法和高效的通信原语（DBuffer），成功解决了现有 FSDP 系统在灵活性、性能和扩展性上的瓶颈，是目前大规模分布式训练领域的一项突破性进展。