fix(ci): resolve 0.8B async OOM on H100 by reducing max-tokens-per-gpu

Root cause: Qwen3.5's 248K vocab produces [T, 248320] fp32 logits tensors.
calculate_log_probs_and_entropy holds 5 copies simultaneously (2 clones +
2 intermediates + original). At max-tokens-per-gpu=9216, each copy is
~8.5 GB → 42.6 GB from logits alone, exceeding H100 80 GB with
activations and reserved pool fragmentation.

Fix: reduce max-tokens-per-gpu from 9216 to 2048. Peak drops from 117.6 GB
to 39.6 GB (measured on H200), well within H100's 80 GB. GSM8K's longest
sequence is ~1200 tokens, so 2048 still fits all samples.

Also removes gsm8k_async_short from SOFT_FAIL_ON_H100 (no longer needed)
and the expandable_segments workaround.

parallel_check remains soft-fail: ~11% flake rate on TP4+per-token-loss,
confirmed same behavior in slime (Megatron FP reduction-order issue).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: aoshen02

.buildkite/gpu_suites.py

@@ -32,32 +32,24 @@
 
 # Known hardware-fit failures on the pool's 80 GB H100s — test-level issues,
 # not pipeline ones (PR #239, builds #6/#7):
-#   * gsm8k_async_short: the actor GPU needs >80 GB as tuned
-#     (--max-tokens-per-gpu 9216 with the 248k vocab; 67 GiB live allocations
-#     after expandable_segments removed fragmentation). The sync twin passes.
-#   * parallel_check: cross-layout grad-norm invariance (CP=2) diverges ~4%
-#     from the same-node baseline recording (likely NVLS reduction ordering).
+#   * gsm8k_async_short: FIXED — max-tokens-per-gpu reduced 9216→2048 (peak
+#     39.6 GB on H200, well within H100 80 GB). Root cause was Qwen3.5 248k
+#     vocab × 5 logits copies in calculate_log_probs_and_entropy.
+#   * parallel_check: cross-layout grad-norm invariance (TP4+per-token-loss)
+#     diverges ~12% on ~11% of rollout data (bimodal: most <1.5%, outliers
+#     10-20%). Confirmed same behavior in slime — Megatron FP reduction-order
+#     non-invariance, not a vime bug.
 # soft_fail keeps them running and visible (orange) without failing the
 # build; the GHA label jobs on the self-hosted boxes remain their
 # authoritative gate.
 SOFT_FAIL_ON_H100 = {
-    "test_qwen3.5_0.8B_gsm8k_async_short.py",
     "test_qwen3_0.6B_parallel_check.py",
 }
 
 # (test_file, num_gpus, extra_args, env overrides)
 SUITES = {
     "short": [
-        # expandable_segments: borderline fit on the pool's 80 GB H100s — OOMed
-        # in compute_log_probs with 7 GiB reserved-but-unallocated (build #6).
-        # Scoped to this test only: vLLM's sleep-mode CuMemAllocator can
-        # conflict with expandable segments, so don't set it pod-wide.
-        (
-            "test_qwen3.5_0.8B_gsm8k_async_short.py",
-            4,
-            "",
-            {"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"},
-        ),
+        ("test_qwen3.5_0.8B_gsm8k_async_short.py", 4, "", {}),
         ("test_qwen3.5_0.8B_gsm8k_short.py", 4, "", {}),
         ("test_qwen2.5_0.5B_ppo_critic_only_short.py", 4, "", {}),
         ("test_qwen2.5_0.5B_fully_async_short.py", 4, "", {}),

tests/test_qwen3.5_0.8B_gsm8k_async_short.py

@@ -57,7 +57,7 @@ def execute():
         "--expert-model-parallel-size 1 "
         "--expert-tensor-parallel-size 1 "
         "--use-dynamic-batch-size "
-        "--max-tokens-per-gpu 9216 "
+        "--max-tokens-per-gpu 2048 "
     )
 
     grpo_args = (