[arXiv 2026] veScale-FSDP：打破分片枷锁，万卡规模下的极致灵活性与性能

TL;DR

在追求大语言模型（LLM）效能的今天，经典的 FSDP（全分片数据并行）正面临现代算法的挑战。字节跳动提出的 veScale-FSDP 通过创新的 RaggedShard 格式，完美解决了非逐元素优化器（如 Muon）和块状量化（Block-wise Quantization）与分布式分片的冲突。它不仅能显著降低 30% 的内存占用，更在万卡规模下展现出近乎完美的线性扩展性。

1. 现状之殇：僵化的分片与灵活的算法

传统的 FSDP 系统（如 PyTorch FSDP1/2, ZeRO-3）为了简化实现，通常强制要求张量按元素或按行均匀分片。这种“一刀切”的做法在面对以下现代训练技术时显得力不从心：

• 矩阵优化器（Muon, Shampoo）：这些优化器需要操作完整的 2D 张量块，而不仅仅是离散的元素。
• 块状量化（Block-wise Quantization）：DeepSeek-V3 等模型采用块量化显著降低显存，但如果分片边界切断了量化块（Quantization Block），则会导致复杂的同步逻辑。
• 性能开销：FSDP2 虽灵活，但由于非对齐的张量拷贝，导致了高达 14% 的额外开销。

2. 核心创新：RaggedShard —— 让分片适应算法

veScale-FSDP 的核心是重新定义了分片抽象。

RaggedShard 格式

不再强制“均分”，RaggedShard 允许每个设备持有不同数量、不同粒度的“块”（Blocks）。

• 粒度感应：开发者可以指定分片的最小原子单位（如 32x32 的量化块）。
• 灵活分布：支持非均匀分布，这对于负载均衡和特定的矩阵运算至关重要。

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

灵活分片往往带来 Padding（填充）过多的隐患。作者将缓冲区布局建模为一个 NP-Hard 优化问题，并利用启发式 DP 算法在毫秒级内找到接近最优的排列组合。

• 目标：最小化全局缓冲区大小，同时满足“块不被切分”和“跨设备负载均衡”。

3. 性能基石：Distributed Buffer (DBuffer)

为了榨干硬件性能，veScale-FSDP 引入了 DBuffer 原语，实现了：

• 零拷贝 (Zero-copy)：通过持久化的地址映射，避免了 AllGather/ReduceScatter 后的冗余数据搬运。
• 内核融合 (Kernel Fusion)：将多个张量的操作（如 Scale, Add）融合为一个扫一遍显存的操作，减少了 CUDA Kernel 发射带来的空隙。

4. 实验战绩：全方位的降本增效

在针对 LLaMA-3 和高性能 MoE 模型的测试中，veScale-FSDP 展现了统治级的表现：

• 吞吐量：在 MoE 模型上比 DeepSpeed 和 PyTorch 原生 FSDP 快 11%~66%。
• 内存消耗：峰值预留内存降低了 16%~30%，这意味着可以在同等硬件上跑更大的 Batch Size。
• 万卡扩展：在 8000B 参数的超大规模场景下，保持了近乎线性的弱扩展（Weak Scaling）和强扩展（Strong Scaling）曲线。

5. 深度洞察：为何 veScale-FSDP 是未来的必经之路？

以往的分布式框架往往试图对算法开发者隐藏物理排布，这导致了“木桶效应”。veScale-FSDP 的成功案例（如原生支持 8-bit Adam 和 Muon）证明了：

1. 抽象的力量：将 RaggedShard 集成进 PyTorch 的 DTensor 体系，保留了 SPMD 编程的简洁性，同时开放了对物理布局的精细控制。
2. 解耦的价值：模型结构（Model Definition）与系统优化（System Optimization）的解耦，让算法研究员可以自由尝试非线性优化器，而不用担心在大规模集群上跑不动。

总结

veScale-FSDP 不仅仅是一个更快的 ZeRO 实现，它代表了一种结构感知的系统设计新思维。通过 RaggedShard 和高效规划，它为下一代万亿参数模型和更复杂的矩阵优化算法铺平了道路。

关键词：FSDP, RaggedShard, 字节跳动, LLM Scaling, Distributed Optimizer, Muon