在人工智能模型训练领域，超大规模语言模型的训练如同搭建精密复杂的科技大厦，需要成千上万个GPU协同工作。其中，“全分片数据并行”技术（FSDP）是关键支撑，但传统FSDP系统存在明显局限，如同固定尺寸的积木难以应对复杂搭建需求。字节跳动研究团队针对这一问题，开发出veScale - FSDP系统，为AI训练带来新突破。
传统FSDP系统问题诸多。当研究人员尝试运用矩阵优化器Muon或8位量化训练等先进技术时，传统系统因只能按固定方式切分数据，无法满足新技术对数据块状结构的特定处理要求，就像刀具只能横切或竖切，难以完成复杂切割。而且，在性能方面，当训练规模扩展到上万个GPU时，系统通信效率和内存使用问题凸显，内存管理缺乏统一规划导致大量内存碎片，通信中的数据拷贝操作也造成显著性能损失。
veScale - FSDP系统的核心创新在于引入“RaggedShard”新型数据分片格式。与传统固定模具式的数据分片不同，RaggedShard如同可随意调整形状的橡皮泥模具，能根据具体需求灵活调整分片大小和形状。以拼图游戏类比，传统系统要求拼图块都是标准长方形，而RaggedShard允许拼图块有各种形状。处理Muon优化器时，它能确保每个设备上的数据块包含完整矩阵，避免矩阵被切断；对于8位量化训练，它能保证每个设备上的数据按量化块边界切分，使每个设备可独立完成量化操作，还能与张量并行、专家并行等现有并行化策略完美融合。
仅有灵活的数据分片还不够，高效通信是关键。veScale - FSDP开发了智能规划算法，解决不规则形状数据块的高效打包和传输问题，类似快递打包中如何将形状各异的物品装入标准快递箱，既要保证物品完好、空间利用率高，又要确保每个快递箱重量平衡。研究团队将此问题数学化为优化问题，在满足“非分片块约束”“连续内存约束”“严格负载平衡约束”三个关键条件下，最小化通信缓冲区大小。基于Transformer模型的结构特性，他们设计了多项式时间的动态规划算法，通过案例分析数据块与分片边界的三种可能关系，找到最佳数据布局方案。
为进一步提升性能，veScale - FSDP引入“分布式缓冲区”（DBuffer）新原语。传统FSDP系统通信时需大量数据拷贝操作，如同搬家时物品需多次搬运，效率低下。DBuffer实现零拷贝通信，像在房间和卡车间建立直通管道，物品可直接传输。它具有全局缓冲区语义、支持组级别操作、提供持久地址映射、支持原地通信和计算等关键特性，在实际测试中，使通信开销降低约7%，规划算法优化带来34%的性能提升，且这些优化自动进行，开发者无需修改模型代码。
veScale - FSDP在各种规模实验中表现卓越。在端到端训练性能对比中，该系统在MoE稀疏模型上比所有基线系统快11%到66%，在LLaMA - 3 - 70B密集模型上有5%的性能提升，这些提升源于优化的通信重叠、基于DBuffer的零拷贝集合通信以及避免填充开销的灵活分片粒度。内存使用方面，它将峰值预留内存降低16%到30%，得益于确定性的批量内存管理策略。扩展性测试显示，它能高效扩展到一万个GPU规模，在弱扩展性测试中保持接近线性扩展性，在强扩展性测试中也有良好表现。
为验证系统灵活性，研究团队测试了veScale - FSDP对8位Adam优化器和分布式Muon优化器的支持能力。8位Adam优化器需支持32×32的块级量化粒度，veScale - FSDP中开发者只需设置参数量化粒度，系统就能自动确保数据分片与量化块边界完美对齐，每个设备可独立量化操作。Muon优化器需在原始2D参数矩阵形状上进行Newton - Schulz迭代，veScale - FSDP通过RaggedShard的不均匀分片能力，让用户用标准SPMD方式编写分布式Muon算法。实验结果显示，该系统能高效支持这两种前沿优化技术，且性能表现良好。
veScale - FSDP的工程实现巧妙扩展了PyTorch的DTensor抽象，RaggedShard作为DTensor的新放置策略，无缝集成现有并行化基础设施。系统包含7600行Python代码，完全兼容标准PyTorch分布式运行时和广泛版本范围，作为FSDP2的即插即用后端，使用相同API，现有训练代码几乎无需修改就能享受性能提升。规划算法实际运行开销极小，时间复杂度为O(|T|?m log(E) log(|T|m))。
研究团队通过消融实验量化各组件贡献，结果显示DBuffer和规划算法贡献大部分性能提升，RaggedShard是支持块级8位Adam的关键抽象。veScale - FSDP已在字节跳动生产环境得到实战验证，成功支持超一万个GPU、2.4万亿参数规模的大规模训练任务。在实际部署中，研究团队总结出小规模实验可准确预测大规模性能、在成熟抽象基础上设计系统抽象、将模型定义与系统优化解耦等重要经验。
veScale - FSDP已开源，为整个AI社区提供了强大工具。无论是学术研究机构还是产业界，都能利用它训练更大、更复杂的AI模型，满足AI模型规模不断增长的需求。
