1 背景
传统GPU资源池采用固定的分配模式，存在明显短板：空闲时段GPU算力闲置，无法转化为实际效率；而任务集中时，资源池又会陷入紧张，高优先级任务因算力不足被迫排队，低优先级任务却占用固定资源，导致整体调度失衡。为破解这一资源利用困境，本文提出弹性训练，核心目标是实现GPU资源池的最大化高效利用。其核心逻辑为”按需伸缩”：当资源池存在空闲GPU时，自动为正在执行的任务扩容分配额外算力，缩短任务完成周期；当资源供不应求时，智能识别并缩容低优先级任务的资源占用，将释放的GPU算力优先供给高优先级任务，实现资源在多任务间的灵活共享与最优配置。
2 问题分析
弹性训练是指根据集群资源情况动态调整训练的数据并行数的训练方式。因此存在如下三个问题：
计算损失
动态扩缩容或软硬件故障，无疑会引发训练任务的重启，因此会导致算力浪费。
准确性
弹性训练无疑会引发工作节点数量发生变化，进而导致batch发的变化，因此可能引发计算不一致。
与调度框架适配
新的架构也需要与调度框架进行适配
接下来依次介绍这三个问题。
3 降低计算损失的影响
分布式训练任务在训练之前需要确定全局进程数WORLD_SIZE，进程编号RANK，本地进程编号LOCAL_RANK等信息。并通过WORLD_SIZE和RANK来给进程分配样本。当出现弹性扩缩容时，WORLD_SIZE会发生变化，需要重新启动任务以配置新的分布式参数。重启的任务需要从上一次checkpoint开始训练，因此出现了计算损失。
软硬件故障也会引发分布式训练任务重启，从而造成计算损失。弹性扩缩容则进一步增加了任务重启的概率，因此意味着弹性训练要求极致的故障恢复。我们对训练架构的改造划分为两个阶段:
- 阶段1: 极致的Checkpoint优化避免计算损失
- 阶段2: 训练框架层实现动态扩缩容
3.1 阶段1: 极致的Checkpoint优化避免计算损失
3.1.1 Checkpiont优化分析
既然重启的训练任务要从上一次checkpoint中恢复，那么最简单的方式就是优化checkpoint的时间，从而提升checkpoint的频率，降低计算损失。
PyTorch默认的checkpoint是同步相对耗时的操作，有很大的优化空间。综合分析Microsoft DeepSpeed的FastPersist、TorchTitan、dlrover的flash checkpoint[3]的方案。主要有如下几种方法:
分布式checkpoint
分布式训练场景下，模型参数存储在多台机器上。因此可以通过多台机器并行checkpoint的方式提高checkpoint的效率。原生pytorch已经支持。
异步checkpoint
先同步地将checkpoint拷贝到CPU内存中。然后再异步地将CPU内存中的checkpoint存储到持久化存储中。这里checkpoint从GPU内存拷贝到CPU内存的过程仍然是同步的。
完全异步checkpoint(pipelined checkpoint write / Zero-Overhead Checkpointing)
FastPersist(pipelined checkpoint write)和TorchTitan(Zero-Overhead Checkpointing)上都提出了完全异步方案。前向计算和反向传播过程中参数的值是不变的，只有在optimizer.step的时候才会更新参数。因此可以新建CUDA流异步地将checkpoint拷贝到pinned memory中。下次optimizer.step更新参数的时候只需保证CUDA流工作完成即可。这样实现了完全异步。
3.1.2 架构
我们选择了基于dlrover flash checkpoint的方案。架构如图所示，每M轮会同步地将checkpoint保存到共享内存，每N轮会同步地将内存持久化到分布式存储中(N>>M)。在指定--save_at_breakpoint的情况下，如果Worker出现故障导致任务重启，会捕获异常并会将共享内存中的checkpoint同步到永久存储中。因此这意味着最多损失M轮的计算。我们只需要将M设置的足够小，就可以实现尽可能小的
通过dlrover架构实现了同步直接写内存，异步写永久存储的方案，降低了checkpoint时间44%-97%。然后我们通过pinned memory和CUDA stream异步写的方式，相比于原生dlrover flash checkpoint, checkpiont时间降低了19%-72%。这样就可以大大提高checkpoint的频率了。
3.1.3 如何进行扩缩容
本方案并没有脱离基础的load/save checkpoint的架构，因此扩缩容方案相对简单，只需要每次启动的时候以不同的WORLD_SIZE启动即可。但是由于DeepSpeed的参数分布在多个进程中，默认情况下每个进程都会存储自己的模型和优化器分片，因此当WORLD_SIZE发生变化的时候，需要使用universal checkpoint来解决该问题。
DeepSpeed架构的扩缩容会引出两个问题:
- 默认情况下，DeepSpeed在所有参与的机器上进行模型分片。对于zero 3, 扩容意味着前向计算和反向传播都可能要依赖于机器间的网络，必然会影响整体性能。
-
WORLD_SIZE发生变化引入的universal checkpoint也会带来性能开线。
这两个问题可以使用misc或zero++的zero_hpz_partition_size进行优化。但考虑到阶段2会解决这些问题，这里就不在详细说明了。
3.2 阶段2: 训练框架层实现动态扩缩容
3.2.1 阶段1的主要问题
阶段1描述的”极致的Checkpoint优化避免计算损失”方法，有如下问题:
异步checkpoint占用了宝贵的GPU带宽。因此可能影响前向计算、反向传播以及offload等过程的性能。
即便更频繁的checkpoint会降低计算损失，任务的启动时间仍然难以极限优化。
以Qwen3 8B的SFT任务为例，num_procs设置为96的情况，也有超过5分钟的启动时间。对于deepspeed，当发生world size变化后，要有长达10分钟以上的universal checkpoint转换时间。
主要问题在于torchrun的rendezvous机制要求在训练之前必须配置好RANK和WORLD_SIZE等参数来进行分布式训练。那么是否有动态加入的机制来避免该问题呢?
3.2.2 架构
正如torchft项目对自己的描述”Easy Per Step Fault Tolerance for PyTorch”那样。torchft实现了一步一容错的方案。
该架构通过Lighthouse服务实现协调功能。训练任务被划分为多个Replica Group，每个Replica Group都维护完整的模型参数。每轮训练开始的时候，每个Replica Group的Manager通过qurom通知Lighthouse自己参与训练，训练完成后通过shouldCommit来决定是否更新参数。每个Replica Group对应一组torchrun任务，在Replica Group内部可能进行allreduce梯度聚集(对于zero, fsdp等模型分片需要该步骤，对于其他不建议使用)，然后在Replica Group间进行allreduce梯度聚集。
3.2.3 扩缩容与故障恢复
当需要对训练任务进行扩容的时候，只需要再启动新的Replica Group，新的Replica Group通过quorum来通知Lighthouse自己即将加入到训练。当Lighthouse检查到当前参与者发生变化的时候，会产生新的quorum id，并通知新的Replica Group从其他Replica Group获取最新的Checkpoint，然后加入到训练过程那种
当训练任务进行缩容的时候，直接关闭部分Replica Group，保留的Replica Group会等待短暂的超时时间后，通过Lighthouse协调新的训练参与者组重新开始训练。
事实上故障的场景与缩容类似。由于这个过程中，Replica Group的加入和退出都是动态进行的，因此扩缩容甚至包括故障都不会重启任务。只要保证同一时间至少有一个Replica Group，任务就不会中断，因此阶段1中强调的极致的checkpoint时间意义似乎就没那么大了。
与Lighthouse的通信都是异步的，对性能影响很小。扩容的时候一般不会带来计算损失(或者说最多一步的损失，详见准确性小节)，缩容或故障的时候一般只有一步的损失。
对于阶段1描述的DeepSpeed的两个问题，我们可以通过把模型分片维护在同一Replica Group下来避免。
3.3 总结
基于torchft的动态扩缩容的方案可以实现接近零损失的弹性训练。另外，阶段1使用的dlrover flash checkpoint方案可以作为训练加速的有效补充。
4 准确性
弹性训练会根据集群资源的情况动态调整运行worker数量，因此可能会导致global_batch的变化。根据附录A: batch变化对训练的影响分析，只需用于梯度聚集的global_batch和样本顺序不变，在不考虑精度损失的情况下，就可以保证计算的一致性。
4.1 batch的保持
用于梯度聚集的batch有如下公式:
1 | global_batch = local_batch * acc_steps * world_size |
其中，global_batch是用于梯度聚集的batch。local_batch表示单个Worker单次训练的batch。acc_steps表示每acc_steps步训练进行一次梯度聚集。world_size表示工作节点数目。
需要调整local_batch，acc_steps，world_size的值，来保证global_batch的不变。考虑到OOM和Batch Normalization的影响，local_batch设置为固定值。而world_size是根据弹性扩缩容实际调整的值。因此需要通过调整acc_steps来保证global_batch的不变。
有的时候难以找到整数的acc_steps来使global_batch不变。那么我们可以通过控制world_size来避免acc_steps不为整数。因此开发expectec_replicas的功能，仅能让指定数量的Replica Group参与训练。
4.2 保证样本顺序
DistributedSampler会根据epoch和自定义的种子来保证Sampler可恢复。但是无论pytorch还是transformers都不支持按照已经训练样本的个数进行恢复训练。因此开发了SkipDistributedSampler来保证Sample的可恢复性。
DistributedSampler默认实现会根据rank和num_replicas来隔项采样分组。这可能导致当replia group数目不同的时候，样本的顺序不一致。因此开发了DistributedBatchSampler配合SkipDistributedSampler来保证顺序不变。
上面两张图介绍了样本分配的流程，可以看出无论有几个运行节点，每个local batch的内容是不会发生变化的，而且可以保证global batch内的样本一致。这有利用降低损失精度带来的影响，而且非常便于调试。
4.3 扩缩容的断点恢复
阶段1中由于使用了checkpoint的机制。因此只要保证3.1和3.2描述的问题，就可以保证训练一致性。而对于阶段2，torchft上不支持训练的一致性，因为在发生扩缩容的时候可能会丢失样本，问题在于关于样本的分配仍然是静态的。因此开发了基于最新quorum动态调整dataloader的功能，来避免样本的丢失。
这么做的唯一问题是当quorum发生变化的时候，会丢弃当前已经训练的梯度，并对当前样本重新训练。因此会损失一步训练。但考虑到quorum发生变化是低概率事件，损失一部训练并没有太大的影响。
5 调度框架
本文使用了volcano框架调度，需要做如下几件事情:
训练框架与调度框架适配。
通过svc,ssh插件配置的环境变量即可实现在volcano运行。
通过抢占和任务资源配置和优先级来实现弹性扩缩容。
根据配置的
minAvailable和replicas控制训练任务的运行副本数。当存在空闲资源的时候，处于pending状态的pod会被调度，以实现扩容。当资源紧张的时候，通过抢占机制杀死低优先级的pod，以实现缩容。根据3.1的描述，有时候需要调整
world_size来保证acc_steps为整数。尽管在训练框架层面可以让多余的工作节点处于等待状态，但这难免浪费很多的资源。因此希望调度框架能够按照指定工作节点数据进行调度，以避免资源浪费。譬如，
globa_batch=4,local_batch=1，为了保证得到整数的acc_steps，我们希望工作节点能够按照[1,2,4]这样调度。如果调度3个节点是没有意义的，为了保证globa_batch的一致，第3个节点不会参与训练，这就做成了资源浪费。传统的
gang调度很难慢匹配这样的需求。因此开发了基于sub job tree的调度方式，实现按指定的工作节点进行调度，以避免资源浪费。如下图所示，通过多级sub job调度，将后两个工作节点的调度绑定在一起，就可以出现工作节点数据为3的中间状态，从而避免资源浪费。
6 效果
6.1 计算一致性
计算一致性是弹性训练的重中之重。在任意工作节点数目下训练都能得到相同的结果，是弹性训练的基本要求。
阶段1的实验:
基于DeepSpeed zero 3在Qwen3-8B模型SFT任务做测试，得到如下loss曲线。其中Base是始终在单机训练的，而Test是使用Flash Checkpint，每100步在1到2台机器上进行切换得到的loss曲线。可以看出弹性扩缩容对训练几乎没有任何影响。
阶段2的实验:
基于DeepSpeed zero3在Qwen3-8B模型SFT任务做测试，得到如下loss曲线。其中Base是始终在单机训练的，而Test则是基于torchft随机对训练任务在到2台机器执行进行切换。可以看出弹性扩缩容对训练几乎没有任何影响。
注: 该实验没有开启grad norm，所以loss曲线不平滑。但这个更能说明计算的一致性效果很高。
6.2 性能
阶段1对任务性能不会产生过大的影响，只是会丢失不必要的训练，因此只需要分析阶段2即可。
先让我们尝试分析一下DeepSpeed zero3 性能测试可能的影响。torchft中manager与ligthouse的通信都是异步的，因此可以忽略。然而通过torchft，我们可以将zero3的前向计算和反向传播过程中的参数收集限定在replica group内，即减少了机器之间的通信量。因此整体上应该会有效果提升。
base为在2机8卡上持续训练。torchft为在2机8卡上通过torchft进行训练。横坐标为步数，纵坐标为单步训练时间。可以看出确实有预期一样有着性能的提升。
事实上对于性能的影响也是分具体类型的任务的。而本例这种比较也略显不公平。如果不考虑zero3的问题或者通过mics和zero++,性能上应该是接近持平的。
7 总结
本文提出的弹性训练解决了如下问题：
可以在”任意”工作节点数目下进行训练，并保持计算的一致性。
这不仅仅对弹性训练意义重大。对于原来必须以固定节点数据进行训练的任务，也有重大的意义。
实现了动态的扩缩容，最多仅有一步的计算损失，并有着良好的性能。
这个意义不仅限于动态扩缩容，对故障恢复的场景也是意义重大。
附录: 相关贡献
参考文献
由于Reduce端内存有限，为了避免在Reduce端进行Merge的时候spill数据到磁盘。Reduce在获取Segment只能读取每个segment的部分buffer，然后对所有buffer进行Merge。然后对然后当某一个segment的部分buffer读取完成，会继续读取这个segment的下一块buffer，将这块buffer继续加到merge过程中。
这样有一个问题，Reduce端从ShuffleServer读取数据的次数大约为为segments_num * (segment_size / buffer_size)，对于大任务这是一个很大的值。过多的RPC意味着性能的下降。
