Skip to content

grpo-colocated-vllm.md

Latest commit

 

History

History
43 lines (26 loc) · 3.46 KB

grpo-colocated-vllm.md

File metadata and controls

43 lines (26 loc) · 3.46 KB
Raw

单卡 vLLM 休眠生成后端

这是原本低显存训练配置的可选生成后端，不修改 GSPO / GRPO 的目标函数、训练 batch、microbatch 或 AdamW 精度，也不需要第二张 GPU 或 HTTP 服务端口。

环境与入口

训练使用仓库根目录下的 .venv/bin/python，vLLM 使用独立环境中的解释器（例如 .venv-vllm/bin/python）。实测版本为 vLLM 0.19.1 / torch 2.10.0+cu129 / Transformers 5.7.0。不要为安装 vLLM 重建或同步已经可用的训练环境；新机器参考安装说明。

在原训练命令中增加：

--low-memory --generation-backend vllm-colocated \
--vllm-python "$VLLM_PYTHON"

单卡后端默认 max_num_seqs=32、max_model_len=1024、gpu_memory_utilization=0.65，采用 eager 执行。这里的 0.65 是 vLLM 活动期间的预算，不是它与训练进程永久各占一块显存。

VLLM_PYTHON 需设置为你自己的独立环境解释器绝对路径。

每轮的执行顺序

  1. 清空梯度，把训练权重和 AdamW 状态显式移到 CPU；保留 Parameter 身份、优化器状态和步数。
  2. 将当前权重暂存至 Linux /dev/shm，启动或复用唯一的私有 vLLM worker。
  3. 唤醒权重区，加载当前训练权重，逐参数校验完整内容 SHA256，清空旧 prefix cache，再唤醒 KV cache。
  4. 一批提交问题并生成每题的 G 个回答，保持问题及组内回答顺序。
  5. sleep(level=2) 丢弃 vLLM 权重和 KV cache，再把训练权重与优化器状态移回 GPU。
  6. 使用原训练循环缓存 old log-probabilities、计算奖励、累积梯度和更新参数。下次生成重新同步最新权重。

level 2 不会保留可直接使用的推理权重，所以每次唤醒必须重新加载；也不会释放所有 CUDA context 占用。这是显式 CPU/GPU 搬运，不是把 Windows 的“共享 GPU 内存”当显存使用。

为兼容 vLLM 的 CuMem 休眠分配器，仅在独立 worker 内移除 PyTorch expandable-segments 环境设置，训练进程的设置保持不变。两进程通过带认证的私有 Unix socket 通信；权重同步不使用同卡双 rank NCCL。

测试与失败处理

scripts/grpo_colocated_benchmark.py 比较相同题目、G=8、最多 512 token 的生成速度。总计时包含 offload、暂存、完整校验、加载、休眠和恢复；第 1 轮冷启动与后续稳态分别记录。

测速里的训练是固定混合奖励的 B8/G8/512-token 更新，确保确实执行 backward 和 AdamW，不能把它的速度倍数当作完整 256-response rollout 的训练倍数。真实 GSM8K pilot 另走 scripts/grpo.py；奖励全相同的组会按原算法跳过实际更新。

同步出错、OOM、非有限训练数值或 worker 异常都应停止，而不是继续用旧权重。SIGTERM 进入 Python cleanup；Linux 父进程死亡保护尽力关闭私有 worker。原生 CUDA 调用未返回时，信号处理可能延迟。训练监督器只管理自己的子进程，不按进程名清理其他任务。

要回到原路径，使用 --generation-backend transformers --low-memory，并选择新的输出目录。已有模型快照没有 Adam/RNG 状态，不能宣称精确断点续训。长实验期间保持电脑及 WSL 运行，避免另开大显存任务。

官方机制说明：vLLM Sleep Mode。公开的诊断摘要不是作业准确率结果；原始本地结果目录不上传 GitHub。