昇腾910B上vLLM-ascend与TRL实现GRPO强化学习训练实践
发布时间:2026/10/19 2:32:12 锦皓数字建站

如果你最近在做大模型的强化学习对齐训练大概率听过“采样瓶颈”这个词。推理场景一次生成几十个token不觉得慢但强化学习里每个训练步都要为一整批prompt各自采样多个response吞吐一上来单靠模型forward硬顶根本扛不住。社区的成熟做法是“训练靠TRL采样靠vLLM”而GRPO这类组内相对奖励的算法又节省了价值网络的显存开销让整条链路在有限算力下更容易落地。这篇博文就是记录我在8卡昇腾910B上把vLLM-ascend和TRL里的GPRO训练完整串起来的全过程包括版本匹配、服务部署、训练接入以及三个折腾到凌晨的性能与稳定性问题。用GPRO还是GRPO只是缩写习惯差异社区里两种写法都在用指的是同一个算法Group Relative Policy Optimization。标题里的GPRO正文统一用GRPO说明但核心思路完全一致。1. 为什么是vLLM-ascend TRL GRPO选型思路与适用场景先说背景。整套项目目标很朴素对某个开源底座模型做数学推理强化学习微调训练数据规模不大但每个prompt需要在线采样多组答案再按答案正确性给奖励。最初选型时我纠结过两条路一条是用训练框架自带的rollout生成能力另一条是把推理引擎独立出来作为采样服务。最后选了后者理由下面会详细说。1.1 这套组合到底解决的是什么问题强化学习和SFT的本质差异是SFT只要“照着标注学”强化学习要“自己生成、被奖励、再更新”。于是每一步都包含三个模型的工作量策略模型生成response、参考模型算KL散度、奖励函数或奖励模型打分。GRPO和PPO最大的区别是不再单独训练一个critic价值网络而是对同一个prompt采样的G条response做组内相对比较。具体做法是每个prompt采样G个输出分别拿到奖励r1到rG然后算组内均值与标准差把每个奖励归一化成advantage。策略更新的方向是让组内高于平均水平的输出概率上升低于平均水平的下降。这个思路特别适合奖励尺度不稳定、或奖励函数本身是规则计算的任务因为它不看绝对分数只看组内相对位置。“小组排名”比“绝对分数线”更抗噪声这是我第一次跑通后体会最深的一点。在显存上GRPO省掉价值网络这一整块负担对NPU这种整体显存规划更“刀尖上跳舞”的环境来说等于直接省了2到4GB的权重和优化器状态可以把budget分给采样batch或上下文长度。这也是我最终确定“vLLM-ascend 独立采样服务 TRL训练端”三角架构的根本原因。1.2 为什么采样必须走vLLM-ascend服务而不是训练进程内生成训练进程内不是不能生成是效果太差。策略模型每更新一步下一轮rollout必须用最新权重重新生成响应。如果在训练进程里用普通forward做beam或multinomial采样一次生成一批G4的response显存峰值可能是训练本身的1.5倍以上而且和训练的前向反向互相抢占算子单元整体吞吐低得没法看。vLLM-ascend的价值在于它把推理部分从“训练进程”中拆出去了而且做了连续批处理、分页KV cache、以及针对昇腾的图模式与算子融合。训练进程只需要把prompt发过去换取completion回来。这里最关键的分工是vLLM-ascend接管全部采样吞吐TRL只负责策略更新和loss计算。两者通过HTTP协议解耦以后才能把训练卡的显存完整留给训练把另外几张卡上的显存完整留给推理。1.3 什么样的项目适合照抄这条路如果你现在的任务是数学、代码、逻辑推理类的强化学习微调奖励以规则判断为主数据量在十万级以内且手里只有昇腾设备那这条路线基本可以直接复用。如果你的任务需要复杂的奖励模型且奖励模型本身很大建议先把奖励模型量化或者拆分再进这条链路。如果只是实验性质、数据量很小、没有吞吐要求那没有必要搭独立vLLM服务直接TRL默认内部采样就能跑。独立服务在显存利用、并发吞吐上的优势必须在大批量训练时才有明显回报。2. 昇腾NPU环境搭建版本矩阵、CANN安装与vLLM-ascend编译昇腾这套东西和NVIDIA生态最大的区别就是“版本匹配极严”。在N卡上你装一个torch 2.1再pip install vllm多半能跑昇腾上不行。驱动版本、CANN版本、torch_npu的wheel版本、vllm-ascend的编译期要求是一整条必须同时锁定的链条。任何一处漂移都会在编译期或运行期指出与表面完全无关的错误。2.1 先把版本组合锁死再谈其他我这次的环境用了某公司提供的昇腾910B设备基础环境是Python 3.10、torch 2.1.0、torch_npu对应的社区wheel包。CANN使用的是8.0系列这是当前昇腾社区较稳定的分支。vLLM-ascend版本选的是0.7系列TRL用支持GRPOTrainer的0.13左右版本transformers要高于4.46否则TRL内部的模型加载逻辑可能不兼容。必须强调不要看着我上面这一串版本号就直接照抄。昇腾社区的兼容矩阵一直在更新更稳妥的做法是到昇腾社区查询你手上固件版本对应的CANN版本根据CANN版本选择匹配的torch_npu wheel再根据torch_npu编译依赖确定vllm-ascend源码版本。如果中间某一步选了“比文档更新”的版本风险反而是最高的。我见过太多Demo先跑通、一进多卡就炸的例子根因全是版本链不齐。2.2 安装中的关键步骤与首个验证程序昇腾侧安装核心流程大致如下# 1. 安装/校验CANN toolkit并确保环境变量已经加载 source /usr/local/Ascend/ascend-toolkit/set_env.sh # 2. 创建conda环境Python 3.10 conda create -n grpo_ascend python3.10 -y conda activate grpo_ascend # 3. 安装torch与torch_nputorch_npu必须使用与CANN匹配的whl包 pip install torch2.1.0 pip install torch_npu2.1.0.post6 # 具体版本以兼容矩阵为准 # 4. 安装vllm-ascend pip install vllm-ascend0.7.x装完后不要急着接训练先用离线生成验证NPU路径是否走通from vllm import LLM, SamplingParams llm LLM( model/path/to/policy-model, devicenpu, tensor_parallel_size2, dtypebfloat16, ) outputs llm.generate( [请计算(23)*4?], SamplingParams(max_tokens64, temperature0.7), ) for item in outputs: print(item.outputs[0].text)如果这一步能正常返回说明CANN、torch_npu、vllm-ascend三层关系没问题。此时再推TRL接入心里就有底了。第一次跑的时候我卡在“torch_npu initialization failed”上翻半天日志才发现是/usr/local/Ascend/ascend-toolkit/latest/lib64没有加入LD_LIBRARY_PATH。昇腾的so库查找路径跟CUDA习惯不一样少加一个目录表现就是运行期符号找不到报错却只在初始化阶段。建议每个终端都先执行source set_env.sh再确认python -c import torch_npu不报错。2.3 多卡通信HCCL验证训练一旦涉及多卡HCCL就是昇腾版的NCCL。GRPO训练里梯度同步每次都会触发全量reduce通信不稳定会直接表现为训练卡住或迭代速度抖动。最基础的自检方式是用一个小脚本在所有可见NPU上做allreduceexport ASCEND_RT_VISIBLE_DEVICES0,1,2,3 python -m torch.distributed.run --nproc_per_node4 \ /path/to/allreduce_check.py脚本里用一个固定的张量在4卡间做sum确认输出符合预期即可。这一步的意义不是性能测试而是提前暴露网卡、路由、共享内存等问题避免训练到一半才出现“HCCL connection timeout”。我在第一次多卡训练时没有做这个自检直接全量启动结果跑到第300多步才出现通信超时。排查了一圈发现梯度过大时某块卡的reduce比其他卡慢整体被拖死。虽然最后通过调大HCCL_CONNECT_TIMEOUT解决了但如果提前做过通信自检至少会少消耗一个晚上的时间。3. vLLM-ascend 推理服务与 TRL rollouts 链路接入真正把推理服务和训练接起来是这套环境里最需要理解“接口契约”的地方。昇腾设备上的TRL并不会自动识别NPU环境尤其当它尝试自己拉起vLLM进程时内部大量默认逻辑都会假设CUDA可用。先明确我的接法训练端TRL只管发起HTTP请求推理端我预先启动好vLLM-ascend服务。TRL把prompt通过OpenAI兼容协议发给服务服务返回completion训练端再计算reward、KL散度、策略损失。这比让TRL自己去管理vLLM进程要稳得多。3.1 首先把vLLM-ascend服务跑成一个纯推理API服务启动命令如下vllm serve /path/to/policy-model \ --served-model-name policy-7b \ --port 8001 \ --max-model-len 2048 \ --gpu-memory-utilization 0.85 \ --tensor-parallel-size 4 \ --dtype bfloat16我特意把--tensor-parallel-size设成4因为我计划用4张卡跑推理剩下4张跑训练。这个分卡策略是后面性能调优的关键一开始我等价于“训练4卡推理4卡”的隔离方案避免了显存互抢。启动后用一条curl验证curl http://127.0.0.1:8001/v1/chat/completions \ -H Content-Type: application/json \ -d {model: policy-7b, messages: [{role: user, content: 你好}], max_tokens: 32}如果返回正常服务侧就绪。值得注意的是--served-model-name和调用时的model字段必须完全一致我在调试过程中因为名称不匹配浪费过不少时间404或400错误一直以为端口没通。3.2 TRL侧关键配置GRPOConfig接入训练时TRL的GRPOTrainer是核心入口。下面是我实际跑通训练的配置骨架from trl import GRPOConfig, GRPOTrainer training_args GRPOConfig( output_dir./grpo_ascend_out, per_device_train_batch_size2, gradient_accumulation_steps4, num_generations4, learning_rate1e-6, beta0.01, # KL散度惩罚系数 epsilon0.2, # 策略裁剪范围 max_prompt_length512, max_completion_length512, use_vllmTrue, # 开启vLLM采样模式 vllm_server_host127.0.0.1, vllm_server_port8001, gradient_checkpointingTrue, logging_steps1, save_steps200, ) trainer GRPOTrainer( model/path/to/policy-model, reward_funcs[accuracy_reward, format_reward], argstraining_args, train_datasetdataset, ) trainer.train()这里每个参数的取舍都值得说。per_device_train_batch_size2和gradient_accumulation_steps4配合相当于每步等效batch是8然后num_generations4会把每个prompt扩展成4条completion等效训练样本量再翻4倍。如果把per_device直接拉到8虽然内存足够但梯度噪声变大训练反而更抖。use_vllmTrue但对TRL来说它只是一个HTTP客户端并不负责启动进程。这正是绕开昇腾环境兼容问题的核心决策。TRL需要vLLM服务时它会按vllm_server_host与vllm_server_port去请求而不会去检查本机是否安装了某个特定vLLM入口。换句话说只要服务端能返回文本训练端根本不关心这块算力是N卡还是NPU。3.3 Reward函数与Reference模型的配合细节GRPO训练如果完全没有reward字段TRL会用reward_funcs逐个对completion打分。需要注意的是不同版本TRL的reward函数签名有差异。我这一版是(prompts, completions, **kwargs)但我后来在另一台机器上遇到过返回参数多了answers的情况建议安装后先看接口签名再写。我用的是两个基础函数一个检查答案是否正确一个检查输出是否遵守指定格式。格式分给0.2正确性分给1.0。为什么要拆开因为格式错误可能纯粹是模型输出不规范不能跟“答错了”混为一谈否则reward信号被拖垮。Reference模型方面GRPOTrainer加载同名权重作为参考模型用来算KL惩罚。注意在昇腾上参考模型和策略模型会同时占用显存所以务必打开梯度检查点。我见过有人只开gradient_checkpointing忘了per_device_train_batch_size结果显存直接溢出得不偿失。跑正式训练前一定要做小batch冒烟测试。我一般取16条样本per_device_train_batch_size1num_generations2跑两三个step确认loss能正常下降、vLLM服务的日志有请求进来。这一步如果不过后面调参时所有问题都会混在一起非常难定位。4. GPRO训练性能瓶颈分卡策略、长尾样本与稳定性调节这部分是全文最实战的部分。从第一轮跑通全量训练到最终达到可接受的吞吐与稳定loss我经历了三轮明显的性能优化。4.1 第一轮瓶颈推理请求并没有跑在全部4张推理卡上全量训练启动后我习惯性地用npu-smi info watch实时盯卡。结果发现一个非常反直觉的现象训练侧4张卡利用率都在80%以上但推理侧4张卡里只有第一张卡是满的其余三张只有10%到20%的波动。我顺着链路排查先看vLLM服务的启动参数确认--tensor-parallel-size4确实生效服务日志里也明确打印了使用4卡再用curl并发压测服务自身吞吐发现服务端自己吃得了并发最后看TRL发出的请求速率发现它根本没把并发打满。结论是TRL最多只有一个rollout线程在按顺序请求vLLM服务导致推理侧除了第一张卡在勤恳工作外其余卡都在“陪跑”。既然服务端支持并发那就让训练端把请求并发度提起来。我调整了TRL侧的数据预取逻辑让每次请求的batch提前组装好、不等前一个step完全结束再发推理侧三张卡利用率很快提高到60%以上。这一轮之后单步耗时从原来的28秒左右降到20秒上下rollout总吞吐提升约40%。4.2 第二轮瓶颈长尾样本拖垮prefill阶段吞吐上来以后我发现一个新问题单步耗时很不稳定上一轮20秒下一轮可能直接跳到45秒。翻看几十个step的耗时分布发现瓶颈集中在prefill阶段。为什么因为数学推理数据里有相当一部分样本答案很长个别样本长度接近512的completion上限。当这批长样本混在一个batch里进入NPU时prefill计算量突然涨一截整个batch都要等它算完表现出来就是单步耗时剧烈抖动。解决思路不是限制所有样本长度而是把长样本隔离到单独的batch。实际操作上并没有专门做长样本分类而是通过两个手段缓解把max_prompt_length从1024压到512prompt大多用不上1024属于白白浪费把max_completion_length从768压到512训练目标本身也不要求过长输出vLLM服务侧调整max_num_batched_tokens让prefill阶段能够切分处理。改完以后单步耗时的方差明显变小稳定在15到18秒之间。4.3 第三轮调整GRPO训练稳定性的参数面调试吞吐稳定之后训练曲线又出问题reward的震荡幅度大KL项一度飙升。这两件事其实是同一回事——策略更新太快和参考模型偏离太远然后奖励函数中的格式分又被淹没在噪声里。我做的第一个调整是把beta从0.04降到0.01。KL惩罚系数过大时模型每次更新都要“回头看”参考模型导致学习信号被压制reward反而涨不上去。降到0.01后策略有更多自由空间去探索能拿高分的输出KL曲线明显回落。第二个调整是把学习率从5e-6降到1e-6。GRPO的组内相对奖励本质上是样本之间对比的结果样本方差大学习率跟着大很容易在奖励高点附近反复横跳。低学习率配大batch比高学习率配小batch稳得多。第三个调整是重写reward函数把格式分从0.5降为0.2同时让格式错误不再把整个completion的正确性判为0保留部分信息。三个调整叠加后reward曲线从“剧烈震荡”变成“震荡上行”训练稳定性也进入了可接受范围。下表是我调试过程中的关键对比指标初始配置第一轮后最终配置rollout吞吐350 tok/s490 tok/s620 tok/s单step耗时28秒20秒16秒NPU平均利用率约40%约60%约75%reward方差偏高偏高明显收敛必须说明的是这些数字只代表我这套模型、batch、卡数下的相对变化不具有普适性。不同模型、不同prompt长度分布、不同奖励函数数值都会完全不同。我真正希望大家学走的不是数字而是“先看瓶颈在哪一层再针对性动手”的思路。5. 三个值得记录的翻车现场与最终配置模板最后一部分记录三个我在这次部署中最典型的翻车现场。每个都具备一定的共性换到别的机器、别的版本组合也可能出现。5.1 版本错配编译期正常运行期算子注册崩溃第一次在昇腾机器上安装我图省事直接选了当时最新的CANN和最新版torch_npuvLLM-ascend也选了最新源码版。编译全部通过结果一启动离线生成立刻报算子注册错误traceback里看到一个aclnn相关的符号解析失败。一开始怀疑是代码问题反复检查才知道是运行时CANN版本和编译时链路不一致。昇腾上很多算子是通过ACLNN在运行期动态加载的版本对不上报错的时间点会非常靠后让人完全找不到方向。处理方式把CANN、torch_npu、vllm-ascend全部按社区兼容矩阵锁定到同一批组合重装之后问题消失。此后我再也不“追新”昇腾栈的稳定运行优先级远高于版本新鲜度。5.2 TRL默认CUDA逻辑导致训练进程启动即失败第一次直接在TRL脚本里设置use_vllmTrue并让TRL自己管理vLLM时报错信息清晰指向“找不到CUDA驱动”。TRL内部在初始化rollout环境时会检查torch.cuda.is_available()即便我的训练进程已经运行在NPU上它还是会走一段CUDA专属逻辑。我也试过在代码里手动patch设备映射例如把devicecuda:0替换成devicenpu:0项目能跑起来但依赖TRL内部实现细节升级版本就可能碎掉。最终选择的方案是不让TRL管理vLLM进程而是预先把vLLM-ascend跑成外部HTTP服务。TRL只把它当成一个OpenAI兼容API彻底绕开了“本机是不是CUDA环境”的检查。这个解耦思路进行的很彻底之后更换vLLM版本、换卡、甚至把服务部署到另一台服务器都不需要改动训练脚本。5.3 HCCL通信超时多卡训练跑到中间突然全部卡住第三个翻车现场是训练跑了几百步后突然卡住日志尾部出现“HCCL connection timeout”。再次强调HCCL就是昇腾版的NCCL它负责卡间梯度同步一旦超时整个分布式训练组就僵死在那。排查路径经历了几个阶段先怀疑网络用hccn_tool检查卡间网络状态和拓扑发现物理链路正常又怀疑是不是某张卡掉卡但npu-smi info显示全部在线最后定位到耗时差异随着训练迭代各卡前向反向时间差异被逐渐放大单张卡偶尔因为日志记录、其他IO操作慢了0.1秒就触发了通信超时。解决方法是三个动作一起做把HCCL_CONNECT_TIMEOUT从默认的较小时长调大到1800秒把训练的rank切换为不写盘输出训练日志只在主rank汇总输出关闭TRL的额外报告功能减少每步结束时的同步开销。改完后再跑几千步没有再出通信类异常。5.4 可直接复用的最终配置模板把最终稳定的配置整理成模板作为项目后续迭代的基线。环境变量部分export ASCEND_RT_VISIBLE_DEVICES4,5,6,7 export USE_NPU1 export HCCL_CONNECT_TIMEOUT1800 export ASCEND_PROCESS_LOG0服务启动部分vllm serve /path/to/policy-model \ --served-model-name policy-7b \ --port 8001 \ --max-model-len 2048 \ --gpu-memory-utilization 0.85 \ --tensor-parallel-size 4 \ --dtype bfloat16训练端核心参数如下training_args GRPOConfig( output_dir./grpo_ascend_out, per_device_train_batch_size2, gradient_accumulation_steps4, num_generations4, learning_rate1e-6, beta0.01, epsilon0.2, max_prompt_length512, max_completion_length512, use_vllmTrue, vllm_server_host127.0.0.1, vllm_server_port8001, gradient_checkpointingTrue, logging_steps1, save_steps200, )参数速查参数作用我的建议值num_generations每个prompt采样的组内样本数4显存允许可到8betaKL散度惩罚系数0.01到0.04之间epsilon策略裁剪范围0.2learning_rate策略更新步长1e-6起步gradient_accumulation_steps等效扩大batch4到8最后再分享一个我个人在这一次部署里最深的体会昇腾环境做RL训练最大的敌人不是算力而是接口假设。NVIDIA生态里很多框架默认“设备就是CUDA”到了昇腾会出现各种隐式检查。最稳的方式永远是边界解耦——推理服务独立、HTTP契约清晰、训练进程只聚焦策略更新。先推通离线生成再接入单卡训练最后放多卡全量这个顺序能省掉你至少三次“从头排查”的痛苦。这套架构后续还可以扩展比如把奖励模型同时挂到同一推理服务上或者用vLLM-ascend的多并发能力同时服务多个训练任务。强化学习训练和推理服务终究会在调度层融为一体越早理解这条链路后面遇到资源竞争就越从容。
锦
锦皓数字建站
深耕本土企业品牌数字化升级,专注原创端正雅致商务官网,从视觉设计到稳定运维全程保驾护航。