slime FSDP 训练后端搭建与测试实战：从 Docker 环境到 Colocated/Disaggregated 部署

• 文档
• 教程
• 人工智能
• 大模型
• RLHF

本篇技术指南以 slime（基于 SGLang 推理与 PyTorch 训练后端的 RL 框架）的 FSDP 训练后端为主线，完整记录从拉取 Docker 镜像、安装 slime、下载模型与数据集，到加载模型配置、运行 FSDP 双卡训练测试脚本的端到端流程，并深入讲解 Colocated（训推一体化）与 Disaggregated（训推分离）两种资源编排方式、SLIME_BACKEND=fsdp后端激活与 GPU 分片机制，以及 Blackwell（B 卡）环境下的适配要点。读完本文，你将能够独立复现 slime 上 FSDP 后端的搭建与测试，理解训练/推理卡数分配的核心参数，并掌握从 Megatron 后端迁移到 FSDP 的关键差异。

本文对应仓库中的 setup_fsdp.md（FSDP 发布配套的搭建与测试文档），并辅以 FSDP2 训练后端介绍 中的架构与参数细节进行纵深展开。

背景：为什么要在 slime 上测试 FSDP

slime 在原有基于 Megatron-LM 的训练后端之外，新增了FSDP（Fully Sharded Data Parallel）作为训练后端，并在精度上完成了与 Megatron 的对齐。文档中提到的 FSDP 均指 PyTorch 原生支持的FSDP2：相比 FSDP1 将所有参数摊平为巨型FlatParameter，FSDP2 引入DTensor（分布式张量），能在保持 Tensor 原始结构（shape、stride）的前提下按指定并行维度切分，并为混合精度训练与 LoRA 提供开箱即用的支持。引入 FSDP 的动机主要包括：

• VLM 架构适配：VLM 的模态交互架构复杂，FSDP 的灵活性使适配远比 Megatron 轻松；
• 架构创新的敏捷性：对 Qwen3-Next/gpt-oss 等快速迭代的新架构，FSDP 能以最快速度接入 RL 流程；
• 低门槛与高易用性：作为 PyTorch 原生后端，无复杂环境依赖，学习曲线与 debug 成本均低于 Megatron；
• 无缝生态兼容：FSDP 直接兼容 HuggingFace 模型格式，无需像 Megatron 那样通过权重转换工具转换权重，社区模型开盒即用。

本文的实战部分正是围绕这一新后端的验证与测试展开，覆盖 H 卡与 B 卡两种硬件形态。

Quick Start：H 卡环境下的完整搭建流程

以下操作在 H 卡（Hopper 架构）上完成，从零开始共分五步：拉取并启动容器、安装 slime、下载模型与数据集、加载模型配置、运行测试脚本。

拉取并启动 Docker 容器

# 拉取最新镜像 # 最新的镜像是 B 卡 H 卡通用的 docker pull slimerl/slime:latest # 启动容器 docker run -d --gpus all --ipc=host --shm-size=16g \ --name slime_wren_fsdp \ -it slimerl/slime:latest /bin/bash

对启动命令中关键参数的理解：

安装 slime

进入 Docker 容器后，克隆 slime 仓库并以可编辑模式安装：

# 路径可根据实际情况调整 cd /root/ git clone https://github.com/THUDM/slime.git cd slime pip install -e .

pip install -e .以 editable 模式安装，便于后续调试与本地修改直接生效。

模型与数据集下载

可从 Hugging Face、ModelScope 等平台下载所需模型与数据集。以下是使用huggingface_hub下载示例资源的命令：

pip install -U huggingface_hub # 下载模型权重 (Qwen3-0.6B) hf download Qwen/Qwen3-0.6B --local-dir /root/Qwen3-0.6B # 下载训练数据集 (dapo-math-17k) hf download --repo-type dataset zhuzilin/dapo-math-17k \ --local-dir /root/dapo-math-17k # 下载评估数据集 (aime-2024) hf download --repo-type dataset zhuzilin/aime-2024 \ --local-dir /root/aime-2024

要点说明：

• hf download是huggingface_hub新版 CLI（pip install -U huggingface_hub确保版本足够新）；
• 模型仓库无需--repo-type，数据集仓库必须显式指定--repo-type dataset；
• --local-dir指定本地保存路径，后续启动脚本中通过--hf-checkpoint、--prompt-data等参数引用这些路径；
• 三个资源分别对应：训练目标模型权重（Qwen3-0.6B）、RL 训练用的数学推理 prompt 数据集（dapo-math-17k）、评估用的 AIME 2024 竞赛题集。

加载目标模型的配置文件

slime/scripts/models目录下包含受支持模型的配置文件。需要source对应模型的脚本，将配置参数加载到当前环境中。以 Qwen3-0.6B 为例：

cd /root/slime source scripts/models/qwen3-0.6B.sh

对于 Qwen3-4B、Qwen3-30B-A3B 是类似的（scripts/models/下分别有对应脚本）。source的作用是把模型相关的公共参数（如 tokenizer、prompt 长度、采样参数、损失相关开关等）注入当前 shell 环境，供后续训练脚本继承使用。

运行训练/测试脚本

完成上述准备工作后，即可运行训练脚本：

cd /root/slime bash tests/test_qwen3-0.6B_fsdp_colocated_2xGPU.sh # 2GPU 协同训练测试

该测试脚本以2 张 GPU、Colocated（训推共卡）方式启动 Qwen3-0.6B 的 FSDP 训练，用于验证 FSDP 后端在 slime 中能够正常完成「Rollout 采样 → 数据切分与打包 → 前向/LogProb → Loss → 反向 → 更新 → 权重回传」的完整 RL 循环。这也是进入大规模训练前最轻量的冒烟测试方式。

除测试脚本外，配套博客还提供了直接可跑的训练脚本 [scripts/run-qwen3-4B-fsdp.sh]（位于 slime 仓库内），在完成权重与数据集下载后执行bash /root/slime/scripts/run-qwen3-4B-fsdp.sh即可一键启动 Qwen3-4B 的 FSDP 训练（colocated 模式 + reference model）。FSDP 通过AutoModelForCausalLM.from_pretrained()自动读取config.json中的所有架构信息，无需手动指定架构参数，也无需权重格式转换——这是与 Megatron 路径最显著的使用差异。

特性介绍：Colocated 与 Disaggregated 两种资源编排

slime 支持两种训练（Actor）与推理（Rollout）的资源编排方式，理解二者的差异是配置 FSDP 测试的前提。

默认的训推分离（Disaggregated）配置

在默认配置下，训练和推理的资源是分开指定的：通过ray给训练部分分配actor_num_nodes * actor_num_gpus_per_node张 GPU，给推理分配actor_num_nodes * rollout_num_gpus张 GPU，也即训推分离。

标准（分离）配置：

ray job submit ... \ -- python3 train.py \ --actor-num-nodes 1 \ --actor-num-gpus-per-node 4 \ --rollout-num-gpus 4 \ ...

上述配置中，Actor使用 4 张卡，Rollout也使用 4 张卡，两者并行运行。

异步训练提示：当进行训推分离时，训练和推理的 GPU 总是相互等待，为了避免这种资源空闲，可以开启异步训练——将启动脚本中的train.py改为train_async.py。这样 slime 就会在进行当前 rollout 的训练时，同时生成下一个 rollout 的数据。

⚠️日志注意：异步训练时，sglang 的性能检测日志与训练日志可能混到一起、不易区分，可通过--sglang-log-level降低 sglang 的日志级别。

训推一体化（Colocated）配置

要将训练和推理部署在同一组 GPU 上，添加--colocate参数即可。开启后 slime 会忽略--rollout-num-gpus，让训练和推理的卡数相等：

ray job submit ... \ -- python3 train.py \ --actor-num-nodes 1 \ --actor-num-gpus-per-node 8 \ --colocate \ ...

此时，训练和推理共享全部 8 张 GPU。测试脚本tests/test_qwen3-0.6B_fsdp_colocated_2xGPU.sh正是这种模式的小规模验证。Colocated 模式下，训练结束调用sleep将模型与优化器 offload 到 CPU，为随后的 rollout 腾出显存；Rollout 完成后通过 weight updater 将最新权重同步回推理引擎。

资源编排的源码视角

从 slime 的 Ray Placement Group 实现（可参见本仓库的 代码走读文档）可以印证上述行为：create_placement_groups中，colocate分支下num_gpus = actor_num_nodes * actor_num_gpus_per_node且rollout_offset = 0（Rollout 与 Actor 完全共用 bundle）；而默认分支下num_gpus = actor_num_nodes * actor_num_gpus_per_node + rollout_num_gpus，rollout_offset = actor_num_nodes * actor_num_gpus_per_node（两者使用不同的 bundle）。异步训练（train_async.py）仅在 Disaggregated 架构下有意义，通过rollout_manager.async_generate()与actor_model.async_train让 rollout 始终领先 train 一个 step（one-step off-policy）。

FSDP 激活机制：后端切换与 GPU 分片

在 slime 中启用 FSDP 后端需要三处配合：指定后端、配置 GPU 分片、选择状态字典模式。

指定 FSDP 后端

# 关键：指定后端为 FSDP "SLIME_BACKEND": "fsdp"

对应命令行参数为--train-backend fsdp（与 Megatron 默认后端区分）。slime 通过「接口标准化 + 物理隔离」的顶层设计同时支持 Megatron 与 FSDP：对外只暴露init / save / sleep / wake_up / train五个核心函数，其余实现采用下划线约定（如_train_core）；利用 Ray Actor 将不同后端封装在独立进程空间中，向上层调度器暴露统一训练原语，上层算法逻辑无需关心底层梯度同步细节。

GPU 分片设置

export CUDA_VISIBLE_DEVICES=1,2 # 使用 GPU 1,2 --actor-num-gpus-per-node 2 # 2 个 GPU 进行模型分片

• CUDA_VISIBLE_DEVICES限定进程可见的物理 GPU，与容器内 GPU 编号对应；
• --actor-num-gpus-per-node决定参与 FSDP 模型分片的 GPU 数量，即 DP（数据并行）分片规模。FSDP 将模型权重、梯度与优化器状态切分并分布在这些 rank 上：前向时通过all-gather临时收集完整参数、计算完立即释放；反向时梯度计算完立即reduce-scatter同步并切分。

FSDP 状态字典模式选择

--fsdp-full-params # 启用 FULL_STATE_DICT 模式 # 注释掉则使用默认的 SHARDED_STATE_DICT 模式

• 默认（注释掉--fsdp-full-params）：使用SHARDED_STATE_DICT，每个 rank 只保存模型参数的一部分与对应元数据，保存 checkpoint 时无需先 all-gather，存储效率更高，读取分片 state_dict 即可直接加载；
• 启用--fsdp-full-params：切换为FULL_STATE_DICT模式，保存/加载完整参数（便于跨环境迁移），但需要汇聚全量参数。

配套的保存行为还受--fsdp-state-dict-cpu-offload（默认 True）控制：保存 checkpoint 时将状态字典 offload 到 CPU，避免峰值显存冲击。

Blackwell（B 卡）GPU 设置

启动容器

B 卡（Blackwell 架构）与 H 卡的操作步骤几乎完全一致，主要差异体现在容器启动参数（这些是环境配置差异，不是硬件差异）：

# 拉取最新镜像，最新的镜像是 B 卡 H 卡通用的 docker pull slimerl/slime:latest # 启动容器 # 这里 GPU 相关参数完全相同，主要差异是镜像版本和挂载目录（这些是环境配置，不是硬件差异） docker run \ -itd \ --shm-size 32g \ --gpus all \ --ipc=host \ --network=host \ --privileged \ -v {your_cache_path}:/root/.cache \ --name slime_fsdp_{your_name} \ slimerl/slime:latest \ /bin/bash

与 H 卡启动命令的差异说明：

• --shm-size 32g：共享内存提升到 32GB，适配 B 卡更大的 batch 与更长的序列；
• --network=host：直接使用宿主机网络栈，便于 Ray 集群、SGLang Router 等跨进程/跨节点通信；
• --privileged：特权模式，规避部分 GPU/NCCL 相关的内核级限制；
• -v {your_cache_path}:/root/.cache：将宿主机的缓存目录（HuggingFace 模型缓存等）挂载到容器内，避免重复下载。

剩余步骤和 H 卡操作步骤完全相同（安装 slime、下载模型与数据集、source 模型配置、运行测试脚本）。

常见问题：NCCL 端口冲突

如果遇到nccl的 error，在ray启动的时候可以指定一个端口：

ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 --port 9987

• --port 9987显式指定 Ray head 的 GCS 服务端口，避免多实例或多节点环境下端口冲突；
• ${MASTER_ADDR}为当前节点的 IP 地址，需要在实际环境中替换；
• --num-gpus 8声明该节点可用的 GPU 数量，需与docker run --gpus all暴露的数量一致；
• --disable-usage-stats关闭遥测上报，--dashboard-host=0.0.0.0 --dashboard-port=8265允许从外部访问 Ray Dashboard。

从 Megatron 迁移到 FSDP：参数对照与能力边界

当你把一套已有的 Megatron 训练脚本迁移到 FSDP 时，可以参考下表（源自 FSDP2 训练后端介绍），快速定位需要替换的参数：

FSDP 目前不支持/不完善的功能

迁移前还需要明确能力边界：

• 并行策略：FSDP 目前仅支持DP + CP，不支持TP、EP、PP；
• CP 实现差异：Megatron Core 的 CP 是原生实现（与 TP/PP 深度集成），而 FSDP 通过外部库 Ring Flash Attention 实现（连续 chunk 切分 + 负载均衡交给 ring flash attn，slime 只需关注输入切分与结果聚合）；
• 重计算：Megatron 的--recompute-granularity(full/selective)、--recompute-method(uniform/block)、--recompute-num-layers均不支持，FSDP 只有简单的--gradient-checkpointing布尔开关；
• 学习率调度：FSDP 优化器的学习率目前仅支持 constant，且没有 warmup 策略。

FSDP 训练后端的运行机制（源码补充）

搭建与测试完成后，了解 FSDP actor 在 slime 内部的运行机制有助于排查问题与调整配置。核心实现位于 slime 仓库的slime/backends/fsdp_utils/actor.py，整体流程仍为标准 RLHF 循环：Rollout → Data Sharding → Packing → Forward/LogProb → Loss → Backward → Update。

初始化与训练主循环

初始化（init）阶段主要完成四类工作：初始化 Actor Model 与 Reference Model（支持从 Checkpoint 恢复，设置true_on_policy_mode与 Optimizer）；初始化 Weight Updater（支持 Colocate 与 Disaggregated 两种权重回传模式）；基于DeviceMesh构建 DP + CP 通信拓扑并调用fully_shard对参数切分；算子优化（enable_batch_invariant_mode强制训练端采用与 SGLang 一致的算子、torch.compile固化 RoPE 实现以确保 True On-Policy 对齐）。

训练主循环（train）分为五步：wake_up（将之前 offload 的 Actor Model 加载回 GPU）→ data preparation（process_rollout_data取当前 DP rank 所需数据，_pack_rollout_data打包为packed_batches，消除 padding 损耗）→ forward & log prob（计算 Actor/Ref 的 log_prob 与 entropy）→ loss calculation（PPO/GRPO loss + importance ratio/clip/KL penalty/entropy bonus，实时计算train_rollout_logprob_abs_diff监控训推数值偏差，可选启用 TIS 重加权）→ update & offload（梯度累积与参数更新，训练结束调用sleep将模型与优化器 offload 到 CPU）。

显存优化：offload 策略

FSDP 训练流程通过以下场景的权重 offload 节省显存：

• Train offload：Colocated 场景下，训练完成后调用sleep将模型 weight 与 optimizer offload 到 CPU，避免在 rollout 阶段占用 GPU 显存；
• Ref model：使用 KL penalty 时，reference model 仅在compute_log_prob时加载到 GPU，计算完成后立即 offload 回 CPU；
• Optimizer offload：训练阶段将不参与计算的 model parameter 与 gradient 都 offload 到 CPU，显著节省训练显存；代价是 optimizer step 在 CPU 上执行，训练时间会明显上升（对应--fsdp-cpu-offload）。

精度对齐验证与 Context Parallel

FSDP 与 Megatron 在训练精度上完成了对齐验证：实验采用单机 H100，对比 Megatron、FSDP colocated w/ ref model、FSDP colocated w/o ref model 三条曲线，收敛效果相近、符合预期。

在 Context Parallelism 方面，理论上max_response_length_with_cp = max_response_length_without_cp * cp_size。在 4 张 B200、global_batch_size = 64的配置下验证结果如下：

即开启 CP=2 后，两种后端都能在 16k response length 下正常训练，实验结果符合预期。

数据打包与训推一致性

为消除 padding 的算力浪费，slime 实现了 Data Packing：按每条序列长度与max_tokens_per_gpu估算 pack（micro-batch）数量，用 Karmarkar-Karp（最大差分法）做负载均衡分配，将长短不一的序列分到不同 pack 中，每个 pack 内拼接为连续 tokens 向量并用cu_seqlens记录边界；训练时cu_seqlens直接传给 Flash Attention 处理变长序列，计算 loss 时unpack_sequences()按边界精确还原每条序列的指标。CP 模式下仅做最多cp_size - 1个 token 的最小对齐 padding，无可见 overhead。

在训推一致性方面，slime 在 FSDP 上实现了 True On-Policy：训练与推理均使用 FlashAttn3 作为 backend 实现 bitwise equal、使用 DeepGEMM 的 Batch-invariant Kernels 实现批次不变性，从系统层面将 training-infer logprob 绝对差收敛到 0（配套博客中有 KL=0 的实测图）。主流程算法若不启用该特性，则通过 TIS（Truncated Importance Sampling，tis = exp(old_log_probs - rollout_log_probs)截断后对pg_loss重加权）来减缓 off-policyness 对训练稳定性的影响。

总结

本文完整复现了 slime 上 FSDP 训练后端的搭建与测试流程：在 H 卡与 B 卡两种环境下，通过 Docker 容器、pip install -e .安装、hf download下载三件套资源、source scripts/models/*.sh加载模型配置，最终运行tests/test_qwen3-0.6B_fsdp_colocated_2xGPU.sh完成 2 GPU 协同训练验证。核心配置要点可归纳为：用--train-backend fsdp（"SLIME_BACKEND": "fsdp"）切换后端；用CUDA_VISIBLE_DEVICES与--actor-num-gpus-per-node控制 FSDP 分片规模；用--colocate开关在训推一体与训推分离之间切换；用--fsdp-full-params在FULL_STATE_DICT与默认SHARDED_STATE_DICT之间选择 checkpoint 模式。若遇到 NCCL 错误，可通过ray start --port显式指定端口解决。进阶阅读可继续参考本仓库的 FSDP2 训练后端介绍、FSDP 训练后端深度解析 与 权重更新机制解析。

• 文档
• 教程
• 人工智能
• 大模型
• RLHF