ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

8 张 A100 训完 13B 对话模型只要 9 小时:DeepSpeed-Chat 的 RLHF 成本账

8 张 A100 训完 13B 对话模型只要 9 小时:DeepSpeed-Chat 的 RLHF 成本账 8 张 A100 训完 13B 对话模型只要 9 小时DeepSpeed-Chat 的 RLHF 成本账【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed先给结论DeepSpeed-Chat 是 DeepSpeed 生态里跑 RLHF 训练的端到端方案把一个 HuggingFace 基座模型完整走完 InstructGPT 式三阶段产出可用的对话模型。它把最贵也最慢的 PPO 阶段第 3 步压进同一个引擎里同时做生成和训练。8 张 A100-80G 训完 OPT-13B 的第 3 步约 9 小时Azure 上账单约 $290一张 48G 消费卡两小时能出 1.3B 的 checkpoint。适合想从零训对话模型、但预算和显卡都有限的人。一、先算钱从单卡到 64 卡第 3 步到底贵在哪RLHF 训练里最烧钱的是第 3 步PPO 强化学习前两步加起来往往不到总耗时的一半。所以下面的账都只算第 3 步口径是 DeepSpeed 官方基准配方135M tokens 训一个 epoch其中 67.5M query tokens 67.5M 生成 tokens每步全局 batch 最多 0.5M tokens1024 组 query-answer 对各 256 长。换数据量或 batch 时这套数字不能直接套用。硬件配置模型规模第 3 步耗时约当成本Azure 估算1× A6000-48GOPT-1.3B1.2 小时端到端约 2.2 小时—8× A100-80G 单节点OPT-13B9 小时$2908× A100-80G 单节点OPT-66B2.1 天$162064× A100-80G 多节点OPT-30B4 小时$102464× A100-80G 多节点OPT-175B20 小时$5120对照一下同样在 8× A100-40G 上6.7B 要 5.7 小时30B 要 1.85 天换 80G 卡后 13B 才能压进 9 小时。可见显存比卡数更敏感——卡数不变显存翻倍就能把 13B 从跑不动变成一天内跑完。二、一条命令走完三阶段1.3B 到 66B 的跑法三阶段沿 InstructGPT 走SFT用人工精选问答对微调基座→奖励模型通常比基座小得多用同一 query 的多个答案排序数据训一个打分模型→RLHFPPO 算法拿奖励反馈继续微调。DeepSpeed-Chat 把两个常被别家省掉的开关内置了EMA 权重平均InstructGPT 的经验是 EMA checkpoint 回答质量更好和混合训练掺入下一词预测目标防止模型在 SQuAD2.0 这类公开基准上退化。跑法上先克隆训练示例仓库再执行train.pygit clone https://gitcode.com/GitHub_Trending/de/DeepSpeed cd DeepSpeedExamples/applications/DeepSpeed-Chat pip install deepspeed0.9.0 pip install -r requirements.txt # 8× A100 单节点训 OPT-13B约半天 python train.py --actor-model facebook/opt-13b --reward-model facebook/opt-350m --deployment-type single_node # 64 卡多节点训 OPT-66B约 9 小时 python train.py --actor-model facebook/opt-66b --reward-model facebook/opt-350m --deployment-type multi_node # 单张消费级卡训 OPT-1.3B约两小时 python train.py --actor-model facebook/opt-1.3b --reward-model facebook/opt-350m --deployment-type single_gpu只换--actor-model和--deployment-type就能覆盖从单卡到集群的所有规模。1.3B 在单张 A6000-48G 上的端到端拆法是SFT 约 2900 秒、奖励模型约 670 秒、RLHF 约 1.2 小时合计约 2.2 小时——真正确认最贵的是最后一步。13B 在 8× A100-40G 节点上则是 2.5 小时 0.25 小时 10.8 小时合计 13.6 小时。三、为什么生成阶段吃掉最多时间1.3B 模型单次 RLHF 迭代的耗时拆开来大头在生成 token训练更新反而是小头。原因不复杂生成阶段是带宽受限每条 256 token 的 prompt 要逐 token 生成 256 个新 token等于把模型前向推几百次每次的算力利用率都很低训练阶段是计算受限参考模型对 512 tokenprompt生成样本做一两次前向/反向GPU 吃得很满。计算量占比上生成只占约两成、训练占八成但墙钟时间反过来。现有 RLHF 系统普遍卡在这里官方对比显示它们在单卡上的 Step 3 吞吐不到峰值的 5%而 DeepSpeed-Chat 在该阶段相对 HuggingFace 方案最高 9 倍吞吐、相对 Colossal-AI 15 倍端到端多卡上则是 6–19 倍对 Colossal-AI和 1.4–10.5 倍对 HuggingFace DDP。模型能跑的上限差距更大同样的单卡 A100-40G 和单节点配置Colossal-AI 到 1.3B / 6.7B 就 OOMDeepSpeed-Chat 能到 6.5B / 50B。四、Hybrid Engine同一个引擎两套执行路径针对生成是瓶颈这个问题Hybrid Engine 的做法是让一个引擎实例同时具备训练和推理两条快路径生成与训练之间反复横跳。落到仓库代码里deepspeed/runtime/hybrid_engine.py的DeepSpeedHybridEngine继承自标准DeepSpeedEngine生成时换内核create_inference_containers()按inference_policies把 Transformer、Linear、Embedding、LayerNorm 等层替换成 DeepSpeed 推理容器并按inference_tp_size张量并行切权重——推理切分方式和训练侧的 ZeRO 分片互不干扰模型没有匹配策略时只打警告回退原生generate()路径LoRA 进出非 ZeRO-3 下生成前fuse_lora_weight()把 LoRA 权重并入推理容器生成完unfuse_lora_weight()还原训练侧拿到的还是独立参数ZeRO-3 下的参数搬运generate()里用GatheredParameters按tp_gather_partition_size默认 8 层分批 gather 非驻留参数避免一次性拉爆显存显存复用开release_inference_cache后retake_inference_cache()重新申请 KV-Cache workspace生成结束workspace.release_workspace()释放并gc.collect()把显存还给训练阶段CUDA Graph开enable_cuda_graph时用DecodeGraphCachehybrid_engine_graph.py缓存 decode 步省掉重复的启动开销构建前会校验当前 ZeRO 阶段是否支持。五、一张卡到底能训多大单卡上限是很多人第一个问题答案是显存决定官方实测V100-32G 到 OPT-2.7BA6000-48G 或 A100-40G 到 OPT-6.7BA100-80G 到 OPT-13B。也就是说没有多卡环境的人也能训出可用而不只是玩具级别的模型。两个边界要注意。第一175B 单卡无解64× A100-80G 下 20 小时约 $512066B 在 64 卡上约 9 小时其中 SFT 82 分钟、奖励模型 5 分钟、RLHF 7.5 小时。第二卡越多越快只在拐点之前成立小规模下 ZeRO 把模型状态摊薄后单卡 batch 可以更大出现超线性加速规模上去后最大全局 batch本配方 1024 组封顶了单卡 batch曲线转为近线性甚至次线性。所以选卡数的依据不是越大越好而是让单卡 batch 尽量接近这个上限的那个点。六、跑通必用的 5 个配置项和几条边界条件hybrid_engine配置块定义在deepspeed/runtime/config.py的HybridEngineConfig写进 DeepSpeed JSON 即可。真正需要动的只有这几个配置项默认说明enabledfalse总开关max_out_tokens512生成最大长度也作为推理容器 workspace 的申请依据inference_tp_size1生成阶段张量并行规模大于 1 且 ZeRO-3 时启用分批 gatherrelease_inference_cachefalse生成后释放 workspace缓解训练/推理交替的显存峰值enable_cuda_graphfalsedecode 阶段 CUDA Graph 缓存最小配置样例在tests/hybrid_engine/hybrid_engine_config.jsontrain_batch_size 32、micro batch 2、ZeRO stage 0 参数 offload 到 CPU、fp16 开、梯度裁剪 1.0。几条边界条件照着对上文所有耗时/成本数字都针对第 3 步且基于 135M tokens 单 epoch 的基准配方跨规模对比时先核口径pin_parametersZeRO-3 下生成前是否 gather 全部非 TP 层默认 true显存特别紧张时才动仓库当前实现比 2023 年博客演进过如 CUDA Graph字段行为以deepspeed/runtime/hybrid_engine.py和tests/hybrid_engine/为准。七、下一步看哪里blogs/deepspeed-chat/README.md官方技术博客全文Figure 3–7 的吞吐与扩展性曲线都在里面tests/hybrid_engine/hybrid_engine_test.py与tests/hybrid_engine/hybrid_engine_config.jsonHybrid Engine 的端到端测试和最小可跑配置deepspeed/runtime/hybrid_engine.py的generate()方法LoRA 融合/还原、ZeRO-3 分批 gather、workspace 回收的实现都集中在这里。【免费下载链接】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),仅供参考
返回列表