GRPO 环境验收指南｜驱动、梯度与显存生命周期

本文以 Qwen3-1.7B 的七卡两步 GRPO 为例，整理四类环境问题的判断方法。读完后，应能检查进程加载的驱动库，解释“跑了两步但参数没变”，估算长回答的 KV 显存，并确认训练退出后的资源状态。

案例来自《Self-Distilled Reasoner》的复现准备。论文中的学生根据题目生成回答，教师还能看到参考解答并提供逐 token 的蒸馏信号。这里先验收同一项目的 GRPO 基线，覆盖生成、正确性奖励、反向传播和优化器更新。正式蒸馏训练与榜单评测需要另外验收。论文与公开代码

1. 先定义环境验收的通过条件

1.1 一次更新经过哪些组件

1
2
3

题目 → vLLM 生成一组回答 → 正确性奖励 → 组内优势
     → 训练模型计算损失 → 反向传播 → 优化器更新
     → 关闭引擎与通信组 → 进程退出 → 外部采样显存

共置模式下，推理引擎和训练模型使用同一组 GPU。vLLM 负责生成，训练框架负责梯度，NCCL 负责多卡通信。某个组件能够导入，只覆盖了链路中的一个环节。

1.2 给每个环节留下独立证据

环节	应保存的证据	本次判定方式
CUDA 初始化	返回码、实际加载的驱动库	初始化成功，再检查数值运算
分布式通信	各进程的集合通信结果	七个进程均得到预期求和 28
生成与奖励	回答长度、结束原因、奖励和优势	检查是否截断、组内是否有差异
优化器更新	步数、梯度、参数前后采样	两步更新均完成，并观测到参数变化
退出	操作系统实际观测的返回码	七个进程全部为 0
显存归还	退出后的 GPU 进程和显存采样	与各卡启动前的空闲基线比较

两步检查可以缩短运行时间，生成、奖励和更新仍走真实入口。本次数据来源为公开数学训练集中的八行题目。

2. CUDA 803：检查进程实际加载了什么

2.1 区分四种版本信息

信息	它描述什么
容器镜像标签	镜像的打包版本
PyTorch 的 CUDA 构建版本	框架构建时对应的 CUDA 运行时版本
宿主 NVIDIA 驱动	GPU 驱动栈的宿主部分
进程加载的 `libcuda.so`	应用实际调用的 CUDA Driver API 用户态库

排查初始化失败时，最后两项尤其需要对照。本例中，镜像默认环境加载了与宿主不匹配的驱动库。取消 LD_LIBRARY_PATH 后仍然选中了相同文件，只有实际加载匹配的库，初始化才成功。

启动环境	加载的库文件名	`cuInit(0)`
默认环境	`libcuda.so.570.124.06`	803
取消 `LD_LIBRARY_PATH`	`libcuda.so.570.124.06`	803
优先加载匹配的驱动库	`libcuda.so.525.125.06`	0

2.2 最小诊断代码

在已分配 GPU 的 Linux 环境中，用计划运行训练的解释器执行。每改变一次库搜索环境，都启动一个新进程。

import ctypes
from pathlib import Path

driver = ctypes.CDLL("libcuda.so.1")
driver.cuInit.argtypes = [ctypes.c_uint]
driver.cuInit.restype = ctypes.c_int
result = driver.cuInit(0)

maps = Path("/proc/self/maps").read_text().splitlines()
libraries = {
    Path(line.split()[-1]).name
    for line in maps if "libcuda.so." in line
}
print("cuInit =", result)
print("loaded =", sorted(libraries))
raise SystemExit(0 if result == 0 else 1)

这段代码检查初始化返回码和已加载文件。成功之后，继续运行浮点运算、多卡通信和模型生成，逐层扩大验证范围。具体库加载策略应按实际镜像与宿主组合确定。

3. GRPO 零优势：为什么步数增加，参数却没变

3.1 从组内归一化解释

同一道题采样 $G$ 个回答，正确性奖励记为 $r_1,\ldots,r_G$ 。采用组内标准差缩放时，优势可写成

\bar r=\frac{1}{G}\sum_{i=1}^{G}r_i, \qquad \hat A_i=\frac{r_i-\bar r}{\operatorname{std}(r_1,\ldots,r_G)+\varepsilon}.

若组内奖励相同，所有分子都为零。在本次 $\beta=0$ 、没有额外 KL 项的设置下，这组样本不会提供非零的优势加权策略梯度。需要同时查看优势、梯度和参数变化，才能判断是否覆盖到了有效更新。

3.2 用八个奖励做一个 CPU 算例

下面使用样本标准差和 $10^{-4}$ 的稳定项，演示锁定版本的组内归一化行为。

from statistics import mean, stdev

for rewards in ([0] * 8, [0] * 4 + [1] * 4):
    center = mean(rewards)
    scale = stdev(rewards) + 1e-4
    advantages = [(reward - center) / scale for reward in rewards]
    print([round(value, 4) for value in advantages])

# 全部为零奖励 → 八个优势都是 0
# 一半正确 → 前四个约 -0.9352，后四个约 +0.9352

零优势本身是合法结果。用于环境验收时，应如实报告该批数据覆盖到的更新情况，不能靠修改奖励制造参数变化。TRL 0.26.0 的优势计算

3.3 优化步数和回答长度分别控制

本例最初把回答上限设为 256 tokens，结果整批被截断，优势全零，抽查参数没有变化。最终保留两步优化，把回答上限恢复到论文 GRPO 配置允许的 16,000 tokens。

参数	本次值	控制的对象
`max_steps`	2	优化器更新次数
`max_prompt_length`	2,048	题目长度上限
`max_completion_length`	16,000	每条回答的生成上限
`num_generations`	8	每道题的回答数
`num_iterations`	2	同一批生成结果用于更新的次数

七卡等量数据并行、每卡 microbatch 1、累积 8 时，全局批量为 $7\times1\times8=56$ 。论文对应设置为 32，因此本次明确记录这个批量差异。论文训练配置

4. KV 缓存预算：采样条数与生成并发分别计算

4.1 先写出显存公式

对于本例使用的 BF16、未量化 GQA，K/V 张量的存储量为

M_{\mathrm{KV}} =2\,N_{\mathrm{layers}}\,H_{\mathrm{KV}}\,D_{\mathrm{head}}\,L\,S\,b.

符号	含义	本次值
$2$	Key 和 Value 两份张量	2
$N_{\mathrm{layers}}$	层数	28
$H_{\mathrm{KV}}$	KV 头数	8
$D_{\mathrm{head}}$	每个头的维度	128
$L$	单条序列的题目与回答总长度	$2048+16000=18048$
$S$	同时处理的序列数	4
$b$	每个元素的字节数	BF16 为 2

这里需要使用 KV 头数。GQA 共享 K/V 头，不能直接把 Query 头数代入。

4.2 用代码复算本次预算

layers = 28
kv_heads = 8
head_dim = 128
sequence_length = 2048 + 16000
concurrent_sequences = 4
bytes_per_element = 2

kv_bytes = (
    2 * layers * kv_heads * head_dim
    * sequence_length * concurrent_sequences * bytes_per_element
)
print(kv_bytes)                     # 8279556096
print(f"{kv_bytes / 1024**3:.2f} GiB")  # 7.71 GiB

本次每卡为 KV 配置 8 GiB。上述公式计算有效 K/V 张量，模型权重、分配块舍入、临时工作区和训练激活还需要另留空间，整卡占用继续通过实时采样检查。

4.3 八个回答可以分批进入引擎

设置	本次值	影响
每题采样条数	8	组内奖励比较的样本数
每卡提交条数	8	本次每卡 rollout 数量
`max_num_seqs`	4	引擎同时处理的序列数，其余排队

降低并发可以减少同时需要容纳的 KV。它也会改变调度与耗时，所以本次把这些值和采样参数一起保存在结果中。

5. 退出与显存：按生命周期判断

5.1 三种显存读数不能混用

读数	主要统计对象	适合回答的问题
`torch.cuda.memory_allocated()`	PyTorch 分配器统计中的活动张量内存	进程内还有多少已分配内存
`torch.cuda.memory_reserved()`	PyTorch 缓存分配器保留的内存	分配器还保留多少缓存空间
驱动侧 GPU 进程与显存采样	进程的 GPU 上下文及设备占用	进程退出后，设备是否回到基线

empty_cache() 用于释放分配器中的空闲缓存，活动张量的引用仍需要正常结束。使用自定义分配器时，也要核对统计接口的覆盖范围。PyTorch 显存管理说明

5.2 共置引擎的 sleep mode 增加了哪些环节

vLLM sleep mode 用于在推理引擎暂时不用时释放或转移一部分显存。固定版本中，启用它会让权重和 KV 使用专用内存池。退出时，模型引用、张量和原生分配器的生命周期需要协调。

本例在这条路径上遇到了内存池析构崩溃。排查确认编译回调能持有模型，注销回调后仍有其他权重引用存活，完整引用链当时没有定位完。随后恢复了仓库 GRPO 入口采用的默认值。

1
2
3

use_vllm: true
vllm_mode: colocate
vllm_enable_sleep_mode: false

这是本次通过验收的版本配置。开启 sleep 的那组结果保留为失败记录。TRL 默认值 · vLLM 专用内存池入口

5.3 把最终验收放在进程外

所有 rank 完成最后一次更新和通信
  → 按当前框架版本关闭引擎、释放训练对象和通信组
  → 监督进程等待每个真实退出码
  → 对同一组设备再次查询计算进程与显存
  → 与开跑前基线比较，写入最终报告

训练函数返回之后，解释器和原生扩展仍会执行析构。因此最终状态需要结合操作系统观测的退出码，不能只依赖训练进程自己写出的完成标记。

6. 本次实测配置与结果

6.1 可复查的版本组合

组件	版本或设置
模型	Qwen3-1.7B
Python／Torch／vLLM	3.10／2.8.0／0.11.0
TRL／Transformers	0.26.0／4.57.1
DeepSpeed／FlashAttention	0.18.2／2.8.3
训练	BF16、LoRA、ZeRO-2、CPUAdam
编译	Inductor 启用，本次关闭 CUDA Graph
推理	共置、TP1、sleep 关闭、KV 每卡 8 GiB

6.2 用计算、退出和设备采样共同判定

观察对象	实测结果
生成	56 条回答全部以 EOS 结束，长度 663 至 12,504 tokens
更新	每个进程完成 16 次 micro-step 和两步 CPUAdam，抽查参数发生变化
通信与编译	七个进程的启动求和均为 28，每个进程记录到 29 次 Inductor 编译
退出	七个训练进程及宿主执行退出码均为 0
进程退出前	每个进程仍有约 11.28 GiB 已分配显存
进程退出后	第一次采样即回到各卡启动前的 0 MiB，计算进程列表为空

本轮约 8 分 37 秒。它验证了七卡两步 GRPO 的计算和退出路径，正式 OPSD 训练、长期稳定性与评测质量仍需后续实验。

7. 排查速查表

现象	下一步读取什么
CUDA 初始化返回 803	新进程的初始化返回码及实际映射的 `libcuda.so`
导入成功，第一次多卡通信失败	当前加载的通信库、实际卡数和首个集合通信的原始错误
步数正常，参数没有变化	回答结束原因、组内奖励标准差、优势和梯度
长回答显存不足	总序列长度、KV 头数、并发数及训练阶段的额外占用
更新完成，退出时 SIGABRT／SIGSEGV	每个进程的原生调用栈、分配器路径及真实退出信号
进程内仍有显存分配	继续等待真实退出，再采设备进程列表与显存基线