Reduce GRPO trainer memory

docs/opencode-grpo.md

@@ -47,7 +47,7 @@ posttrainarena-train model-bridge \
   --tokenizer-revision <immutable-sha> \
   --max-tokens 4096 \
   --max-context-tokens 49152 \
-  --max-logprob-context-tokens 24576 \
+  --max-logprob-context-tokens 16384 \
   --max-sidecar-entries 2048 \
   --port 8001
 ```
@@ -81,9 +81,11 @@ configured 49,152-token context window. If tool output would overflow the
 prompt budget, it preserves system/user messages and truncates the oldest tool
 outputs with an explicit marker. Non-tool context overflow fails before calling
 the TRL server.
-Sampled-logprob GRPO requests use a stricter 24,576-token context cap so TRL can
+Sampled-logprob GRPO requests use a stricter 16,384-token context cap so TRL can
 recompute policy logprobs on one H100 without materializing 49k-token logits.
 Ordinary baseline, SFT, and final evaluation keep the full context window.
+The CLI also defaults `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` to
+reduce fragmentation across repeated variable-length GRPO batches.
 
 ## Per-update lifecycle
 

docs/training-pipeline.md

@@ -222,6 +222,8 @@ The bridge fits each served prompt to the configured model context by truncating
 oldest tool outputs only; system and user instructions are never truncated.
 GRPO sampled-logprob requests use a smaller context cap than evaluation requests
 to keep trainer-side policy-logprob recomputation within GPU memory.
+The CLI enables PyTorch expandable CUDA segments unless the operator already
+supplied a custom allocator configuration.
 
 ## Execute and resume
 

pipelines/benchflow-task-posttrain/README.md

@@ -168,6 +168,8 @@ The bridge normalizes OpenCode follow-up tool arguments and token-fits oversized
 tool results to the server context without truncating system or user messages.
 Its sampled-logprob path uses a stricter context cap for trainer memory while
 ordinary evaluation retains the full model context.
+CLI runs also enable expandable CUDA segments by default to reduce GRPO memory
+fragmentation.
 
 The final contract is:
 

pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/cli.py

@@ -9,6 +9,8 @@
 import sys
 from pathlib import Path
 
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
 from .config import load_config
 from .pipeline import Pipeline
 
@@ -179,7 +181,7 @@ def build_parser() -> argparse.ArgumentParser:
     bridge.add_argument("--api-key-env", default="BENCHFLOW_PROVIDER_API_KEY")
     bridge.add_argument("--max-tokens", type=int, default=4096)
     bridge.add_argument("--max-context-tokens", type=int, default=49152)
-    bridge.add_argument("--max-logprob-context-tokens", type=int, default=24576)
+    bridge.add_argument("--max-logprob-context-tokens", type=int, default=16384)
     bridge.add_argument("--max-sidecar-entries", type=int, default=2048)
     bridge.add_argument("--host", default="0.0.0.0")
     bridge.add_argument("--port", type=int, default=8001)

pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/model_bridge.py

@@ -35,7 +35,7 @@ class ModelBridgeConfig:
     api_key: str | None = None
     max_tokens_per_call: int = 4096
     max_context_tokens: int = 49152
-    max_logprob_context_tokens: int = 24576
+    max_logprob_context_tokens: int = 16384
     timeout_seconds: float = 900.0
     max_sidecar_entries: int = 2048
 

pipelines/benchflow-task-posttrain/tests/test_cli.py

@@ -46,7 +46,7 @@ def test_model_bridge_cli_contract() -> None:
             "--max-context-tokens",
             "32768",
             "--max-logprob-context-tokens",
-            "16384",
+            "12288",
             "--max-sidecar-entries",
             "256",
             "--port",
@@ -59,7 +59,7 @@ def test_model_bridge_cli_contract() -> None:
     assert args.api_key_env == "BENCHFLOW_PROVIDER_API_KEY"
     assert args.max_tokens == 2048
     assert args.max_context_tokens == 32768
-    assert args.max_logprob_context_tokens == 16384
+    assert args.max_logprob_context_tokens == 12288
     assert args.max_sidecar_entries == 256
     assert args.port == 9001