DeepSpeed-Chat单卡训 13B RLHF 模型1 条命令跑完三阶段【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeedDeepSpeed-Chat 让 RLHF基于人类反馈的强化学习三阶段训练可以一条命令跑完单张 A100-80G 就能训 13B 模型。读完你能带走三样东西一条可直接复制的 train.py 命令、hybrid_engine 配置块的全部 7 个字段及默认值、一组实测数字8×A100-80G 上 9 小时训完 OPT-13BAzure 成本约 $290来源官方博客 Table 1。核心引擎叫 Hybrid Engine同一个模型在训练和生成两条执行路径间切换。三层结构数据、训练和一个会换路的引擎整个系统分三层。数据层统一不同数据源的格式再切分混合进三个阶段训练层按 InstructGPT 流程走Step 1 监督微调SFT、Step 2 奖励模型微调、Step 3 用 PPO一种限制策略更新幅度的强化学习算法微调 actor 模型另附 EMA checkpoint 与混合训练两个可选特性。前两步是普通微调由 ZeRO把优化器状态切到多卡以省显存驱动第三步每次迭代要做两件事——先生成回答、再训练 actor 和奖励模型。Hybrid Engine 专门接管第三步训练时它就是普通 DeepSpeed 引擎生成时切换成推理内核像同一台发动机在烧油和发电两种模式间切换不需要重建模型或搬运权重。最小路径2 小时跑通 1.3B 示例第一步克隆本仓库并装依赖约 5 分钟任意 CPU 机器即可git clone https://gitcode.com/GitHub_Trending/de/DeepSpeed.git cd DeepSpeed pip install -r requirements.txt python setup.py install装完应看到deepspeed --version有正常版本号输出。第二步先跑仓库自带的 Hybrid Engine 测试验证训练—生成切换可用单卡 10 分钟内结束跑完应看到1 passedpython -m pytest tests/hybrid_engine/hybrid_engine_test.py第三步正式训练。train.py 示例脚本位于官方 DeepSpeedExamples 仓库的 DeepSpeed-Chat 应用目录克隆后装依赖单卡 48G 显存约 2.2 小时出 checkpoint来源官方博客 Table 6python train.py --actor-model facebook/opt-1.3b \ --reward-model facebook/opt-350m --deployment-type single_gpu跑完后用推理 API 交互多轮对话输出长这样Human: Do you know Microsoft? Assistant: Yes, Im a big fan of Microsoft. What about them? Human: Can you explain it to a 6-year old child? Assistant: Sure... Microsoft is a company that makes computers, and they make a program called Windows...不同规模与资源的需求来源官方博客 Table 1/4/6模型硬件三阶段总耗时OPT-1.3B单张 A6000 48G约 2.2 小时OPT-6.7B8×A100-80G4.1 小时$132OPT-13B8×A100-80G13.6 小时Step 3 占 10.8 小时OPT-66B64×A100-80G约 9 小时想自定义训练策略时示例仓库给出的 API 入口只有两个动作engine DeepSpeedRLHFEngine(actor_model_name_or_pathpath, tokenizertokenizer, argsargs) trainer DeepSpeedPPOTrainer(engineengine, argsargs) for prompt_batch in prompt_dataloader: out trainer.generate_experience(prompt_batch) # 生成阶段推理加速路径 actor_loss, critic_loss trainer.train_rlhf(out) # 训练阶段PPO 更新该 API 在 DeepSpeedExamples 的 DeepSpeed-Chat 应用中本仓库提供其背后的引擎实现。Hybrid Engine 里值得看的四处实现generate()同一模型上训练与生成的切换开关它做什么同一个引擎实例在不重建模型、不复制权重的前提下切换执行路径。怎么做到的初始化时预建推理容器并保存原始 forwardgenerate()按 ZeRO 阶段把分片参数 gather 进显存、融合 LoRA、走推理内核结束后解融合并回收推理 workspace。源码落点deepspeed/runtime/hybrid_engine.pyDeepSpeedHybridEngine.generate类定义在 L31-L75。这段代码的关键行为是L243-317 是 ZeRO-3 分支GatheredParameters逐层 gather 参数后推理L318-329 是非 ZeRO-3 分支直接 fuse LoRA → 生成 → unfuseL332-335 在release_inference_cache开启时调用workspace.release_workspace()把显存还回训练阶段。一个具体例子配置为 ZeRO-3 且inference_tp_size2时生成前引擎按每 8 层一组tp_gather_partition_size默认 8gather 非 TP 参数生成后对每层推理容器调用release_memory()L297-298释放 KV-Cache 与中间结果占用的显存。容易踩的坑LoRA 的融合与还原必须配对。ZeRO-3 下还原走的是unfuse_lora_weight_non_pinned()L170-177只对已 gather 的参数逐层 unfuse若误用非 pin 路径会残留未还原的参数。hybrid_engine配置块的全部字段定义在 deepspeed/runtime/config.py可直接写进 DeepSpeed JSON 配置仓库样例见 tests/hybrid_engine/hybrid_engine_config.json字段类型/默认值实际作用enabledbool /False开启 Hybrid Enginemax_out_tokensint /512生成最大长度同时作为推理容器的 workspace 上限inference_tp_sizeint /1推理阶段张量并行规模1时按此规模分组release_inference_cachebool /False生成后释放 workspace训练前重新申请pin_parametersbool /TrueZeRO-3 下生成前 gather 全部非 TP 层参数tp_gather_partition_sizeint /8ZeRO-3 TP 场景按每 8 层分组 gather 的步长enable_cuda_graphbool /False启用 decode 阶段 CUDA Graph 缓存推理容器与张量并行生成阶段为什么快它做什么把训练层的 forward 换成带 KV-Cache 的推理内核让内存带宽受限的生成阶段跑得更快。怎么做到的create_inference_containers()递归遍历模型子模块命中策略的层nn.Linear、nn.Embedding、nn.LayerNorm、OPT 位置编码等替换为推理容器并记录原始 forwardinference_tp_size 1时create_inference_module()把全部 rank 按 tp_size 张卡一组建 TP 组每组持有完整模型batch 在各组间用 all_gather 分配。源码落点deepspeed/runtime/hybrid_engine.py。关键行为是 L386-394 用global_rank // inference_tp_size算出组号、dist.new_group(ranks)建组生成时 L279-289 把输入张量扩inference_tp_size倍并跨组 all_gather即每个 TP 组只处理全局 batch 的一部分。一个具体例子8 卡 inference_tp_size2得到 4 组双卡每组承担 1/4 的 batch。此时生成吞吐取决于组内 TP 通信而非 ZeRO 的梯度 allreduce——博客实测生成阶段相对 HuggingFace 最高 9 倍、相对 Colossal-AI 15 倍来源官方博客 Figure 5。CUDA Graph把 decode 步骤录下来重放它做什么消除每个 decode 步骤重复的 kernel launch 开销。怎么做到的enable_cuda_graph开启后初始化时用DecodeGraphCache包裹原始 forward并先校验 ZeRO 阶段兼容性generate()入口调用begin_sequence()绑定本次 decode 步数图被录制后按步重放。源码落点deepspeed/runtime/hybrid_engine.py缓存本体在 deepspeed/runtime/hybrid_engine_graph.py。关键行为是 L75 用max_out_tokens作为 decode 步数上限建缓存L240-241 在每次生成前重置序列校验不通过时 L70-71 打印警告并自动退回无图路径不会报错中断。一个具体例子生成 256 token 的回答时max_out_tokens512即可让图缓存命中若把生成长度调到 1024缓存录制成本上升且 ZeRO-3 等不支持的阶段会看到 running without CUDA graphs 警告。调优提示CUDA Graph 只覆盖 decode 阶段开启后先看日志确认没有触发警告回退否则白录一遍图。边界与选型单卡上限 13B64 卡到 175B天花板。单卡支持的最大模型A100-80G 为 OPT-13BA100-40G 为 OPT-6.7BV100 32G 只有 2.7B来源官方博客 Table 3。64×A100-80G 训 175B 需 20 小时约 $5120但单卡吞吐回落到 1.3B 模型的 1.2 倍原因是显存撑不起大 batch来源官方博客 Figure 6。以上数据均基于 2023 年基准135M tokens、单 epoch、每步全局 batch 0.5M tokens做成本对比时以此规格为准。怎么调。三个影响最大的旋钮inference_tp_size模型 6.7B 及以下用 130B 以上设 2 或 4。生成阶段是内存带宽受限的TP 通信量远小于 ZeRO allreduce比加卡分片更划算。release_inference_cache训练显存紧张时设true。代价是每次生成前重新申请 workspacedeepspeed/runtime/hybrid_engine.py 里有一次失败重试逻辑申请不到会直接抛RuntimeError。max_out_tokens对齐实际回答长度。它同时决定推理 workspace 和 CUDA Graph 缓存上限设小会截断输出设大则浪费显存。选型判断。当你做 OPT 系模型的 PPO 训练、需要 ZeRO LoRA 且希望一个引擎同时负责训练与生成时选 DeepSpeed-Chat 的 Hybrid Engine当你要频繁更换 RL 算法或模型家族时考虑 HuggingFace 生态的 RLHF 工具链——本仓库的推理容器策略按有限模型结构匹配见create_inference_containers的策略表匹配逻辑非覆盖范围回退到模型原生生成路径当你需要 66B 以上规模时直接按 64 卡 8 节点规划单节点跑不了。下一步建议先跑通tests/hybrid_engine/hybrid_engine_test.py然后在自己的迭代里观察generate()记录的_generate_latency与_gather_latencyL337-338确认你的瓶颈在生成还是参数 gather再决定动哪个旋钮。【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考