Qwen3.5-4B / 4×A800 训练 OOM 备忘

2026-09-05 在 run_full_train.sh 200-step 作业上连续三次 torch.OutOfMemoryError，都发生在 Megatron 训练 forward 的 vocab softmax / log-prob，不是 rollout、不是 ShopSimulator。本文记录现象、原因和最终旋钮，避免把几个看起来都像「显存」的开关当成一回事。

硬件：4×A800-SXM 80GB。模型：Qwen3.5-4B，text_config.vocab_size=248320（约 3/4 层 linear attention）。训练 colocated：Megatron 与 SGLang 分时占用同一组 GPU。

1. 三次崩溃分别卡在哪

次序	当时配置	报错点	额外申请	当时占用
1	`TP=1`，`MAX_TOKENS_PER_GPU=32768`	logits `float16_to_fp32`	~29.5 GiB	单卡空闲不够一整份 `[T,V]` fp32
2	`TP=1`，打包降到 8192，仍无 chunk	softmax 复制整份 `[T,V]`	~13 GiB	单条超长 episode 独占 micro-batch，打包上限管不到
3	`TP=2`，`--log-probs-chunk-size 1024`	仍在 softmax	464 MiB	77.3 / 79.3 GiB，空闲 207 MiB

第三次说明：分块只压瞬时峰值，不压 backward 保留量。 slime ppo_utils._VocabParallelLogProbEntropy 每个 chunk 算完 softmax 都 save_for_backward(log_prob_softmax, …)。8 个 chunk 各存 1/8，加起来仍是完整的 [T, V/TP]。TP=2 把这份从约 8 GiB 砍到 4 GiB，chunk 把单次分配从 4 GiB 砍到 ~0.5 GiB，但 4 GiB 一直挂到反传结束。

对症开关是 --recompute-loss-function（slime/backends/megatron_utils/loss.py）：对整个 loss 段做梯度检查点，softmax 不进 saved_tensors，backward 时重算。梯度精确，代价是 loss 段多跑一遍 forward。不要和「截断序列」混为一谈。

不要开 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True 来赌碎片。sglang torch_memory_saver 同样走 CUDA VMM，两者有已知冲突；第三次 reserved-未分配只有 ~600 MiB，瓶颈不是碎片。

2. 三个长度旋钮不是一件事

一条购物 episode 在训练里会变成多条 slime Sample，再被打进若干 micro-batch。三个旋钮作用在不同层：

旋钮	作用对象	会不会截断单条 Sample	决定什么显存
`ROLLOUT_MAX_CONTEXT_LEN`	生成侧 context + finalize `tokens[:cap]`	会。触顶且解析不出 tool call → `generation_length`，环境零分，`Sample.status=TRUNCATED`，仍进 GRPO	单条 Sample 的上限，也是峰值上限
`MAX_TOKENS_PER_GPU`	first-fit 把多条 Sample 装进同一 micro-batch	不会。单件比箱子大就独占一箱、允许超载	短样本箱子的容量。峰值 = `max(箱子, 本步最长 Sample)`
`LOG_PROBS_CHUNK_SIZE`	softmax 沿 T 切片	不会	瞬时分配；有 recompute 时不再决定保留量

dp_schedule 的不变量：除「单条超过容量」外，每个 micro-batch ≤ max_tokens_per_gpu * cp_size。训练用 THD packed sequence（qkv_format=thd），注意力按各条 cu_seqlens 算，不是把拼箱当成一条超长序列做平方注意力。

Qwen3.5-4B、TP=2 时一份 fp32 [T, V/2] 大约：

T	一份 `[T, 124160]` fp32
4096	2.0 GiB
8192	4.0 GiB
16384	7.8 GiB
24576	11.6 GiB
32768	15.5 GiB

无 recompute 时 backward 还要再挂一份同尺寸的 softmax。有 recompute 后这份不再常驻。

3. 截断、fork、打包：报表会骗人

ShopSim 适配器 fork_threshold_tokens=0：工具 schema / chat template / reasoning 回显造成 token drift 时直接切段，保留每一轮真实采样的 logprob。每条 fork 出来的 Sample 带的是 从该段起点起的完整上下文，不是增量。

n_samples_per_prompt=8 是同一题 8 条独立 episode（GRPO 组）。fork 是 一条 episode 内部再切成多条训练 Sample。两件事不要混。

rollout/truncated_ratio 是 mean(status == TRUNCATED)，分母是 Sample 数。episode 因 generation_length 结束时，runtime 给 这条轨迹的每一段都打上 TRUNCATED，包括前面没触顶的短段。触顶 episode 往往最长、fork 最多，所以 5% 的 episode 触顶可以显示成 10–25% 的 truncated_ratio。它仍是可评分的零分策略失败，不是被 filter 丢掉；只有 ABORTED + remove_sample 才整组丢弃。

round-000 的 256 条 episode（Qwen tokenizer，当时 context 还是 32K）：

区间	n	reward>0
0–8,192	191	68%
8,192–12,288	32	84%
12,288–16,384	12	17%
16,384–24,576	8	62%
24,576+	13	8%

约 1100–1200 token/环境步；4–6 步中位 6.5k（80% 样本）。把 context 砍到 8192 会切在成功率最高的带上。24576 落在垃圾带边缘：牺牲约 0.6% 成功 episode，触顶率约 5%。generation_length 与单 turn rollout-max-response-len=2048 在 ShopSim 里不区分，统一零分进训练。

wandb run jin9wfo9 前 8 步 total_lengths 均值 6.7k–8.4k，已经大于当时的 4096 箱子。first-fit 对「单件 > 容量」是独占一箱，所以 4096 对大多数 Sample 是空转：峰值仍是 24k 独苗，短样本也没捆上。真正的分配混合是 7k 独苗 + 24k 独苗，不是「4k 满箱」。

4. 训练 TP 和推理 TP

训练 OOM 的主因是 lm_head 词表维。4 卡上 TP_SIZE=2（DP=2）把每卡词表切半，并打开 sequence parallel，够用。TP_SIZE=4（DP=1）再砍一半 logits，但每卡总 FLOPs 不变、丢掉 DP、GEMM 更瘦、all-reduce 更宽，4B 模型通常更慢。只在打算把 context 加回 32K 时才值得考虑。

Rollout / val 没有这份反传 logits。保持 ROLLOUT_TP_SIZE=1，4 个独立引擎。评测同样 4 卡 tp=1 dp=4 起 sglang，不要为训练 OOM 给推理加 TP。

5. 最终配置（2026-09-06）

scripts/run_full_train.sh / scripts/run_shopsimrl_slime.sh 默认：

NUM_GPUS=4
TP_SIZE=2                          # 仅 Megatron
ROLLOUT_TP_SIZE=1                  # rollout / val
MAX_TOKENS_PER_GPU=16384           # 打包；峰值仍由 24k 独苗决定
ROLLOUT_MAX_CONTEXT_LEN=24576      # 单条 Sample 硬顶
LOG_PROBS_CHUNK_SIZE=1024
--recompute-loss-function
--recompute-granularity full
--recompute-method uniform
--recompute-num-layers 1

MAX_TOKENS_PER_GPU=16384 的理由：峰值公式是 max(箱子, 最长 Sample)。箱子只要 不超过 24576，就不会创造比已跑通的 24k 独苗更大的 micro-batch。16384 让两条典型 ~7k Sample 进同一箱，箱数大约减半，THD 下注意力仍按短序列算。提到 24576 也能保持峰值不变，但会把「截断率低的好步骤」的峰值从 ~12–16k 抬到 24k，边际吞吐收益更小。

不要把 MAX_TOKENS_PER_GPU 加到 24576 以上：first-fit 可能把 20k+8k 捆成 28k，峰值高于独苗。

--balance-by-flops 按 FLOPs 切 micro-batch，不保证 token 上限，箱子已经很满时不要开。

改这些环境变量只在下一轮 run_shopsimrl_slime.sh 启动时生效（外循环 round 边界）。正在跑的 slime job 不会热更新。

6. 若再次 OOM

先看堆栈，不要先改 TP。

仍在 softmax / log-prob，额外只要几百 MB：多半是碎片或 24k logits 余量见底。先确认 --recompute-loss-function 还在；再把 ROLLOUT_MAX_CONTEXT_LEN 降到 16384（会真的截断长 episode，零分进训练）。不要优先降 MAX_TOKENS_PER_GPU——它已经不决定峰值。
一次申请十几 GiB：封顶没生效，或某条 Sample 绕过了 max_sample_tokens。查 rollout/total_lengths 的最大值，不是均值。
推理阶段 / SGLang scheduler：与本文路径无关，不要用训练 TP 去「修」它。

现场可看：rollout/truncated_ratio（Sample 口径，有 fork 放大）、perf/actor_train_time（24k 独苗多的步骤会到 500s+）、perf/train_wait_time（colocated 下这是 rollout，GPU 通常仍在忙）。wait_time_ratio≈30% 不是 GPU 空闲率。