字节跳动发布分布式训练技术：让模型训练像搭积木一样灵活高效

这项由字节跳动种子实验室主导的研究发表于2026年2月的arXiv预印本论文库，论文编号为arXiv:2602.22437v1。有兴趣深入了解的读者可以通过该编号查询完整论文内容。

在人工智能快速发展的今天，训练超大规模语言模型就像建造摩天大楼一样复杂。每当我们想要训练一个拥有数千亿参数的AI模型时，就需要成千上万个GPU协同工作，这个过程中最关键的技术之一就是"全分片数据并行"技术，英文名叫FSDP。不过，现有的FSDP系统就像一套死板的积木玩具，每个积木的形状都是固定的，无法根据具体需求进行调整。字节跳动的研究团队发现了这个问题，并开发出了一种全新的解决方案——veScale-FSDP系统。

现有FSDP系统的问题就像用固定尺寸的积木搭建复杂模型一样困难。当研究人员想要使用一些先进的训练技术时，比如矩阵优化器Muon或者8位量化训练时，传统系统就显得力不从心了。这些新技术需要将数据按照特定的块状结构进行处理，但传统FSDP系统只能按照固定的方式切分数据，就像只能横切或竖切的刀具，无法满足复杂切割需求。

更糟糕的是，传统系统在性能上也存在明显不足。当训练规模扩展到上万个GPU时，系统的通信效率和内存使用都会出现问题，就像一个城市的交通系统在高峰期出现拥堵一样。研究人员发现，现有系统在内存管理方面缺乏统一规划，导致大量内存碎片，同时通信过程中的数据拷贝操作也造成了显著的性能损失。

**一、革命性的灵活切分技术**

veScale-FSDP的核心创新在于引入了一种叫做"RaggedShard"的新型数据分片格式。如果把传统的数据分片比作用固定模具制作饼干，那么RaggedShard就像是可以随意调整形状的橡皮泥模具，能够根据具体需求灵活调整分片的大小和形状。

这种灵活性的实现原理可以用拼图游戏来理解。传统FSDP系统就像强制要求所有拼图块都必须是标准的长方形，而RaggedShard允许拼图块有各种不同的形状。当我们需要处理Muon优化器时，它需要完整的2D矩阵结构，RaggedShard就可以确保每个设备上的数据块恰好包含完整的矩阵，避免了矩阵被切断的问题。

对于8位量化训练来说，这种灵活性更是至关重要。量化训练需要将参数按照32×32的块进行处理，就像制作瓷砖时需要按照固定尺寸切割一样。传统系统由于无法保证切分边界与量化块边界对齐，往往需要额外的通信来收集完整的量化块，这就像为了得到完整的瓷砖而需要从多个工厂收集碎片一样低效。RaggedShard则能确保每个设备上的数据恰好按照量化块的边界进行切分，每个设备都能独立完成量化操作。

更重要的是，RaggedShard还能与现有的并行化策略完美融合。现代大模型训练通常需要同时使用多种并行化技术，包括张量并行、专家并行等。RaggedShard的设计使得它能够与这些技术无缝配合，就像一个万能适配器，可以连接各种不同的设备一样。

**二、智能规划算法解决通信难题**

仅有灵活的数据分片还不够，如何高效地进行通信才是关键。veScale-FSDP开发了一套智能规划算法，专门解决如何将不规则形状的数据块高效地打包和传输。

这个问题的复杂性可以用快递打包来类比。当我们需要邮寄一堆形状各异的物品时，如何将它们装入标准尺寸的快递箱中，既要确保物品不被损坏，又要最大化空间利用率，同时还要保证每个快递箱的重量平衡？这正是veScale-FSDP面临的挑战。

研究团队将这个问题数学化为一个优化问题。他们需要在满足三个关键约束条件的前提下，最小化通信缓冲区的大小。第一个约束是"非分片块约束"，确保数据块不会在通信过程中被意外切断；第二个是"连续内存约束"，保证数据在内存中的连续性以避免碎片化的拷贝操作；第三个是"严格负载平衡约束"，确保每个设备的通信负载相等。

理论上，这是一个NP困难问题，意味着寻找最优解需要指数级的时间。但研究团队发现了一个巧妙的解决方案。他们观察到Transformer模型具有高度规律的结构特性：线性层权重占据了参数总数的绝大部分，而且分片块的大小在不同层之间往往是一致的。基于这个观察，他们设计了一个多项式时间的动态规划算法，能够在实际应用中找到接近最优的解决方案。

算法的核心思想是案例分析。对于每个数据块，算法分析它与分片边界的三种可能关系：完全位于单个分片内、跨越两个相邻分片但不包含完整分片、完全包含至少一个分片。通过这种分类，算法能够有效地搜索解空间，找到最佳的数据布局方案。

**三、分布式缓冲区实现零拷贝通信**

为了进一步提升性能，veScale-FSDP引入了一个叫做"分布式缓冲区"（DBuffer）的新原语。这个组件可以理解为一个智能的内存管理器，专门为分布式通信优化。

传统FSDP系统在通信时需要进行大量的数据拷贝操作，就像搬家时需要把物品从房间搬到走廊，再搬到卡车上一样低效。DBuffer则实现了零拷贝通信，就像在房间和卡车之间建立了一条直通管道，物品可以直接传输。

DBuffer的设计有四个关键特性。首先，它提供了全局缓冲区语义，抽象掉了N维设备拓扑的复杂性，让开发者可以像操作单机内存一样操作分布式内存。其次，它支持组级别的操作，能够将多个张量的相同操作（如加法、缩放、清零）融合成一个批量操作，减少了内核启动的开销。第三，它通过RaggedShard的规划算法提供了持久的地址映射，每个张量的数据指针都有固定的地址，避免了通信前后的数据拷贝。最后，它支持原地通信和计算，进一步减少了内存占用。

这种设计的效果非常显著。在实际测试中，DBuffer使得通信开销降低了约7%，而规划算法的优化则带来了34%的性能提升。更重要的是，这些优化是自动进行的，开发者无需修改任何模型代码就能享受到性能提升。

**四、突破性的性能表现**

veScale-FSDP在各种规模的实验中都展现出了卓越的性能。在端到端训练性能对比中，该系统在MoE稀疏模型上比所有基线系统快11%到66%，在LLaMA-3-70B密集模型上也有5%的性能提升。这些提升主要来自于三个方面的优化：优化的通信重叠、基于DBuffer的零拷贝集合通信、以及避免填充开销的灵活分片粒度。

内存使用方面的改善同样令人印象深刻。veScale-FSDP在各个测试中都将峰值预留内存降低了16%到30%。这种内存节省源于确定性的批量内存管理策略：系统显式管理流依赖关系以实现可预测的内存释放，并且通过批量分配来减少碎片化。相比之下，传统系统由于隐式的record_stream机制导致不确定性的内存释放，往往会阻止缓存分配器重用缓冲区，使峰值预留内存增加20%。

扩展性测试显示，veScale-FSDP能够高效地扩展到一万个GPU的规模。在弱扩展性测试中，系统保持了接近线性的扩展性，这是因为FSDP的通信成本和每GPU的计算成本都与GPU数量无关，只依赖于模型和输入大小。在强扩展性测试中，系统在128M token的全局批处理规模下能够线性扩展到一万个GPU，即使在16M token的较小批处理下，从1K到8K GPU也能实现3.4倍的吞吐量提升。

**五、支持前沿优化技术的灵活性验证**

为了验证系统的灵活性，研究团队特别测试了veScale-FSDP对两种前沿优化技术的支持能力：8位Adam优化器和分布式Muon优化器。

8位Adam优化器通过将梯度统计信息量化为INT8格式来大幅减少优化器状态的内存占用。要实现这种量化，系统需要支持32×32的块级量化粒度。在veScale-FSDP中，开发者只需要设置参数的量化粒度，系统就会自动确保每个设备上的数据分片与量化块边界完美对齐。这样，每个设备都可以独立进行量化操作，无需任何额外的通信。相比之下，现有的FSDP系统由于无法原生地跟踪块边界，要么需要侵入性的系统修改，要么需要手动实现复杂的集合通信来交换量化元数据。

Muon优化器的情况更加复杂。这种矩阵符号预调节器需要在原始2D参数矩阵形状上进行Newton-Schulz迭代。veScale-FSDP通过RaggedShard的不均匀分片能力，让用户可以用标准的SPMD方式编写分布式Muon算法。具体来说，系统通过负载平衡选择一个根设备，然后使用DTensor的redistribute操作将完整的2D参数收集到根设备上，在那里执行Newton-Schulz更新，最后再将更新后的参数分发回原始设备。整个过程对其他设备来说Newton-Schulz更新变成了无操作，实现了清晰的SPMD语义。

实验结果显示，无论是8位Adam还是分布式Muon，veScale-FSDP都能在几乎不修改代码的情况下实现高效支持。8位Adam的损失曲线与传统DDP实现高度一致，只是偶尔出现一些量化精度导致的小波动。分布式Muon的收敛速度明显快于AdamW，在训练约800亿token后稳定在比AdamW低0.01的损失水平，这与之前的研究结果一致。

**六、工程实现的精巧设计**

veScale-FSDP的工程实现体现了"站在巨人肩膀上"的设计哲学。系统没有重新发明轮子，而是巧妙地扩展了PyTorch的DTensor抽象。RaggedShard被实现为DTensor的一个新的放置策略，这使得它能够无缝集成现有的并行化基础设施，包括张量并行、专家并行，以及成熟的训练工具如分布式检查点。

系统的实现包含了7600行Python代码，完全兼容标准的PyTorch分布式运行时和广泛的PyTorch版本范围。更重要的是，它作为FSDP2的即插即用后端，使用相同的PyTorch原生fully_shard API，这意味着现有的训练代码几乎无需修改就能享受到性能提升。

规划算法的实际运行开销极小，在所有实验中算法运行时间都少于0.3秒，这是一次性的初始化成本，在分布式训练的整个生命周期中可以忽略不计。算法的时间复杂度为O(|T|?m log(E) log(|T|m))，其中T是张量数量，m是设备数量，E是总元素数量。

**七、深度性能分析与组件贡献**

研究团队进行了详细的消融实验来量化各个组件的贡献。结果显示，DBuffer和规划算法贡献了大部分的性能提升：禁用DBuffer会使吞吐量降低7.2%，禁用规划算法则会导致34.6%的性能下降。

DBuffer的性能提升主要来自于消除了集合通信周围的拷贝输入/拷贝输出开销。当通信缓冲区需要拷贝时，这些操作会成为性能瓶颈。规划算法的影响更为显著，因为当禁用规划时，量化块无法保证完全包含在设备的本地分片内，系统必须回退到DTensor重分布来在每块量化之前组装所需的优化器状态，这会产生大量额外的通信开销。

至于RaggedShard本身，它不仅仅是一个优化——它是使得块级8位Adam在不进行侵入性模型/优化器修改或手写集合通信的情况下可用的关键抽象。禁用RaggedShard会使系统变得几乎不可用，用户必须要么仔细修改每个模型和优化器张量使得32×32块边界与分片边界对齐，要么手动实现复杂的集合通信来恢复块级语义。

**八、生产环境的实战验证**

veScale-FSDP不仅在研究环境中表现出色，更重要的是它已经在字节跳动的生产环境中得到了实战验证。系统成功支持了超过一万个GPU的大规模训练任务，训练的模型参数规模达到了2.4万亿。这种规模的训练对系统的稳定性、效率和可维护性都提出了极高的要求。

在实际部署中，研究团队总结了三个重要的经验教训。首先，小规模实验能够准确预测大规模性能。FSDP工作负载的性能可以通过各层的计算时间和FSDP通信时间来准确估算，因为计算完全在GPU内部进行，而FSDP通信时间在GPU数量增加时基本保持不变。这使得团队可以在64个GPU上进行性能分析，然后外推到数千个GPU的规模。

其次，在成熟抽象基础上设计系统抽象的重要性。DTensor提供了一个强大的抽象层，已经支持广泛的并行化技术。通过将RaggedShard设计为DTensor的扩展，veScale-FSDP能够无缝集成现有的并行化策略，最小化工程投入的同时为更广泛的社区做出贡献。

第三，将模型定义与系统优化解耦的价值。模型架构的快速演进要求频繁更新模型定义，但像Megatron-LM这样的框架将系统级并行化优化与模型代码紧密耦合，使得研究人员难以修改或扩展架构。veScale-FSDP通过将模型定义从系统框架中解耦出来，让研究人员能够专注于模型设计，同时保持在上万个GPU上的线性扩展性。

说到底，veScale-FSDP的意义远超一个技术系统的范畴。它代表了一种新的设计哲学：在追求极致性能的同时，不牺牲灵活性和易用性。就像一个优秀的工具应该既强大又好用一样，veScale-FSDP让研究人员能够专注于AI模型本身的创新，而不用被复杂的分布式训练细节所束缚。

这项技术的开源释放，为整个AI社区提供了一个强大的工具。无论是学术研究机构还是产业界，都可以利用这项技术来训练更大、更复杂的AI模型。而随着AI模型规模的不断增长，类似veScale-FSDP这样的技术创新将变得越来越重要，它们正在默默地为AI技术的发展提供坚实的基础设施支撑。

Q&A

Q1：RaggedShard是什么技术？

A：RaggedShard是veScale-FSDP系统的核心创新，它是一种新型的数据分片格式。传统FSDP系统只能按照固定方式切分数据，而RaggedShard可以根据具体需求灵活调整分片的大小和形状，就像可调节的模具一样。这种灵活性让它能够完美支持Muon优化器和8位量化训练等先进技术。

Q2：veScale-FSDP比传统系统性能提升多少？

A：veScale-FSDP在不同测试中表现出显著的性能优势。在MoE稀疏模型上，吞吐量比其他系统高11%到66%，在密集模型上也有5%的提升。同时，内存使用降低了16%到30%。这些提升主要来自优化的通信重叠、零拷贝集合通信和智能的内存管理策略。

Q3：普通研究团队能使用veScale-FSDP吗？

A：完全可以。veScale-FSDP已经开源，包含7600行Python代码，完全兼容标准PyTorch环境。它作为即插即用的模块，使用与PyTorch FSDP2相同的API，现有训练代码几乎无需修改就能享受性能提升。无论是学术研究还是产业应用，都可以直接使用这项技术。