尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

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

slime FSDP 训练后端搭建与测试实战:从 Docker 环境到 Colocated/Disaggregated 部署 文档教程人工智能大模型RLHF【免费下载链接】Awesome-ML-SYS-TutorialMy learning notes for ML SYS.项目地址https://gitcode.com/gh_mirrors/aw/Awesome-ML-SYS-Tutorial点击查看免费下载本篇技术指南以 slime基于 SGLang 推理与 PyTorch 训练后端的 RL 框架的 FSDP 训练后端为主线完整记录从拉取 Docker 镜像、安装 slime、下载模型与数据集到加载模型配置、运行 FSDP 双卡训练测试脚本的端到端流程并深入讲解 Colocated训推一体化与 Disaggregated训推分离两种资源编排方式、SLIME_BACKENDfsdp后端激活与 GPU 分片机制以及 BlackwellB 卡环境下的适配要点。读完本文你将能够独立复现 slime 上 FSDP 后端的搭建与测试理解训练/推理卡数分配的核心参数并掌握从 Megatron 后端迁移到 FSDP 的关键差异。本文对应仓库中的 setup_fsdp.mdFSDP 发布配套的搭建与测试文档并辅以 FSDP2 训练后端介绍 中的架构与参数细节进行纵深展开。背景为什么要在 slime 上测试 FSDPslime 在原有基于 Megatron-LM 的训练后端之外新增了FSDPFully Sharded Data Parallel作为训练后端并在精度上完成了与 Megatron 的对齐。文档中提到的 FSDP 均指 PyTorch 原生支持的FSDP2相比 FSDP1 将所有参数摊平为巨型FlatParameterFSDP2 引入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 StartH 卡环境下的完整搭建流程以下操作在 H 卡Hopper 架构上完成从零开始共分五步拉取并启动容器、安装 slime、下载模型与数据集、加载模型配置、运行测试脚本。拉取并启动 Docker 容器# 拉取最新镜像 # 最新的镜像是 B 卡 H 卡通用的 docker pull slimerl/slime:latest # 启动容器 docker run -d --gpus all --ipchost --shm-size16g \ --name slime_wren_fsdp \ -it slimerl/slime:latest /bin/bash对启动命令中关键参数的理解参数作用--gpus all将宿主机全部 GPU 暴露给容器供 Ray 调度与训练/推理使用--ipchost共享宿主机 IPC 命名空间避免进程间通信如共享内存队列受限--shm-size16g将/dev/shm设为 16GB满足 PyTorch DataLoader 多进程与分布式训练对共享内存的需求--name slime_wren_fsdp指定容器名可按需修改后续进入容器时使用该名称-it ... /bin/bash以交互模式启动并进入 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新版 CLIpip 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 0Rollout 与 Actor 完全共用 bundle而默认分支下num_gpus actor_num_nodes * actor_num_gpus_per_node rollout_num_gpusrollout_offset actor_num_nodes * actor_num_gpus_per_node两者使用不同的 bundle。异步训练train_async.py仅在 Disaggregated 架构下有意义通过rollout_manager.async_generate()与actor_model.async_train让 rollout 始终领先 train 一个 stepone-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_DEVICES1,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避免峰值显存冲击。BlackwellB 卡GPU 设置启动容器B 卡Blackwell 架构与 H 卡的操作步骤几乎完全一致主要差异体现在容器启动参数这些是环境配置差异不是硬件差异# 拉取最新镜像最新的镜像是 B 卡 H 卡通用的 docker pull slimerl/slime:latest # 启动容器 # 这里 GPU 相关参数完全相同主要差异是镜像版本和挂载目录这些是环境配置不是硬件差异 docker run \ -itd \ --shm-size 32g \ --gpus all \ --ipchost \ --networkhost \ --privileged \ -v {your_cache_path}:/root/.cache \ --name slime_fsdp_{your_name} \ slimerl/slime:latest \ /bin/bash与 H 卡启动命令的差异说明--shm-size 32g共享内存提升到 32GB适配 B 卡更大的 batch 与更长的序列--networkhost直接使用宿主机网络栈便于 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-host0.0.0.0 --dashboard-port8265 --port 9987--port 9987显式指定 Ray head 的 GCS 服务端口避免多实例或多节点环境下端口冲突${MASTER_ADDR}为当前节点的 IP 地址需要在实际环境中替换--num-gpus 8声明该节点可用的 GPU 数量需与docker run --gpus all暴露的数量一致--disable-usage-stats关闭遥测上报--dashboard-host0.0.0.0 --dashboard-port8265允许从外部访问 Ray Dashboard。从 Megatron 迁移到 FSDP参数对照与能力边界当你把一套已有的 Megatron 训练脚本迁移到 FSDP 时可以参考下表源自 FSDP2 训练后端介绍快速定位需要替换的参数配置类别Megatron 参数FSDP 参数说明模型加载--load(Megatron checkpoint) 架构参数 (--num-layers,--hidden-size等) 或--use-hf-config-for-megatron--hf-checkpoint(必需)FSDP直接使用 HuggingFace 格式无需转换权重通过AutoConfig自动推断架构张量并行--tensor-model-parallel-sizeComing Soon流水线并行--pipeline-model-parallel-sizeComing Soon专家并行--expert-model-parallel-sizeComing Soon上下文并行--context-parallel-size--context-parallel-size两者都支持 CP初始学习率--lr--lr参数相同学习率衰减--lr-decay-style(linear/cosine)--lr-decay-style(仅 constant)Warmup--lr-warmup-iters(步数)Coming Soon最小学习率--min-lrComing Soon优化器类型--optimizer(adam/sgd 等)--optimizer(默认 adam)基本相同分布式优化器--use-distributed-optimizer内置于 FSDPFSDP 默认使用分布式优化器梯度检查点--recompute-granularity,--recompute-method--gradient-checkpointingFSDP简化为布尔开关CPU Offload通过分布式优化器实现--fsdp-cpu-offloadFSDP将参数/梯度/优化器状态卸载到 CPUAttention 后端由 Megatron Core 决定--attn-implementation(flash_attention_2/sdpa/eager)FSDP直接透传给 HuggingFace混合精度--fp16或--bf16--fp16(bf16 自动推断)基本相同保存时 Offload---fsdp-state-dict-cpu-offload(默认 True)FSDP保存 checkpoint 时 offload 到 CPU训练后端默认或--train-backend megatron--train-backend fsdp(必需)用于切换后端FSDP 目前不支持/不完善的功能迁移前还需要明确能力边界并行策略FSDP 目前仅支持DP CP不支持TP、EP、PPCP 实现差异Megatron Core 的 CP 是原生实现与 TP/PP 深度集成而 FSDP 通过外部库 Ring Flash Attention 实现连续 chunk 切分 负载均衡交给 ring flash attnslime 只需关注输入切分与结果聚合重计算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 preparationprocess_rollout_data取当前 DP rank 所需数据_pack_rollout_data打包为packed_batches消除 padding 损耗→ forward log prob计算 Actor/Ref 的 log_prob 与 entropy→ loss calculationPPO/GRPO loss importance ratio/clip/KL penalty/entropy bonus实时计算train_rollout_logprob_abs_diff监控训推数值偏差可选启用 TIS 重加权→ update offload梯度累积与参数更新训练结束调用sleep将模型与优化器 offload 到 CPU。显存优化offload 策略FSDP 训练流程通过以下场景的权重 offload 节省显存Train offloadColocated 场景下训练完成后调用sleep将模型 weight 与 optimizer offload 到 CPU避免在 rollout 阶段占用 GPU 显存Ref model使用 KL penalty 时reference model 仅在compute_log_prob时加载到 GPU计算完成后立即 offload 回 CPUOptimizer offload训练阶段将不参与计算的 model parameter 与 gradient 都 offload 到 CPU显著节省训练显存代价是 optimizer step 在 CPU 上执行训练时间会明显上升对应--fsdp-cpu-offload。精度对齐验证与 Context ParallelFSDP 与 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的配置下验证结果如下response_length 8kresponse_length 16kFSDP, cp 1workOOMFSDP, cp 2workworkMegatron(TP 1), cp 1workOOMMegatron(TP 1), cp 2workwork即开启 CP2 后两种后端都能在 16k response length 下正常训练实验结果符合预期。数据打包与训推一致性为消除 padding 的算力浪费slime 实现了 Data Packing按每条序列长度与max_tokens_per_gpu估算 packmicro-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配套博客中有 KL0 的实测图。主流程算法若不启用该特性则通过 TISTruncated Importance Samplingtis 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 fsdpSLIME_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【免费下载链接】Awesome-ML-SYS-TutorialMy learning notes for ML SYS.项目地址https://gitcode.com/gh_mirrors/aw/Awesome-ML-SYS-Tutorial点击查看免费下载相关推荐slime 8xH100 训练 GLM4-9B 实战从环境搭建到 dynamic sampling 进阶采样slime 8xH100 训练 GLM4 9B 实战从环境搭建到 dynamic sampling 进阶采样 本篇技术指南围绕 slime 官方示例 scri人工智能大模型强化学习RLHF分布式训练Megatron-LM Docker部署容器化训练环境搭建Megatron LM Docker部署容器化训练环境搭建 概述 还在为大规模Transformer模型训练的环境配置而头疼吗复杂的依赖关系、版本冲突、CU人工智能大模型预训练分布式训练深度学习强化学习数据库变更管理实战指南如何用Liquibase像管理代码一样管好你的库结构数据库变更管理实战指南如何用Liquibase像管理代码一样管好你的库结构 你的数据库结构是不是已经失控了开发环境加了一个字段、测试库改了字段类型、生产库却数据库开发工具后端上一篇Protobuf 打包与分发体系深入解读 protobuf pkg 目录的 Bazel 打包规则下一篇Karpenter v1 FAQ 深度解析从 NodePool 供给到调度、中断与升级的实战指南创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表