From d624a35692c39e1374e487f8af0d837089d365bb Mon Sep 17 00:00:00 2001 From: Bingran You Date: Tue, 14 Jul 2026 14:04:09 -0700 Subject: [PATCH] Cap GRPO logprob context --- docs/opencode-grpo.md | 4 ++ docs/training-pipeline.md | 2 + pipelines/benchflow-task-posttrain/README.md | 2 + .../posttrainarena/benchflow_pipeline/cli.py | 2 + .../benchflow_pipeline/model_bridge.py | 17 ++++- .../tests/test_cli.py | 3 + .../tests/test_model_bridge.py | 70 +++++++++++++++++++ 7 files changed, 99 insertions(+), 1 deletion(-) diff --git a/docs/opencode-grpo.md b/docs/opencode-grpo.md index 7a85d03..cb08b70 100644 --- a/docs/opencode-grpo.md +++ b/docs/opencode-grpo.md @@ -47,6 +47,7 @@ posttrainarena-train model-bridge \ --tokenizer-revision \ --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 diff --git a/docs/training-pipeline.md b/docs/training-pipeline.md index a13f53f..c8ddad8 100644 --- a/docs/training-pipeline.md +++ b/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 diff --git a/pipelines/benchflow-task-posttrain/README.md b/pipelines/benchflow-task-posttrain/README.md index 346fbb4..9fd9aef 100644 --- a/pipelines/benchflow-task-posttrain/README.md +++ b/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: diff --git a/pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/cli.py b/pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/cli.py index 423ec4e..07bc4b0 100644 --- a/pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/cli.py +++ b/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, diff --git a/pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/model_bridge.py b/pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/model_bridge.py index c784226..51d8cb5 100644 --- a/pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/model_bridge.py +++ b/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, ) ) diff --git a/pipelines/benchflow-task-posttrain/tests/test_cli.py b/pipelines/benchflow-task-posttrain/tests/test_cli.py index a415aa2..7f0fb76 100644 --- a/pipelines/benchflow-task-posttrain/tests/test_cli.py +++ b/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 diff --git a/pipelines/benchflow-task-posttrain/tests/test_model_bridge.py b/pipelines/benchflow-task-posttrain/tests/test_model_bridge.py index 8ad30ff..b7a3008 100644 --- a/pipelines/benchflow-task-posttrain/tests/test_model_bridge.py +++ b/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, + )