训推同卡：colocate 的显存怎么分

两种模式

RL 训练需要一个推理引擎持续产出数据。它和训练进程的关系有两种安排：

server 模式（多卡）
┌────────────┐   HTTP   ┌────────────┐
│  训练进程   │ ──────→ │ vLLM 服务   │
│  卡 4-7    │ ←────── │  卡 0-3     │
└────────────┘          └────────────┘
两边同时工作，采样和训练可以重叠

colocate 模式（单卡）
┌──────────────────────────────┐
│      同一张卡，交替执行         │
│  采样 ████    训练 ████        │
│      ████        ████         │
└──────────────────────────────┘
一个时刻只有一件事在跑

单卡只能用 colocate。 这不是配置选择，是物理约束——一张卡不可能同时跑训练 kernel 和生成 kernel。

核心参数：显存怎么划

vllm_gpu_memory_utilization（slime 里叫 --sglang-mem-fraction-static）决定给 vLLM 多少。

关键认知：这是预先占用，不是按需申请。 vLLM 启动时就把这个比例的显存吃掉，用来放权重和 KV cache 池。

所以单卡上你必须手工划一条线：

32 GiB
├── vLLM        util × 32     ← 你划的
└── 训练侧      剩下的

怎么定这条线？用 L0 的显存计算器：

1. 算出训练侧需要多少（权重 + 梯度 + 优化器 + 激活 + logits）
2. 留 1.5 GiB 给 CUDA context 和碎片
3. 剩下的除以 32，就是 util 的上限
4. 从这个上限往下减 5%，作为起始值

以 0.6B 全参为例：

训练侧 ≈ 10.4 GiB   (权重1.1 + 梯度1.1 + AdamW6.7 + 激活0.2 + logits1.2)
context ≈ 1.5 GiB
────────────────────
可给 vLLM ≈ 32 − 10.4 − 1.5 = 20 GiB  →  util 上限 ≈ 0.62
起始值取 0.55

sleep mode：把显存借来借去

vLLM 的 sleep mode 是 colocate 能舒服工作的关键：

engine = LLM(..., enable_sleep_mode=True)

engine.wake_up()      # 采样前：申请 KV cache
samples = rollout(...)
engine.sleep(level=1) # 采样后：归还 KV cache 显存
# ↓ 这段时间训练侧可以用这块显存
train_step(...)

两个 level：

TRL 里对应 vllm_enable_sleep_mode=True（部分版本），Axolotl 的文档也推荐开启。

OOM 排查顺序

colocate 下 OOM 有两种，先分清是哪一种：

报错发生在 采样阶段  →  vLLM 侧不够   →  降 max_model_len / 降并发 / 提 util
报错发生在 训练阶段  →  训练侧不够     →  降 micro batch / 降 util / 开检查点

看 traceback 里出现的是 vLLM 的栈还是 transformers 的栈就能区分。

然后按这个顺序动，一次只动一个：

时间去哪了

记一次实测的分解（0.6B、8 prompt × 8 条、平均 152 token）：

采样            28.4 s   62%   ████████████████████████
前向 + 反向     12.1 s   26%   ██████████
权重同步         2.2 s    5%   ██
sleep/wake      3.1 s    7%   ███
─────────────────────────────
单步            45.8 s

结论很清楚：优化的重点在采样。

按性价比排序的四个动作：

1. 提高 util（如果显存允许）—— KV cache 池更大，能容纳更多并发，吞吐直接上去
2. 降 max_completion_length —— 采样时间对长度基本是线性的。但要先确认截断率不高
3. 确认 CUDA graph 开着 —— 关掉能损失几倍吞吐
4. 减小固定开销的占比 —— 每步 batch 调大，让同步和切换的 5 秒摊薄

colocate 调优演练。目标是在不 OOM 的前提下把单步时间压下来。
goals 看目标，hint 要提示。

这节课的落点

• 单卡只能 colocate，代价约 1.5 倍（不能重叠 + 切换开销），这是物理约束
• vllm_gpu_memory_utilization 是预占不是按需，必须手工划线
• 理论上限用显存计算器算，但实践从 0.3 起步，跑通再往上加
• sleep mode level=1 归还 KV cache 显存，等价于让 util 能设得更高
• OOM 先看 traceback 在哪一侧，再按「不影响效果」的顺序动刀
• 不要先降 num_generations，它是 GRPO 的基线来源，一般不低于 4
• 采样占 60%+，优化重点：提 util、控长度、保住 CUDA graph、摊薄固定开销