<aside> ✅

TLDR：我们在 slime 中新增了 FSDP 作为更为灵活的训练框架，并与 Megatron 完成对齐。FSDP 能够更加灵活支持诸如 Qwen3-Next 等架构创新的模型，并且有助于我们进一步支持 VLM RL。

</aside>

背景

什么是 FSDP？

FSDP (Fully Sharded Data Parallel) 继承了 DeepSpeed ZeRO Stage 3 的设计哲学，可以被视为是对传统 DDP (Distributed Data Parallel) 的强力优化。

从 Replicate 到 Shard

在传统的 DDP 中，每个 GPU 都维护一份完整的模型权重、梯度和优化器状态（Replication），通过 all-reduce 同步梯度。而在 FSDP 中，我们转向了 Sharding（切分） 模式：上述所有数据都被切分并分布在不同的 GPU rank 上。

• 前向传播：需要计算某层时，通过 all-gather 临时收集完整参数，计算完立即释放。
• 反向传播：梯度计算完成后，立即进行 reduce-scatter 同步并切分，随即释放完整梯度。

FSDP1 vs FSDP2

相比于 FSDP1 将所有参数摊平成一个巨大的 FlatParameter，FSDP2 引入了 DTensor (Distributed Tensor)。它能够在保持 Tensor 原始结构（如 shape, stride）的前提下，在指定的并行维度上进行更优的切分。这不仅解决了 FSDP1 中元数据易失和 padding 复杂的痛点，更为 MixedPrecision Training 和 LoRA 提供了开箱即用的支持；本文中提到的 FSDP 均指 PyTorch 原生支持的 FSDP2。

<aside> ✅

关于 FSDP 的更多内容可以查阅 SGLang RL team 以往的博客：RL System Deep Dive: FSDP Training Backend

Awesome-ML-SYS-Tutorial/rlhf/sys-design/readme-2-en.md at main · zhaochenyang20/Awesome-ML-SYS-Tutorial

</aside>

为什么 slime 需要 FSDP？

熟悉 slime 的朋友都知道，我们已经拥有了基于 Megatron-LM 的成熟训练引擎。考虑到引入新后端带来的显著维护成本，为什么我们还要坚定地支持 FSDP？

1. VLM 架构适配：VLM 的模态交互架构复杂，FSDP 的灵活性使其在适配上远比 Megatron 轻松。因此，我们选择 FSDP 作为 VLM RL 训练的首选路径（当然，Megatron 版本的适配也在计划中）。
2. 架构创新的敏捷性：对于 Qwen3-Next 这类处于快速迭代中的新架构，FSDP 能让我们以最快速度支持 RL 流程。
3. 低门槛与高易用性：作为 PyTorch native 的训练后端，FSDP 没有复杂的环境依赖和安装过程。无论是学习曲线还是 debug 成本都显著低于 Megatron。
4. 无缝生态兼容：FSDP 能直接兼容 HuggingFace Model 格式。这意味着我们无需像使用 Megatron 那样通过 mbridge 进行繁琐的权重转换(备注: 部分模型 Megatron 现在也无需手动权重转换了，内部会自动转换 PR)，社区模型开盒即用。

FSDP in slime：架构设计

要在 slime 中同时支持 Megatron 和 FSDP 两种截然不同的分布式后端，应该如何避免底层冲突并保持代码整洁？我们采用了 "接口标准化 + 物理隔离" 的顶层设计, 也就是说向外只暴露 FSDP 的核心函数, init, save, sleep, wake_up, train , 其他函数会尽量用下划线约定, 也就是像_train_core , 具体来说如下:

利用 Ray Actor 机制将不同后端封装在独立的进程空间中，向上层调度器暴露统一的训练原语（如 train），从而使上层算法逻辑无需关注底层的梯度同步细节。这种设计大幅消除了全局变量冲突并降低了条件分支复杂度，允许我们针对 FSDP2 的 Sharding 机制和 DTensor 结构进行深度优化。核心实现位于 slime/backends/fsdp_utils/actor.py，我们在保持对外业务逻辑（如 Data Packing、Context Parallel）与 Megatron 高度一致的同时，在内核实现上重构了数据流转路径，确保在享受 FSDP 灵活性的同时，最大化训练效率并维持数值精度。

完善的 FSDP 设计让顶层架构未受影响，整体流程仍旧是标准的 RLHF 循环：Rollout → Data Sharding → Packing → Forward/LogProb → Loss → Backward → Update。在此基础上，我们针对 FSDP 做了多项优化，包括 Data Packing、True On-Policy 模式、CUDA Graph Aware Weight Wake Up 以及 Training-Inference Mismatch 的众多缓解机制。我们接着讨论最顶层的 init 和 train 函数入口。

初始化

在 init 阶段，主要完成以下工作：

FSDP actor init 流程

• 模型与优化器：初始化 Actor Model 和 Reference Model，并支持从 Checkpoint 恢复；设置 true_on_policy_mode 和 Optimizer。
• Weight Updater：支持 Colocate（训练任务和推理任务放在同一组 GPU）和 Disaggregated（训练任务和推理任务放在不同 GPU）两种模式，用于将训练后的权重同步回 Inference Engine。
• Device Mesh：基于 DeviceMesh 构建 DP + CP 的通信拓扑，并调用 fully_shard 对参数进行切分。
• 算子优化：
  • 通过 enable_batch_invariant_mode 强制训练端采用与 SGLang 一致的算子，消除 batch size 对计算结果的影响。
  • 利用 torch.compile 固化 RoPE 实现，底层消除算子行为差异，确保 True On-Policy 的对齐。

训练流程

train 函数作为训练主入口：

FSDP actor train 流程

1. wake up：将之前被 Offload 的 Actor Model 加载回到 GPU。
2. data preparation：
   • 通过 process_rollout_data 获取当前 DP rank 所需数据。
   • 调用 _pack_rollout_data 将数据打包成 packed_batches（详见附录 Data Packing），消除 Padding 带来的性能损耗。
3. forward & log prob：
   • 计算 Actor 和 Ref 的 log_prob 与 entropy。
4. loss calculation：
   • 计算 PPO/GRPO loss（importance ratio, clip, KL penalty, entropy bonus）。
   • mismatch feature：实时计算 train_rollout_logprob_abs_diff 监控训练与推理的数值偏差。启用 TIS (Truncated Importance Sampling)[Source]，对 policy gradient loss 进行重加权，防止因为训推差异带来的 off policyness 造成模型崩塌。
5. update & offload：
   • 进行梯度累积和参数更新。
   • offload 策略：训练结束后调用 sleep 将模型和优化器 offload 到 CPU（colocated 模式）；Ref model 仅在计算 log prob 时加载，用完即 offload。

FSDP in slime 特性 & 优化

在架构设计基础上，我们进一步剖析目前做出的优化。

Data Prepare And Packing

每一轮训练开始时，FSDP actor (也就是这个 actor class) 首先从 rollout 拿到一批 balance 之后的 rollout sequence，然后按 DP rank 做简单的样本拆分，这一步和常规实现没有差别。为了极致效率，在这里我们实现了 数据打包（data packing） [PR Link]**。**简单来说，在 slime/backends/fsdp_utils/data_packing.py 中处理全部的 pack_sequences，对于输入的一批序列，根据每条的长度和 max_tokens_per_gpu 估算需要多少个 pack，即 micro-batch 的数量。接下来把长短不一的 sequence 分到不同 pack 中，使每个 pack 的 token 总数尽量接近。在每个 pack 内，将多条序列摊平成一条长的 tokens 向量，并构建 cu_seqlens 记录各条序列的起止位置。这种策略确保了每个 Pack 的 Token 总量高度一致，消除了传统 Padding 带来的算力浪费。具体细节可以参考附录  数据打包

严格训推一致

完成 Data Packing 后，actor 会对 packed micro‑batch 计算 ref/actor 的 log‑prob 和 entropy。我们在 FSDP 上实现了 True On Policy。也即对于近期非常火爆的 training inference mismatch 问题，我们给出了最为严格的答案，实现了同一个 policy model 在 training backend 和 inference backend 的 logp rob 绝对一致，从系统层面上解决了 training-infer mismatch。

<aside> ✅

简单说一下 training-infer kl = 0 的实现和思想如下:

• Training 和 Inference 都使用 FlashAttn3 当作 backend，来实现bitwise equal
• 使用 DeepGEMM进行矩阵乘法, Batch-invariant Kernels 实现批次不变性

具体细节在 slime 的 Doc 里有更详细的记载, 主要实现的 PR 是 [PR link1], [PR link2]

</aside>

training-rollout logprob diff = 0

我们更进一步优化 true on policy 情况下的性能。get_logprob_and_entropy_with_cp 直接复用了 Rollout 传入的 temperature，并关闭了可能引入偏差的 allow_compile , disable compile 会禁止compile selective_log_softmax_raw，防止因为编译带来的和 batch invariant‣中.确保训练端重算的 log‑prob 能精准还原 Rollout 时的数值表现，杜绝因计算路径不同而产生的估算偏差

<aside> ⚠️

</aside>

Algorithms Mitigation For Mismatch

数据打包