最近在昇腾NPU上折腾 vLLM-ascend 配合 GRPO 强化学习训练也就是很多人常挂在嘴边的“GPRO”训练。实际上 TRL 框架里对应的是 GRPO群体相对策略优化。这段链路跑通之后采样速度比直接用原生生成方法快了五六倍整个训练迭代节奏完全不一样了。先说明这套东西是干什么用的一个 LLM 在 SFT 之后通常还要用强化学习做一轮偏好对齐。GRPO 这类轻量算法不需要像 PPO 那样额外维护价值网络而是通过对同一个 prompt 采样一组回答用组内相对奖励来更新策略。TRL 已经实现了 GRPO 的训练流程但默认采样还是走 transformers 的 generate。对于大规模 RL 任务采样本身是最大的开销所以这里引入 vLLM-ascend 作为生成引擎利用它的连续批处理和 PagedAttention 来加速采样。这篇内容主要面向两类人一是想在昇腾 NPU 上跑 LLM 对齐训练的同学二是在其他国产芯片上搞推理优化、需要理解 vLLM 与训练框架如何衔接的朋友。我会把环境搭建、版本匹配、训练代码、模型部署、性能调优和踩坑记录都讲清楚尽量给出能直接照抄的配置。1. 整体方案设计与技术选型1.1 这条链路到底要解决什么问题在昇腾 NPU 上做 LLM 强化学习常规思路是直接用 PyTorch transformers 训练采样也靠 transformers 的 generate 方法。但 GRPO 有个特点每个 prompt 要采样多组输出如果组大小 n 设为 4意味着模型离线生成的数量是训练更新数量的 4 倍。如果采样过程太慢训练会一直等生成NPU 利用率上不去。同时transformers 原生的 generate 在 NPU 上对动态 batch 和长序列的支持没有专门优化很容易出现显存碎片和低吞吐。vLLM-ascend 是 vLLM 在昇腾上的适配版本它把连续批处理和 paged KV cache 带到了 NPU。这个方案要解决的就是在单卡 NPU 上把“采样”和“更新”两个环节串起来互不拖后腿。实话说一开始我把问题想简单了。我以为只要把 TRL 的 device 参数改成 npu 就能跑结果发现 TRL 的生成后端还是走 transformers 的接口采样速度根本没上去。直到我把use_vllmTrue打开并把vllm_device指到npu:0才真正把 vLLM-ascend 的能力接进训练流程。1.2 为什么选 TRL 的 GRPO 而不是其他方案RLHF 类方案里PPO、DPO、GRPO 是最常被提到的三个。PPO 需要额外的 critic 模型除了策略模型还要在每一步估计价值函数显存占用高在昇腾上的算子适配也更复杂。DPO 虽然简单但它是离线偏好优化没有在线探索通常放在 RLHF 之前做轻量对齐。GRPO 介于两者之间它在线采样但没有 critic 模型显存占用显著降低训练稳定性也不错。我最终选 TRL 的 GRPOTrainer还有一个现实考虑它把 vLLM 生成后端集成得比较干净。在GRPOConfig里开启use_vllmTruetrainer 内部就会创建一个 vLLM 引擎来采样。这个设计让我不需要魔改训练主循环只要把 vLLM 的 device 指到 NPU并把模型路径传给 vLLM 就行。相比自己写一个采样进程再和训练进程通信这个方案更省事。当然TRL 的 GRPOTrainer 也不是没有缺点。它内部封装了很多细节报错信息比较抽象一旦 vLLM 和 trainer 之间的接口版本不匹配问题很难定位。所以选择它之前必须接受“版本锁死”这件事。1.3 vLLM-ascend 在训练链路里的定位一个完整 GRPO step 包含采样、打分、更新三步。vLLM-ascend 只负责第一步采样。它会常驻在训练进程内接收 prompt 列表返回采样完的文本列表。TRL 拿到这些文本后再交给奖励函数打分。vLLM-ascend 带来的提升主要来自三个方面连续批处理能动态拼接长短不同的请求减少等待浪费PagedAttention 把 KV cache 拆成固定大小的 page显存利用率比原生 generate 高很多昇腾上有针对 flash attention 和全连接算子的融合优化长序列生成时开销更稳定。在训练链路里vLLM-ascend 不是额外起一个服务进程而是以库的形式被 import 进来。这么做的好处是数据和参数都在同一个 Python 进程内交互不需要把文本序列化成 JSON 走网络省掉通信开销。训练进程和 vLLM 引擎共享同一块 NPU 显存这是它最有价值也最容易翻车的地方。1.4 版本匹配的教训这条链路最大的坑就是版本组合。vLLM-ascend 是紧接着某个 vLLM 版本发布适配包的TRL 内部对 vLLM 的调用接口又会随版本变化。我一开始装的是比较新的 TRL 开发版结果它对 vLLM 的初始化参数变了vLLM-ascend 还没同步导致启动即报错。这个问题的本质是三方项目版本联动的滞后。要等 TRL 先适配新参数然后 vLLM-ascend 再跟上来中间往往有窗口期。所以版本不是越新越好稳定才是第一优先级。我建议固定一个“经过验证的组合”。比如Python 3.10 torch 2.1 对应版本的 torch_npu CANN 7.x vLLM-ascend 某稳定版 TRL 0.10 到 0.14 之间。在跑真实训练前先用一个小模型跑 20 步小实验确认生成和训练都没问题再上大规模任务。版本冲撞的时候优先看 vLLM-ascend 的 release notes它支持的 vLLM 版本范围是硬边界。2. 昇腾NPU上的环境搭建与验证2.1 基础环境准备CANN、torch_npu和Python环境昇腾 NPU 上的 PyTorch 应用依赖两层底层组件CANN 提供算子库和 runtimetorch_npu 是 PyTorch 的昇腾扩展插件。安装顺序不能乱先装 CANN再装 torch_npu。如果你用的是 Docker 镜像从官方镜像启动会比较快但版本也要和宿主机驱动匹配。我习惯用 conda 创建一个专门的 Python 3.10 环境不污染系统环境。依赖装完后先做一个最小验证python -c import torch, torch_npu; print(torch.npu.device_count())输出不为 0说明驱动和插件基本可用。然后手动把一个随机模型放到npu:0上跑一次 forward 和 backward确认梯度能正常回传。这一步虽然基础但能避免后续把环境问题误判成代码问题。我见过不少人在同一个脚本里既跑 transformers 又跑 vLLM最后报错是 torch_npu 和 PyTorch 版本冲突白白排查一整天。2.2 安装 vLLM-ascend 与 TRLvLLM-ascend 的安装有两种方式预编译 wheel 包和源码编译。如果你的 CANN 版本和 Python 版本都在它支持的列表里优先用预编译包省时间。源码编译需要 cmake、gcc 等工具链时间较长但遇到 wheel 没有对应架构时就得走这条路。pip install vllm-ascendTRL 的安装也很简单但版本一定要提前锁定。我在实验中用稳定版然后强制把 transformers 退回到一个和 torch_npu 兼容的版本pip install trl0.10,0.15 pip install transformers4.40.0为什么这么在意 transformers因为 TRL 依赖 transformers而 transformers 新版本可能会调用 torch_npu 还没实现的新接口。我遇到过几次莫名其妙的 SEGV最后定位都是 transformers 版本太新。所以锁版本不是保守是在给后续调试减少变量。2.3 用一个小模型验证链路环境装好后先别急着写 GRPO 训练代码而是先用一个小模型单独验证 vLLM-ascend 能否正常加载和生成。我用的是某个开源 2B instruct 模型tokenizer 和权重都从本地目录读取from vllm import LLM, SamplingParams llm LLM( modelpath/to/2b-model, devicenpu, gpu_memory_utilization0.3, max_model_len2048, trust_remote_codeTrue, ) params SamplingParams(n4, max_tokens256, temperature0.6) outputs llm.generate([请用一两句话解释什么是相对策略优化], params) for o in outputs: for out in o.outputs: print(out.text)这里n4是 GRPO 里最常用的做法对同一个 prompt 采样 4 条回答。如果这段代码能在 NPU 上顺利跑完说明 vLLM-ascend 的生成引擎没问题。这一步最容易翻车的地方是device参数。vLLM-ascend 早期版本要求写devicenpu新版又改成了通过环境变量VLLM_TARGET_DEVICEnpu来声明。不同版本行为不一样一定以你本地版本的文档为准。如果两个都写上通常会报重复配置。2.4 环境变量与运行时配置昇腾环境下有几个环境变量会影响整条链路。ASCEND_RT_VISIBLE_DEVICES用于指定当前 Python 进程可见的 NPU 卡号类似 CUDA_VISIBLE_DEVICES。PYTORCH_NPU_ALLOC_CONF可以控制显存分配策略例如expandable_segments:True能减少碎片化。在启动训练脚本之前我通常会在 shell 里做一套统一设置export ASCEND_RT_VISIBLE_DEVICES0 export PYTORCH_NPU_ALLOC_CONFexpandable_segments:True export VLLM_TARGET_DEVICEnpuexpandable_segments:True这个参数我用下来比较关键。GRPO 训练时张量频繁申请和释放显存会逐渐碎片化最终导致明明还有富余却申请失败。开启后页面碎片减少但也会带来一点额外内存开销需要自己权衡。3. GRPO训练与vLLM-ascend的联合实现3.1 准备训练数据与奖励函数GRPO 训练首先需要一组训练数据格式可以很简单每条数据包含一个 prompt 和对应的参考答案。奖励函数我一般先用规则版因为在一个开源 2B 模型上验证链路不需要加载额外的奖励模型跑通了再换真实 RM。示例奖励函数判断回答长度是否在合理区间以及是否包含关键标志词。奖励分数取 0 到 1 之间的浮点数然后 GRPO 会在组内做归一化。from datasets import Dataset train_data [ {prompt: 一列火车每小时行驶120公里行驶360公里需要多少小时, answer: 3小时}, {prompt: 如果一个矩形的长是6米宽是4米面积是多少, answer: 24平方米}, ] dataset Dataset.from_list(train_data) def reward_func(prompts, completions, **kwargs): rewards [] golden kwargs.get(answer) for completion in completions: score 0.0 if 10 len(completion) 200: score 0.5 if golden and golden in completion: score 0.5 rewards.append(score) return rewards有几个细节容易被忽略奖励函数接收的prompts和completions都是 list必须返回同样长度的 list。如果用了kwargs.get(answer)训练数据里必须有answer字段并且奖励函数要在初始化时允许接收额外参数。3.2 初始化策略模型、参考模型和奖励模型GRPO 里用于更新的策略模型和用于约束结果的参考模型初始权重是同一个模型所以代码里要加载两份权重。昇腾上加载 transformers 模型时我把它们放到 NPU 上from transformers import AutoModelForCausalLM, AutoTokenizer model_path path/to/2b-model tokenizer AutoTokenizer.from_pretrained(model_path, trust_remote_codeTrue) policy_model AutoModelForCausalLM.from_pretrained( model_path, device_mapauto, trust_remote_codeTrue ) ref_model AutoModelForCausalLM.from_pretrained( model_path, device_mapauto, trust_remote_codeTrue )这里有个容易忽略的点TRL 内部更新策略模型时会保留一份“冻结”的参考模型所以即使初始权重一样两个对象必须独立不能用同一个对象反复赋值否则参考模型会被意外更新。如果你的 NPU 显存比较紧张参考模型可以包在torch.inference_mode(modeTrue)里运行减少梯度相关开销。不过 TRL 的 GRPOTrainer 内部也会对 ref_model 做冻结处理所以大部分情况直接用就行。3.3 配置 GRPOTrainer 和 vLLM 生成后端接下来是核心配置。TRL 的GRPOConfig支持很多参数这里只列重点from trl import GRPOConfig, GRPOTrainer training_args GRPOConfig( output_dir./grpo_ascend_logs, per_device_train_batch_size4, gradient_accumulation_steps2, max_steps80, learning_rate5e-6, adam_beta10.9, adam_beta20.99, warmup_steps5, use_vllmTrue, vllm_devicenpu:0, vllm_gpu_memory_utilization0.35, vllm_max_model_len2048, vllm_dtypeauto, logging_steps1, ) trainer GRPOTrainer( modelpolicy_model, ref_modelref_model, reward_funcsreward_func, argstraining_args, train_datasetdataset, processing_classtokenizer, ) trainer.train()use_vllmTrue这一步就是让 TRL 不再走原生 generate而是创建一个 vLLM-ascend 引擎。vllm_device必须写成npu:0如果写成 cpu 或 cuda会直接报错。vllm_gpu_memory_utilization是关键设太大会导致后面训练进程 OOM设太小采样会很慢。如果你的 TRL 版本里没有use_vllm这个参数说明版本偏旧需要升级或者改用vllm_config的方式。升级前先看 vLLM-ascend 支持的版本范围别盲目升到最新。3.4 训练循环里发生了什么采样、奖励归一化、策略更新跑起来之后训练器每步大致做这些事从训练集取一个 batch 的 prompt送到 vLLM-ascend 的引擎对每个 prompt 采样 n 条回答拿到 completions 后奖励函数批量计算分数然后在组内做零均值化得到每个 token 位置的优势值最后用策略梯度更新策略模型。组内归一化是 GRPO 的核心。假设一个 prompt 采样了 4 条回答奖励分数分别是 0.6、0.5、0.5、0.2归一化后就是相对奖励。这样训练信号来自“这条回答比同组的其他回答好多少”不依赖绝对分数也不需要额外的基线模型。这个过程里vLLM 生成的 tokens 数量会直接影响每一步的时间。我自己实测在 2B 模型、max_new_tokens256、单卡昇腾上vLLM-ascend 采样 128 个 prompt 平均耗时约 3 秒而原生 transformers generate 大概要 16 秒。换用 vLLM 后训练瓶颈从生成变成了反向传播这更符合预期。3.5 踩坑记录vLLM与Trainer之间的设备冲突联调阶段最典型的问题是 vLLM 引擎和训练进程出现设备争抢。vLLM-ascend 启动时会申请一部分 NPU 显存作为 KV cache这部分显存不会立刻释放。如果训练时再加载优化器显存很容易超。我第一次跑的时候把vllm_gpu_memory_utilization设成 0.6结果训练到第 3 步直接 OOM。后来调成 0.35并把训练 batch 从 8 改成 4才稳定下来。另一个坑是 tokenizer 不一致。TRL trainer 用自己的 tokenizer 对 prompt 做 encodevLLM 生成时也会有自己的 tokenizer 配置。如果两个 tokenizer 并不是同一个 checkpoint 的同一个版本生成结果里会出现乱 token。解决办法是显式给 GRPOTrainer 传入和 vLLM 同样的 tokenizer 路径并确保 trust_remote_code 一致。4. 训练产物部署从checkpoint到推理服务4.1 checkpoint的保存与导出训练结束后TRL 会在输出目录保存 checkpoint。如果想用 vLLM-ascend 做线上推理最好把权重导出成 vLLM 能直接加载的格式比如 safetensors并确保 tokenizer 文件在同一目录。通常trainer.save_model()会把模型和 tokenizer 都保存到指定目录。保存完之后检查一下目录结构确认有config.json、tokenizer.json和*.safetensors。如果没有 tokenizer 文件vLLM 启动时会尝试从 model 目录找找不到就报错。我一般还会单独做一次本地加载测试用 transformers 加载保存出来的模型跑一次生成确认权重没有损坏。这一步虽然多花几分钟但能避免部署阶段才发现 checkpoint 不完整。4.2 用 vLLM-ascend 启动推理服务线上推理和训练时的 vLLM 配置不一样线上没有训练进程抢显存可以把gpu_memory_utilization调高给 KV cache 留更多空间。我用下面的命令启动vllm serve /path/to/saved-model \ --device npu \ --max-model-len 2048 \ --gpu-memory-utilization 0.85 \ --enable-prefix-caching \ --trust-remote-code如果希望用 Python 代码控制服务生命周期也可以直接写from vllm import LLM, SamplingParams llm LLM( model/path/to/saved-model, devicenpu, gpu_memory_utilization0.85, max_model_len2048, enable_prefix_cachingTrue, trust_remote_codeTrue, )启动后可以用 curl 测试接口确认模型能正常响应curl http://localhost:8000/v1/completions \ -H Content-Type: application/json \ -d {model: saved-model, prompt: 你好介绍一下你自己, max_tokens: 64}4.3 部署时的性能配置与注意事项线上推理的显存分配逻辑和训练完全不同但有几个点值得专门提醒。max_model_len不能拍脑袋。它决定了 KV cache 的上限设太大会浪费显存设太小会把长 prompt 截断。我会先统计线上请求的序列长度分布取 P99 值加上生成的最大 tokens 数再留一点余量。max_num_seqs控制同一时刻能并发几条序列。这个值太大会导致 GPU 算不过来反而增加排队延迟太小又浪费显存。建议从小往大调观察请求的 end-to-end 时延。enable_prefix_caching对多轮对话场景很有效。当用户发新消息时历史消息的前缀 KV cache 可以复用减少重复计算。但在纯单轮生成场景中收益不明显。5. 性能优化与问题排查实操5.1 先量化瓶颈生成吞吐和训练吞吐怎么测性能调优不能靠感觉建议先把两个指标拆开生成吞吐tokens/s和训练吞吐samples/s。生成吞吐可以通过 vLLM 的 log 直接看或者用llm.generate跑一批固定长度的生成长度然后用总 token 数除以耗时。训练吞吐可以用trainer.state里的 global_step 和运行时间来算或者直接盯日志里的step打印。我一般会单独测两步先测 vLLM 生成 200 条 completion 的耗时再测训练端一个 step 更新权重的耗时。谁的耗时高瓶颈就在谁身上。举个例子如果生成 200 条 completion 花了 8 秒而一个 step 更新权重只花 2 秒那生成端就是瓶颈优先调 vLLM 参数。反过来如果生成只要 1 秒但 step 更新要 5 秒那就要从训练侧优化比如梯度累积和混合精度。5.2 生成端调优max_model_len、gpu_memory_utilization、前缀缓存生成端最重要的参数是max_model_len。它决定了 KV cache 能覆盖的最大序列长度。我的做法是先统计训练数据里最长的 prompt 长度加上 max_tokens取 1.1 倍余量再向上取整。gpu_memory_utilization则是给 vLLM 划多少显存。在训练链路里它不是越高越好。我最终的取值是 0.35相比默认的 0.9 会损失一部分并发能力但换来了训练侧的稳定。如果你的训练 batch 很小可以适当提高到 0.5。vLLM 的 prefix caching 也很值得开。GRPO 里同一个 prompt 会反复采样多个 completion前面前缀的 KV cache 可以复用。开启之后在固定 prompt 集合的反复采样场景下生成延迟能再降 20% 左右。对应 vLLM 启动参数是--enable-prefix-caching在 GRPOConfig 里可以传一个 dict 到vllm_extra_args。5.3 训练端调优批量大小、梯度累积、混合精度训练端一个常见误区是把per_device_train_batch_size调得很大来提升吞吐。但在 GRPO 训练里batch 大会导致 vLLM 一次性要生成大量 completions显存压力反而更大。建议保持 batch 在 4 到 8 之间再用gradient_accumulation_steps把有效 batch 撑大。混合精度上昇腾 NPU 对 FP16 和 BF16 支持比较完整。策略更新阶段用 BF16 比 FP16 更稳避免 loss 突然 NaN。设置 transformers 的torch_dtypebfloat16同时把GRPOConfig.bf16True打开。如果某些算子不支持 BF16可以给关键模块单独保持 FP16但这会增加代码复杂度。还有一个容易忽略的点学习率的设定。GRPO 训练中策略模型很容易被奖励信号推偏学习率不要太大。我在 2B 模型上用 5e-6配合 warmup效果比较稳。如果奖励函数本身噪声大可以再降到 1e-6。5.4 显存分配策略vLLM与训练进程如何共存一张 NPU 里跑两个进程不现实因为 TRL 的训练进程和 vLLM 引擎在同一个进程内它们共享同一块显存。显存分配的核心是给 vLLM 的 KV cache 留一个合理上限同时给训练权重、梯度、优化器状态预留出明确空间。以 2B 模型、BF16 单卡为例模型权重约 4GBAdamW 优化器状态约 8GB梯度约 4GB一共 16GB如果 NPU 总显存是 32GB那 vLLM 最多用 10GB大约 0.3 到 0.4 的 util。这个比例不是死的要按模型大小重新算。我建议写一个简单的脚本统计每部分峰值。可以先关掉 vLLM只跑训练看训练侧的显存峰值再单独跑 vLLM看生成侧的峰值。两边相加小于总显存才算留出了安全余量。很多 OOM 都是因为低估了训练侧的 AdamW 状态这部分往往比模型权重还大。5.5 常见问题速查表现象可能原因解决办法vLLM 启动报 device 不匹配未设置devicenpu或环境变量检查 vLLM-ascend 文档统一用devicenpu训练到第几步就 OOMvLLM 显存分配太大调低vllm_gpu_memory_utilization减小训练 batch生成结果乱码tokenizer 不一致给 vLLM 和 Trainer 传同一个 tokenizer 路径算子不支持报错transformers 或 torch_npu 版本过新回退到 vLLM-ascend release 锁定的版本loss 出现 NaNFP16 溢出切换 BF16检查 reward 是否出现 infvLLM 和训练速度差距过大prefix caching 未启用开启缓存并复用同 prompt 的采样显存碎片严重反复调整 util 导致设置PYTORCH_NPU_ALLOC_CONFexpandable_segments:True部署接口响应慢max_num_seqs 或 gpu_memory_utilization 不合理调大gpu-memory-utilization控制并发序列数这些是我实际调试中高频出现的问题。如果你也在昇腾 NPU 上跑这条链路大概率会碰到其中几个。6. 实测效果与优化心得6.1 性能数据对比用某开源 2B 模型单一昇腾 NPUprompt 长度为 128max_new_tokens256GRPO 组大小 n4我这边实测的数据大致如下配置采样耗时128 prompt一个训练 step 总耗时原生 generate约 16s约 22svLLM-ascend 默认约 5.2s约 8.1svLLM-ascend 开 prefix cache约 3.9s约 6.7s数据不算惊艳但对一个 2B 规模的实验来说采样部分提速约 4 倍整个训练 step 提速约 3 倍。如果你的 prompt 更长、采样数量更多提速会更明显。6.2 可复用的优化经验第一版本锁死比追新重要。vLLM-ascend 和 TRL 的组合每换一个版本都可能带来接口变化我建议把整套依赖版本写进 requirements 或 docker 镜像不要散落在个人环境里。第二显存分配要有预算表。先算训练侧峰值再给 vLLM 定 util别让 vLLM 一上来吃掉大块显存。你会发现显存预算表比调参日志更有用。第三奖励函数尽量先用规则版跑通链路再换奖励模型。规则函数简单、可控能快速暴露链路问题一上来就用大 RM遇到问题会分不清是采样问题还是打分问题。6.3 后续扩展与个人体会如果你打算把这个方案往更大模型和更长序列上迁移还有几个方向可以尝试用托管内存做更大 KV cache把 vLLM 引擎放到独立进程通过跨进程通信传 completion这样训练进程和设备分配更干净或者考虑多卡版本用数据并行把采样任务切分到多个 NPU 上。我个人接下来准备把奖励模型替换成可学习的 RM并对比不同采样组大小对训练稳定性的影响。最后再分享一个小技巧无论你用什么版本组合第一步永远先跑通 100 步小实验不要上来就全量训练。小步实验能帮你把显存、梯度、采样三块都看到峰值再放大规模心里就有底了。