最近在昇腾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_vllm=True打开，并把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_vllm=True，trainer 内部就会创建一个 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 提供算子库和 runtime，torch_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 与 TRL

vLLM-ascend 的安装有两种方式：预编译 wheel 包和源码编译。如果你的 CANN 版本和 Python 版本都在它支持的列表里，优先用预编译包，省时间。源码编译需要 cmake、gcc 等工具链，时间较长，但遇到 wheel 没有对应架构时就得走这条路。

pip install vllm-ascend

TRL 的安装也很简单，但版本一定要提前锁定。我在实验中用稳定版，然后强制把 transformers 退回到一个和 torch_npu 兼容的版本：

pip install "trl>=0.10,<0.15" pip install "transformers==4.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( model="path/to/2b-model", device="npu", gpu_memory_utilization=0.3, max_model_len=2048, trust_remote_code=True, ) params = SamplingParams(n=4, max_tokens=256, temperature=0.6) outputs = llm.generate(["请用一两句话解释什么是相对策略优化"], params) for o in outputs: for out in o.outputs: print(out.text)

这里n=4是 GRPO 里最常用的做法，对同一个 prompt 采样 4 条回答。如果这段代码能在 NPU 上顺利跑完，说明 vLLM-ascend 的生成引擎没问题。

这一步最容易翻车的地方是device参数。vLLM-ascend 早期版本要求写device="npu"，新版又改成了通过环境变量VLLM_TARGET_DEVICE=npu来声明。不同版本行为不一样，一定以你本地版本的文档为准。如果两个都写上，通常会报重复配置。

2.4 环境变量与运行时配置

昇腾环境下有几个环境变量会影响整条链路。ASCEND_RT_VISIBLE_DEVICES用于指定当前 Python 进程可见的 NPU 卡号，类似 CUDA_VISIBLE_DEVICES。PYTORCH_NPU_ALLOC_CONF可以控制显存分配策略，例如expandable_segments:True能减少碎片化。

在启动训练脚本之前，我通常会在 shell 里做一套统一设置：

export ASCEND_RT_VISIBLE_DEVICES=0 export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export VLLM_TARGET_DEVICE=npu

expandable_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_code=True) policy_model = AutoModelForCausalLM.from_pretrained( model_path, device_map="auto", trust_remote_code=True ) ref_model = AutoModelForCausalLM.from_pretrained( model_path, device_map="auto", trust_remote_code=True )

这里有个容易忽略的点：TRL 内部更新策略模型时会保留一份“冻结”的参考模型，所以即使初始权重一样，两个对象必须独立，不能用同一个对象反复赋值，否则参考模型会被意外更新。

如果你的 NPU 显存比较紧张，参考模型可以包在torch.inference_mode(mode=True)里运行，减少梯度相关开销。不过 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_size=4, gradient_accumulation_steps=2, max_steps=80, learning_rate=5e-6, adam_beta1=0.9, adam_beta2=0.99, warmup_steps=5, use_vllm=True, vllm_device="npu:0", vllm_gpu_memory_utilization=0.35, vllm_max_model_len=2048, vllm_dtype="auto", logging_steps=1, ) trainer = GRPOTrainer( model=policy_model, ref_model=ref_model, reward_funcs=reward_func, args=training_args, train_dataset=dataset, processing_class=tokenizer, ) trainer.train()

use_vllm=True这一步就是让 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_tokens=256、单卡昇腾上，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 做 encode，vLLM 生成时也会有自己的 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", device="npu", gpu_memory_utilization=0.85, max_model_len=2048, enable_prefix_caching=True, trust_remote_code=True, )

启动后可以用 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_dtype="bfloat16"，同时把GRPOConfig.bf16=True打开。如果某些算子不支持 BF16，可以给关键模块单独保持 FP16，但这会增加代码复杂度。

还有一个容易忽略的点：学习率的设定。GRPO 训练中策略模型很容易被奖励信号推偏，学习率不要太大。我在 2B 模型上用 5e-6，配合 warmup，效果比较稳。如果奖励函数本身噪声大，可以再降到 1e-6。

5.4 显存分配策略：vLLM与训练进程如何共存

一张 NPU 里跑两个进程不现实，因为 TRL 的训练进程和 vLLM 引擎在同一个进程内，它们共享同一块显存。显存分配的核心是给 vLLM 的 KV cache 留一个合理上限，同时给训练权重、梯度、优化器状态预留出明确空间。

以 2B 模型、BF16 单卡为例：模型权重约 4GB，AdamW 优化器状态约 8GB，梯度约 4GB，一共 16GB；如果 NPU 总显存是 32GB，那 vLLM 最多用 10GB，大约 0.3 到 0.4 的 util。这个比例不是死的，要按模型大小重新算。

我建议写一个简单的脚本统计每部分峰值。可以先关掉 vLLM，只跑训练，看训练侧的显存峰值；再单独跑 vLLM，看生成侧的峰值。两边相加小于总显存，才算留出了安全余量。很多 OOM 都是因为低估了训练侧的 AdamW 状态，这部分往往比模型权重还大。

5.5 常见问题速查表

这些是我实际调试中高频出现的问题。如果你也在昇腾 NPU 上跑这条链路，大概率会碰到其中几个。

6. 实测效果与优化心得

6.1 性能数据对比

用某开源 2B 模型，单一昇腾 NPU，prompt 长度为 128，max_new_tokens=256，GRPO 组大小 n=4，我这边实测的数据大致如下：

数据不算惊艳，但对一个 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 步小实验，不要上来就全量训练。小步实验能帮你把显存、梯度、采样三块都看到峰值，再放大规模，心里就有底了。