重新审视参数服务器在大语言模型后训练中的应用 Revisiting Parameter Server in LLM Post-Training
用点对点通信替代FSDP集合通信，消除逐层同步屏障，LLM后训练提速36%
前置知识
Fully Sharded Data Parallel (FSDP)
FSDP 是 PyTorch 实现的一种分布式数据并行训练策略，源自 ZeRO（Zero Redundancy Optimizer）的思想。其核心是将模型的参数、梯度和优化器状态分片（shard）存储在不同设备上，每个设备只保存 1/N 的模型状态。在前向传播时，通过 all-gather 操作从所有设备收集完整参数用于计算，计算完毕后立即丢弃以节省内存；在反向传播后，通过 reduce-scatter 操作聚合梯度并分发给各设备。这种设计使得万亿参数模型也能在有限显存的 GPU 集群上训练，是当前 LLM 后训练的标准并行方案。
本文的核心就是改进 FSDP 的通信模式，理解 FSDP 的逐层 all-gather 和 reduce-scatter 机制是理解本文动机和方法的基础。
集合通信 (Collective Communication)
集合通信是分布式训练中的通信范式，要求所有参与设备同时参与通信操作。典型的集合操作包括 all-gather（所有设备收集所有其他设备的数据片段以重建完整数据）、reduce-scatter（所有设备将各自的数据进行归约后分发）、all-reduce 等。这些操作通过 NCCL 等库高度优化，支持 Ring AllReduce 等高效拓扑。集合通信的效率建立在所有设备工作负载均衡的假设之上——因为所有设备必须同时完成当前步骤才能进入下一步，最快完成的设备必须等待最慢的设备。
本文指出集合通信在负载不均衡场景下的同步屏障是性能瓶颈的根本原因，理解这一点才能理解为什么需要回归参数服务器架构。
参数服务器 (Parameter Server, PS)
参数服务器是早期分布式深度学习的通信架构，由 Dean 等人在 DistBelief（2012）中提出。其设计是：一组服务器节点负责存储模型参数和优化器状态，一组工作节点从服务器拉取参数、在本地数据上执行前向和反向计算、然后将梯度推送给服务器进行聚合和更新。关键特点是工作节点之间相互独立，不需要同步推进——快的节点可以继续处理下一批数据，慢的节点不会拖累整体进度。这种异步或松同步的设计天然容忍工作负载不均衡和硬件异构性。
本文的核心创新就是将经典参数服务器的思想重新引入现代分片数据并行训练，理解 PS 的架构特点是理解 ODC 方法的关键。
序列打包 (Sequence Packing)
序列打包是处理 LLM 训练中变长序列的一种技术。由于文本数据的序列长度差异巨大（本文数据集平均长度从 9K 到 35K tokens 不等），简单的填充（padding）策略会浪费大量计算资源。序列打包将多个短样本拼接成一个长序列，并使用注意力掩码（attention mask）防止不同样本之间的信息泄露。FlashAttention 等库提供了对打包序列的高效支持。然而，打包只能缓解而不能消除负载不均衡——因为每个微批次（microbatch）的大小受设备显存限制，且注意力计算的 $O(s^2)$ 复杂度与激活内存的 $O(s)$ 复杂度之间存在根本性不匹配。
序列打包是当前缓解负载不均衡的主要手段，本文指出其根本局限性，并提出 ODC 从通信层面解决问题。
微批次与梯度累积 (Microbatch & Gradient Accumulation)
在 LLM 训练中，一个优化步（minibatch）的样本量往往超出单次前向-反向传播的显存容量。因此将 minibatch 拆分为 M 个 microbatch，每个 microbatch 独立执行前向和反向计算得到梯度 $g^{(m)}$，然后累积梯度 $\bar{g} = \sum_{m=1}^{M} w_m g^{(m)}$ 后再执行优化器更新。在标准 FSDP 中，每个 microbatch 的处理都包含逐层的集合通信，这意味着 M 个 microbatch 就产生 $M \times L$ 个同步点（L 为层数），大幅放大了负载不均衡的影响。
ODC 的一个核心设计就是将同步粒度从逐层放松到逐 minibatch，理解 microbatch 和梯度累积机制才能理解这一放松的意义。
研究动机
现代 LLM 后训练（包括监督微调 SFT 和强化学习 RL）面临一个被长期忽视的根本性问题：FSDP 所依赖的集合通信机制隐含地假设各设备工作负载均衡，但这一假设在后训练场景中被严重破坏。真实世界的文本语料库包含长度差异巨大的序列，本文使用的三个数据集就是典型例子：LongAlign 数据集平均序列长度 16,540 tokens，最长超过 64K；SWE-Smith 数据集平均长度 34,724 tokens，最长超过 60K；AIME 数据集平均长度 9,066 tokens。由于注意力计算的计算复杂度随序列长度呈 $O(s^2)$ 增长，而激活内存呈 $O(s)$ 增长，这种长度差异直接转化为各设备之间持久的计算不平衡。尽管序列打包（Sequence Packing）等技术可以在一定程度上缓解这个问题，但打包策略受设备显存限制，无法完全消除不平衡。论文的评估显示，即使使用最先进的打包策略，在长序列 SFT 训练中设备空闲时间仍可高达 50%（见 Table 6，1.5B 模型在 LongAlign 数据集上，minibatch size=1 时 bubble rate 达 66.81%）。在 FSDP 中，每一层的前向传播和反向传播都需要集合通信操作，形成了逐层的同步屏障——所有设备必须完成当前层的 all-gather 或 reduce-scatter 才能进入下一层。这意味着最快完成的设备必须等待最慢的设备，产生大量的设备空闲时间（bubble）。当 microbatch 数量增加时，同步点数量也随之线性增长，进一步放大了这一低效性。
本文的目标是本文的目标是从通信机制层面解决 FSDP 在 LLM 后训练中的负载不均衡问题，而不是继续在打包策略上做文章。具体而言，作者希望设计一种新的通信方案，将 FSDP 的同步粒度从逐层放松到逐 minibatch，从而消除逐层同步屏障导致的设备空闲时间。该方案需要保留 FSDP 的核心优势——显存效率、去中心化、可扩展性和实现简单性——同时获得参数服务器架构对异构工作负载的天然容忍能力。在性能目标上，论文希望在多种 LLM 后训练任务（SFT 和 RL）、多种模型规模（1.5B 到 32B）上实现一致的吞吐量提升。
与已有工作不同的是，本文的独特切入角度在于：它没有像之前的工作那样继续在「如何更好地平衡工作负载」这个方向上努力（即优化打包策略），而是退后一步，重新审视通信机制本身。论文提出了一个关键洞察——FSDP 中的同步屏障是通信模型的产物，而非训练算法本身的要求。分布式数据并行的第一性原理是：各设备的计算是相互独立的。标准 FSDP 通过集合通信违反了这种独立性，强制所有设备同步推进。本文指出，如果将集合通信替换为点对点通信，就可以恢复这种独立性，让快的设备不用等待慢的设备。这一视角将经典参数服务器的思想引入现代分片数据并行范式，是一个被忽视但本质性的解决思路。
核心方法
ODC 的核心思路可以用一个简单的类比来理解：想象一个班级一起做作业。FSDP 的方式是「大家一起翻到第一页，做完后一起翻到第二页」——如果有人做得慢，其他人只能干等。ODC 的方式则是「每个人按自己的节奏做，做完后统一交作业」——在交作业之前（minibatch 结束），每个人可以自由控制进度。具体技术路线是：保留 FSDP 的参数分片和显存布局不变，但将逐层的集合通信操作（all-gather 和 reduce-scatter）替换为点对点通信操作（gather 和 scatter-accumulate）。每个设备独立地从其他设备获取所需参数、独立地执行计算、独立地推送梯度，只在 minibatch 结束时进行同步（optimizer step）。这种设计将 FSDP 重新框架化为一个去中心化的参数服务器：每个设备同时扮演服务器（拥有和管理一部分参数/优化器状态）和 worker（在本地数据上执行计算）的角色。
ODC 的核心创新在于一个关键的通信原语替换：将 FSDP 中的集合操作分解为点对点操作。在标准 FSDP 中，all-gather 要求所有设备同时参与以重建完整参数；ODC 将其分解为一系列独立的 gather 请求，每个设备只从需要的 peer 拉取特定的参数分片。类似地，reduce-scatter 被分解为 scatter-accumulate 操作，每个设备将计算好的梯度直接推送给拥有对应梯度分片的设备。这一替换的本质区别在于：集合通信是「全员同步」的——所有设备必须同时到达通信点才能开始；而点对点通信是「按需异步」的——每个设备准备好就可以发起通信，不需要等待其他设备。更重要的是，这些点对点数据传输是非侵入性的（non-intrusive）：当一个设备向另一个设备发起 gather 或 scatter-accumulate 请求时，不会打断目标设备正在进行的计算。这是通过 RDMA（远程直接内存访问）技术实现的——RDMA 允许在不消耗目标设备 GPU 计算资源的情况下直接读写其内存。同步屏障从每层一次放松到每个 minibatch 一次，设备之间的执行完全解耦。
方法步骤详情
ODC 的完整流程可以分为以下几个步骤：第一步，参数获取（Gather）。在前向传播开始前，每个设备独立地向其他设备发送 gather 请求，拉取当前层所需的参数分片。与 FSDP 的 all-gather 不同，设备无需等待其他设备也完成参数获取——只要拿到自己需要的参数就可以开始计算。第二步，前向计算。每个设备使用获取的参数在自己的本地数据上执行前向传播。不同设备可以同时处理不同层的计算（因为它们可能处于不同的进度）。第三步，反向计算与梯度推送（Scatter-accumulate）。反向传播完成后，每个设备将计算好的梯度通过 scatter-accumulate 操作直接推送给拥有对应梯度分片的设备。接收端运行一个轻量级守护进程（daemon），通过轮询（polling）检测到梯度到达后进行累积。由于轮询不占用 GPU SM，不会影响接收端正在进行的计算。第四步，梯度累积。各设备独立地累积从其他设备收到的梯度，累积过程在整个 minibatch 的所有 microbatch 完成后结束。第五步，优化器更新（Optimizer Step）。在 minibatch 结束时，所有设备同步执行优化器更新——这是 ODC 唯一保留的同步点。通信原语的实现上，节点内通信使用 CUDA IPC（直接读写远程 GPU 张量），节点间通信使用基于 NVSHMEM 的自定义 Triton kernel。整个 ODC 集成到 FSDP 中非常简单：只需将集合通信调用替换为 ODC 原语，并在 minibatch 结束时获取累积的梯度。
技术新颖性
ODC 的技术新颖性体现在三个层面。首先，在架构层面，它首次将经典参数服务器的「解耦设备进度」思想直接融入现代 FSDP 的分片机制中。虽然「colocated server+worker」的设计在一些 PS 系统中已有先例（如 Jiang et al., 2020），但 ODC 是首次将其与 FSDP 的分片机制直接集成。其次，在通信层面，ODC 使用 RDMA 原生接口实现非侵入式的点对点通信。CUDA IPC 实现节点内通信（无需自定义 GPU kernel），NVSHMEM + Triton-Distributed 实现节点间通信。关键的 scatter-accumulate 操作设计巧妙：客户端通过 put_mem 推送数据后通过同一 RDMA 通道通知服务端，服务端的轻量级 daemon 进程通过轮询检测并执行梯度累积，轮询不占用 GPU SM。为每个客户端分配独立 buffer 以支持并行数据传输，每台服务器的总 buffer 内存为 M（M 为一层 Transformer 的元素数）。第三，在负载均衡层面，ODC 解除了「所有设备必须处理相同数量 microbatch」的约束，允许将平衡目标从微批次级别提升到 minibatch 级别。这不仅简化了打包算法，还通过在更大、更灵活的样本集合上操作实现了更优的负载均衡。
实验结果
论文在三个数据集（LongAlign、SWE-Smith、AIME）和四个模型规模（1.5B、7B、14B、32B）上进行了全面评估，涵盖 SFT 和 RL 两类后训练任务。核心发现如下：在 SFT 任务上，ODC 配合 LB-Mini 负载均衡策略实现了最高 36% 的吞吐量提升。具体来看，14B 模型在 LongAlign 数据集上，minibatch size=4 时，ODC+LB-Mini 达到 61.4 samples/sec，而 Collective+LB-Micro 仅为 45.1 samples/sec，提升 36%。32B 模型在同一数据集上，minibatch size=4 时提升 31%。值得注意的是，ODC 在未打包（LocalSort）的场景下也能带来提升（最高 12%），说明其优势不仅来自更好的负载均衡，还来自消除同步屏障本身。在 RL 任务上，ODC 实现了最高 10% 的提速。RL 任务的提升幅度小于 SFT，主要原因是：（a）verl 框架的实现约束要求各设备处理相同数量的样本，限制了 LB-Mini 的效果；（b）AIME 数据集的序列长度分布比 SFT 数据集的长尾程度更轻。Bubble rate 数据直观地展示了 ODC 的效果：在 1.5B 模型、SWE-Smith 数据集、minibatch size=4 的设置下，Collective+LB-Micro 的 bubble rate 为 20%，而 ODC+LB-Mini 降至 0.71%；在 14B 模型、LongAlign 数据集、minibatch size=4 的设置下，bubble rate 从 37.75% 降至 14.43%。当 minibatch size=8 时，ODC 的 bubble rate 可以降至接近零（0.01-0.08%）。参数研究表明：加速比在中等 minibatch size 时达到峰值；随序列长度增加而增大（更长序列放大了 $O(s^2)$ 的计算不平衡）；随 packing ratio 增加而降低（更好的打包策略为基线提供了更多灵活性）；随设备数量增加而增长（更多设备引入更大的异构性）。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| SFT-LongAlign (14B, minibatch=4) | samples/sec | 61.4 (ODC+LB-Mini) | 45.1 (Collective+LB-Micro) | +36% |
| SFT-LongAlign (32B, minibatch=4) | samples/sec | 33.6 (ODC+LB-Mini) | 25.6 (Collective+LB-Micro) | +31% |
| SFT-SWE-Smith (1.5B, minibatch=4) | samples/sec | 172.0 (ODC+LB-Mini) | 140.9 (Collective+LB-Micro) | +22% |
| RL-AIME (7B, minibatch=4) | samples/sec | 302.5 (ODC+LB-Mini) | 273.0 (Collective+LB-Micro) | +11% |
| SFT-LongAlign (14B, minibatch=8) | samples/sec | 69.0 (ODC+LB-Mini) | 53.9 (Collective+LB-Micro) | +28% |
局限与改进
论文坦诚地指出了几个重要的局限性。首先，节点间通信效率是 ODC 的主要短板。通信原语基准测试（Figure 11）显示，节点内 ODC 的带宽与集合通信相当，但跨节点时 ODC 显著慢于集合通信。这是因为集合通信可以通过层次化优化（如先节点间广播再节点内广播）来最小化跨节点流量，而 ODC 的点对点 RDMA 通信无法利用这些优化。Table 2 的通信量分析证实，虽然 ODC 不增加总通信量（均为 $(D-1) \times K$），但增加了跨节点通信量。其次，在 RL 任务上的加速效果不如 SFT 显著（最高 10% vs 36%），主要受限于 verl 框架的实现约束（要求各设备处理相同数量样本）和 AIME 数据集较轻的长尾分布。第三，当 minibatch size=1 时，ODC 与集合通信性能完全相同，因为此时 ODC 也是每个样本同步一次，退化为与集合通信相同的同步模式。第四，论文没有在超过 32 个 GPU 的更大规模集群上验证 ODC 的效果，而更大规模通常意味着更大的异构性和更明显的 straggler 效应。最后，ODC 的 hybrid sharding 策略虽然能缓解跨节点通信问题，但会增加显存使用（Figure 13），需要在通信效率和显存容量之间权衡。
独立分析的弱点
从独立分析的角度，ODC 存在以下几个值得关注的弱点。第一，跨节点通信效率问题的解决方案尚不成熟。论文提出的 hybrid sharding（类似 ZeRO++，节点内分片参数和梯度、跨节点分片优化器状态）虽然有效，但本质上是用显存换通信效率——在 32B 模型上，hybrid sharding 的显存使用比 full sharding 高出约 40%（Figure 13），对于显存已经捉襟见肘的大模型训练来说这是一个不小的代价。一个可能的改进方向是设计层次化的点对点通信策略：让设备优先从同一节点内已缓存了目标参数的 peer 获取数据，形成类似集合通信的层次化路径。第二，ODC 的 RDMA daemon 虽然声称不占用 GPU SM，但在高并发场景下（如 32 个设备同时向同一台服务器推送梯度），轮询和梯度累积的开销可能会累积。论文没有报告在大规模场景下的 daemon 开销分析。改进方向是将 daemon 的梯度累积操作 offload 到 CPU 或专用硬件上。第三，ODC 要求每个客户端分配独立 buffer（总大小为 M per server），虽然论文声称这是可管理的，但在节点数很多时这个 buffer 的显存开销也需要评估。第四，论文只验证了同步更新语义——minibatch 结束时的 barrier 是硬性的。在高度异构的环境中（如不同代 GPU 混合部署），这个 barrier 仍然可能造成显著等待。可以探索 bounded-staleness 异步更新来进一步放松同步约束。
未来方向
论文提出了几个明确的未来研究方向，基于这些方向还可以做更多延伸。第一，ODC 特定的通信优化：当前实现使用直接点对点通信，通信图可以进一步优化，例如设备可以从同一节点内已缓存参数的 peer 拉取数据，创建层次化通信路径。第二，放松同步保证：当前设计在 minibatch 边界保留同步更新以保证训练语义一致性，但可以扩展支持有界陈旧度（bounded-staleness）异步 SGD，在高度异构环境中进一步减少空闲时间——这需要仔细分析对 LLM 训练收敛性的影响。第三，弹性与容错：PS 架构天然支持弹性和容错（节点可以动态加入/退出），而集合通信系统在这方面非常脆弱。将这些能力集成到 ODC 中可以提高大规模长时间 LLM 训练任务的韧性。第四，将 ODC 应用到更多场景：如 MoE（混合专家）模型的训练（天然存在负载不均衡）、多模态模型训练（不同模态的计算量差异大）、以及长上下文推理的 KV-cache 管理。
复现评估
论文在可复现性方面做得很好。作者已经开源了 ODC 的核心通信库以及集成到 FSDP 的代码补丁，地址为 https://github.com/sail-sg/odc。实验使用的模型（DeepSeek-R1-Distill-Qwen 系列）和数据集（LongAlign、SWE-Smith、AIME/NuminaMath）都是公开可获取的。RL 实验使用了开源框架 verl。训练硬件为 NVIDIA A100 80G GPU（最多 32 卡），使用 NVSwitch 节点内通信和 RoCE RDMA（800 Gbps/节点）节点间通信。需要注意的是，ODC 的实现依赖 Triton-Distributed 和 RDMA 接口（CUDA IPC 节点内、NVSHMEM 节点间），对硬件和软件环境有一定要求——不支持没有 RDMA 能力的普通以太网集群。整体复现难度中等：代码已开源且集成简单（只需替换集合通信调用），但需要特定的硬件配置。
论文图表
展示了标准 FSDP 中两个设备的计算时间线。每个设备交替执行前向传播（F）和反向传播（B），每层之间都有同步屏障（用竖线表示）。Device 0 的计算完成后必须等待 Device 1 完成才能进入下一层，导致 Device 0 产生大量空闲时间。
这张图直观地展示了问题的根源：逐层同步屏障导致快设备等待慢设备，是理解整篇论文动机的关键。
展示了三个评估数据集的序列长度分布直方图。LongAlign 的平均长度为 16,540 tokens，分布较为均匀地覆盖 0-64K；SWE-Smith 的平均长度为 34,724 tokens，呈现出明显的长尾分布，大部分样本在 10K-60K 之间；AIME 的平均长度为 9,066 tokens，分布相对集中，主要在 0-10K 范围内，但也有一些长尾样本超过 20K。
这张图直接支撑了论文的核心论点：LLM 后训练数据的序列长度差异巨大，是导致负载不均衡的根本原因。
四个子图分别展示了 1.5B、7B、14B、32B 模型一层参数的通信带宽对比。节点内（2-8 设备）ODC 的 gather 和 scatter-accumulate 带宽与集合通信的 all-gather 和 reduce-scatter 相当；但跨节点（16-32 设备）时 ODC 显著慢于集合通信。例如 32B 模型在 32 设备时，集合通信的 all-gather 带宽约 150 GB/s，而 ODC 的 gather 约 80 GB/s。
这张图揭示了 ODC 的主要技术短板——跨节点通信效率，是理解论文局限性的关键。