diff --git a/docs/opencode-grpo.md b/docs/opencode-grpo.md index cb08b70..bddefd8 100644 --- a/docs/opencode-grpo.md +++ b/docs/opencode-grpo.md @@ -47,7 +47,7 @@ posttrainarena-train model-bridge \ --tokenizer-revision \ --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 diff --git a/docs/training-pipeline.md b/docs/training-pipeline.md index c8ddad8..fb527fe 100644 --- a/docs/training-pipeline.md +++ b/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 diff --git a/pipelines/benchflow-task-posttrain/README.md b/pipelines/benchflow-task-posttrain/README.md index 9fd9aef..d37c285 100644 --- a/pipelines/benchflow-task-posttrain/README.md +++ b/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: 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 07bc4b0..36b3ba6 100644 --- a/pipelines/benchflow-task-posttrain/src/posttrainarena/benchflow_pipeline/cli.py +++ b/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) 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 51d8cb5..9db4b14 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,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 diff --git a/pipelines/benchflow-task-posttrain/tests/test_cli.py b/pipelines/benchflow-task-posttrain/tests/test_cli.py index 7f0fb76..734aa3b 100644 --- a/pipelines/benchflow-task-posttrain/tests/test_cli.py +++ b/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