Cap GRPO sampled-logprob context

docs/opencode-grpo.md

@@ -47,6 +47,7 @@ posttrainarena-train model-bridge \
   --tokenizer-revision <immutable-sha> \
   --max-tokens 4096 \
   --max-context-tokens 49152 \
+  --max-logprob-context-tokens 24576 \
   --max-sidecar-entries 2048 \
   --port 8001
 ```
@@ -80,6 +81,9 @@ 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
+recompute policy logprobs on one H100 without materializing 49k-token logits.
+Ordinary baseline, SFT, and final evaluation keep the full context window.
 
 ## Per-update lifecycle
 

docs/training-pipeline.md

@@ -220,6 +220,8 @@ It also converts OpenCode's stringified follow-up tool arguments back to JSON
 objects before Qwen3.5 chat-template rendering.
 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.
 
 ## Execute and resume
 

pipelines/benchflow-task-posttrain/README.md

@@ -166,6 +166,8 @@ before SFT evaluation, the current GRPO policy before each rollout batch, and
 final weights before the held-out evaluation.
 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.
 
 The final contract is:
 

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

@@ -179,6 +179,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-sidecar-entries", type=int, default=2048)
     bridge.add_argument("--host", default="0.0.0.0")
     bridge.add_argument("--port", type=int, default=8001)
@@ -402,6 +403,7 @@ def main(argv: list[str] | None = None) -> int:
             api_key=os.environ.get(args.api_key_env),
             max_tokens_per_call=args.max_tokens,
             max_context_tokens=args.max_context_tokens,
+            max_logprob_context_tokens=args.max_logprob_context_tokens,
             max_sidecar_entries=args.max_sidecar_entries,
             host=args.host,
             port=args.port,

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

@@ -35,6 +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
     timeout_seconds: float = 900.0
     max_sidecar_entries: int = 2048
 
@@ -53,6 +54,15 @@ def __post_init__(self) -> None:
             raise ValueError(
                 "max_context_tokens must be an integer greater than max_tokens_per_call"
             )
+        if (
+            not isinstance(self.max_logprob_context_tokens, int)
+            or isinstance(self.max_logprob_context_tokens, bool)
+            or self.max_logprob_context_tokens <= self.max_tokens_per_call
+        ):
+            raise ValueError(
+                "max_logprob_context_tokens must be an integer greater than "
+                "max_tokens_per_call"
+            )
         if (
             not isinstance(self.max_sidecar_entries, int)
             or isinstance(self.max_sidecar_entries, bool)
@@ -401,12 +411,15 @@ def _trl_request(
     )
     normalized_messages = normalize_tool_call_arguments(messages)
     tools = body.get("tools")
+    context_tokens = config.max_context_tokens
+    if capture_logprobs:
+        context_tokens = min(context_tokens, config.max_logprob_context_tokens)
     fitted_messages, original_tokens, fitted_tokens, truncated_messages = (
         fit_messages_to_context(
             tokenizer=tokenizer,
             messages=normalized_messages,
             tools=tools,
-            max_prompt_tokens=config.max_context_tokens - max_tokens,
+            max_prompt_tokens=context_tokens - max_tokens,
         )
     )
     if truncated_messages:
@@ -616,6 +629,7 @@ def serve_model_bridge(
     api_key: str | None,
     max_tokens_per_call: int,
     max_context_tokens: int,
+    max_logprob_context_tokens: int,
     max_sidecar_entries: int,
     host: str,
     port: int,
@@ -630,6 +644,7 @@ def serve_model_bridge(
             api_key=api_key,
             max_tokens_per_call=max_tokens_per_call,
             max_context_tokens=max_context_tokens,
+            max_logprob_context_tokens=max_logprob_context_tokens,
             max_sidecar_entries=max_sidecar_entries,
         )
     )

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

@@ -45,6 +45,8 @@ def test_model_bridge_cli_contract() -> None:
             "2048",
             "--max-context-tokens",
             "32768",
+            "--max-logprob-context-tokens",
+            "16384",
             "--max-sidecar-entries",
             "256",
             "--port",
@@ -57,6 +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_sidecar_entries == 256
     assert args.port == 9001
 

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

@@ -473,6 +473,68 @@ async def fake_chat(payload: dict[str, Any]) -> dict[str, Any]:
     assert len(tool_content) < 1000
 
 
+def test_model_bridge_uses_stricter_context_for_logprob_requests() -> None:
+    captured: dict[str, Any] = {}
+
+    async def fake_chat(payload: dict[str, Any]) -> dict[str, Any]:
+        captured["payload"] = payload
+        return _upstream("OK")
+
+    tokenizer = FakeTokenizer()
+    app = create_model_bridge_app(
+        ModelBridgeConfig(
+            upstream_url="http://127.0.0.1:8000",
+            tokenizer_id="Qwen/Qwen3-4B",
+            max_tokens_per_call=64,
+            max_context_tokens=512,
+            max_logprob_context_tokens=448,
+        ),
+        tokenizer=tokenizer,
+        chat_call=fake_chat,
+    )
+    response = TestClient(app).post(
+        "/v1/chat/completions",
+        json={
+            "messages": [
+                {"role": "user", "content": "inspect"},
+                {
+                    "role": "assistant",
+                    "content": "",
+                    "tool_calls": [
+                        {
+                            "id": "call-1",
+                            "type": "function",
+                            "function": {
+                                "name": "read",
+                                "arguments": '{"filePath":"/tmp/data.csv"}',
+                            },
+                        }
+                    ],
+                },
+                {
+                    "role": "tool",
+                    "tool_call_id": "call-1",
+                    "content": "x" * 1000,
+                },
+            ],
+            "logprobs": True,
+        },
+    )
+
+    assert response.status_code == 200
+    payload = captured["payload"]
+    prompt_tokens = len(
+        tokenizer.apply_chat_template(
+            payload["messages"][0],
+            tools=payload["tools"],
+            tokenize=True,
+            add_generation_prompt=True,
+        )["input_ids"][0]
+    )
+    assert prompt_tokens <= 384
+    assert payload["messages"][0][2]["content"].endswith(TOOL_OUTPUT_TRUNCATION_MARKER)
+
+
 def test_model_bridge_caps_tokens_per_call() -> None:
     captured: dict[str, Any] = {}
 
@@ -598,3 +660,11 @@ def test_model_bridge_rejects_invalid_token_cap() -> None:
             max_tokens_per_call=64,
             max_context_tokens=64,
         )
+
+    with pytest.raises(ValueError, match="max_logprob_context_tokens"):
+        ModelBridgeConfig(
+            upstream_url="http://127.0.0.1:8000",
+            tokenizer_id="Qwen/Qwen3-4B",
+            max_tokens_per_call=64,
+            max_logprob_context_tokens=64,
+        )