diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 38f96c73c..41433015f 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -228,7 +228,7 @@ jobs: strategy: fail-fast: false matrix: - info: [{"num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_glm4.7_30B_A3B_pd_mooncake.py"}, {"num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"num_gpus": 8, "test_file": "test_qwen3.6_35B_A3B_pd_mooncake.py", "use_deepep": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo_disaggregate.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo_train_critic_only.py"}, {"num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_vllm.py"}] + info: [{"num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_glm4.7_30B_A3B_pd_mooncake.py"}, {"num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"num_gpus": 8, "test_file": "test_qwen3.6_35B_A3B_pd_mooncake.py", "use_deepep": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo_disaggregate.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo_train_critic_only.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_streaming_partial_rollout.py"}, {"num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_vllm.py"}] defaults: run: working-directory: ${{ github.workspace }} @@ -482,7 +482,7 @@ jobs: strategy: fail-fast: false matrix: - info: [{"num_gpus": 0, "test_file": "test_megatron_argument_validation.py"}, {"num_gpus": 0, "test_file": "test_rollout_validation.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_rollout_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_runtime_hook_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_path_loading_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_generate_contracts.py"}, {"num_gpus": 0, "test_file": "test_rm_deepscaler.py"}, {"num_gpus": 0, "test_file": "test_rm_f1.py"}, {"num_gpus": 0, "test_file": "test_rm_gpqa.py"}, {"num_gpus": 0, "test_file": "test_rm_math.py"}, {"num_gpus": 0, "test_file": "test_rm_math_dapo.py"}] + info: [{"num_gpus": 0, "test_file": "test_megatron_argument_validation.py"}, {"num_gpus": 0, "test_file": "test_rollout_validation.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_rollout_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_runtime_hook_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_path_loading_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_generate_contracts.py"}, {"num_gpus": 0, "test_file": "test_rm_deepscaler.py"}, {"num_gpus": 0, "test_file": "test_rm_f1.py"}, {"num_gpus": 0, "test_file": "test_rm_gpqa.py"}, {"num_gpus": 0, "test_file": "test_rm_math.py"}, {"num_gpus": 0, "test_file": "test_rm_math_dapo.py"}, {"num_gpus": 0, "test_file": "test_dp_schedule.py"}, {"num_gpus": 0, "test_file": "test_cp_utils.py"}, {"num_gpus": 0, "test_file": "test_metric_report.py"}, {"num_gpus": 0, "test_file": "test_metric_report_dist.py"}, {"num_gpus": 0, "test_file": "test_loss_cp_invariance.py"}, {"num_gpus": 0, "test_file": "test_sample.py"}, {"num_gpus": 0, "test_file": "utils/test_hf_checkpoint_saver.py"}] defaults: run: working-directory: ${{ github.workspace }} @@ -511,6 +511,76 @@ jobs: pip install torch --index-url https://download.pytorch.org/whl/cpu pip install pytest numpy packaging pyyaml omegaconf tqdm httpx pybase64 pylatexenc sympy aiohttp pillow + + - name: Install + shell: bash + run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps + + + - name: Execute + shell: bash + run: | + + TEST_PATH="${{ matrix.info.test_file }}" + if [[ "$TEST_PATH" != tests/* ]]; then + TEST_PATH="tests/$TEST_PATH" + fi + TEST_ARGS="${{ matrix.info.test_args || '' }}" + if [[ -n "$TEST_ARGS" ]]; then + read -r -a TEST_ARGS_ARRAY < <(printf '%s\n' "$TEST_ARGS") + else + TEST_ARGS_ARRAY=() + fi + if [ "${{ matrix.info.num_gpus }}" = "0" ]; then + python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" + else + python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" + fi + + + agent-adapter-test: + needs: pre-commit + + if: github.event_name == 'pull_request' || github.event_name == 'workflow_dispatch' + + + runs-on: ubuntu-latest + + strategy: + fail-fast: false + matrix: + info: [{"num_gpus": 0, "test_file": "test_agent_trajectory.py"}, {"num_gpus": 0, "test_file": "test_agent_adapters.py"}, {"num_gpus": 0, "test_file": "test_agent_sdk_adapters.py"}] + defaults: + run: + working-directory: ${{ github.workspace }} + env: + GITHUB_COMMIT_NAME: ${{ github.sha }}_${{ github.event.pull_request.number || 'non-pr' }} + WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }} + SLIME_TEST_ENABLE_INFINITE_RUN: ${{ (github.event_name == 'workflow_dispatch' && github.event.inputs.infinite_run) || 'false' }} + SLIME_TEST_USE_DEEPEP: ${{ matrix.info.use_deepep || '0' }} + SLIME_TEST_USE_FP8_ROLLOUT: ${{ matrix.info.use_fp8_rollout || '0' }} + SLIME_TEST_ENABLE_EVAL: ${{ matrix.info.enable_eval || '1' }} + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.10' + cache: 'pip' + + - name: Install dependencies + shell: bash + run: | + pip install torch --index-url https://download.pytorch.org/whl/cpu + pip install pytest numpy packaging pyyaml omegaconf tqdm httpx pybase64 pylatexenc sympy aiohttp pillow + + pip install openai openai-agents anthropic + + - name: Install shell: bash run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index 0687a0cbb..cd097b6f1 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.github/workflows/pr-test.yml.j2 @@ -29,6 +29,7 @@ {'test_file': 'test_qwen3_4B_ppo.py', 'num_gpus': 8}, {'test_file': 'test_qwen3_4B_ppo_disaggregate.py', 'num_gpus': 8}, {'test_file': 'test_qwen3_4B_ppo_train_critic_only.py', 'num_gpus': 8}, + {'test_file': 'test_qwen3_4B_streaming_partial_rollout.py', 'num_gpus': 8}, {'test_file': 'test_moonlight_16B_A3B.py', 'num_gpus': 8}, {'test_file': 'test_moonlight_16B_A3B_r3.py', 'num_gpus': 8, 'enable_eval': '0'}, {'test_file': 'test_qwen2.5_0.5B_debug_rollout_then_train.py', 'num_gpus': 8}, @@ -65,6 +66,25 @@ {'test_file': 'test_rm_gpqa.py', 'num_gpus': 0}, {'test_file': 'test_rm_math.py', 'num_gpus': 0}, {'test_file': 'test_rm_math_dapo.py', 'num_gpus': 0}, + {'test_file': 'test_dp_schedule.py', 'num_gpus': 0}, + {'test_file': 'test_cp_utils.py', 'num_gpus': 0}, + {'test_file': 'test_metric_report.py', 'num_gpus': 0}, + {'test_file': 'test_metric_report_dist.py', 'num_gpus': 0}, + {'test_file': 'test_loss_cp_invariance.py', 'num_gpus': 0}, + {'test_file': 'test_sample.py', 'num_gpus': 0}, + {'test_file': 'utils/test_hf_checkpoint_saver.py', 'num_gpus': 0}, + ], + }, + + 'agent-adapter-test': { + 'label': 'run-ci-agent-adapter', + 'always': True, + 'cpu': True, + 'extra_pip_deps': 'openai openai-agents anthropic', + 'tests': [ + {'test_file': 'test_agent_trajectory.py', 'num_gpus': 0}, + {'test_file': 'test_agent_adapters.py', 'num_gpus': 0}, + {'test_file': 'test_agent_sdk_adapters.py', 'num_gpus': 0}, ], }, @@ -182,6 +202,9 @@ jobs: run: | pip install torch --index-url https://download.pytorch.org/whl/cpu pip install pytest numpy packaging pyyaml omegaconf tqdm httpx pybase64 pylatexenc sympy aiohttp pillow +<% if config.get('extra_pip_deps') %> + pip install << config.extra_pip_deps >> +<% endif %> - name: Install shell: bash diff --git a/examples/coding_agent_rl/README.md b/examples/coding_agent_rl/README.md new file mode 100644 index 000000000..61c81e036 --- /dev/null +++ b/examples/coding_agent_rl/README.md @@ -0,0 +1,179 @@ +# Coding-Agent RL + +This directory provides an example of running end-to-end **SWE (Software-Engineering) coding-agent RL** with vime: a real coding agent (claude-code CLI) drives `Read/Edit/Grep/Bash/Agent` tools inside a fresh sandbox per sample, the model produces a `git diff`, and the diff is graded against the dataset's test harness in a second clean sandbox (no test-cheating). + +Two example files and one shared adapter implement the loop: + +- `generate.py` — per-sample `generate()` registered via `--custom-generate-function-path`. Boots the sandbox, runs claude-code, captures the diff, scores it, and emits one or more `Sample`s back to vime. +- `vime.agent.adapters.AnthropicAdapter` — the shared Anthropic Messages adapter. claude-code talks to it as if it were Anthropic; the adapter tokenizes the current message history each turn, records prompt/output token snapshots, preserves model-generated tokens (`loss_mask=1`) only while later prompts stitch onto them, masks template/observation tokens (`0`), and emits **three kinds of segments** per trajectory: `subagent` (completed `Task/Agent` dispatch), `wipe` (chain frozen by auto-compact), `final` (tail of the main chain). +- `sandbox.py` — coding-agent/SWE helpers built on `vime.agent.sandbox`: install bootstraps, spawn claude-code, capture patches, and run the fresh-sandbox evaluator. The shared sandbox contract lives in `vime.agent.sandbox.Sandbox`. + +`generate.py` owns one `AnthropicAdapter` instance. For each sample it calls +`adapter.open_session(...)` before starting claude-code, serves `adapter.app` as +the Anthropic-compatible endpoint, and drains trainable `TokenSegment`s with +`await adapter.finish_session(...)` when the trajectory ends. + +## Environment Setup + +The vime training stack itself follows the standard setup. On top of that you need: + +1. **An E2B-compatible sandbox cluster** (or any provider that speaks the E2B SDK). Configure via `E2B_API_KEY` (e.g. the standard `e2b_xxx` key from https://e2b.dev, or any internal endpoint that accepts the same SDK). The official SDK validates this value locally, so internal gateways that ignore auth still need a syntactically valid `e2b_` + 40 hex-character placeholder. +2. **Host-side tarballs** that get uploaded into each sandbox at boot: + - Node 22 (`node-v22.x-linux-x64.tar.xz`) — exported as `SWE_HOST_NODE_TARBALL`. + - Claude Code CLI npm tarball (`anthropic-ai-claude-code-local-linux-x64.tgz`) — exported as `SWE_HOST_CC_TARBALL`. +3. **A sandbox metadata file** (`SWE_SANDBOX_METADATA_FILE`, or the generic `VIME_AGENT_SANDBOX_METADATA_FILE`) — JSON dict whose keys are passed as routing tags when booting an E2B sandbox. Must contain the image key referenced by `SWE_SANDBOX_IMAGE_METADATA_KEY` / `VIME_AGENT_SANDBOX_IMAGE_METADATA_KEY` (e.g. `image`). +4. **Network reachability**: each sandbox dials back to the vime head node's Anthropic adapter over `http://${VIME_HEAD_HOST}:${SHIM_PORT}`. The head host must be reachable from inside the sandboxes (set `VIME_HEAD_HOST` to a routable IP, not `127.0.0.1`). + +## Dataset Format + +Standard vime JSONL with three keys: + +```jsonc +{ + "prompt": "", + "label": "", + "metadata": { + "image": "swedev/scaleswe.oh.34:", // sandbox image reference + "workdir": "/workspace/", // repo path inside the sandbox + "problem_statement": "", + // exactly one of the following two graders: + "swepro": { /* SWE-bench Pro test harness — preferred */ }, + "eval_cmd": "pytest -x tests/..." // last-resort: exit 0 = solved + // sweb-style rows: metadata.remote_env_info.f2p_script (Python file + // ending in `sys.exit(pytest.main(...))`) is auto-wrapped into eval_cmd. + } +} +``` + +Wire it up with `--input-key prompt --label-key label --metadata-key metadata`. + +## Running the Script + +Override the paths at the top of the launcher, then run from a long-lived shell on the Ray head node (do **not** wrap in `nohup` — Ray child processes get cleaned up with it): + +```bash +cd vime/ + +export HF_CHECKPOINT=/path/to/Qwen3.6-35B-A3B +export REF_MODEL_PATH=/path/to/Qwen3.6-35B-A3B_torch_dist +export PROMPT_DATA=/path/to/swe_train.jsonl +export SANDBOX_METADATA_FILE=/path/to/sandbox_metadata.json +export SWE_HOST_NODE_TARBALL=/path/to/node-v22.20.0-linux-x64.tar.xz +export SWE_HOST_CC_TARBALL=/path/to/anthropic-ai-claude-code-local-linux-x64.tgz + +bash examples/coding_agent_rl/run_qwen36_35b_a3b_swe_8nodes.sh +``` + +The launcher brings up Ray across all hosts in `/root/mpi_rack_hostfile`, dumps every rollout to `runs/${EXP_TAG}_${STAMP}/rollout_dumps/`, and tees stdout into `runs/${EXP_TAG}_${STAMP}/run.log`. + +## New Arguments + +`generate.py` is wired in through vime's standard custom-generate hook: + +```bash +ROLLOUT_ARGS=( + --custom-generate-function-path examples.coding_agent_rl.generate.generate + --prompt-data "${PROMPT_DATA}" + --input-key prompt + --label-key label + --metadata-key metadata + --rollout-batch-size 8 + --n-samples-per-prompt 8 + --rollout-max-context-len 96000 + --rollout-max-response-len 32768 + --rollout-stop-token-ids 248046 248044 + --save-debug-rollout-data "${RUN_ROOT}/rollout_dumps/rollout_{rollout_id}.pt" +) +``` + +claude-code's tool invocations are parsed from the model output. By default the +adapter uses the built-in XML tool-call fallback (`parse_xml_tool_uses`), which +handles the `` text Qwen3-Coder emits, so no extra +engine-side parser configuration is required. (The adapter can optionally +delegate to SGLang's reasoning/function-call parsers when +`vllm_tool_call_parser` / `vllm_reasoning_parser` are set on `args`, but this is +not used by this example and is not wired into vime's vLLM arguments.) + +## SWE-specific Environment Knobs + +All set in the launcher; tune per cluster. + +| Variable | Default | Meaning | +| --- | --- | --- | +| `VIME_HEAD_HOST` | `${MASTER_ADDR}` | Public IP the sandbox uses to reach the Anthropic adapter. **Must be routable from inside the sandbox.** | +| `SHIM_BIND_HOST` / `SHIM_PORT` | `0.0.0.0` / `18001` | Bind address of the adapter shim on the head node. | +| `E2B_API_KEY` | — | E2B (or compatible) API key. | +| `SWE_SANDBOX_METADATA_FILE` / `VIME_AGENT_SANDBOX_METADATA_FILE` | — | JSON dict of routing metadata passed at sandbox boot. | +| `SWE_SANDBOX_IMAGE_METADATA_KEY` / `VIME_AGENT_SANDBOX_IMAGE_METADATA_KEY` | — | Which key in the metadata file holds the image reference (e.g. `image`). | +| `SWE_HOST_NODE_TARBALL` | — | Host path to Node 22 tarball uploaded into each sandbox. | +| `SWE_HOST_CC_TARBALL` | — | Host path to the Claude Code CLI npm tarball. | +| `SWE_TIME_BUDGET_SEC` | `1800` | Wallclock budget for one agent run. | +| `SWE_EVAL_TIMEOUT_SEC` | `600` | Wallclock cap on the evaluator sandbox. | +| `SWE_BOOT_CONCURRENCY` | `6` | Cap on simultaneous sandbox boots (eases h2/SSL long-tail). | +| `SWE_CLAUDE_EXTRA_ARGS` | (see launcher) | Extra flags appended to the `claude` CLI invocation — registers the read-only `investigator` sub-agent, disables `WebFetch`/`WebSearch`, disables slash commands. | +| `SWE_CC_PROMPT` | unset | Optional override for the user-turn prompt. Setting this to require sub-agent dispatch is the most reliable way to maximize fan-out. | + +`--rollout-max-response-len` is the per-turn generation cap passed to each vLLM +`/inference/v1/generate` call as the sampling-params `max_tokens`. +`--rollout-max-context-len` is the multi-turn prompt+response budget enforced +only during generation: each turn clamps the generation length to the remaining +context. Trajectory merge/export keeps the emitted segments and does not drop +them for length. + +## String-in, Token-out Trajectories + +The coding-agent environment is string/message based: claude-code sends +Anthropic Messages requests, receives streamed text/thinking/tool-use blocks, +and later sends back rendered tool observations. Training, however, must stay +token based. A trajectory is only a valid RL target when the optimized tokens +are the same tokens the rollout model actually sampled. + +The Anthropic adapter therefore follows a **string in, token out** contract: + +- Each incoming message history is rendered with the served model's chat + template and sent to vLLM as `token_ids`. +- vLLM's `/inference/v1/generate` is called with `logprobs` set; the adapter + records the exact `prompt_ids`, sampled `output_ids` (`choices[0].token_ids`), + and per-token rollout logprobs (`choices[0].logprobs.content[i].logprob`) for + that turn. +- At training export time, samples are assembled from those saved token ids. + The decoded `response` field is only a readable sidecar; it is not + re-tokenized to recover the training sequence. + +Multi-turn agents still force the adapter to tokenize later message +histories, because tool observations and claude-code's own compacted messages +arrive as strings. `vime.agent.trajectory.merge_turns` stitches those later +prompts against the saved token stream: + +- New prompt suffixes that are tool/user/environment context are appended with + `loss_mask=0`. +- Fresh model outputs from vLLM are appended with `loss_mask=1`. +- If a later prompt no longer token-matches an earlier sampled output, the + unmatched suffix is dropped. If the drift cuts through the middle of a + previous model output, the retained prefix of that whole output turn is also + assigned `loss_mask=0`. + +That last case is the important correctness guard. A re-tokenization mismatch +can make a string-level conversation look continuous while token-level +provenance is broken. vime keeps the context needed to continue the agent, but +does not backprop through tokens whose sampled origin can no longer be proven. + +## Fan-out Semantics + +- `generate()` returns `list[Sample]` — one Sample per trajectory **segment** (`subagent` / `wipe` / `final`). +- Per-trajectory reward is split as `reward / K` across segments; `rollout_id` is shared so the per-rollout-mean loss reducer still counts the trajectory once. +- Sub-agent dispatch increases `K` (each completed `Agent` turn block becomes its own segment), so the effective batch after flatten can be much larger than `rollout_batch_size * n_samples_per_prompt`. + +## Porting to a New Sandbox Backend + +`vime.agent.sandbox.Sandbox` exposes the shared sandbox contract, and +`vime.agent.sandbox.E2BSandbox` is the E2B implementation: + +```python +await sb.exec(cmd, user=..., check=..., timeout=...) +await sb.write_file(sandbox_path, content_or_host_path, user=...) +await sb.read_file(sandbox_path, user=...) +async with E2BSandbox(...) as sb: ... +``` + +Reimplement those on Docker / Modal / a local VM and everything in `generate.py` keeps working unchanged. diff --git a/examples/coding_agent_rl/aiohttp_threaded.py b/examples/coding_agent_rl/aiohttp_threaded.py new file mode 100644 index 000000000..a5a17652d --- /dev/null +++ b/examples/coding_agent_rl/aiohttp_threaded.py @@ -0,0 +1,90 @@ +"""Run an ``aiohttp.web.Application`` in a background daemon thread.""" + +from __future__ import annotations + +import asyncio +import threading +from dataclasses import dataclass +from typing import Any + +from aiohttp import web + + +@dataclass +class AppHandle: + host: str + port: int + thread: threading.Thread + loop: asyncio.AbstractEventLoop + runner: web.AppRunner + + @property + def url(self) -> str: + return f"http://{self.host}:{self.port}" + + def stop(self) -> None: + async def _shutdown() -> None: + await self.runner.cleanup() + + try: + fut = asyncio.run_coroutine_threadsafe(_shutdown(), self.loop) + fut.result(timeout=10) + except Exception: + pass + self.loop.call_soon_threadsafe(self.loop.stop) + self.thread.join(timeout=5) + + +def run_app_in_thread( + app: web.Application, + *, + host: str = "0.0.0.0", + port: int = 0, + thread_name: str = "aiohttp-app", + start_timeout_sec: float = 15.0, + runner_kwargs: dict[str, Any] | None = None, +) -> AppHandle: + """Spin up ``app`` on a daemon thread; block until it is listening. + + ``runner_kwargs`` is forwarded to ``web.AppRunner`` (e.g. pass + ``{"handler_cancellation": True}`` to make a client disconnect cancel + the in-flight handler coroutine). + """ + started = threading.Event() + err_box: list[BaseException] = [] + box: dict[str, Any] = {} + + def _run() -> None: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + runner = web.AppRunner(app, **(runner_kwargs or {})) + loop.run_until_complete(runner.setup()) + site = web.TCPSite(runner, host, port) + loop.run_until_complete(site.start()) + actual_port = port + for sock in site._server.sockets: # type: ignore[attr-defined] + actual_port = sock.getsockname()[1] + break + box["loop"] = loop + box["runner"] = runner + box["port"] = actual_port + started.set() + loop.run_forever() + except BaseException as e: # pragma: no cover + err_box.append(e) + started.set() + raise + + thread = threading.Thread(target=_run, name=thread_name, daemon=True) + thread.start() + started.wait(timeout=start_timeout_sec) + if err_box: + raise err_box[0] + return AppHandle( + host=host, + port=int(box["port"]), + thread=thread, + loop=box["loop"], + runner=box["runner"], + ) diff --git a/examples/coding_agent_rl/generate.py b/examples/coding_agent_rl/generate.py new file mode 100644 index 000000000..b048cff52 --- /dev/null +++ b/examples/coding_agent_rl/generate.py @@ -0,0 +1,364 @@ +"""Coding-Agent RL: per-sample generate() function for vime. + +Wire-up: + + --custom-generate-function-path examples.coding_agent_rl.generate.generate + +``generate()`` is intentionally a small four-stage orchestrator: + + 1. ``sandbox.run_claude_code`` prepares the agent sandbox and runs claude-code. + 2. ``sandbox.git_diff`` captures the model-produced patch. + 3. ``sandbox.evaluate`` scores that patch in a second clean sandbox. + 4. ``_merge_samples`` combines reward + adapter ``TokenSegment``s, + delegating segment-to-``Sample`` fan-out to ``vime.agent.trajectory``. + +All sandbox-side details live in ``sandbox.py``; the LLM plumbing +(Anthropic <-> vLLM ``/inference/v1/generate``, token capture, 3-kind segment +split) uses ``vime.agent.adapters.AnthropicAdapter``. + +Dataset row ``metadata`` schema:: + + image: str # sandbox image + workdir: str # repo path inside the sandbox + problem_statement: str # issue body (falls back to sample.prompt) + swepro: dict|None # SWE-bench Pro test harness (preferred) + eval_cmd: str|None # last-resort: shell command (exit 0 = solved) + +Also accepted (sweb-style rows): ``metadata.remote_env_info.f2p_script`` — +a self-contained Python test file ending in ``sys.exit(pytest.main(...))``. +When ``eval_cmd`` is absent, ``_metadata`` wraps this script into a base64 +materialize-and-run shell command so the existing eval path stays unchanged. + +Env knobs (set in run.sh): + + SWE_HOST_NODE_TARBALL host path to a Node 22 tarball (REQUIRED) + SWE_HOST_CC_TARBALL host path to the Claude Code npm tarball (REQUIRED) + SWE_TIME_BUDGET_SEC 1800 per agent run, wallclock + SWE_EVAL_TIMEOUT_SEC 600 per eval test execution + SHIM_BIND_HOST 0.0.0.0 + SHIM_PORT 18001 + VIME_HEAD_HOST public host the sandboxes use to reach the adapter (REQUIRED) +""" + +from __future__ import annotations + +import asyncio +import base64 +import logging +import os +import secrets +import time +import traceback +from dataclasses import dataclass +from typing import Any + +from vime.agent.adapters import AnthropicAdapter +from vime.agent.trajectory import TokenSegment, fan_out_sample_segments +from vime.utils.misc import SingletonMeta +from vime.utils.processing_utils import load_tokenizer +from vime.utils.types import Sample + +from . import sandbox +from .aiohttp_threaded import run_app_in_thread + +logger = logging.getLogger(__name__) + + +SWE_TIME_BUDGET_SEC = int(os.environ.get("SWE_TIME_BUDGET_SEC", "1800")) +SWE_EVAL_TIMEOUT_SEC = int(os.environ.get("SWE_EVAL_TIMEOUT_SEC", "600")) +# Wall-clock guard for the entire generate() call. Defaults to +# SWE_TIME_BUDGET_SEC + SWE_EVAL_TIMEOUT_SEC + 180 (buffer for sandbox boot, +# diff capture, etc). When exceeded, the in-flight sample is aborted with +# reason `wall_clock_timeout` and the rest of the rollout continues -- this +# isolates a single hung trajectory (e.g. stuck in sandbox.evaluate) so it +# does not kill the whole training step. +SWE_GENERATE_GUARD_SEC = int(os.environ.get("SWE_GENERATE_GUARD_SEC", "0") or 0) or ( + SWE_TIME_BUDGET_SEC + SWE_EVAL_TIMEOUT_SEC + 180 +) +SHIM_BIND_HOST = os.environ.get("SHIM_BIND_HOST", "0.0.0.0") +SHIM_PORT = int(os.environ.get("SHIM_PORT", "18001")) + + +# --------------------------------------------------------------------------- +# Singleton: tokenizer + in-process Anthropic adapter + reducer +# --------------------------------------------------------------------------- +class _State(metaclass=SingletonMeta): + def __init__(self, args) -> None: + self.tokenizer = load_tokenizer(args.hf_checkpoint, trust_remote_code=True) + self.max_context_len = int(getattr(args, "rollout_max_context_len", 0) or 0) + self.tool_parser = getattr(args, "vllm_tool_call_parser", None) or None + self.reasoning_parser = getattr(args, "vllm_reasoning_parser", None) or None + vllm_url = f"http://{args.vllm_router_ip}:{args.vllm_router_port}" + public_host = os.environ.get("VIME_HEAD_HOST") + if not public_host: + raise RuntimeError( + "VIME_HEAD_HOST is not set. Export it to the host IP that " + "sandboxes can reach for reverse-connection to the Anthropic adapter. " + "Without it the sandbox cannot dial back and the rollout will " + "silently abort." + ) + self.adapter = AnthropicAdapter( + tokenizer=self.tokenizer, + vllm_url=vllm_url, + model=args.hf_checkpoint, + tool_parser=self.tool_parser, + reasoning_parser=self.reasoning_parser, + ) + # handler_cancellation=True so a client disconnect cancels the handler + # coroutine, tearing down the in-flight engine ``/inference/v1/generate`` + # request. Without it a cancelled client leaves an inflight generate that + # races with the next release_memory_occupation. + self.app_handle = run_app_in_thread( + self.adapter.app, + host=SHIM_BIND_HOST, + port=SHIM_PORT, + thread_name="anthropic-adapter", + runner_kwargs={"handler_cancellation": True}, + ) + self.adapter_url = f"http://{public_host}:{self.app_handle.port}" + logger.info( + "[coding_agent_rl] tokenizer=%s adapter=%s max_context_len=%s tool_parser=%s reasoning_parser=%s", + args.hf_checkpoint, + self.adapter_url, + self.max_context_len, + self.tool_parser, + self.reasoning_parser, + ) + + +# --------------------------------------------------------------------------- +# Trajectory -> Sample conversion +# adapter.finish_session() returns TokenSegments. One trajectory yields >=1 +# segments because the agent may compact + reset mid-run; trajectory.py handles +# the mechanical segment -> Sample fan-out. +# --------------------------------------------------------------------------- +@dataclass(frozen=True) +class RewardResult: + reward: float + is_solved: bool + applied_cleanly: bool + + +def _start_session( + state: _State, + sample: Sample, + md: dict[str, Any], + sampling_params: dict[str, Any], +) -> str: + # claude-code inside the sandbox dials back to the adapter with this + # session_id (passed as the Bearer token) so its turns are grouped under + # one chain history. Build from (instance_id, index, group_index) when + # possible; fall back to random hex if either index is missing. + if sample.session_id: + session_id = sample.session_id + elif sample.index is not None and sample.group_index is not None: + session_id = f"cagent-{md['instance_id']}-{sample.index}-{sample.group_index}" + else: + session_id = f"cagent-{md['instance_id']}-{secrets.token_hex(8)}" + sample.session_id = session_id + state.adapter.open_session( + session_id, + sampling_defaults=sampling_params, + max_context_tokens=state.max_context_len, + ) + return session_id + + +def _merge_samples( + *, + sample: Sample, + state: _State, + segments: list[TokenSegment], + reward_result: RewardResult, + elapsed_sec: float, + instance_id: str, +): + if not segments: + return _abort_result(sample, "adapter_session_empty") + + trajectory_metadata = { + **(sample.metadata or {}), + "instance_id": instance_id, + "is_solved": reward_result.is_solved, + "applied_cleanly": reward_result.applied_cleanly, + "elapsed_sec": elapsed_sec, + } + + # All K samples share rollout_id so the loss reducer counts this + # trajectory once. + fanned = fan_out_sample_segments( + sample, + segments, + reward_result.reward, + state.tokenizer, + metadata=trajectory_metadata, + ) + if not fanned: + raise ValueError("fan-out produced no samples") + + logger.info( + "[coding_agent_rl] %s: reward=%.2f solved=%s applied=%s elapsed=%.1fs segments=%d", + instance_id, + reward_result.reward, + reward_result.is_solved, + reward_result.applied_cleanly, + elapsed_sec, + len(fanned), + ) + return fanned + + +# --------------------------------------------------------------------------- +# Main per-sample agent function +# +# The four calls inside the timeout are the high-level rollout recipe: +# run_claude_code -> git_diff -> sandbox.evaluate -> merge_samples. +# --------------------------------------------------------------------------- +async def generate(args, sample: Sample, sampling_params: dict[str, Any]): + """Per-sample agent function with wall-clock guard. See + SWE_GENERATE_GUARD_SEC docstring above.""" + state = _State(args) + md = _metadata(sample) + if not md["image"] or not md["workdir"]: + return _abort_result(sample, "missing_image_or_workdir") + + instance_id = md["instance_id"] + session_id = _start_session(state, sample, md, sampling_params) + t0 = time.time() + try: + async with asyncio.timeout(SWE_GENERATE_GUARD_SEC): + async with sandbox.boot_agent_sandbox(md["image"]) as sb: + await sandbox.run_claude_code( + sb, + workdir=md["workdir"], + session_id=session_id, + adapter_url=state.adapter_url, + time_budget_sec=SWE_TIME_BUDGET_SEC, + problem_statement=md["problem_statement"], + swepro=md["swepro"], + pre_commands=md["pre_commands"], + ) + diff_text = await sandbox.git_diff(sb, md["workdir"]) + + reward, is_solved, applied_cleanly = await sandbox.evaluate( + image=md["image"], + workdir=md["workdir"], + diff_text=diff_text, + swepro=md["swepro"], + eval_cmd=md["eval_cmd"], + pre_commands=md["pre_commands"], + timeout_sec=SWE_EVAL_TIMEOUT_SEC, + ) + reward_result = RewardResult( + reward=float(reward), + is_solved=bool(is_solved), + applied_cleanly=bool(applied_cleanly), + ) + segments = await state.adapter.finish_session(session_id) + return _merge_samples( + sample=sample, + state=state, + segments=segments, + reward_result=reward_result, + elapsed_sec=time.time() - t0, + instance_id=instance_id, + ) + + except asyncio.TimeoutError: + _log_timeout_diagnostic(t0) + return _abort_result(sample, "wall_clock_timeout") + except Exception as e: + logger.error( + "[coding_agent_rl] %s: rollout failed: %s\n%s", + instance_id, + e, + traceback.format_exc(), + ) + return _abort_result(sample, f"exception:{type(e).__name__}") + finally: + # Close the sid before next train step's release_memory_occupation; + # stragglers from this trajectory would otherwise race its idle assert. + await state.adapter.finish_session(session_id) # idempotent + + +def _log_timeout_diagnostic(t0: float) -> None: + """Dump pending-task names when the wall-clock guard fires so future + debugging can see which await was stuck. Must never crash.""" + try: + elapsed = time.time() - t0 + pending = [t for t in asyncio.all_tasks() if not t.done()] + stuck = [] + for t in pending[:5]: # cap to avoid log spam + coro = getattr(t, "_coro", None) + stuck.append(getattr(coro, "__qualname__", repr(coro))) + logger.warning( + "[coding_agent_rl] generate() wall_clock_timeout after %.1fs " + "(guard=%ds); %d tasks pending; sample of stuck: %s", + elapsed, + SWE_GENERATE_GUARD_SEC, + len(pending), + stuck, + ) + except Exception: # pragma: no cover - diag must never crash + pass + + +# --------------------------------------------------------------------------- +# Metadata helpers +# --------------------------------------------------------------------------- +def _wrap_f2p_script(script: str | None) -> str | None: + # Materialize a self-contained pytest script (typical sweb f2p_script: + # ends with `sys.exit(pytest.main([...]))`) into the sandbox via base64 + # so we sidestep all shell quoting; python's exit code carries the + # pytest pass/fail signal that `_run_eval_cmd` turns into reward. + if not script: + return None + b64 = base64.b64encode(script.encode("utf-8")).decode("ascii") + return f"echo {b64} | base64 -d > /tmp/vime_f2p.py && python /tmp/vime_f2p.py" + + +def _metadata(sample: Sample) -> dict[str, Any]: + """Normalize the two dataset schemas (flat vs ``remote_env_info``).""" + m = sample.metadata or {} + rem = m.get("remote_env_info") or {} + label = sample.label if (isinstance(sample.label, str) and len(sample.label) < 256) else None + return { + "instance_id": m.get("instance_id") or rem.get("instance_id") or label or "unknown", + "image": m.get("image") or rem.get("image_url"), + "workdir": m.get("workdir") or rem.get("workdir"), + "problem_statement": m.get("problem_statement") or _coerce_prompt(sample.prompt), + "swepro": m.get("swepro"), + "eval_cmd": m.get("eval_cmd") or _wrap_f2p_script(rem.get("f2p_script")), + "pre_commands": m.get("pre_commands") or rem.get("pre_commands"), + } + + +def _coerce_prompt(prompt) -> str: + if isinstance(prompt, str): + return prompt + if isinstance(prompt, list): + for m in prompt: + if isinstance(m, dict) and m.get("role") == "user": + c = m.get("content") + if isinstance(c, str): + return c + if isinstance(c, list): + return "\n".join(p.get("text", "") for p in c if isinstance(p, dict) and p.get("type") == "text") + return "" + + +def _abort(sample: Sample, reason: str) -> Sample: + sample.tokens = [0, 0] + sample.response = "" + sample.response_length = 1 + sample.loss_mask = [0] + sample.reward = 0.0 + sample.status = Sample.Status.ABORTED + sample.metadata = {**(sample.metadata or {}), "abort_reason": reason} + logger.warning("[coding_agent_rl] aborted: %s", reason) + return sample + + +def _abort_result(sample: Sample, reason: str): + """Return a uniform list shape for this fan-out generate function.""" + return [_abort(sample, reason)] diff --git a/examples/coding_agent_rl/run_qwen36_35b_a3b_swe_8nodes.sh b/examples/coding_agent_rl/run_qwen36_35b_a3b_swe_8nodes.sh new file mode 100644 index 000000000..ef14087c0 --- /dev/null +++ b/examples/coding_agent_rl/run_qwen36_35b_a3b_swe_8nodes.sh @@ -0,0 +1,374 @@ +#!/usr/bin/env bash +# End-to-end SWE coding-agent RL on 8 nodes. +# +# Same model and training loop as run_qwen36_35b_a3b_swe_8node.sh, with three +# extra layers that actively encourage the rollout to dispatch sub-agents. +# Trajectory trees produced by this script show real `sibling` branches: +# +# (1) An `investigator` sub-agent is registered via claude-code's --agents +# flag (Grep/Read/Glob only, no edits) — a concrete, narrowly-scoped +# dispatch target. +# (2) SWE_CC_PROMPT requires the model to dispatch the investigator before +# any edit, naming the exact call form (Agent tool with +# subagent_type=investigator). +# (3) Agent/Task tools stay in the allowed set; WebFetch/WebSearch are +# disabled (sandbox has no outbound internet); --disable-slash-commands +# removes /compact as a competing branching pathway. +# +# Fan-out semantics: +# * generate() returns list[Sample] (one Sample per trajectory segment); +# the per-trajectory reward is split as reward/K across segments. +# * Sub-agent dispatch increases K (each sub-agent turn block becomes its +# own segment), so the effective batch after flatten can be much larger +# than rollout_batch_size * n_samples_per_prompt. If pinned-memory or +# GPU wake_up OOM appears, lower rollout_batch_size or n_samples_per_prompt +# first — not max-tokens-per-gpu. +# Run from a long-lived shell / tmux session on the Ray head node; do not wrap +# in a short-lived nohup launcher or Ray child processes get cleaned up with it. + +# Best-effort cleanup so a rerun does not collide with stale workers. +pkill -9 -f "vllm serve" || true +sleep 3 +ray stop --force || true +pkill -9 ray || true +sleep 3 +pkill -9 ray || true + +set -ex + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +VIME_DIR="${VIME_DIR:-$(cd "${SCRIPT_DIR}/../.." && pwd)}" + +# ============ model parallelism ============ +# CP=8 (higher than the baseline script's CP=2) gives more rank-local context +# room for the longer per-segment payloads typical under sub-agent dispatch. +export TP_SIZE="${TP_SIZE:-2}" +export PP_SIZE="${PP_SIZE:-1}" +export CP_SIZE="${CP_SIZE:-8}" +export EP_SIZE="${EP_SIZE:-8}" +export ETP_SIZE="${ETP_SIZE:-1}" + +# ============ rollout engine ============ +ROLLOUT_TP_SIZE="${ROLLOUT_TP_SIZE:-8}" +ROLLOUT_DP_SIZE="${ROLLOUT_DP_SIZE:-8}" +ROLLOUT_MEM_UTILIZATION="${ROLLOUT_MEM_UTILIZATION:-0.75}" + +# ============ Qwen3.5-35B-A3B architecture ============ +NLAYERS=40 +FIRST_K_DENSE_REPLACE=0 + +arr=() +for ((i=0; ivLLM arg map (see translation_guide.md "Arg map"): +# --sglang-mem-fraction-static -> --vllm-gpu-memory-utilization +# --sglang-dp-size N -> --vllm-data-parallel-size N +# --sglang-ep-size N -> --vllm-enable-expert-parallel (boolean) +# --sglang-{enable-dp-attention,enable-dp-lm-head,moe-dense-tp-size, +# mamba-scheduler-strategy} -> sglang-only, no vLLM equivalent (dropped) +# --sglang-{tool-call,reasoning}-parser -> no vLLM arg; coding_agent_rl uses the +# XML tool-call fallback (parsers=None), so these are dropped. +# --sglang-speculative-* per-field flags -> a single --vllm-speculative-config +# JSON dict (commented below; supply your EAGLE draft model to enable). +VLLM_ARGS=( + --rollout-num-gpus 64 + --rollout-num-gpus-per-engine ${ROLLOUT_TP_SIZE} + --vllm-gpu-memory-utilization ${ROLLOUT_MEM_UTILIZATION} + --vllm-data-parallel-size ${ROLLOUT_DP_SIZE} + --vllm-enable-expert-parallel + --prefill-num-servers 1 +) +# Speculative decoding (EAGLE): vLLM takes one JSON dict rather than the +# sglang per-field flags. Uncomment and point "model" at your draft checkpoint. +# VLLM_ARGS+=( +# --vllm-speculative-config '{"method":"eagle","num_speculative_tokens":4,"model":"/path/to/eagle_draft"}' +# ) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --moe-token-dispatcher-type flex + --moe-enable-deepep + --colocate +) + +# ============ ray cluster network ============ +# Set MASTER_ADDR before the SWE block: VIME_HEAD_HOST below falls back to it. +export MASTER_ADDR="${MASTER_ADDR:-${MLP_WORKER_0_HOST:-$(hostname -I | awk '{print $1}')}}" +export MASTER_PORT="${MASTER_PORT:-${MLP_WORKER_0_PORT:-6379}}" +export GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-${MLP_SOCKET_IFNAME:-eth0}}" +export NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-${MLP_SOCKET_IFNAME:-eth0}}" + +# ============ SWE / claude-code rollout knobs ============ + +# --- sandbox provisioning (E2B) --- +# The E2B SDK validates the API key format locally before it reaches +# E2B-compatible gateways. If your internal gateway ignores auth, use any +# syntactically valid e2b_... hex placeholder. +E2B_DUMMY_API_KEY="e2b_0000000000000000000000000000000000000000" +E2B_API_KEY="${E2B_API_KEY:-${E2B_DUMMY_API_KEY}}" +if [[ ! "${E2B_API_KEY}" =~ ^e2b_[0-9a-fA-F]{40}$ ]]; then + echo "WARN: E2B_API_KEY does not pass local E2B SDK format validation; using dummy key." >&2 + E2B_API_KEY="${E2B_DUMMY_API_KEY}" +fi +export E2B_API_KEY +export SWE_SANDBOX_METADATA_FILE="${SANDBOX_METADATA_FILE}" +export SWE_SANDBOX_IMAGE_METADATA_KEY="${SWE_SANDBOX_IMAGE_METADATA_KEY:-glm-platform/image}" +# Host-side tarballs injected into each sandbox at boot. +export SWE_HOST_NODE_TARBALL="${SWE_HOST_NODE_TARBALL:-/path/to/node-v22.x-linux-x64.tar.xz}" +export SWE_HOST_CC_TARBALL="${SWE_HOST_CC_TARBALL:-/path/to/anthropic-ai-claude-code-local-linux-x64.tgz}" + +# --- reply path (sandbox -> host shim) --- +export VIME_HEAD_HOST="${VIME_HEAD_HOST:-${MASTER_ADDR:-${MLP_WORKER_0_HOST:-127.0.0.1}}}" +export SHIM_BIND_HOST="${SHIM_BIND_HOST:-0.0.0.0}" +export SHIM_PORT="${SHIM_PORT:-18001}" + +# --- per-trajectory time / concurrency budgets --- +# Time budget 1800s (vs baseline 1200): sub-agent dispatch on large repos blows +# past a tighter budget — investigator passes are the long tail. +# Boot concurrency 6 (vs baseline 8) eases h2/SSL long-tail stalls under +# heavier sub-agent dispatch. +export SWE_TIME_BUDGET_SEC="${SWE_TIME_BUDGET_SEC:-1800}" +export SWE_EVAL_TIMEOUT_SEC="${SWE_EVAL_TIMEOUT_SEC:-600}" +export SWE_BOOT_CONCURRENCY="${SWE_BOOT_CONCURRENCY:-6}" + +# --- trajectory fan-out --- +# generate() emits one Sample per segment (reducer splits reward/K); +# group_id is shared so the per-rollout-mean loss reducer still counts +# the trajectory once. +# --rollout-max-response-len caps one model turn. The custom generate function +# uses --rollout-max-context-len as the multi-turn prompt+response budget. + +# --- claude-code CLI extras --- +# SETTINGS_JSON: autoCompactWindow (80k) < MAX_CONTEXT_LEN (96k) so the CLI +# compacts before any segment crosses the training-side cap. +# AGENTS_JSON: register a read-only `investigator` sub-agent (Grep/Read/Glob) +# as a concrete, narrowly-scoped dispatch target. +# SWE_CLAUDE_EXTRA_ARGS: WebFetch/WebSearch are off (sandbox has no outbound +# internet); --disable-slash-commands keeps the model from emitting /compact +# as a competing branching pathway. +SETTINGS_JSON='{"permissions":{"defaultMode":"bypassPermissions"},"autoCompactEnabled":true,"autoCompactWindow":80000}' +AGENTS_JSON='{"investigator":{"description":"Searches the repo for relevant files before any edit","prompt":"You are an investigator sub-agent. Use Grep/Read/Glob to find every file relevant to the user task, then return a short bulleted summary. Do NOT edit anything.","tools":["Grep","Read","Glob"]}}' +export SWE_CLAUDE_EXTRA_ARGS="--settings '${SETTINGS_JSON}' --disable-slash-commands --agents '${AGENTS_JSON}' --disallowedTools WebFetch WebSearch" + +# Optional: bias the model to dispatch the investigator before any edit. +# Uncomment to maximize sub-agent dispatch — naming the exact call form +# (Agent tool with subagent_type=investigator) is what reliably triggers it. +# export SWE_CC_PROMPT="Read PROBLEM_STATEMENT.md. BEFORE editing any file, dispatch the 'investigator' sub-agent (via the Agent tool with subagent_type=investigator) to locate every file relevant to the issue. Then fix the issue and run the tests." + +# ============ proxy bypass for in-cluster traffic ============ +export no_proxy="127.0.0.1,${MASTER_ADDR},${VIME_HEAD_HOST}" +export NO_PROXY="${no_proxy}" + +cd "${VIME_DIR}" + +# ============ bring up ray cluster ============ +HOSTFILE="${HOSTFILE:-/root/mpi_rack_hostfile}" +ACTOR_NUM_NODES="${ACTOR_NUM_NODES:-${MLP_WORKER_NUM:-8}}" +ACTOR_NUM_GPUS_PER_NODE="${ACTOR_NUM_GPUS_PER_NODE:-8}" + +ray start --head --node-ip-address "${MASTER_ADDR}" --num-gpus "${ACTOR_NUM_GPUS_PER_NODE}" \ + --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +if [[ -f "${HOSTFILE}" ]]; then + for WORKER_IP in $(awk '{print $1}' "${HOSTFILE}"); do + [[ -z "${WORKER_IP}" ]] && continue + [[ "${WORKER_IP}" == "${MASTER_ADDR}" ]] && continue + echo "Starting Ray worker on ${WORKER_IP}" + ssh -o StrictHostKeyChecking=no "root@${WORKER_IP}" \ + "pkill -9 -f 'vllm serve' ; ray stop --force ; pkill -9 python ; \ + ray start --address=${MASTER_ADDR}:6379 --num-gpus ${ACTOR_NUM_GPUS_PER_NODE} \ + --node-ip-address ${WORKER_IP} --disable-usage-stats" & + done + wait +fi + +echo "Waiting for Ray cluster to stabilize..." +sleep 30 +ray status + +# ============ runtime env propagated to ray workers ============ +export VIME_DIR +RUNTIME_ENV_JSON=$(python3 - <&1 | tee "${LOG_FILE}" + +echo "RUN_ROOT=${RUN_ROOT}" diff --git a/examples/coding_agent_rl/sandbox.py b/examples/coding_agent_rl/sandbox.py new file mode 100644 index 000000000..cd9fbe48b --- /dev/null +++ b/examples/coding_agent_rl/sandbox.py @@ -0,0 +1,400 @@ +"""Coding-agent sandbox helpers. + +The provider-agnostic sandbox contract and E2B backend live in +``vime.agent.sandbox``. This module keeps the coding-agent/SWE-specific +bootstrap, Claude Code runner, diff capture, and fresh-sandbox evaluator. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import lzma +import os +import shlex +import shutil +import tempfile +import time +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from pathlib import Path + +from vime.agent.sandbox import E2BSandbox, Sandbox + + +logger = logging.getLogger(__name__) + +# Paths inside the sandbox (avoid clashes with image-shipped paths). +_PATCH = "/workspace/__cagent_patch__.diff" +_PRE = "/workspace/__cagent_pre__.sh" +_SWEPRO_DIR = "/workspace/swepro_eval" + +SWE_HOST_NODE_TARBALL = Path( + os.environ.get( + "SWE_HOST_NODE_TARBALL", + "/path/to/node-v22.20.0-linux-x64.tar.xz", + ) +) +SWE_HOST_CC_TARBALL = Path( + os.environ.get( + "SWE_HOST_CC_TARBALL", + "/path/to/anthropic-ai-claude-code.tgz", + ) +) +SWE_BOOT_CONCURRENCY = int(os.environ.get("SWE_BOOT_CONCURRENCY", "16")) +SWE_BOOT_RETRIES = int(os.environ.get("SWE_BOOT_RETRIES", "2")) +CC_PROMPT = os.environ.get( + "SWE_CC_PROMPT", + "Read PROBLEM_STATEMENT.md in the current directory and resolve the issue. " + "Edit source files only (do NOT touch tests). After editing, run the relevant " + "tests to verify your fix passes. Do NOT modify PROBLEM_STATEMENT.md and do " + "NOT commit. When finished, print a one-line summary and exit.", +) + +_BOOT_SEM: asyncio.Semaphore | None = None + + +# --------------------------------------------------------------------------- +# Sandbox bootstrap (Node + Claude Code + agent user) +# --------------------------------------------------------------------------- +@asynccontextmanager +async def boot_agent_sandbox(image: str) -> AsyncIterator[E2BSandbox]: + """Boot a fresh E2B sandbox and install the Claude Code toolchain. + + This is the provisioning wrapper for the work sandbox: create the sandbox + from the dataset image, install Node 22 + Claude Code CLI from host + tarballs, retry transient boot/install failures, and close the sandbox when + the caller leaves the context. + """ + global _BOOT_SEM + if _BOOT_SEM is None: + _BOOT_SEM = asyncio.Semaphore(SWE_BOOT_CONCURRENCY) + + sb = None + last_err: Exception | None = None + for attempt in range(SWE_BOOT_RETRIES): + cand = E2BSandbox(image) + try: + async with _BOOT_SEM: + await cand.__aenter__() + try: + await install_node22(cand, SWE_HOST_NODE_TARBALL) + await install_claude_code(cand, SWE_HOST_CC_TARBALL) + except BaseException: + await cand.__aexit__(None, None, None) + raise + sb = cand + break + except Exception as e: + last_err = e + logger.warning( + "[coding_agent_rl] provision attempt %d/%d failed: %s: %s", + attempt + 1, + SWE_BOOT_RETRIES, + type(e).__name__, + str(e)[:200], + ) + await asyncio.sleep(1 + attempt) + if sb is None: + assert last_err is not None + raise last_err + try: + yield sb + finally: + await sb.__aexit__(None, None, None) + + +async def install_node22(sb: Sandbox, host_tarball: Path) -> None: + """Node 22 over the base image (Debian 12 ships 16; cli.js needs >= 20). + Decompresses .xz on the host (cached) so sandboxes without xz-utils can + still run plain `tar xf`. npm prefix=/usr/local required for sweap-images.""" + host_tarball = Path(host_tarball) + if host_tarball.suffix == ".xz": + plain = Path(tempfile.gettempdir()) / f"coding_agent_rl.{host_tarball.stem}.tar" + if not plain.exists(): + tmp = plain.with_suffix(".tar.partial") + with lzma.open(host_tarball, "rb") as src, open(tmp, "wb") as dst: + shutil.copyfileobj(src, dst) + os.replace(tmp, plain) + host_tarball = plain + await sb.write_file("/tmp/node22.tar", host_tarball) + await sb.exec( + "set -e && mkdir -p /opt/node22 && " + "tar xf /tmp/node22.tar -C /opt/node22 --strip-components=1 && " + "ln -sf /opt/node22/bin/node /usr/local/bin/node && " + "ln -sf /opt/node22/bin/npm /usr/local/bin/npm && " + "ln -sf /opt/node22/bin/npx /usr/local/bin/npx && " + "hash -r 2>/dev/null || true && node --version && npm --version", + user="root", + timeout=180, + check=True, + ) + + +async def install_claude_code(sb: Sandbox, host_tarball: Path) -> None: + await sb.write_file("/tmp/claude-code.tgz", host_tarball) + await sb.exec( + "npm install -g --prefix=/usr/local --no-audit --no-fund /tmp/claude-code.tgz " + "&& ls -la /usr/local/bin/claude && /usr/local/bin/claude --version", + user="root", + timeout=300, + check=True, + ) + + +async def ensure_agent_user(sb: Sandbox, workdir: str) -> None: + """Create the unprivileged 'agent' user that owns workdir + can git diff. + Settings file pre-acks bypass-permissions so claude-code starts headless.""" + await sb.exec( + f"id agent >/dev/null 2>&1 || useradd -m -s /bin/bash agent && " + f"chown -R agent:agent /home/agent {workdir} && " + f"git config --system --add safe.directory '*' && id agent && " + f"mkdir -p /home/agent/.claude && " + f'echo \'{{"hasCompletedOnboarding": true, "bypassPermissionsModeAccepted": true}}\' ' + f"| tee /home/agent/.claude.json /home/agent/.claude/settings.json > /dev/null && " + f"chown -R agent:agent /home/agent/.claude /home/agent/.claude.json", + user="root", + check=True, + timeout=60, + ) + + +async def apply_before_repo_set_cmd(sb: Sandbox, workdir: str, swepro: dict) -> None: + """Run swepro['before_repo_set_cmd'] in the sandbox if present (no-op if not).""" + before = swepro.get("before_repo_set_cmd") if swepro else None + if not before: + return + payload = f"set -e\ncd {workdir}\n{before}\n" + await sb.exec( + "mkdir -p /workspace/swepro_setup && chown agent:agent /workspace/swepro_setup", user="root", check=True + ) + await sb.write_file("/workspace/swepro_setup/before.sh", payload, user="agent") + await sb.exec("bash /workspace/swepro_setup/before.sh", user="agent", check=False, timeout=600) + + +# --------------------------------------------------------------------------- +# Agent run (workspace prep + claude-code spawn + done-marker poll) +# --------------------------------------------------------------------------- +async def run_claude_code( + sb: Sandbox, + *, + workdir: str, + session_id: str, + adapter_url: str, + time_budget_sec: int, + problem_statement: str = "", + swepro: dict | None = None, + pre_commands: list[str] | str | None = None, + prompt: str | None = None, +) -> int: + """Prepare the SWE workspace, write PROBLEM_STATEMENT.md, then run CC.""" + await ensure_agent_user(sb, workdir) + if swepro: + await apply_before_repo_set_cmd(sb, workdir, swepro) + if pre_commands: + await apply_pre_commands(sb, workdir, pre_commands) + await sb.write_file( + f"{workdir}/PROBLEM_STATEMENT.md", + problem_statement or "", + user="agent", + ) + return await _spawn_claude_code( + sb, + workdir=workdir, + session_id=session_id, + adapter_url=adapter_url, + prompt=prompt or CC_PROMPT, + time_budget_sec=time_budget_sec, + ) + + +async def _spawn_claude_code( + sb: Sandbox, + *, + workdir: str, + session_id: str, + adapter_url: str, + prompt: str, + time_budget_sec: int, +) -> int: + """Spawn claude-code detached + poll a done-marker file. + + E2B's gateway resets HTTP/2 around 6.5 min, so we can't keep a long-lived + foreground exec. The launcher writes the exit code into a marker file + and we poll it every 5s via short RPCs (which also keeps the sandbox + alive against idle GC).""" + done = f"{workdir}/.cagent_done" + launcher = f"{workdir}/.cagent_run.sh" + traj = f"{workdir}/claude_code_trajectory.jsonl" + + launcher_body = ( + "#!/bin/bash\n" + f"cd {workdir}\n" + "export HOME=/home/agent\n" + f"/usr/local/bin/claude -p {json.dumps(prompt)} " + f"--permission-mode bypassPermissions " + f"--output-format stream-json --include-partial-messages " + f"--include-hook-events --verbose " + f"{os.environ.get('SWE_CLAUDE_EXTRA_ARGS', '').strip()} " + f"2>&1 | tee {shlex.quote(traj)}\n" + f"echo $? > {done}\n" + ) + await sb.write_file(launcher, launcher_body, user="agent") + await sb.exec(f"chmod +x {launcher}", user="agent", timeout=30) + + env = { + "ANTHROPIC_BASE_URL": adapter_url, + "ANTHROPIC_AUTH_TOKEN": session_id, + "ANTHROPIC_MODEL": "vime-actor", + "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1", + "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1", + "CLAUDE_CODE_ATTRIBUTION_HEADER": "0", + } + env_keys = ",".join(env.keys()) + await sb.exec( + f"runuser -u agent --whitelist-environment={env_keys}" + f" -- bash -c 'setsid {launcher} < /dev/null > /dev/null 2>&1 &'", + user="root", + env=env, + timeout=30, + check=True, + ) + + deadline = time.time() + time_budget_sec + exit_code = -2 # convention: -2 = budget exceeded + while time.time() < deadline: + await asyncio.sleep(5) + ec, out, _ = await sb.exec( + f"test -f {done} && cat {done}", + user="agent", + timeout=15, + check=False, + ) + if ec == 0: + try: + exit_code = int((out or "").strip() or "-1") + except ValueError: + exit_code = -1 + break + return exit_code + + +async def git_diff(sb: Sandbox, workdir: str) -> str: + cmd = ( + f"cd {workdir} && git add -N . && " + f"git diff -- . ':(exclude)PROBLEM_STATEMENT.md' " + f"':(exclude)claude_code_trajectory.jsonl' " + f"':(exclude).cagent_done' ':(exclude).cagent_run.sh'" + ) + _, out, _ = await sb.exec(cmd, user="agent", timeout=120) + return out + + +# --------------------------------------------------------------------------- +# Eval (fresh sandbox, apply diff, run dataset tests) +# --------------------------------------------------------------------------- +async def evaluate( + *, + image: str, + workdir: str, + diff_text: str, + swepro: dict | None = None, + eval_cmd: str | None = None, + pre_commands: list[str] | str | None = None, + timeout_sec: int = 600, +) -> tuple[float, bool, bool]: + """Returns (reward, solved, applied_cleanly). + + No-test-cheating guarantee: the eval sandbox is built from the same image + but starts CLEAN, so only the model-produced diff affects reward.""" + if not (swepro or eval_cmd): + logger.warning("[e2b.evaluate] no swepro/eval_cmd; reward=0") + return 0.0, False, True + + async with E2BSandbox(image) as ev: + await ensure_agent_user(ev, workdir) + if swepro: + await _setup_swepro_assets(ev, swepro) + await apply_before_repo_set_cmd(ev, workdir, swepro) + if pre_commands: + await apply_pre_commands(ev, workdir, pre_commands) + + applied = await _apply_diff(ev, workdir, diff_text) + if not applied: + return 0.0, False, False + + if swepro: + r, s = await _run_swepro(ev, workdir, swepro, timeout_sec) + return r, s, True + r, s = await _run_eval_cmd(ev, workdir, eval_cmd, timeout_sec) + return r, s, True + + +async def _setup_swepro_assets(ev: Sandbox, swepro: dict) -> None: + await ev.exec(f"mkdir -p {_SWEPRO_DIR} && chmod 777 {_SWEPRO_DIR}", user="root", check=True) + for k, dst in [("run_script_path", "run_script.sh"), ("parser_script_path", "parser.py")]: + host_p = swepro.get(k) + if host_p: + text = Path(host_p).read_text() + await ev.write_file(f"{_SWEPRO_DIR}/{dst}", text, user="root") + await ev.exec(f"chmod 755 {_SWEPRO_DIR}/* && chown -R agent:agent {_SWEPRO_DIR}", user="root", check=True) + + +async def apply_pre_commands(ev: Sandbox, workdir: str, pre: list[str] | str) -> None: + # Public: also called by generate.py to keep the work sandbox baseline + # aligned with eval (sweb-style pre_commands typically `git checkout + # -f`, so skipping in work sandbox makes the model's diff + # context mismatch the eval base -> 100% apply failure). + if isinstance(pre, str): + body = pre.replace("\\n", "\n") + else: + body = "\n".join(c for c in (pre or []) if c) + await ev.write_file(_PRE, "set -e\n" + body, user="agent") + await ev.exec(f"chmod 755 {_PRE} && cd {workdir} && bash {_PRE}", user="agent", check=False, timeout=600) + + +async def _apply_diff(ev: Sandbox, workdir: str, diff_text: str) -> bool: + if not diff_text.strip(): + return True + await ev.write_file(_PATCH, diff_text, user="agent") + for cmd in [ + f"cd {workdir} && git apply --3way --whitespace=nowarn {_PATCH}", + f"cd {workdir} && git apply --whitespace=nowarn {_PATCH}", + f"cd {workdir} && patch -p1 --no-backup-if-mismatch < {_PATCH}", + ]: + ec, _, _ = await ev.exec(cmd, user="agent", check=False, timeout=120) + if ec == 0: + return True + return False + + +async def _run_swepro(ev: Sandbox, workdir: str, swepro: dict, timeout: int) -> tuple[float, bool]: + test_arg = ",".join(swepro.get("selected_test_files") or []) + stdout_f = f"{_SWEPRO_DIR}/stdout.log" + stderr_f = f"{_SWEPRO_DIR}/stderr.log" + result_f = f"{_SWEPRO_DIR}/result.json" + await ev.exec( + f"cd {workdir} && bash {_SWEPRO_DIR}/run_script.sh " + f"{json.dumps(test_arg)} > {stdout_f} 2> {stderr_f} || true", + user="agent", + check=False, + timeout=timeout, + ) + await ev.exec( + f"python3 {_SWEPRO_DIR}/parser.py {stdout_f} {stderr_f} {result_f}", + user="agent", + check=False, + timeout=120, + ) + raw = await ev.read_file(result_f, user="agent") + parsed = json.loads(raw) if raw else {"tests": []} + passed = {t["name"] for t in parsed.get("tests", []) if t.get("status") == "PASSED"} + required = set(swepro.get("fail_to_pass") or []) | set(swepro.get("pass_to_pass") or []) + solved = bool(required) and required.issubset(passed) + return (1.0 if solved else 0.0), solved + + +async def _run_eval_cmd(ev: Sandbox, workdir: str, cmd: str, timeout: int) -> tuple[float, bool]: + ec, _, _ = await ev.exec(f"cd {workdir} && {cmd}", user="agent", check=False, timeout=timeout) + return (1.0 if ec == 0 else 0.0), ec == 0 diff --git a/examples/retool/generate_with_retool.py b/examples/retool/generate_with_retool.py new file mode 100644 index 000000000..f0506db66 --- /dev/null +++ b/examples/retool/generate_with_retool.py @@ -0,0 +1,461 @@ +# Adapted from https://github.com/volcengine/verl/blob/cb809d66e46dfd3342d008628891a14a054fa424/recipe/retool/retool.py +# +# vime/vLLM counterpart of slime's ``examples/retool/generate_with_retool.py``. +# The only structural change versus slime is the rollout engine: requests go to +# vLLM's ``/inference/v1/generate`` (token-only body, ``token_ids`` not +# ``input_ids``) and the response ``choices[0]`` is normalized into a +# sglang-shaped ``meta_info`` (``finish_reason.type`` + ``output_token_logprobs``) +# via the canonical helpers in :mod:`vime.rollout.vllm_rollout`, so the +# multi-turn tool-call control flow below is identical to slime's. +import re +from typing import Any + +try: + from jinja2 import Template +except ImportError as e: + raise ImportError("Jinja2 is required. Please install it with: pip install jinja2") from e + +from vime.rollout.vllm_rollout import ( + GenerateState, + _build_inference_sampling_params, + _coerce_flat_int_token_ids, + _inference_generate_tokens_and_logprobs, + _vllm_meta_from_generate_choice, +) +from vime.utils.http_utils import post +from vime.utils.types import Sample + +# Import reward models +try: + from vime.rollout.rm_hub.math_dapo_utils import compute_score as math_dapo_compute_score +except ImportError as e: + raise ImportError("MathDapo is not installed") from e + +# Import tool sandbox functionality +from tool_sandbox import SEMAPHORE, TOOL_CONFIGS, tool_registry + +# Jinja2 template for tool-enabled conversations +TOOL_TEMPLATE = """<|im_start|>system +{%- if messages[0]['role'] == 'system' %} +{{- messages[0]['content'] }} +{%- else %} +You are a helpful assistant. +{%- endif %} +{%- if tools %} +# Tools + +You may call one or more functions to assist with the user query. + +You are provided with function signatures within XML tags: + +{%- for tool in tools %} +{{- tool | tojson }} +{%- endfor %} + + +For each function call, return a json object with function name and arguments within XML tags: + +{"name": , "arguments": } + +{%- endif %} +<|im_end|> +{%- for message in messages %} +{%- if message['role'] == 'user' %} +<|im_start|>user +{{- message['content'] }}<|im_end|> +{%- elif message['role'] == 'assistant' %} +<|im_start|>assistant +{{- message['content'] }}<|im_end|> +{%- endif %} +{%- endfor %} +<|im_start|>assistant +""" + + +def format_conversation_with_tools( + prompt: str, tools: list[dict[str, Any]] = None, system_prompt: str = None, messages: list[dict[str, Any]] = None +) -> str: + """Format conversation using Jinja2 template with tool support""" + template = Template(TOOL_TEMPLATE) + + # Prepare messages + messages_to_render = [] + + # Always add system message - use provided one or default + if system_prompt: + system_content = system_prompt + else: + system_content = ( + "You are a helpful assistant that can use Python " + "tools to solve mathematical problems. When you need " + "to perform calculations, use the code_interpreter " + "tool to execute code and get results." + ) + + messages_to_render.append({"role": "system", "content": system_content}) + + # Add user message if provided + if prompt: + messages_to_render.append({"role": "user", "content": prompt}) + + # Add assistant responses from previous turns if provided + if messages: + messages_to_render.extend(messages) + + # Render template + formatted_text = template.render(messages=messages_to_render, tools=tools or []) + + return formatted_text + + +def postprocess_predictions(prediction: str): + """Extract action and content from prediction string""" + # Check for Answer: \boxed{...} format (only format we need for math_dapo) + # Use a more robust regex that handles nested braces + answer_pattern = r"Answer:\s*\\boxed\{((?:[^{}]|\{[^{}]*\})*)\}" + answer_match = re.search(answer_pattern, prediction, re.DOTALL) + if answer_match: + content = answer_match.group(1).strip() + return "answer", content + + # Then check for tags (new format from Jinja2 template) + tool_call_pattern = r"\s*(\{.*?\})\s*" + tool_call_match = re.search(tool_call_pattern, prediction, re.DOTALL) + if tool_call_match: + try: + import json + + # Clean up the JSON string by removing newlines and extra + # whitespace + json_str = tool_call_match.group(1) + # Replace newlines in string values with \n + json_str = json_str.replace("\n", "\\n") + tool_call_data = json.loads(json_str) + tool_name = tool_call_data.get("name") + arguments = tool_call_data.get("arguments", {}) + + if tool_name == "code_interpreter": + code = arguments.get("code", "") + if code.strip(): + return "code", code + except (json.JSONDecodeError, KeyError, AttributeError): + pass + + # Then check for tags + code_pattern = r"(.*?)" + code_match = re.search(code_pattern, prediction, re.DOTALL) + if code_match: + content = code_match.group(1).strip() + return "code", content + + # Finally check for ```python code blocks (lowest priority) + python_code_pattern = r"```python\s*(.*?)\s*```" + python_code_match = re.search(python_code_pattern, prediction, re.DOTALL) + if python_code_match: + content = python_code_match.group(1).strip() + return "code", content + + return None, "" + + +def postprocess_responses(resp: str) -> str: + """Post-process response to ensure tag completeness""" + # Handle tags (new format from Jinja2 template) + if "" in resp: + # Find the last occurrence of ... + tool_call_pattern = r"\s*\{.*?\}\s*" + matches = list(re.finditer(tool_call_pattern, resp, re.DOTALL)) + if matches: + last_match = matches[-1] + return resp[: last_match.end()] + + # Handle tags + if "" in resp: + return resp.split("")[0] + "" + + # Handle ```python code blocks + if "```python" in resp: + # Find the last occurrence of ```python...``` + python_pattern = r"```python\s*.*?```" + matches = list(re.finditer(python_pattern, resp, re.DOTALL)) + if matches: + last_match = matches[-1] + return resp[: last_match.end()] + + # Handle Answer: \boxed{...} format (only format we need for math_dapo) + if "Answer:" in resp and "\\boxed{" in resp: + # Find the last occurrence of Answer: \boxed{...} with nested braces support + answer_pattern = r"Answer:\s*\\boxed\{((?:[^{}]|\{[^{}]*\})*)\}" + matches = list(re.finditer(answer_pattern, resp, re.DOTALL)) + if matches: + last_match = matches[-1] + return resp[: last_match.end()] + + return resp + + +async def execute_predictions(prediction: str) -> str: + """Execute predictions and return results""" + action, content = postprocess_predictions(prediction) + + if action == "code": + # Content is already the Python code (extracted by + # postprocess_predictions) + code = content.strip() + if code: + async with SEMAPHORE: + result = await tool_registry.execute_tool("code_interpreter", {"code": code}) + next_obs = f"\n\n\n{result}\n\n\n" + done = False + else: + next_obs = "\n\n\nError: No Python code found" "\n\n\n" + done = False + elif action == "answer": + next_obs = "" + done = True + else: + next_obs = ( + "\nMy previous action is invalid. " + "If I want to execute code, I should put the code between " + " and . " + "If I want to give the final answer, I should use the format " + "'Answer: \\boxed{answer}'. Let me try again.\n" + ) + done = False + + return next_obs, done + + +async def generate(args, sample: Sample, sampling_params) -> Sample: + """Custom generation function supporting tool calls""" + assert not args.partial_rollout, "Partial rollout is not supported for " "this function at the moment." + + # Retried samples (previously aborted / partial) arrive here with stale + # rollout state from the first attempt. Clear it so this generation starts + # clean; otherwise the concatenation below appends new tokens to old ones + # and downstream `slice_log_prob_with_cp` sees a length mismatch. + sample.rollout_log_probs = None + sample.response = "" + sample.response_length = 0 + sample.loss_mask = None + + state = GenerateState(args) + url = f"http://{args.vllm_router_ip}:{args.vllm_router_port}/inference/v1/generate" + + # Set up the initial prompt with system prompt and tools (outside the loop) + tool_specs = tool_registry.get_tool_specs() + prompt = format_conversation_with_tools(prompt=sample.prompt, tools=tool_specs) + + prompt_tokens_ids = state.tokenizer(prompt, add_special_tokens=False)["input_ids"] + response = "" + response_token_ids = [] + loss_masks = [] + tool_call_count = 0 # Track actual tool call rounds + + if args.rollout_max_context_len is not None: + max_context_length = args.rollout_max_context_len + else: + max_context_length = args.context_parallel_size * args.max_tokens_per_gpu + + # Track the finish reason of the most recent turn so the post-loop status + # mapping mirrors slime (which read it back off the last response). + finish_type = "stop" + + for turn in range(TOOL_CONFIGS["max_turns"]): + # Check if total length exceeds max context length + total_length = len(prompt_tokens_ids) + len(response_token_ids) + if total_length >= max_context_length: + sample.status = Sample.Status.TRUNCATED + break + + # Clamp per-turn max_new_tokens to the remaining context budget so a + # single turn cannot push total_length past max_context_length. Without + # this, a turn can append up to rollout_max_response_len tokens on top + # of a total that was just barely under the cap, producing samples + # that exceed the training-side max_tokens_per_gpu * cp_size budget + # and crash the partition/batch code (asserts or OOMs on an oversized + # partition). + remaining_budget = max_context_length - total_length + per_turn_sampling_params = dict(sampling_params) + per_turn_sampling_params["max_new_tokens"] = min( + sampling_params.get("max_new_tokens", remaining_budget), + remaining_budget, + ) + + # Use token IDs instead of text. vLLM's /inference/v1/generate is + # token-only; sampling params are nested under their own key and built + # by the canonical helper (logprobs enabled there). + current_token_ids = _coerce_flat_int_token_ids(prompt_tokens_ids + response_token_ids) + payload = { + "model": args.hf_checkpoint, + "token_ids": current_token_ids, + "sampling_params": _build_inference_sampling_params(per_turn_sampling_params), + } + + # Log payload to wandb for debugging + try: + import wandb + + if wandb.run is not None: + # Count available tools (from tool_specs) + available_tools = len(tool_specs) + # Count tools used in the current response + tools_used = response.count("") + + wandb.log( + { + "debug/payload_length": len(prompt + response), + "debug/available_tools": available_tools, + "debug/tools_used": tools_used, + "debug/turn": turn, + } + ) + except ImportError: + pass # wandb not available + + output = await post(url, payload) + + # Normalize the vLLM choice into a sglang-shaped meta_info: finish + # reason wrapped as {"type": str} and per-token (logprob, token_id) + # pairs under output_token_logprobs, exactly what the non-streaming + # vllm_rollout.generate produces. The retool control flow below then + # reads meta_info the same way slime did. + choice = output["choices"][0] + meta_info = _vllm_meta_from_generate_choice(args, choice, output.get("usage")) + cur_token_ids, cur_logprobs = _inference_generate_tokens_and_logprobs(choice) + if cur_token_ids: + meta_info["output_token_logprobs"] = [ + [float(lp), int(tid)] for lp, tid in zip(cur_logprobs, cur_token_ids, strict=False) + ] + finish_type = meta_info["finish_reason"]["type"] + + # Handle abort + if finish_type == "abort": + sample.status = Sample.Status.ABORTED + return sample + + if not meta_info.get("output_token_logprobs"): + # vLLM returned a choice but no per-token logprobs — we cannot + # recover per-token logprobs for this turn, which would desync + # rollout_log_probs from response_token_ids and blow up + # `slice_log_prob_with_cp` downstream. Abort the sample so the + # fully_async rollout manager returns the whole group to the + # buffer for retry instead of poisoning the trainer. + sample.status = Sample.Status.ABORTED + return sample + + cur_response_token_ids = [item[1] for item in meta_info["output_token_logprobs"]] + cur_response = state.tokenizer.decode(cur_response_token_ids) + cur_log_probs = [item[0] for item in meta_info["output_token_logprobs"]] + if sample.rollout_log_probs is None: + sample.rollout_log_probs = [] + sample.rollout_log_probs += cur_log_probs + + response += cur_response + response_token_ids += cur_response_token_ids + loss_masks += [1] * len(cur_response_token_ids) + + # Check length limit + if finish_type == "length": + break + + next_obs, done = await execute_predictions(cur_response) + if done: + break + + # Count tool calls (when we get interpreter output, it means a tool + # was called) + if "" in next_obs: + tool_call_count += 1 + + assert next_obs != "", "Next observation should not be empty." + obs_tokens_ids = state.tokenizer(next_obs, add_special_tokens=False)["input_ids"] + response += next_obs + response_token_ids += obs_tokens_ids + loss_masks += [0] * len(obs_tokens_ids) + + # Add dummy log probs for observation tokens (they won't be used due to loss_mask=0) + # Check if maximum tool call count reached + if sample.rollout_log_probs is not None: + sample.rollout_log_probs += [0.0] * len(obs_tokens_ids) + + assert len(response_token_ids) == len( + sample.rollout_log_probs + ), f"Token/logp length mismatch at turn {turn}: {len(response_token_ids)} tokens vs {len(sample.rollout_log_probs)} logps" + + # Tool output is appended verbatim and can push total_length past + # max_context_length (the per-turn generation was clamped to the + # remaining budget, but tool output is unconstrained). Trim tail + # tokens so the final sample fits the training budget exactly. + overflow = len(prompt_tokens_ids) + len(response_token_ids) - max_context_length + if overflow > 0: + response_token_ids = response_token_ids[:-overflow] + loss_masks = loss_masks[:-overflow] + if sample.rollout_log_probs is not None: + sample.rollout_log_probs = sample.rollout_log_probs[:-overflow] + # Resync the text field from the trimmed token list so + # reward_func's `sample.prompt + sample.response` matches what + # the model was actually trained on. decode(tokenize(text)) can + # be lossy on some tokenizers (whitespace / special-token + # collapse), but reward_func's regex is whitespace-robust and + # the trainer sees tokens, not text — so the drift is safe. + response = state.tokenizer.decode(response_token_ids) + sample.status = Sample.Status.TRUNCATED + break + + if tool_call_count >= TOOL_CONFIGS["max_tool_calls"]: + break + + # Set sample attributes + sample.tokens = prompt_tokens_ids + response_token_ids + sample.response_length = len(response_token_ids) + sample.response = response + sample.loss_mask = loss_masks + + # Store payload information for wandb logging + sample.payload_text = prompt + response + sample.payload_has_system = "<|im_start|>system" in prompt + response + sample.payload_has_tools = "# Tools" in prompt + response + + # Store tool call count for reward calculation + sample.tool_call_count = tool_call_count + + # Set status (skip if a context-budget branch above already set it) + if sample.status not in (Sample.Status.TRUNCATED, Sample.Status.ABORTED): + match finish_type: + case "length": + sample.status = Sample.Status.TRUNCATED + case "abort": + sample.status = Sample.Status.ABORTED + case "stop": + sample.status = Sample.Status.COMPLETED + + return sample + + +async def reward_func(args, sample, **kwargs): + """Tool call reward function using math_dapo as primary reward model""" + if not isinstance(sample, Sample): + raise TypeError("Sample must be an instance of Sample class.") + + # Build complete solution string + solution_str = sample.prompt + sample.response + + # Get ground truth answer - label is a string, not a dict + ground_truth = sample.label if sample.label is not None else "" + + # Get tool call count as num_turns + num_turns = getattr(sample, "tool_call_count", 0) + + # use \\boxed{...} answer + result = math_dapo_compute_score(solution_str, ground_truth, strict_box_verify=True) + + # encourage model to call tools + if result["score"] < 0: + tool_call_reward = (num_turns - 2) / 2 * 0.1 + result["score"] = min(-0.6, result["score"] + tool_call_reward) + + if result["pred"] is None: + result["pred"] = "" + + return result diff --git a/examples/retool/tool_sandbox.py b/examples/retool/tool_sandbox.py new file mode 100644 index 000000000..cdf68aa02 --- /dev/null +++ b/examples/retool/tool_sandbox.py @@ -0,0 +1,360 @@ +""" +Tool sandbox module for safe code execution and tool management. + +This module provides: +- PythonSandbox: Safe Python code execution environment +- ToolRegistry: Tool registration and execution management +- Memory management utilities +""" + +import asyncio +import gc +import os +import re +import subprocess +import tempfile +from contextlib import contextmanager +from typing import Any + +import psutil + +# Configuration for tool execution +TOOL_CONFIGS = { + "max_turns": 16, + "max_tool_calls": 16, + "tool_concurrency": 32, # Aggressive: 32 concurrent processes + # Python interpreter settings + "python_timeout": 120, # 2 minutes for complex calculations + "python_memory_limit": "4GB", # 4GB per Python process + "python_cpu_limit": 1, + # Memory management settings + "max_memory_usage": 12288, # 12GB total (75% of 16GB) + "cleanup_threshold": 6144, # 6GB + "aggressive_cleanup_threshold": 3072, # 3GB + "force_cleanup_threshold": 9216, # 9GB +} + +# Global semaphore for controlling concurrent tool executions +SEMAPHORE = asyncio.Semaphore(TOOL_CONFIGS["tool_concurrency"]) + + +def get_memory_usage() -> float: + """Get current memory usage in MB""" + process = psutil.Process() + return process.memory_info().rss / 1024 / 1024 + + +def cleanup_memory(): + """Force garbage collection to free memory""" + gc.collect() + + +def aggressive_cleanup_memory(): + """More aggressive memory cleanup""" + # Force multiple garbage collection cycles + for _ in range(3): + gc.collect() + + # Clear Python's internal caches + import sys + + # Note: sys.intern doesn't have a clear method, so we skip this + # Clear module cache if possible + if hasattr(sys, "modules"): + # Don't clear all modules, but clear some common ones that might cache data + modules_to_clear = ["numpy", "pandas", "matplotlib", "scipy"] + for module_name in modules_to_clear: + if module_name in sys.modules: + module = sys.modules[module_name] + if hasattr(module, "clear_cache"): + module.clear_cache() + + +def check_and_cleanup_memory(): + """Check memory usage and perform appropriate cleanup""" + current_memory = get_memory_usage() + + if current_memory > TOOL_CONFIGS["force_cleanup_threshold"]: + # Force aggressive cleanup + aggressive_cleanup_memory() + return f"Warning: High memory usage ({current_memory:.1f}MB), performed aggressive cleanup" + elif current_memory > TOOL_CONFIGS["cleanup_threshold"]: + # Normal cleanup + cleanup_memory() + return f"Info: Memory usage ({current_memory:.1f}MB), performed cleanup" + elif current_memory > TOOL_CONFIGS["aggressive_cleanup_threshold"]: + # Light cleanup + gc.collect() + return f"Info: Memory usage ({current_memory:.1f}MB), performed light cleanup" + + return None + + +class PythonSandbox: + """Python code sandbox, provides safe code execution environment""" + + def __init__(self, timeout: int = 10, memory_limit: str = "100MB"): + self.timeout = timeout + self.memory_limit = memory_limit + self.allowed_modules = { + "math", + "random", + "datetime", + "collections", + "itertools", + "functools", + "operator", + "statistics", + "decimal", + "fractions", + } + + def _check_code_safety(self, code: str) -> tuple[bool, str]: + """Check code safety by scanning for dangerous patterns""" + # Check for dangerous operations + dangerous_patterns = [ + r"import\s+os", + r"import\s+sys", + r"import\s+subprocess", + r"import\s+shutil", + r"import\s+glob", + r"import\s+pathlib", + r"__import__", + r"eval\s*\(", + r"exec\s*\(", + r"open\s*\(", + r"file\s*\(", + r"input\s*\(", + r"raw_input\s*\(", + r"compile\s*\(", + r"execfile\s*\(", + r"getattr\s*\(", + r"setattr\s*\(", + r"delattr\s*\(", + r"hasattr\s*\(", + r"globals\s*\(", + r"locals\s*\(", + r"vars\s*\(", + r"dir\s*\(", + r"type\s*\(", + r"isinstance\s*\(", + r"issubclass\s*\(", + r"super\s*\(", + r"property\s*\(", + r"staticmethod\s*\(", + r"classmethod\s*\(", + r"__\w+__", # double underscore methods + ] + + for pattern in dangerous_patterns: + if re.search(pattern, code, re.IGNORECASE): + return False, f"Code contains dangerous pattern: {pattern}" + + # Check imported modules + import_pattern = r"import\s+(\w+)" + from_pattern = r"from\s+(\w+)" + + imports = re.findall(import_pattern, code) + froms = re.findall(from_pattern, code) + + all_imports = set(imports + froms) + for imp in all_imports: + if imp not in self.allowed_modules: + return False, f"Import of '{imp}' is not allowed" + + return True, "Code is safe" + + @contextmanager + def _create_safe_environment(self): + """Create safe execution environment with temporary directory""" + # Create temporary directory + temp_dir = tempfile.mkdtemp(prefix="python_sandbox_") + + try: + # Create safe Python script + script_path = os.path.join(temp_dir, "code.py") + + # Set environment variables + env = os.environ.copy() + env["PYTHONPATH"] = temp_dir + env["PYTHONUNBUFFERED"] = "1" + + yield script_path, env, temp_dir + + finally: + # Clean up temporary directory + try: + import shutil + + shutil.rmtree(temp_dir) + except Exception: + pass + + async def execute_code(self, code: str) -> str: + """Execute Python code in sandbox with safety checks""" + # Check memory usage before execution + current_memory = get_memory_usage() + if current_memory > TOOL_CONFIGS["max_memory_usage"]: + aggressive_cleanup_memory() + return "Error: Memory usage too high, please try again" + + # Check code safety + is_safe, message = self._check_code_safety(code) + if not is_safe: + return f"Error: {message}" + + # Add necessary wrapper code with memory limits + # Properly indent the user code within the try block + # Handle indentation properly by adding 4 spaces to each line + indented_code = "\n".join(" " + line for line in code.split("\n")) + + wrapped_code = f"""import sys +import traceback +from io import StringIO +import resource + +# Set memory limit (4GB) +try: + resource.setrlimit(resource.RLIMIT_AS, (4 * 1024 * 1024 * 1024, -1)) +except Exception: + pass + +# Redirect stdout and stderr +old_stdout = sys.stdout +old_stderr = sys.stderr +stdout_capture = StringIO() +stderr_capture = StringIO() +sys.stdout = stdout_capture +sys.stderr = stderr_capture + +try: + # User code +{indented_code} + + # Get output + stdout_output = stdout_capture.getvalue() + stderr_output = stderr_capture.getvalue() + + # Restore standard output + sys.stdout = old_stdout + sys.stderr = old_stderr + + # Return result + result = "" + if stdout_output: + result += f"Output:\\n{{stdout_output}}" + if stderr_output: + result += f"\\nErrors:\\n{{stderr_output}}" + + print(result) + +except Exception as e: + # Restore standard output + sys.stdout = old_stdout + sys.stderr = old_stderr + + # Return error information + error_msg = f"Error: {{str(e)}}\\nTraceback:\\n{{traceback.format_exc()}}" + print(error_msg)""" + + with self._create_safe_environment() as (script_path, env, temp_dir): + # Write code to file + with open(script_path, "w") as f: + f.write(wrapped_code) + + try: + # Use subprocess to run code + process = subprocess.Popen( + ["python3", script_path], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=env, + cwd=temp_dir, + text=True, + ) + + # Set timeout + try: + stdout, stderr = process.communicate(timeout=self.timeout) + + if process.returncode == 0: + result = stdout.strip() + else: + result = f"Error: Process exited with code {process.returncode}\n{stderr}" + + except subprocess.TimeoutExpired: + process.kill() + result = f"Error: Code execution timed out after {self.timeout} seconds" + + except Exception as e: + result = f"Error: Failed to execute code: {str(e)}" + + # Check memory usage after execution and cleanup if needed + cleanup_message = check_and_cleanup_memory() + if cleanup_message: + print(f"Memory cleanup: {cleanup_message}") + + return result + + +class ToolRegistry: + """Tool registry, manages available tools and their execution""" + + def __init__(self): + self.tools = {} + self.python_sandbox = PythonSandbox( + timeout=TOOL_CONFIGS["python_timeout"], memory_limit=TOOL_CONFIGS["python_memory_limit"] + ) + self._register_default_tools() + + def _register_default_tools(self): + """Register default tools in the registry""" + # Python code interpreter + self.register_tool( + "code_interpreter", + { + "type": "function", + "function": { + "name": "code_interpreter", + "description": "A tool for executing Python code in a safe sandbox environment.", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string", "description": "The Python code to execute"}}, + "required": ["code"], + }, + }, + }, + ) + + def register_tool(self, name: str, tool_spec: dict[str, Any]): + """Register a new tool in the registry""" + self.tools[name] = tool_spec + + def get_tool_specs(self) -> list[dict[str, Any]]: + """Get all tool specifications as a list""" + return list(self.tools.values()) + + async def execute_tool(self, tool_name: str, arguments: dict[str, Any]) -> str: + """Execute a tool call with the given arguments""" + if tool_name not in self.tools: + return f"Error: Tool '{tool_name}' not found" + + async with SEMAPHORE: + if tool_name == "code_interpreter": + return await self._execute_python(arguments) + else: + return f"Error: Tool '{tool_name}' not implemented" + + async def _execute_python(self, arguments: dict[str, Any]) -> str: + """Execute Python code using the sandbox""" + code = arguments.get("code", "") + if not code.strip(): + return "Error: No code provided" + + # Execute code in sandbox + result = await self.python_sandbox.execute_code(code) + return result + + +# Global tool registry instance +tool_registry = ToolRegistry() diff --git a/requirements.txt b/requirements.txt index 7db56918c..aa8cb6d32 100644 --- a/requirements.txt +++ b/requirements.txt @@ -13,6 +13,7 @@ pyyaml qwen_vl_utils # for VLM ray[default] ring_flash_attn +safetensors tensorboard transformers vllm-router>=0.1.14 diff --git a/tests/_cp_dist_helpers.py b/tests/_cp_dist_helpers.py new file mode 100644 index 000000000..1382094fe --- /dev/null +++ b/tests/_cp_dist_helpers.py @@ -0,0 +1,167 @@ +"""Shared infrastructure for the CP-related multi-process CPU tests. + +Why this module exists +---------------------- +The CP / metric-report / backward-grad-norm tests all want to: + +1. Stub ``megatron.core.mpu`` *before* importing + ``vime.backends.megatron_utils.cp_utils`` (the CPU CI image has no real + megatron). +2. Spawn ``dp_size * cp_size`` workers with real ``torch.distributed`` and + exercise the actual production helpers (``get_sum_of_sample_mean``, + ``reduce_train_step_metrics``, ``gather_and_reduce_log_dict``, + ``rollout_log_metric_contribution``). +3. Chunk each sample's response tensor across CP ranks the same way the + real forward pass does — using + ``get_logits_and_tokens_offset_with_cp`` so the slicing stays in lock- + step with the production reducer. + +Putting that here keeps the per-feature test files focused on the +behaviour they check (numerics / report formulas / backward) rather than +on plumbing. + +Mapping to Megatron +------------------- +- ``mp.spawn(...)`` + gloo backend mirrors the per-rank entry-point that + ``torch.distributed.run`` would create for a real launch. +- ``dp_cp_group = new_group(range(world_size))`` matches + ``parallel_state.get_data_parallel_group(with_context_parallel=True)`` + (Megatron-LM ``finalize_model_grads.py:437``). In the no-TP / no-PP + CPU test setup the whole world *is* that group. +- The per-rank CP chunking mirrors what the attention layer feeds into + the loss in Megatron: each CP rank only sees its 2-chunk slice of the + response tokens (cf. ``cp_utils.get_logits_and_tokens_offset_with_cp``, + the same helper used by the real forward pass). +""" + +from __future__ import annotations + +import os +import socket +import sys +import types + + +# --- Stub ``megatron.core.mpu`` (must run before cp_utils is imported) --- +# +# Both this module and any test file that imports it should *import this +# helper first*. Doing so installs the stub at import time so that the +# subsequent ``from vime.backends.megatron_utils.cp_utils import ...`` in +# the test file binds ``cp_utils.mpu`` to this stub. +# +# In spawned workers, ``mp.spawn`` re-imports the test module fresh, which +# re-runs this stub installation; then the worker mutates the stub's +# ``get_context_parallel_*`` attributes via ``_stub_megatron_in_worker`` +# below to pin (cp_size, cp_rank) for that worker. +_fake_mpu = types.ModuleType("megatron.core.mpu") +_fake_mpu.get_context_parallel_world_size = lambda: 1 +_fake_mpu.get_context_parallel_rank = lambda: 0 +_fake_core = types.ModuleType("megatron.core") +_fake_core.mpu = _fake_mpu +_fake_megatron = types.ModuleType("megatron") +_fake_megatron.core = _fake_core +sys.modules.setdefault("megatron", _fake_megatron) +sys.modules.setdefault("megatron.core", _fake_core) +sys.modules.setdefault("megatron.core.mpu", _fake_mpu) + + +def stub_megatron_in_worker(cp_size: int, cp_rank: int) -> None: + """Override ``mpu.get_context_parallel_*`` inside an ``mp.spawn`` worker. + + ``mp.spawn`` pickles the worker function by name and re-imports the + test module in the child — that re-runs the top-of-file stub install + with ``cp_size=1``. By the time the worker runs, ``cp_utils`` has + already bound its module-level ``mpu`` reference to the stub. + + So we must MUTATE the stub module's attributes in place rather than + replace ``sys.modules['megatron.core.mpu']`` — replacing the module + would leave ``cp_utils.mpu`` pointing at the now-shadowed stub. + """ + from megatron.core import mpu # the stub installed at import time + + mpu.get_context_parallel_world_size = lambda: cp_size + mpu.get_context_parallel_rank = lambda: cp_rank + + +def free_port() -> int: + """Pick an unused TCP port for ``init_process_group``'s rendezvous. + + Equivalent to what ``torchrun`` does when ``--master-port`` is not + set; we just need a port nothing else is bound to so multiple + parametrized test cases can spawn without colliding. + """ + s = socket.socket() + s.bind(("", 0)) + port = s.getsockname()[1] + s.close() + return port + + +def init_worker_process_group(rank: int, world_size: int, master_port: int): + """Stand up gloo ``torch.distributed`` and return the DP*CP group. + + The CPU CI image ships gloo but not NCCL; in the no-TP / no-PP setup + the DP-with-CP group is the whole world, mirroring + ``parallel_state.get_data_parallel_group(with_context_parallel=True)`` + in Megatron-LM ``finalize_model_grads.py:437``. + """ + import torch.distributed as _dist + + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(master_port) + _dist.init_process_group(backend="gloo", rank=rank, world_size=world_size) + return _dist.new_group(ranks=list(range(world_size))) + + +def cp_chunk_response_tensor(x, total_length: int, response_length: int): + """Slice a sample's response tensor to what the current CP rank sees. + + Mirrors the real forward pass: at CP > 1 each rank's attention only + consumes the two response-token chunks selected by + ``get_logits_and_tokens_offset_with_cp`` (the same helper used by the + production reducer in ``cp_utils.get_sum_of_sample_mean``). So the + "x" we feed into the reducer on a CP rank must be sliced the same + way to keep the numbers honest. + + Importing locally so callers don't pay the import cost before + ``stub_megatron_in_worker`` has had a chance to pin (cp_size, cp_rank). + """ + import torch + + from vime.backends.megatron_utils.cp_utils import get_logits_and_tokens_offset_with_cp + + prompt_length = total_length - response_length + _, _, _, tokens_offset = get_logits_and_tokens_offset_with_cp(total_length, response_length) + c0 = x[tokens_offset[0][0] - prompt_length : tokens_offset[0][1] - prompt_length] + c1 = x[tokens_offset[1][0] - prompt_length : tokens_offset[1][1] - prompt_length] + return torch.cat([c0, c1]) + + +# --------------------------------------------------------------------------- +# Shared four-rollout fixture, used by both the metric-report distributed +# tests and the backward-grad-norm test. Keeping the data in one place so +# the "train report matches rollout report matches grad-norm baseline" +# contract is anchored on the same numbers everywhere. +# +# Four samples (1 rollout each), total_length=12 (4 prompt + 8 response), +# loss_mask=all-ones. x values differ by orders of magnitude so any cross- +# rank summation bug shows up as a visibly wrong number. +# +# Per-sample token-mean: 4.5 / 45 / 450 / 4500. +# Per-rollout-mean report (sum / num_rollouts): +# (4.5 + 45 + 450 + 4500) / 4 = 1249.875 +# Per-token-loss report (sum_x / total_tokens): +# (36 + 360 + 3600 + 36000) / 32 = 1249.875 +# (the two paths agree by construction so the test expectations stay +# simple — the *report formulas* are still distinct as exercised inside +# ``reduce_train_step_metrics``.) +# --------------------------------------------------------------------------- +FOUR_ROLLOUT_TOTAL_LENGTHS = [12, 12, 12, 12] +FOUR_ROLLOUT_RESPONSE_LENGTHS = [8, 8, 8, 8] +FOUR_ROLLOUT_X_VALUES = [ + [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], + [10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0], + [100.0, 200.0, 300.0, 400.0, 500.0, 600.0, 700.0, 800.0], + [1000.0, 2000.0, 3000.0, 4000.0, 5000.0, 6000.0, 7000.0, 8000.0], +] +FOUR_ROLLOUT_EXPECTED_REPORT = 1249.875 diff --git a/tests/test_agent_adapters.py b/tests/test_agent_adapters.py new file mode 100644 index 000000000..94e055f67 --- /dev/null +++ b/tests/test_agent_adapters.py @@ -0,0 +1,853 @@ +import asyncio +import json +import sys +from pathlib import Path + +import pytest +from aiohttp import web +from aiohttp.test_utils import TestClient, TestServer + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from vime.agent.adapters import anthropic, openai +from vime.agent.adapters.common import VLLM_URL_KEY +from vime.agent.trajectory import TurnRecord + + +NUM_GPUS = 0 + + +class ToyTokenizer: + def __init__(self, outputs: dict[tuple[int, ...], str] | None = None) -> None: + self.outputs = outputs or {} + self.rendered: list[tuple[list[dict], list[dict] | None]] = [] + + def apply_chat_template(self, messages, tools=None, tokenize=True, add_generation_prompt=True): + self.rendered.append((list(messages), tools)) + return list(range(1, len(messages) + 2)) + + def decode(self, ids, skip_special_tokens=False): + return self.outputs.get(tuple(ids), "") + + +class ScriptedTokenizer(ToyTokenizer): + def __init__(self, prompts: list[list[int]], outputs: dict[tuple[int, ...], str]) -> None: + super().__init__(outputs) + self.prompts = [list(prompt) for prompt in prompts] + + def apply_chat_template(self, messages, tools=None, tokenize=True, add_generation_prompt=True): + self.rendered.append((list(messages), tools)) + assert self.prompts, "unexpected chat-template render" + return self.prompts.pop(0) + + +class FakeVLLM: + def __init__(self, turns: list[list[tuple[float, int]]]) -> None: + self.turns = [list(turn) for turn in turns] + self.requests: list[dict] = [] + self.routing_keys: list[str | None] = [] + + async def handle_generate(self, request): + self.routing_keys.append(request.headers.get("x-session-id")) + self.requests.append(await request.json()) + assert self.turns, "unexpected /inference/v1/generate call" + turn = self.turns.pop(0) + token_ids = [token_id for _logprob, token_id in turn] + content = [{"logprob": logprob} for logprob, _token_id in turn] + return web.json_response( + { + "choices": [ + { + "token_ids": token_ids, + "logprobs": {"content": content}, + "finish_reason": "stop", + } + ] + } + ) + + +class FakeRequest: + def __init__(self, headers: dict[str, str]) -> None: + self.headers = headers + + +def _parse_sse(raw: str) -> list[tuple[str, object]]: + events: list[tuple[str, object]] = [] + event_name = "message" + data_lines: list[str] = [] + + def flush() -> None: + nonlocal event_name, data_lines + if not data_lines: + event_name = "message" + return + data = "\n".join(data_lines) + payload: object + if data == "[DONE]": + payload = data + else: + payload = json.loads(data) + events.append((event_name, payload)) + event_name = "message" + data_lines = [] + + for line in raw.splitlines(): + if not line: + flush() + elif line.startswith("event:"): + event_name = line.removeprefix("event:").strip() + elif line.startswith("data:"): + data_lines.append(line.removeprefix("data:").strip()) + flush() + return events + + +@pytest.mark.unit +def test_session_id_comes_from_protocol_fields_not_custom_header(): + assert ( + openai._request_session_id( + FakeRequest({"X-Slime-Session-Id": "custom"}), + {"metadata": {"session_id": "meta-session"}, "user": "body-user"}, + ) + == "meta-session" + ) + assert ( + openai._request_session_id(FakeRequest({"X-Slime-Session-Id": "custom"}), {"user": "body-user"}) == "body-user" + ) + assert ( + anthropic._request_session_id(FakeRequest({"X-Slime-Session-Id": "custom", "X-Api-Key": "anthropic-key"})) + == "anthropic-key" + ) + assert ( + anthropic._request_session_id( + FakeRequest({"Authorization": "Bearer bearer-session", "X-Api-Key": "anthropic-key"}) + ) + == "bearer-session" + ) + + +@pytest.mark.unit +def test_anthropic_translation_keeps_tool_results_and_tool_schema(): + messages = [ + {"role": "user", "content": [{"type": "text", "text": "hi"}]}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "plan"}, + {"type": "text", "text": "ok"}, + {"type": "tool_use", "name": "lookup", "input": {"q": "slime"}}, + ], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "u1", "content": "result"}]}, + ] + + translated = anthropic._translate_anthropic(messages, system="sys") + tools = anthropic._anthropic_tools_to_chat_tools( + [{"name": "lookup", "description": "search", "input_schema": {"type": "object"}}] + ) + + assert translated == [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": "ok", + "reasoning_content": "plan", + "tool_calls": [{"function": {"name": "lookup", "arguments": {"q": "slime"}}}], + }, + {"role": "tool", "content": "result"}, + ] + assert tools == [ + { + "type": "function", + "function": { + "name": "lookup", + "description": "search", + "parameters": {"type": "object"}, + }, + } + ] + + +@pytest.mark.unit +def test_openai_translation_and_responses_input_shapes(): + chat_messages = openai._translate_chat_messages( + [ + {"role": "developer", "content": "rules"}, + {"role": "user", "content": [{"type": "text", "text": "hello"}]}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": {"q": "slime"}}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "found"}, + ] + ) + response_messages = openai._responses_input_to_messages( + [ + {"role": "user", "content": [{"type": "input_text", "text": "question"}]}, + {"type": "function_call_output", "call_id": "call_1", "output": "answer"}, + ], + instructions="be brief", + ) + + assert chat_messages == [ + {"role": "system", "content": "rules"}, + {"role": "user", "content": "hello"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": '{"q": "slime"}'}, + } + ], + }, + {"role": "tool", "content": "found", "tool_call_id": "call_1"}, + ] + assert response_messages == [ + {"role": "system", "content": "be brief"}, + {"role": "user", "content": [{"type": "input_text", "text": "question"}]}, + {"role": "tool", "tool_call_id": "call_1", "content": "answer"}, + ] + + +@pytest.mark.unit +def test_openai_chat_completion_endpoint_records_token_segments(monkeypatch): + async def fake_generate(prompt_ids, session, body, app, **kwargs): + return TurnRecord(prompt_ids=list(prompt_ids), output_ids=[101], finish_reason="stop", output_log_probs=[-0.1]) + + async def run_case(): + monkeypatch.setattr(openai, "_generate", fake_generate) + tokenizer = ToyTokenizer({(101,): "hello"}) + adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") + adapter.open_session("sid-chat", sampling_defaults={"max_new_tokens": 8}) + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + resp = await client.post( + "/v1/chat/completions", + headers={"Authorization": "Bearer sid-chat"}, + json={ + "model": "actor", + "messages": [{"role": "user", "content": "hello?"}], + "max_tokens": 4, + }, + ) + data = await resp.json() + finally: + await client.close() + + segments = await adapter.finish_session("sid-chat") + assert resp.status == 200 + assert data["object"] == "chat.completion" + assert data["choices"][0]["message"] == {"role": "assistant", "content": "hello"} + assert data["usage"] == {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3} + assert segments[0].prompt_ids == [1, 2] + assert segments[0].response_ids == [101] + assert segments[0].loss_mask == [1] + + asyncio.run(run_case()) + + +@pytest.mark.unit +def test_openai_chat_completion_streaming_returns_sse_chunks_and_records_segments(monkeypatch): + async def fake_generate(prompt_ids, session, body, app, **kwargs): + return TurnRecord(prompt_ids=list(prompt_ids), output_ids=[401], finish_reason="stop", output_log_probs=[-0.4]) + + async def run_case(): + monkeypatch.setattr(openai, "_generate", fake_generate) + tokenizer = ToyTokenizer({(401,): "streamed text"}) + adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") + adapter.open_session("sid-chat-stream", sampling_defaults={"max_new_tokens": 8}) + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + resp = await client.post( + "/v1/chat/completions", + headers={"Authorization": "Bearer sid-chat-stream"}, + json={ + "model": "actor", + "stream": True, + "messages": [{"role": "user", "content": "hello?"}], + }, + ) + raw = await resp.text() + finally: + await client.close() + + events = _parse_sse(raw) + chunks = [payload for _, payload in events if isinstance(payload, dict)] + segments = await adapter.finish_session("sid-chat-stream") + assert resp.status == 200 + assert chunks[0]["object"] == "chat.completion.chunk" + assert chunks[0]["choices"][0]["delta"] == {"role": "assistant"} + assert any(c["choices"][0]["delta"] == {"content": "streamed text"} for c in chunks) + assert chunks[-1]["choices"][0]["finish_reason"] == "stop" + assert chunks[-1]["usage"] == {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3} + assert events[-1] == ("message", "[DONE]") + assert segments[0].prompt_ids == [1, 2] + assert segments[0].response_ids == [401] + assert segments[0].rollout_log_probs == [-0.4] + + asyncio.run(run_case()) + + +@pytest.mark.unit +def test_openai_chat_completion_streaming_returns_tool_call_delta(monkeypatch): + async def fake_generate(prompt_ids, session, body, app, **kwargs): + return TurnRecord( + prompt_ids=list(prompt_ids), output_ids=[451], finish_reason="stop", output_log_probs=[-0.45] + ) + + async def run_case(): + monkeypatch.setattr(openai, "_generate", fake_generate) + raw = "use it slime" + tokenizer = ToyTokenizer({(451,): raw}) + adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") + adapter.open_session("sid-chat-tool-stream", sampling_defaults={"max_new_tokens": 8}) + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + resp = await client.post( + "/v1/chat/completions", + headers={"Authorization": "Bearer sid-chat-tool-stream"}, + json={ + "model": "actor", + "stream": True, + "messages": [{"role": "user", "content": "call lookup"}], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "description": "search", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}, + }, + } + ], + }, + ) + raw_sse = await resp.text() + finally: + await client.close() + + chunks = [payload for _, payload in _parse_sse(raw_sse) if isinstance(payload, dict)] + tool_delta = next(c["choices"][0]["delta"] for c in chunks if "tool_calls" in c["choices"][0]["delta"]) + segments = await adapter.finish_session("sid-chat-tool-stream") + assert resp.status == 200 + assert any(c["choices"][0]["delta"] == {"content": "use it"} for c in chunks) + assert tool_delta["tool_calls"][0]["index"] == 0 + assert tool_delta["tool_calls"][0]["function"]["name"] == "lookup" + assert tool_delta["tool_calls"][0]["function"]["arguments"] == '{"query": "slime"}' + assert chunks[-1]["choices"][0]["finish_reason"] == "tool_calls" + assert segments[0].response_ids == [451] + + asyncio.run(run_case()) + + +@pytest.mark.unit +def test_openai_responses_endpoint_returns_function_calls(monkeypatch): + async def fake_generate(prompt_ids, session, body, app, **kwargs): + return TurnRecord(prompt_ids=list(prompt_ids), output_ids=[301], finish_reason="stop", output_log_probs=[-0.3]) + + async def run_case(): + monkeypatch.setattr(openai, "_generate", fake_generate) + raw = "look slime" + tokenizer = ToyTokenizer({(301,): raw}) + adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") + adapter.open_session("sid-responses", sampling_defaults={"max_new_tokens": 8}) + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + resp = await client.post( + "/v1/responses", + headers={"Authorization": "Bearer sid-responses"}, + json={ + "model": "actor", + "input": "find it", + "tools": [ + { + "type": "function", + "name": "lookup", + "description": "search", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}, + } + ], + }, + ) + data = await resp.json() + finally: + await client.close() + + output_types = [item["type"] for item in data["output"]] + function_call = next(item for item in data["output"] if item["type"] == "function_call") + segments = await adapter.finish_session("sid-responses") + assert resp.status == 200 + assert data["object"] == "response" + assert output_types == ["message", "function_call"] + assert data["output"][0]["content"][0]["text"] == "look" + assert function_call["name"] == "lookup" + assert function_call["arguments"] == '{"query": "slime"}' + assert segments[0].response_ids == [301] + + asyncio.run(run_case()) + + +@pytest.mark.unit +def test_openai_responses_streaming_preserves_function_call_output(monkeypatch): + async def fake_generate(prompt_ids, session, body, app, **kwargs): + return TurnRecord( + prompt_ids=list(prompt_ids), output_ids=[551], finish_reason="stop", output_log_probs=[-0.55] + ) + + async def run_case(): + monkeypatch.setattr(openai, "_generate", fake_generate) + raw = "slime" + tokenizer = ToyTokenizer({(551,): raw}) + adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") + adapter.open_session("sid-responses-tool-stream", sampling_defaults={"max_new_tokens": 8}) + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + resp = await client.post( + "/v1/responses", + headers={"Authorization": "Bearer sid-responses-tool-stream"}, + json={ + "model": "actor", + "stream": True, + "input": "call lookup", + "tools": [ + { + "type": "function", + "name": "lookup", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}, + } + ], + }, + ) + raw_sse = await resp.text() + finally: + await client.close() + + events = _parse_sse(raw_sse) + created = next(payload for name, payload in events if name == "response.created") + completed = next(payload for name, payload in events if name == "response.completed") + completed_call = next(item for item in completed["response"]["output"] if item["type"] == "function_call") + segments = await adapter.finish_session("sid-responses-tool-stream") + assert resp.status == 200 + assert created["type"] == "response.created" + assert [item["type"] for item in created["response"]["output"]] == ["function_call"] + assert completed_call["name"] == "lookup" + assert completed_call["arguments"] == '{"query": "slime"}' + assert segments[0].response_ids == [551] + + asyncio.run(run_case()) + + +@pytest.mark.unit +def test_openai_responses_streaming_returns_sse_events_and_records_segments(monkeypatch): + async def fake_generate(prompt_ids, session, body, app, **kwargs): + return TurnRecord(prompt_ids=list(prompt_ids), output_ids=[501], finish_reason="stop", output_log_probs=[-0.5]) + + async def run_case(): + monkeypatch.setattr(openai, "_generate", fake_generate) + tokenizer = ToyTokenizer({(501,): "response text"}) + adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") + adapter.open_session("sid-responses-stream", sampling_defaults={"max_new_tokens": 8}) + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + resp = await client.post( + "/v1/responses", + headers={"Authorization": "Bearer sid-responses-stream"}, + json={ + "model": "actor", + "stream": True, + "instructions": "be brief", + "input": [{"role": "user", "content": [{"type": "input_text", "text": "hello?"}]}], + }, + ) + raw = await resp.text() + finally: + await client.close() + + events = _parse_sse(raw) + event_names = [name for name, _ in events] + text_delta = next(payload for name, payload in events if name == "response.output_text.delta") + completed = next(payload for name, payload in events if name == "response.completed") + segments = await adapter.finish_session("sid-responses-stream") + assert resp.status == 200 + assert event_names == ["response.created", "response.output_text.delta", "response.completed"] + assert text_delta == {"type": "response.output_text.delta", "delta": "response text"} + assert completed["response"]["status"] == "completed" + assert completed["response"]["usage"] == {"input_tokens": 3, "output_tokens": 1, "total_tokens": 4} + assert segments[0].prompt_ids == [1, 2, 3] + assert segments[0].response_ids == [501] + + asyncio.run(run_case()) + + +@pytest.mark.unit +def test_anthropic_messages_endpoint_returns_non_stream_json_and_records_segments(monkeypatch): + async def fake_generate(prompt_ids, session, body, app, **kwargs): + return TurnRecord( + prompt_ids=list(prompt_ids), output_ids=[581], finish_reason="stop", output_log_probs=[-0.58] + ) + + async def run_case(): + monkeypatch.setattr(anthropic, "_generate", fake_generate) + tokenizer = ToyTokenizer({(581,): "plain response"}) + adapter = anthropic.AnthropicAdapter(tokenizer=tokenizer, vllm_url="http://unused") + adapter.open_session("sid-anthropic-json", sampling_defaults={"max_new_tokens": 8}) + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + resp = await client.post( + "/v1/messages", + headers={"Authorization": "Bearer sid-anthropic-json"}, + json={ + "model": "actor", + "system": "be useful", + "max_tokens": 4, + "messages": [{"role": "user", "content": [{"type": "text", "text": "solve"}]}], + }, + ) + data = await resp.json() + finally: + await client.close() + + segments = await adapter.finish_session("sid-anthropic-json") + assert resp.status == 200 + assert data["type"] == "message" + assert data["model"] == "actor" + assert data["content"] == [{"type": "text", "text": "plain response"}] + assert data["stop_reason"] == "end_turn" + assert data["usage"] == {"input_tokens": 3, "output_tokens": 1} + assert segments[0].prompt_ids == [1, 2, 3] + assert segments[0].response_ids == [581] + + asyncio.run(run_case()) + + +@pytest.mark.unit +def test_anthropic_messages_endpoint_streams_blocks_and_records_segments(monkeypatch): + async def fake_generate(prompt_ids, session, body, app, **kwargs): + return TurnRecord(prompt_ids=list(prompt_ids), output_ids=[601], finish_reason="stop", output_log_probs=[-0.6]) + + async def run_case(): + monkeypatch.setattr(anthropic, "_generate", fake_generate) + raw_output = ( + "delegate inspect" + ) + tokenizer = ToyTokenizer({(601,): raw_output}) + adapter = anthropic.AnthropicAdapter(tokenizer=tokenizer, vllm_url="http://unused") + adapter.open_session("sid-anthropic-stream", sampling_defaults={"max_new_tokens": 8}) + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + resp = await client.post( + "/v1/messages", + headers={"Authorization": "Bearer sid-anthropic-stream"}, + json={ + "model": "actor", + "system": "be useful", + "stream": True, + "max_tokens": 4, + "messages": [{"role": "user", "content": [{"type": "text", "text": "solve"}]}], + "tools": [ + { + "name": "Task", + "description": "spawn subagent", + "input_schema": {"type": "object", "properties": {"description": {"type": "string"}}}, + } + ], + }, + ) + raw = await resp.text() + finally: + await client.close() + + events = _parse_sse(raw) + names = [name for name, _ in events] + starts = [payload for name, payload in events if name == "content_block_start"] + deltas = [payload for name, payload in events if name == "content_block_delta"] + message_delta = next(payload for name, payload in events if name == "message_delta") + segments = await adapter.finish_session("sid-anthropic-stream") + assert resp.status == 200 + assert names[0] == "message_start" + assert names[-1] == "message_stop" + assert any(s["content_block"]["type"] == "text" for s in starts) + assert any(s["content_block"]["type"] == "tool_use" and s["content_block"]["name"] == "Task" for s in starts) + assert any(d["delta"].get("text") == "delegate" for d in deltas) + assert any(json.loads(d["delta"].get("partial_json", "{}")) == {"description": "inspect"} for d in deltas) + assert message_delta["delta"]["stop_reason"] == "tool_use" + assert message_delta["usage"] == {"input_tokens": 3, "output_tokens": 1} + assert segments[0].metadata["segment_kind"] == "final" + assert segments[0].prompt_ids == [1, 2, 3] + assert segments[0].response_ids == [601] + assert segments[0].loss_mask == [1] + + asyncio.run(run_case()) + + +@pytest.mark.unit +def test_openai_responses_multiturn_uses_vllm_tokens_for_training_segment(): + async def run_case(): + upstream = FakeVLLM( + [ + [(-0.20, 20), (-0.21, 21)], + [(-0.40, 40)], + ] + ) + upstream_app = web.Application() + upstream_app.router.add_post("/inference/v1/generate", upstream.handle_generate) + upstream_server = TestServer(upstream_app) + await upstream_server.start_server() + + tool_raw = "slime" + tokenizer = ScriptedTokenizer( + prompts=[ + [10, 11], + [10, 11, 20, 21, 30, 31], + ], + outputs={ + (20, 21): tool_raw, + (40,): "done", + }, + ) + adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url=str(upstream_server.make_url("")).rstrip("/")) + adapter.open_session("sid-openai-token", sampling_defaults={"max_new_tokens": 99}) + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + first = await client.post( + "/v1/responses", + headers={"Authorization": "Bearer sid-openai-token"}, + json={ + "model": "actor", + "input": "find slime", + "max_output_tokens": 5, + "tools": [ + { + "type": "function", + "name": "lookup", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}, + } + ], + }, + ) + first_data = await first.json() + function_call = next(item for item in first_data["output"] if item["type"] == "function_call") + + second = await client.post( + "/v1/responses", + headers={"Authorization": "Bearer sid-openai-token"}, + json={ + "model": "actor", + "input": [ + {"role": "user", "content": "find slime"}, + function_call, + { + "type": "function_call_output", + "call_id": function_call["call_id"], + "output": "found slime", + }, + ], + "max_output_tokens": 7, + "tools": [ + { + "type": "function", + "name": "lookup", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}, + } + ], + }, + ) + second_data = await second.json() + finally: + await client.close() + await upstream_server.close() + + segments = await adapter.finish_session("sid-openai-token") + assert first.status == 200 + assert second.status == 200 + assert function_call["name"] == "lookup" + assert function_call["arguments"] == '{"query": "slime"}' + assert second_data["output"][0]["content"][0]["text"] == "done" + assert [req["token_ids"] for req in upstream.requests] == [[10, 11], [10, 11, 20, 21, 30, 31]] + assert upstream.routing_keys == ["sid-openai-token", "sid-openai-token"] + assert upstream.requests[0]["sampling_params"]["max_tokens"] == 5 + assert upstream.requests[1]["sampling_params"]["max_tokens"] == 7 + assert segments[0].prompt_ids == [10, 11] + assert segments[0].response_ids == [20, 21, 30, 31, 40] + assert segments[0].loss_mask == [1, 1, 0, 0, 1] + assert segments[0].rollout_log_probs == [-0.20, -0.21, 0.0, 0.0, -0.40] + + asyncio.run(run_case()) + + +@pytest.mark.unit +def test_anthropic_messages_multiturn_uses_vllm_tokens_for_training_segment(): + async def run_case(): + upstream = FakeVLLM( + [ + [(-1.20, 120), (-1.21, 121)], + [(-1.40, 140), (-1.41, 141)], + ] + ) + upstream_app = web.Application() + upstream_app.router.add_post("/inference/v1/generate", upstream.handle_generate) + upstream_server = TestServer(upstream_app) + await upstream_server.start_server() + + tool_raw = "slime" + tokenizer = ScriptedTokenizer( + prompts=[ + [110, 111], + [110, 111, 120, 121, 130], + ], + outputs={ + (120, 121): tool_raw, + (140, 141): "anthropic done", + }, + ) + adapter = anthropic.AnthropicAdapter( + tokenizer=tokenizer, + vllm_url=str(upstream_server.make_url("")).rstrip("/"), + ) + adapter.open_session("sid-anthropic-token", sampling_defaults={"max_new_tokens": 99}) + client = TestClient(TestServer(adapter.app)) + await client.start_server() + try: + first = await client.post( + "/v1/messages", + headers={"Authorization": "Bearer sid-anthropic-token"}, + json={ + "model": "actor", + "max_tokens": 5, + "messages": [{"role": "user", "content": [{"type": "text", "text": "find slime"}]}], + "tools": [ + { + "name": "lookup", + "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, + } + ], + }, + ) + first_data = await first.json() + tool_use = next(block for block in first_data["content"] if block["type"] == "tool_use") + + second = await client.post( + "/v1/messages", + headers={"Authorization": "Bearer sid-anthropic-token"}, + json={ + "model": "actor", + "max_tokens": 7, + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "find slime"}]}, + {"role": "assistant", "content": first_data["content"]}, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": tool_use["id"], + "content": "found slime", + } + ], + }, + ], + "tools": [ + { + "name": "lookup", + "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, + } + ], + }, + ) + second_data = await second.json() + finally: + await client.close() + await upstream_server.close() + + segments = await adapter.finish_session("sid-anthropic-token") + assert first.status == 200 + assert second.status == 200 + assert tool_use["name"] == "lookup" + assert tool_use["input"] == {"query": "slime"} + assert second_data["content"] == [{"type": "text", "text": "anthropic done"}] + assert [req["token_ids"] for req in upstream.requests] == [[110, 111], [110, 111, 120, 121, 130]] + assert upstream.routing_keys == ["sid-anthropic-token", "sid-anthropic-token"] + assert upstream.requests[0]["sampling_params"]["max_tokens"] == 5 + assert upstream.requests[1]["sampling_params"]["max_tokens"] == 7 + assert segments[0].prompt_ids == [110, 111] + assert segments[0].response_ids == [120, 121, 130, 140, 141] + assert segments[0].loss_mask == [1, 1, 0, 1, 1] + assert segments[0].rollout_log_probs == [-1.20, -1.21, 0.0, -1.40, -1.41] + + asyncio.run(run_case()) + + +@pytest.mark.unit +def test_openai_generate_posts_token_ids_and_extracts_logprobs(): + async def run_case(): + captured = {} + captured_headers = {} + + async def handle_generate(request): + captured_headers.update(request.headers) + captured.update(await request.json()) + return web.json_response( + { + "choices": [ + { + "token_ids": [701, 702], + "logprobs": {"content": [{"logprob": -0.7}, {"logprob": -0.8}]}, + "finish_reason": "stop", + } + ] + } + ) + + upstream_app = web.Application() + upstream_app.router.add_post("/inference/v1/generate", handle_generate) + server = TestServer(upstream_app) + await server.start_server() + try: + session = openai.Session(sampling_defaults={"max_new_tokens": 9}) + turn = await openai._generate( + [11, 12], + session, + {"max_tokens": 3, "temperature": 0.25, "stop": [""]}, + {VLLM_URL_KEY: str(server.make_url("")).rstrip("/")}, + ) + finally: + await server.close() + + assert captured["token_ids"] == [11, 12] + assert captured["sampling_params"]["logprobs"] == 1 + assert captured_headers.get("x-session-id") is None + assert captured["sampling_params"]["max_tokens"] == 3 + assert captured["sampling_params"]["temperature"] == 0.25 + assert captured["sampling_params"]["stop"] == [""] + assert turn.prompt_ids == [11, 12] + assert turn.output_ids == [701, 702] + assert turn.output_log_probs == [-0.7, -0.8] + + asyncio.run(run_case()) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_agent_sdk_adapters.py b/tests/test_agent_sdk_adapters.py new file mode 100644 index 000000000..a4ad3d4ce --- /dev/null +++ b/tests/test_agent_sdk_adapters.py @@ -0,0 +1,418 @@ +import asyncio +import sys +from pathlib import Path + +import httpx +import pytest +from aiohttp.test_utils import TestClient, TestServer + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from vime.agent.adapters import anthropic, openai +from vime.agent.trajectory import TurnRecord + + +NUM_GPUS = 0 + + +agents = pytest.importorskip("agents") +anthropic_sdk = pytest.importorskip("anthropic") +openai_sdk = pytest.importorskip("openai") + + +class SDKTokenizer: + def __init__(self, outputs: list[str]) -> None: + self.outputs = outputs + self.rendered: list[tuple[list[dict], list[dict] | None]] = [] + + def apply_chat_template(self, messages, tools=None, tokenize=True, add_generation_prompt=True): + self.rendered.append((list(messages), tools)) + return list(range(1, len(messages) + 2)) + + def decode(self, ids, skip_special_tokens=False): + return self.outputs[ids[0] - 1] + + +@pytest.mark.integration +def test_openai_agents_sdk_responses_runs_tool_loop_against_adapter(monkeypatch): + async def run_case(): + calls = [] + + async def fake_generate(prompt_ids, session, body, app, **kwargs): + calls.append({"prompt_ids": list(prompt_ids), "body": body}) + return TurnRecord( + prompt_ids=list(prompt_ids), + output_ids=[len(calls)], + finish_reason="stop", + output_log_probs=[-0.1 * len(calls)], + ) + + monkeypatch.setattr(openai, "_generate", fake_generate) + tokenizer = SDKTokenizer( + [ + "slime", + "final after tool", + ] + ) + adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") + client = TestClient(TestServer(adapter.app)) + await client.start_server() + base_url = str(client.make_url("/v1/")) + http_client = httpx.AsyncClient(trust_env=False) + oai = openai_sdk.AsyncOpenAI( + api_key="sdk-openai", + base_url=base_url, + max_retries=0, + http_client=http_client, + ) + + @agents.function_tool + def lookup(query: str) -> str: + return f"found {query}" + + agents.set_tracing_disabled(True) + model = agents.OpenAIResponsesModel(model="actor", openai_client=oai) + agent = agents.Agent( + name="sdk-responses", + instructions="Use lookup.", + model=model, + tools=[lookup], + model_settings=agents.ModelSettings(max_tokens=4), + ) + try: + result = await agents.Runner.run(agent, "find slime") + finally: + await client.close() + await oai.close() + + segments = await adapter.finish_session("sdk-openai") + assert result.final_output == "final after tool" + assert len(calls) == 2 + assert calls[0]["body"]["max_output_tokens"] == 4 + assert calls[0]["body"]["tools"][0]["name"] == "lookup" + assert calls[1]["body"]["input"][-1] == { + "call_id": calls[1]["body"]["input"][-2]["call_id"], + "output": "found slime", + "type": "function_call_output", + } + assert tokenizer.rendered[0][0] == [ + {"role": "system", "content": "Use lookup."}, + {"role": "user", "content": "find slime"}, + ] + assert tokenizer.rendered[1][0][-1] == { + "role": "tool", + "content": "found slime", + "tool_call_id": calls[1]["body"]["input"][-2]["call_id"], + } + assert segments[0].metadata["segment_kind"] == "final" + assert segments[0].response_ids[-1] == 2 + assert segments[0].loss_mask[-1] == 1 + + asyncio.run(run_case()) + + +@pytest.mark.integration +def test_openai_agents_sdk_chat_completions_runs_against_adapter(monkeypatch): + async def run_case(): + calls = [] + + async def fake_generate(prompt_ids, session, body, app, **kwargs): + calls.append({"prompt_ids": list(prompt_ids), "body": body}) + return TurnRecord( + prompt_ids=list(prompt_ids), output_ids=[1], finish_reason="stop", output_log_probs=[-0.2] + ) + + monkeypatch.setattr(openai, "_generate", fake_generate) + tokenizer = SDKTokenizer(["chat final"]) + adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") + client = TestClient(TestServer(adapter.app)) + await client.start_server() + base_url = str(client.make_url("/v1/")) + http_client = httpx.AsyncClient(trust_env=False) + oai = openai_sdk.AsyncOpenAI( + api_key="sdk-openai-chat", + base_url=base_url, + max_retries=0, + http_client=http_client, + ) + + agents.set_tracing_disabled(True) + model = agents.OpenAIChatCompletionsModel(model="actor", openai_client=oai) + agent = agents.Agent( + name="sdk-chat", + instructions="Be short.", + model=model, + model_settings=agents.ModelSettings(max_tokens=5), + ) + try: + result = await agents.Runner.run(agent, "say hi") + finally: + await client.close() + await oai.close() + + segments = await adapter.finish_session("sdk-openai-chat") + assert result.final_output == "chat final" + assert calls[0]["body"]["max_tokens"] == 5 + assert calls[0]["body"]["messages"] == [ + {"content": "Be short.", "role": "system"}, + {"role": "user", "content": "say hi"}, + ] + assert segments[0].prompt_ids == [1, 2, 3] + assert segments[0].response_ids == [1] + assert segments[0].loss_mask == [1] + + asyncio.run(run_case()) + + +@pytest.mark.integration +def test_openai_sdk_chat_completion_streaming_runs_against_adapter(monkeypatch): + async def run_case(): + calls = [] + + async def fake_generate(prompt_ids, session, body, app, **kwargs): + calls.append({"prompt_ids": list(prompt_ids), "body": body}) + return TurnRecord( + prompt_ids=list(prompt_ids), output_ids=[1], finish_reason="stop", output_log_probs=[-0.25] + ) + + monkeypatch.setattr(openai, "_generate", fake_generate) + tokenizer = SDKTokenizer( + ["streamed via sdk slime"] + ) + adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") + client = TestClient(TestServer(adapter.app)) + await client.start_server() + base_url = str(client.make_url("/v1/")) + http_client = httpx.AsyncClient(trust_env=False) + oai = openai_sdk.AsyncOpenAI( + api_key="sdk-openai-chat-stream", + base_url=base_url, + max_retries=0, + http_client=http_client, + ) + + try: + stream = await oai.chat.completions.create( + model="actor", + messages=[{"role": "user", "content": "call lookup"}], + tools=[ + { + "type": "function", + "function": { + "name": "lookup", + "description": "search", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}, + }, + } + ], + stream=True, + ) + text_parts = [] + tool_names = [] + tool_arguments = [] + finish_reasons = [] + usages = [] + async for chunk in stream: + choice = chunk.choices[0] + if choice.delta.content: + text_parts.append(choice.delta.content) + if choice.delta.tool_calls: + for tool_call in choice.delta.tool_calls: + tool_names.append(tool_call.function.name) + tool_arguments.append(tool_call.function.arguments) + if choice.finish_reason: + finish_reasons.append(choice.finish_reason) + if chunk.usage: + usages.append(chunk.usage) + finally: + await client.close() + await oai.close() + + segments = await adapter.finish_session("sdk-openai-chat-stream") + assert "".join(text_parts) == "streamed via sdk" + assert tool_names == ["lookup"] + assert tool_arguments == ['{"query": "slime"}'] + assert finish_reasons == ["tool_calls"] + assert usages[-1].prompt_tokens == 2 + assert usages[-1].completion_tokens == 1 + assert calls[0]["body"]["stream"] is True + assert segments[0].response_ids == [1] + + asyncio.run(run_case()) + + +@pytest.mark.integration +def test_openai_sdk_responses_streaming_runs_against_adapter(monkeypatch): + async def run_case(): + calls = [] + + async def fake_generate(prompt_ids, session, body, app, **kwargs): + calls.append({"prompt_ids": list(prompt_ids), "body": body}) + return TurnRecord( + prompt_ids=list(prompt_ids), output_ids=[1], finish_reason="stop", output_log_probs=[-0.35] + ) + + monkeypatch.setattr(openai, "_generate", fake_generate) + tokenizer = SDKTokenizer(["response stream via sdk"]) + adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") + client = TestClient(TestServer(adapter.app)) + await client.start_server() + base_url = str(client.make_url("/v1/")) + http_client = httpx.AsyncClient(trust_env=False) + oai = openai_sdk.AsyncOpenAI( + api_key="sdk-openai-responses-stream", + base_url=base_url, + max_retries=0, + http_client=http_client, + ) + + try: + stream = await oai.responses.create( + model="actor", + instructions="Be brief.", + input="say hi", + stream=True, + ) + event_types = [] + deltas = [] + completed_response = None + async for event in stream: + event_types.append(event.type) + if event.type == "response.output_text.delta": + deltas.append(event.delta) + if event.type == "response.completed": + completed_response = event.response + finally: + await client.close() + await oai.close() + + segments = await adapter.finish_session("sdk-openai-responses-stream") + assert event_types == ["response.created", "response.output_text.delta", "response.completed"] + assert "".join(deltas) == "response stream via sdk" + assert completed_response.status == "completed" + assert completed_response.usage.input_tokens == 3 + assert completed_response.usage.output_tokens == 1 + assert calls[0]["body"]["stream"] is True + assert segments[0].response_ids == [1] + + asyncio.run(run_case()) + + +@pytest.mark.integration +def test_anthropic_sdk_non_streaming_messages_runs_against_adapter(monkeypatch): + async def run_case(): + calls = [] + + async def fake_generate(prompt_ids, session, body, app, **kwargs): + calls.append({"prompt_ids": list(prompt_ids), "body": body}) + return TurnRecord( + prompt_ids=list(prompt_ids), output_ids=[1], finish_reason="stop", output_log_probs=[-0.28] + ) + + monkeypatch.setattr(anthropic, "_generate", fake_generate) + tokenizer = SDKTokenizer(["anthropic json"]) + adapter = anthropic.AnthropicAdapter(tokenizer=tokenizer, vllm_url="http://unused") + client = TestClient(TestServer(adapter.app)) + await client.start_server() + base_url = str(client.make_url("/")) + http_client = httpx.AsyncClient(trust_env=False) + anth = anthropic_sdk.AsyncAnthropic( + api_key="sdk-anthropic-json", + base_url=base_url, + max_retries=0, + http_client=http_client, + ) + + try: + message = await anth.messages.create( + model="actor", + max_tokens=6, + system="Be direct.", + messages=[{"role": "user", "content": "say hi"}], + ) + finally: + await client.close() + await anth.close() + + segments = await adapter.finish_session("sdk-anthropic-json") + assert message.type == "message" + assert message.content[0].type == "text" + assert message.content[0].text == "anthropic json" + assert message.stop_reason == "end_turn" + assert message.usage.input_tokens == 3 + assert message.usage.output_tokens == 1 + assert calls[0]["body"]["max_tokens"] == 6 + assert segments[0].response_ids == [1] + + asyncio.run(run_case()) + + +@pytest.mark.integration +def test_anthropic_sdk_streaming_messages_runs_against_adapter(monkeypatch): + async def run_case(): + calls = [] + + async def fake_generate(prompt_ids, session, body, app, **kwargs): + calls.append({"prompt_ids": list(prompt_ids), "body": body}) + return TurnRecord( + prompt_ids=list(prompt_ids), output_ids=[1], finish_reason="stop", output_log_probs=[-0.3] + ) + + monkeypatch.setattr(anthropic, "_generate", fake_generate) + tokenizer = SDKTokenizer(["anthropic final"]) + adapter = anthropic.AnthropicAdapter(tokenizer=tokenizer, vllm_url="http://unused") + client = TestClient(TestServer(adapter.app)) + await client.start_server() + base_url = str(client.make_url("/")) + http_client = httpx.AsyncClient(trust_env=False) + anth = anthropic_sdk.AsyncAnthropic( + api_key="sdk-anthropic", + base_url=base_url, + max_retries=0, + http_client=http_client, + ) + + try: + stream = await anth.messages.create( + model="actor", + max_tokens=6, + system="Be direct.", + messages=[{"role": "user", "content": "say hi"}], + stream=True, + ) + text_parts = [] + event_types = [] + async for event in stream: + event_types.append(event.type) + if event.type == "content_block_delta" and getattr(event.delta, "text", None): + text_parts.append(event.delta.text) + finally: + await client.close() + await anth.close() + + segments = await adapter.finish_session("sdk-anthropic") + assert "".join(text_parts) == "anthropic final" + assert event_types == [ + "message_start", + "content_block_start", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", + ] + assert calls[0]["body"]["max_tokens"] == 6 + assert tokenizer.rendered[0][0] == [ + {"role": "system", "content": "Be direct."}, + {"role": "user", "content": "say hi"}, + ] + assert segments[0].prompt_ids == [1, 2, 3] + assert segments[0].response_ids == [1] + assert segments[0].loss_mask == [1] + + asyncio.run(run_case()) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_agent_trajectory.py b/tests/test_agent_trajectory.py new file mode 100644 index 000000000..3abf88610 --- /dev/null +++ b/tests/test_agent_trajectory.py @@ -0,0 +1,143 @@ +import sys +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from vime.agent.trajectory import TurnRecord, TurnSegment, merge_turn_segments, merge_turns + + +NUM_GPUS = 0 + + +def _turn(prompt_ids: list[int], output_ids: list[int], output_log_probs: list[float] | None = None) -> TurnRecord: + return TurnRecord( + prompt_ids=prompt_ids, + output_ids=output_ids, + finish_reason="stop", + output_log_probs=( + output_log_probs if output_log_probs is not None else [-token_id / 100 for token_id in output_ids] + ), + ) + + +@pytest.mark.unit +def test_merge_turns_preserves_matched_prefix_on_prompt_drift(): + segment = merge_turns( + [ + _turn([10], [11]), + _turn([10, 11, 21], [12]), + _turn([10, 11, 21, 12, 31], [13]), + _turn([10, 11, 21, 12, 22], [14]), + ] + ) + + assert segment is not None + assert segment.prompt_ids == [10] + assert segment.response_ids == [11, 21, 12, 22, 14] + assert segment.loss_mask == [1, 0, 1, 0, 1] + assert segment.rollout_log_probs == [-0.11, 0.0, -0.12, 0.0, -0.14] + + +@pytest.mark.unit +def test_merge_turns_drops_middle_turn_when_next_prompt_skips_it(): + segment = merge_turns( + [ + _turn([10], [11]), + _turn([10, 11, 21], [12]), + _turn([10, 11, 22], [13]), + _turn([10, 11, 22, 13, 31], [14]), + ] + ) + + assert segment is not None + assert segment.prompt_ids == [10] + assert segment.response_ids == [11, 22, 13, 31, 14] + assert segment.loss_mask == [1, 0, 1, 0, 1] + assert segment.rollout_log_probs == [-0.11, 0.0, -0.13, 0.0, -0.14] + + +@pytest.mark.unit +def test_merge_turns_handles_consecutive_prompt_drifts(): + segment = merge_turns( + [ + _turn([10], [11]), + _turn([10, 11, 21], [12]), + _turn([10, 11, 22], [13]), + _turn([10, 11, 23], [14]), + _turn([10, 11, 23, 14, 31], [15]), + ] + ) + + assert segment is not None + assert segment.prompt_ids == [10] + assert segment.response_ids == [11, 23, 14, 31, 15] + assert segment.loss_mask == [1, 0, 1, 0, 1] + assert segment.rollout_log_probs == [-0.11, 0.0, -0.14, 0.0, -0.15] + + +@pytest.mark.unit +def test_merge_turns_masks_whole_output_when_prompt_drift_splits_it(): + segment = merge_turns( + [ + _turn([10], [11, 12, 13, 14]), + _turn([10, 11, 12, 99, 14], [15]), + ] + ) + + assert segment is not None + assert segment.prompt_ids == [10] + assert segment.response_ids == [11, 12, 99, 14, 15] + assert segment.loss_mask == [0, 0, 0, 0, 1] + assert segment.rollout_log_probs == [0.0, 0.0, 0.0, 0.0, -0.15] + + +@pytest.mark.unit +def test_merge_turns_masks_whole_output_when_prompt_drift_changes_token_count(): + segment = merge_turns( + [ + _turn([10], [11, 12, 13, 14]), + _turn([10, 11, 12, 99, 100, 14], [15]), + ] + ) + + assert segment is not None + assert segment.prompt_ids == [10] + assert segment.response_ids == [11, 12, 99, 100, 14, 15] + assert segment.loss_mask == [0, 0, 0, 0, 0, 1] + assert segment.rollout_log_probs == [0.0, 0.0, 0.0, 0.0, 0.0, -0.15] + + +@pytest.mark.unit +def test_merge_turns_restarts_when_prompt_base_changes(): + segment = merge_turns( + [ + _turn([10], [11]), + _turn([20, 21], [22]), + _turn([20, 21, 22, 23], [24]), + ] + ) + + assert segment is not None + assert segment.prompt_ids == [20, 21] + assert segment.response_ids == [22, 23, 24] + assert segment.loss_mask == [1, 0, 1] + assert segment.rollout_log_probs == [-0.22, 0.0, -0.24] + + +@pytest.mark.unit +def test_merge_turn_segments_keeps_oversized_segments(): + segments = [TurnSegment(turns=[_turn([10, 11, 12], [13, 14])])] + + merged = merge_turn_segments(segments) + + assert len(merged) == 1 + assert merged[0].prompt_ids == [10, 11, 12] + assert merged[0].response_ids == [13, 14] + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_cp_utils.py b/tests/test_cp_utils.py new file mode 100644 index 000000000..d37e870a7 --- /dev/null +++ b/tests/test_cp_utils.py @@ -0,0 +1,177 @@ +"""CPU unit tests for ``vime.backends.megatron_utils.cp_utils.get_sum_of_sample_mean``. + +Pins the per-rollout reducer contract: a rollout split into N training +samples (compact / subagent) must contribute exactly one token-weighted +mean to the sum, even when first-fit packing puts those siblings into +different micro-batches at training time. + +The CPU-only CI image does not ship megatron — ``_cp_dist_helpers`` +stubs ``megatron.core.mpu`` at import time so the subsequent +``cp_utils`` import binds against the stub. + +End-to-end report-formula invariance and multi-process distributed +checks live in ``test_metric_report.py`` and ``test_metric_report_dist.py``. +""" + +from __future__ import annotations + +# Import the helpers BEFORE the slime imports so the megatron stub lands +# in sys.modules first. pytest's prepend importmode puts this file's +# directory (``tests/``) on sys.path, which is what makes the bare-name +# import work without an ``__init__.py``. +import _cp_dist_helpers # noqa: F401 +import pytest +import torch + +from vime.backends.megatron_utils.cp_utils import ( # noqa: E402 + get_logits_and_tokens_offset_with_cp, + get_sum_of_sample_mean, +) + + +def _make_inputs(per_sample_lengths: list[int]): + """Build (total_lengths, response_lengths, loss_masks) for samples of the given lengths. + + Each sample has loss_mask = all-ones (so mask sum == length); total length + is response length + 4 fake prompt tokens (unused by the reducer in + cp_size==1 mode). + """ + response_lengths = list(per_sample_lengths) + total_lengths = [r + 4 for r in response_lengths] + loss_masks = [torch.ones(r, dtype=torch.float32) for r in response_lengths] + return total_lengths, response_lengths, loss_masks + + +def _denoms(*values: int) -> torch.Tensor: + """Wrap per-sample denoms as the float tensor that the actor side promotes + them to before calling the reducer.""" + return torch.tensor(values, dtype=torch.float32) + + +@pytest.mark.unit +def test_default_reduces_to_per_sample_mean(): + """``sample_denoms=None`` reproduces the legacy per-sample-mean.""" + total_lengths, response_lengths, loss_masks = _make_inputs([3, 3, 3]) + reducer = get_sum_of_sample_mean(total_lengths, response_lengths, loss_masks) + x = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]) + # per-sample means: 2, 5, 8 → sum = 15 + assert reducer(x).item() == pytest.approx(15.0) + + +@pytest.mark.unit +def test_per_rollout_denom_collapses_siblings_into_one_mean(): + """Pre-computed per-rollout mask sums make N sibling samples contribute one + token-weighted mean instead of N per-sample means.""" + # 4 samples: rollout R0 owns indices 0,1,2 (mask sums 3+3+3=9); rollout R1 + # owns index 3 (mask sum 3). Pre-computed per-sample denom = group sum. + total_lengths, response_lengths, loss_masks = _make_inputs([3, 3, 3, 3]) + sample_denoms = _denoms(9, 9, 9, 3) + reducer = get_sum_of_sample_mean(total_lengths, response_lengths, loss_masks, sample_denoms) + x = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0]) + # R0 token-mean: (1+2+...+9)/9 = 5. R1 token-mean: (10+11+12)/3 = 11. Sum = 16. + assert reducer(x).item() == pytest.approx(16.0) + + +@pytest.mark.unit +def test_split_across_mbs_recovers_full_per_rollout_mean(): + """The critical contract: when a rollout's samples land in different mbs, + summing each mb's reducer output equals one whole-step reducer call with + the same pre-computed denominators. This is exactly the bug that motivated + the precomputation — if the denom were computed per-mb (partial mask sum), + the two halves wouldn't add up.""" + # 4 samples (same as above). Whole-step denoms = [9, 9, 9, 3]. + total_lengths, response_lengths, loss_masks = _make_inputs([3, 3, 3, 3]) + sample_denoms = _denoms(9, 9, 9, 3) + x = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0]) + + whole = get_sum_of_sample_mean(total_lengths, response_lengths, loss_masks, sample_denoms) + whole_value = whole(x).item() + + # mb_a holds samples 0, 1 of R0; mb_b holds sample 2 of R0 and sample 3 (R1). + # Each mb carries the SAME per-sample denoms (precomputed at step level) + # — that's what makes the split safe. + mb_a = get_sum_of_sample_mean(total_lengths[:2], response_lengths[:2], loss_masks[:2], sample_denoms[:2]) + mb_b = get_sum_of_sample_mean(total_lengths[2:], response_lengths[2:], loss_masks[2:], sample_denoms[2:]) + split_value = mb_a(x[:6]).item() + mb_b(x[6:]).item() + + assert split_value == pytest.approx(whole_value) + + +@pytest.mark.unit +def test_split_with_per_mb_denom_would_be_wrong(): + """Sanity-check the bug we're guarding against: if the caller naively + computes per-rollout denoms from each mb's own samples (the local mask + sum, NOT the precomputed whole-rollout sum), the two halves DON'T add up + to the whole-step value. This pins down WHY the precomputation must + happen at the step level.""" + total_lengths, response_lengths, loss_masks = _make_inputs([3, 3, 3, 3]) + x = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0]) + + whole = get_sum_of_sample_mean(total_lengths, response_lengths, loss_masks, _denoms(9, 9, 9, 3)) + whole_value = whole(x).item() + + # Wrong denom: each mb only sees its own samples of R0. + # mb_a's "rollout mask sum" for R0 would be 3+3=6 (instead of 9). mb_b's + # would be 3. Different from the true whole-rollout total. + mb_a_wrong = get_sum_of_sample_mean(total_lengths[:2], response_lengths[:2], loss_masks[:2], _denoms(6, 6)) + mb_b_wrong = get_sum_of_sample_mean(total_lengths[2:], response_lengths[2:], loss_masks[2:], _denoms(3, 3)) + wrong_total = mb_a_wrong(x[:6]).item() + mb_b_wrong(x[6:]).item() + + assert wrong_total != pytest.approx(whole_value), ( + "Expected the per-mb denom path to produce a different (incorrect) value; " + "if these match, the regression test is no longer guarding the precomputation contract." + ) + + +@pytest.mark.unit +def test_cp_chunking_preserves_per_rollout_mean_report(monkeypatch): + """Turning CP on must not change the reducer's output. + + Real flow: each CP rank only sees its chunk of the response tokens; the + reducer's CP>1 branch slices ``loss_mask`` to match. Summing each CP + rank's reducer output across CP ranks reproduces the cp=1 result, which + is what train_one_step then divides by ``step_global_batch_size``. + """ + from megatron.core import mpu as _mpu + + # Use lengths that line up cleanly with the CP chunking + # (chunk_size = ceil(total_length / (2*cp_size))). + total_lengths = [12, 12] # 2 samples + response_lengths = [8, 8] # 4 prompt + 8 response each + loss_masks = [torch.ones(r, dtype=torch.float32) for r in response_lengths] + sample_denoms = torch.tensor([16.0, 16.0], dtype=torch.float32) # = sum of both mask totals (one rollout) + x_full = [ + torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]), + torch.tensor([10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0]), + ] + x_concat = torch.cat(x_full) + + # --- cp=1 baseline --- + monkeypatch.setattr(_mpu, "get_context_parallel_world_size", lambda: 1) + monkeypatch.setattr(_mpu, "get_context_parallel_rank", lambda: 0) + reducer_cp1 = get_sum_of_sample_mean(total_lengths, response_lengths, loss_masks, sample_denoms) + baseline = reducer_cp1(x_concat).item() + + # --- cp=2: sum partial reducer outputs across the two CP ranks --- + monkeypatch.setattr(_mpu, "get_context_parallel_world_size", lambda: 2) + cp_total = 0.0 + for cp_rank in range(2): + monkeypatch.setattr(_mpu, "get_context_parallel_rank", lambda r=cp_rank: r) + # Slice each sample's response-token tensor to the chunks this CP + # rank owns, mirroring what the forward pass would feed in. + x_chunks_per_sample = [] + for tl, rl, x in zip(total_lengths, response_lengths, x_full, strict=True): + prompt_length = tl - rl + _, _, _, tokens_offset = get_logits_and_tokens_offset_with_cp(tl, rl) + chunk_0 = x[tokens_offset[0][0] - prompt_length : tokens_offset[0][1] - prompt_length] + chunk_1 = x[tokens_offset[1][0] - prompt_length : tokens_offset[1][1] - prompt_length] + x_chunks_per_sample.append(torch.cat([chunk_0, chunk_1])) + x_for_rank = torch.cat(x_chunks_per_sample) + reducer_cp2 = get_sum_of_sample_mean(total_lengths, response_lengths, loss_masks, sample_denoms) + cp_total += reducer_cp2(x_for_rank).item() + + assert cp_total == pytest.approx(baseline) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_dp_schedule.py b/tests/test_dp_schedule.py new file mode 100644 index 000000000..6ce03ce2b --- /dev/null +++ b/tests/test_dp_schedule.py @@ -0,0 +1,288 @@ +"""CPU unit tests for vime.utils.dp_schedule.build_dp_schedule. + +The tests assert the invariants documented at the top of dp_schedule.py against +a range of static / dynamic / VPP / oversize / balance / uneven scenarios. +""" + +from types import SimpleNamespace + +import pytest + +from vime.utils.dp_schedule import build_dp_schedule + + +def make_args( + *, + micro_batch_size=1, + use_dynamic_batch_size=False, + max_tokens_per_gpu=None, + balance_data=False, +): + return SimpleNamespace( + micro_batch_size=micro_batch_size, + use_dynamic_batch_size=use_dynamic_batch_size, + max_tokens_per_gpu=max_tokens_per_gpu, + balance_data=balance_data, + ) + + +def make_tp(dp_size=1, cp_size=1, vpp_size=1, microbatch_group_size_per_vp_stage=1): + return { + "dp_size": dp_size, + "cp_size": cp_size, + "vpp_size": vpp_size, + "microbatch_group_size_per_vp_stage": microbatch_group_size_per_vp_stage, + } + + +def assert_invariants( + partitions, + micro_batch_indices, + num_microbatches, + *, + dp_size, + expected_global_sample_indices, + total_lengths, + max_per_bin=None, +): + """Check the invariants documented at the top of dp_schedule.py. + + ``expected_global_sample_indices`` is the set of global sample indices + that should end up covered (after trim). Trailing rollouts that don't + fit are excluded. + """ + seen_global: set[int] = set() + for r in range(dp_size): + partition = partitions[r] + mbi = micro_batch_indices[r] + + # Same num_mbs per rank (PP sync). + assert len(mbi) == sum(num_microbatches), f"rank {r}: mbs count mismatch" + + # Flattened micro_batch_indices == range(len(partition)). + flat = [i for mbs in mbi for i in mbs] + assert flat == list(range(len(partition))), f"rank {r}: micro_batch_indices don't tile [0, n)" + + # Disjoint partitions whose union covers every kept sample. + assert seen_global.isdisjoint(partition), f"rank {r}: overlap with other ranks" + seen_global.update(partition) + assert seen_global == set(expected_global_sample_indices), "covered sample set mismatch" + + if max_per_bin is None: + return + + # Every mbs <= max_per_bin tokens, EXCEPT a singleton bin holding an oversized sample. + for r in range(dp_size): + partition = partitions[r] + for mbs in micro_batch_indices[r]: + bin_total = sum(total_lengths[partition[i]] for i in mbs) + if bin_total > max_per_bin: + assert len(mbs) == 1, f"rank {r}: mbs sum {bin_total} > {max_per_bin} but contains {len(mbs)} samples" + + +@pytest.mark.unit +def test_static_stride_single_step(): + """Static + strided DP split, single step (1 rollout = 1 sample).""" + total_lengths = [10] * 16 + rollout_indices = list(range(16)) + args = make_args(micro_batch_size=2) + tp = make_tp(dp_size=4) + + partitions, mbi, nmb, gbs_per_step = build_dp_schedule( + args, tp, total_lengths, global_batch_size=16, rollout_indices=rollout_indices + ) + + assert nmb == [2] + assert gbs_per_step == [16] + assert_invariants( + partitions, + mbi, + nmb, + dp_size=4, + expected_global_sample_indices=range(16), + total_lengths=total_lengths, + ) + + +@pytest.mark.unit +def test_static_balance_multi_step(): + """Static + balance_data + 2 training steps.""" + total_lengths = [1, 2, 3, 4, 5, 6, 7, 8, 8, 7, 6, 5, 4, 3, 2, 1] + rollout_indices = list(range(16)) + args = make_args(micro_batch_size=2, balance_data=True) + tp = make_tp(dp_size=2) + + partitions, mbi, nmb, gbs_per_step = build_dp_schedule( + args, tp, total_lengths, global_batch_size=8, rollout_indices=rollout_indices + ) + + assert nmb == [2, 2] + assert gbs_per_step == [8, 8] + assert_invariants( + partitions, + mbi, + nmb, + dp_size=2, + expected_global_sample_indices=range(16), + total_lengths=total_lengths, + ) + + +@pytest.mark.unit +def test_dynamic_uniform(): + """Dynamic mbs on uniform-length samples.""" + total_lengths = [5] * 8 + rollout_indices = list(range(8)) + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=10) + tp = make_tp(dp_size=2) + + partitions, mbi, nmb, gbs_per_step = build_dp_schedule( + args, tp, total_lengths, global_batch_size=8, rollout_indices=rollout_indices + ) + + assert gbs_per_step == [8] + assert_invariants( + partitions, + mbi, + nmb, + dp_size=2, + expected_global_sample_indices=range(8), + total_lengths=total_lengths, + max_per_bin=10, + ) + + +@pytest.mark.unit +def test_dynamic_oversized_sample_lands_alone(): + """A sample larger than max_per_bin must end up alone in its mbs.""" + total_lengths = [15, 3, 3, 3, 3, 3, 3, 3] + rollout_indices = list(range(8)) + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=10) + tp = make_tp(dp_size=2) + + partitions, mbi, nmb, gbs_per_step = build_dp_schedule( + args, tp, total_lengths, global_batch_size=8, rollout_indices=rollout_indices + ) + + assert_invariants( + partitions, + mbi, + nmb, + dp_size=2, + expected_global_sample_indices=range(8), + total_lengths=total_lengths, + max_per_bin=10, + ) + oversize_idx = total_lengths.index(15) + found = False + for r in range(2): + if oversize_idx not in partitions[r]: + continue + local = partitions[r].index(oversize_idx) + for mbs in mbi[r]: + if local in mbs: + assert mbs == [local], f"oversized sample shares an mbs: {mbs}" + found = True + assert found + + +@pytest.mark.unit +def test_dynamic_with_vpp_rounds_to_mb_group(): + """num_microbatches per rank should be a multiple of mb_group when vpp_size > 1.""" + total_lengths = [4] * 32 + rollout_indices = list(range(32)) + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=8) + tp = make_tp(dp_size=2, vpp_size=2, microbatch_group_size_per_vp_stage=2) + + partitions, mbi, nmb, gbs_per_step = build_dp_schedule( + args, tp, total_lengths, global_batch_size=16, rollout_indices=rollout_indices + ) + + for n in nmb: + assert n % 2 == 0, f"num_microbatches {n} is not a multiple of mb_group=2" + assert_invariants( + partitions, + mbi, + nmb, + dp_size=2, + expected_global_sample_indices=range(32), + total_lengths=total_lengths, + max_per_bin=8, + ) + + +@pytest.mark.unit +def test_rollout_grouping_keeps_samples_together(): + """compact / subagent simulation: rollout 0 emits 3 samples, rollout 1 emits 2, + rollout 2 emits 4. Splitter keeps every rollout's samples in a single step.""" + rollout_indices = [0, 0, 0, 1, 1, 2, 2, 2, 2] + total_lengths = [3] * 9 + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=12) + tp = make_tp(dp_size=1) + + partitions, mbi, nmb, gbs_per_step = build_dp_schedule( + args, tp, total_lengths, global_batch_size=1, rollout_indices=rollout_indices + ) + + # 3 rollouts / 1 per step → 3 steps, gbs constant. + assert gbs_per_step == [1, 1, 1] + # For each step, collect the samples (global indices) that landed in that step's mbs + # on rank 0, then verify they exactly equal the rollout's sample positions. + expected_per_step = [[0, 1, 2], [3, 4], [5, 6, 7, 8]] + rank0_partition = partitions[0] + mbs_cursor = 0 + for step_i, n_mbs in enumerate(nmb): + step_locals = sorted(j for mbs in mbi[0][mbs_cursor : mbs_cursor + n_mbs] for j in mbs) + step_globals = [rank0_partition[j] for j in step_locals] + assert ( + sorted(step_globals) == expected_per_step[step_i] + ), f"step {step_i} samples = {step_globals}, expected {expected_per_step[step_i]}" + mbs_cursor += n_mbs + assert_invariants( + partitions, + mbi, + nmb, + dp_size=1, + expected_global_sample_indices=range(9), + total_lengths=total_lengths, + max_per_bin=12, + ) + + +@pytest.mark.unit +def test_trims_trailing_rollouts_that_dont_fill_a_step(): + """5 rollouts, gbs=2 → 2 steps × 2 rollouts; trailing rollout 4 (sample positions 6, 7) + is dropped.""" + rollout_indices = [0, 0, 1, 2, 2, 3, 4, 4] + total_lengths = [3] * 8 + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=12) + tp = make_tp(dp_size=1) + + partitions, mbi, nmb, gbs_per_step = build_dp_schedule( + args, tp, total_lengths, global_batch_size=2, rollout_indices=rollout_indices + ) + + assert gbs_per_step == [2, 2] + # Sample positions 6 and 7 belong to the trimmed rollout 4 and must be absent. + assert_invariants( + partitions, + mbi, + nmb, + dp_size=1, + expected_global_sample_indices=range(6), + total_lengths=total_lengths, + max_per_bin=12, + ) + + +@pytest.mark.unit +def test_rejects_when_fewer_rollouts_than_gbs(): + """gbs=4 with only 3 distinct rollouts → cannot form one step.""" + args = make_args(use_dynamic_batch_size=True, max_tokens_per_gpu=12) + tp = make_tp(dp_size=1) + with pytest.raises(AssertionError, match="num_rollouts"): + build_dp_schedule(args, tp, [3] * 6, global_batch_size=4, rollout_indices=[0, 0, 1, 1, 2, 2]) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_loss_cp_invariance.py b/tests/test_loss_cp_invariance.py new file mode 100644 index 000000000..998ba1d6c --- /dev/null +++ b/tests/test_loss_cp_invariance.py @@ -0,0 +1,253 @@ +"""End-to-end backward / gradient-norm CP-invariance check on CPU. + +This is the closest thing to a real training-step backward we can run on +the CPU CI image without standing up Megatron, FlashAttention, or NCCL. +The goal: prove that for the same training samples, *the gradient norm +after the optimizer-side all-reduce is identical regardless of CP size*. + +Why this matters +---------------- +Slime's loss prescaling + Megatron's per-mb scaling + DDP's grad +averaging compose into one big formula. Any time we touch any one of +those three layers the numbers should land in the same place. Until +this test existed we only had end-to-end report-formula checks +(`test_metric_report_dist.py`); none of them ran a real ``backward()``, +so a sign or factor error in the prescaling would slip through. + +Mapping to Megatron source +-------------------------- +We reproduce, for each spawned rank, the exact sequence Megatron applies +when a 3-tuple ``(loss, num_tokens, log)`` comes back from the loss +function with ``calculate_per_token_loss=False`` — slime's per-rollout- +mean path: + + 1. Loss function pre-scales:: + loss *= num_microbatches / step_global_batch_size * (dp * cp) + See ``slime/backends/megatron_utils/loss.py:1209-1215``. + 2. Megatron divides by ``clamp(num_tokens, 1)`` then by + ``num_microbatches``:: + output_tensor /= torch.clamp(num_tokens, min=1) # num_tokens=1 → no-op + output_tensor /= num_microbatches + See ``Megatron-LM/megatron/core/pipeline_parallel/schedules.py:258-264`` + (the ``len(outputs) == 3`` branch with ``not calculate_per_token_loss``). + 3. Backward fills grad buffers; per-mb contributions sum on each rank. + 4. DDP grad sync averages across the DP-with-CP group:: + grad_sum_across_dp_cp_world / (dp * cp) + See ``Megatron-LM/megatron/core/distributed/distributed_data_parallel.py:283-290`` + (``average_in_collective=False``, ``gradient_scaling_factor = 1.0 / dp_cp_group.size()``). + +Composing 1-4 collapses to + final_grad = total_sum_of_rollout_means / step_global_batch_size, +i.e. the gradient of ``mean_of_per_rollout_means(x)``. That doesn't +contain ``cp`` anywhere, so the grad norm must be identical for any +(dp, cp) factorization of the same world size. + +What this test does NOT exercise: the actual Megatron model classes, the +real DDP buffer code, fused optimizers, mixed-precision. We use a plain +``nn.Linear`` with manual all-reduce-average to simulate steps 1-4 above. +The contract here is on *our* scaling math (steps 1 + 4 are slime's; +step 2 is what Megatron does to our 3-tuple). If Megatron later changes +step 2 — e.g. drops the ``/= num_microbatches`` — this test won't catch +it, but the real GPU integration suite (``test_qwen2.5_0.5B_short.py``) +will. +""" + +from __future__ import annotations + +# Megatron stub must land in sys.modules first; the slime imports inside +# the worker pick it up via this same module. pytest's prepend importmode +# puts ``tests/`` on sys.path so the bare-name import works without an +# ``__init__.py``; mp.spawn children inherit the parent's sys.path. +import _cp_dist_helpers +import pytest +import torch +from _cp_dist_helpers import ( + FOUR_ROLLOUT_EXPECTED_REPORT, + FOUR_ROLLOUT_RESPONSE_LENGTHS, + FOUR_ROLLOUT_TOTAL_LENGTHS, + FOUR_ROLLOUT_X_VALUES, + cp_chunk_response_tensor, + free_port, + init_worker_process_group, + stub_megatron_in_worker, +) + + +def _grad_norm_worker( + rank: int, + world_size: int, + cp_size: int, + dp_size: int, + seed: int, + master_port: int, + result_path: str, +) -> None: + """One spawned rank. + + Builds a tiny ``nn.Linear`` model (deterministic init via ``seed``), + runs slime's per-rollout-mean loss reducer with the rank's share of + the four-rollout fixture, applies the slime-side prescaling, then + Megatron's per-mb scaling, then ``.backward()``, then a manual + all-reduce-average across the dp-with-cp group (mirroring DDP's + ``average_in_collective=False`` path with + ``gradient_scaling_factor = 1 / dp_cp_world_size``). Rank 0 writes the + final ``grad_norm`` to ``result_path``. + """ + import torch.distributed as _dist + + cp_rank = rank % cp_size + dp_rank = rank // cp_size + stub_megatron_in_worker(cp_size, cp_rank) + + dp_cp_group = init_worker_process_group(rank, world_size, master_port) + try: + from vime.backends.megatron_utils.cp_utils import get_sum_of_sample_mean + + # Same init across all (dp, cp) configs so the grad we backprop + # into is comparable. ``manual_seed`` is enough on CPU because we + # only do one forward/backward and no dropout. + torch.manual_seed(seed) + model = torch.nn.Linear(1, 1, bias=False) + # Force a known weight value to keep the math hand-checkable: with + # weight = 1.0 and input = x, the linear output equals x, and the + # grad of (output * x).sum() wrt weight equals (x*x).sum(). That + # makes the gradient a pure function of the fixture's x values, + # independent of the random init draw. + with torch.no_grad(): + model.weight.fill_(1.0) + + all_total_lengths = FOUR_ROLLOUT_TOTAL_LENGTHS + all_response_lengths = FOUR_ROLLOUT_RESPONSE_LENGTHS + all_loss_masks = [torch.ones(r, dtype=torch.float32) for r in all_response_lengths] + all_x = [torch.tensor(v) for v in FOUR_ROLLOUT_X_VALUES] + step_global_batch_size = 4 # 4 rollouts in the step + num_microbatches = 1 # this CPU model does the whole rank-share in one mb + + my_indices = [i for i in range(4) if i % dp_size == dp_rank] + my_tl = [all_total_lengths[i] for i in my_indices] + my_rl = [all_response_lengths[i] for i in my_indices] + my_masks = [all_loss_masks[i] for i in my_indices] + my_x_full = [all_x[i] for i in my_indices] + # Pre-computed per-rollout denoms = each sample's own mask sum + # (each rollout in the fixture has exactly one sample, so the + # per-rollout denom collapses to the per-sample denom). + my_denoms = torch.tensor([float(m.sum().item()) for m in my_masks], dtype=torch.float32) + + if cp_size == 1: + x_for_rank = torch.cat(my_x_full) + else: + x_for_rank = torch.cat( + [cp_chunk_response_tensor(x, tl, rl) for tl, rl, x in zip(my_tl, my_rl, my_x_full, strict=True)] + ) + + # === Forward path ===================================================== + # Tiny "model": output[i] = x[i] * weight. We treat the linear + # output as the per-token quantity the loss is computed over — + # this stands in for the (logits @ token_emb) the policy loss + # consumes in real training. + x_input = x_for_rank.unsqueeze(-1) # shape [T, 1] + output = model(x_input).squeeze(-1) # shape [T] + + reducer = get_sum_of_sample_mean(my_tl, my_rl, my_masks, my_denoms) + loss = reducer(output) + + # === Step 1: slime's per-rollout-mean prescaling ====================== + # loss.py:1209-1215. ``mpu.get_data_parallel_world_size(with_context_parallel=True)`` + # is the dp-with-cp world size, which is ``world_size`` in this setup. + loss = loss * num_microbatches / step_global_batch_size * world_size + + # === Step 2: Megatron's forward_step_calc_loss scaling ================ + # schedules.py:258-264 — for the 3-tuple, not-per-token-loss path: + # output_tensor /= torch.clamp(num_tokens, min=1) + # output_tensor /= num_microbatches + # slime passes num_tokens=1 in this path (loss.py:1221), so the + # first divide is a no-op; we keep it explicit to mirror the + # source faithfully. + num_tokens_for_scaling = torch.tensor(1.0) # slime's placeholder + loss = loss / torch.clamp(num_tokens_for_scaling, min=1.0) + loss = loss / num_microbatches + + # === Step 3: backward fills per-rank grad ============================= + loss.backward() + + # === Step 4: DDP all-reduce-average across dp-with-cp world =========== + # distributed_data_parallel.py:283-290, ``average_in_collective=False`` + # case: ``gradient_scaling_factor = 1.0 / dp_cp_group.size()`` is + # baked into the buffer, so the all-reduce is a SUM and the + # 1/world_size scaling pre-applies. We do the equivalent here by + # all-reducing then dividing. + grad = model.weight.grad.detach() + _dist.all_reduce(grad, group=dp_cp_group) + grad = grad / world_size + + # The norm of a 1-element gradient is its absolute value. We + # report ``grad.item()`` directly so the assertion side can also + # eyeball the sign, which is more useful than a strict norm when + # debugging a regression. + grad_value = grad.item() + + if rank == 0: + with open(result_path, "w") as f: + f.write(repr(grad_value)) + finally: + _dist.destroy_process_group() + + +def _run_grad_norm_worker(dp_size: int, cp_size: int, tmp_path) -> float: + """Spawn ``dp_size * cp_size`` workers and return rank-0's final grad.""" + import torch.multiprocessing as mp + + world_size = dp_size * cp_size + result_path = str(tmp_path / f"grad_dp{dp_size}_cp{cp_size}.txt") + mp.spawn( + _grad_norm_worker, + args=(world_size, cp_size, dp_size, 0, free_port(), result_path), + nprocs=world_size, + join=True, + ) + with open(result_path) as f: + return float(f.read()) + + +# Subset of (dp, cp) configs to keep runtime down; covers the four +# qualitatively distinct cases: +# - (1, 1) baseline (no parallelism) +# - (2, 1) DP-only +# - (1, 2) CP-only +# - (2, 2) DP + CP combined +# - (1, 4) deeper CP-only +# - (4, 1) deeper DP-only +# The full 3*3 matrix lives in test_metric_report_dist.py — here we just +# want enough coverage to catch a sign/factor regression in the slime +# prescaling math. +_PARALLELISM_CASES = [(1, 1), (2, 1), (1, 2), (2, 2), (1, 4), (4, 1)] + + +@pytest.mark.unit +@pytest.mark.parametrize("dp_size,cp_size", _PARALLELISM_CASES) +def test_backward_grad_is_cp_invariant(dp_size, cp_size, tmp_path): + """The post-DDP-average gradient must be identical across all + (dp, cp) configurations of the same global batch. + + Hand-derivable expectation: with weight = 1.0 and the fixture above, + the gradient of ``mean_of_per_rollout_means(model(x))`` wrt weight is + the same quantity the rollout-report tests pin + (FOUR_ROLLOUT_EXPECTED_REPORT = 1249.875), because for each rollout + the per-token mean of ``x * weight`` differentiates to the per-token + mean of ``x``. + """ + grad = _run_grad_norm_worker(dp_size=dp_size, cp_size=cp_size, tmp_path=tmp_path) + # Tolerance: float32 / multi-rank sums introduce ~1e-3 relative error + # on numbers up to ~1250; that's still 5+ digits of agreement. Each + # (dp, cp) case is pinned to the same hand-derived value, so a sign + # or factor regression in the prescaling math will fail the whole + # matrix uniformly — easy to spot in CI logs. + assert grad == pytest.approx(FOUR_ROLLOUT_EXPECTED_REPORT, rel=1e-4) + + +# Keep the helpers import load-bearing (it installs the megatron stub). +_ = _cp_dist_helpers + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_metric_report.py b/tests/test_metric_report.py new file mode 100644 index 000000000..98d7d69c6 --- /dev/null +++ b/tests/test_metric_report.py @@ -0,0 +1,320 @@ +"""Single-process metric-report invariance tests. + +Pins train-side / rollout-side report formulas implemented in +``vime.backends.megatron_utils.cp_utils.reduce_train_step_metrics`` and +``rollout_log_metric_contribution``: the reported number for a given set +of samples must be the same regardless of + + - how samples are distributed across micro-batches / DP ranks + - whether context parallelism is on or off + - whether the path is per-rollout-mean or per-token-loss + +Single-process variants use a mock dp-with-cp group + a no-op +``dist.all_reduce`` to keep things lightweight; the multi-process +end-to-end variants (real torch.distributed) live in +``test_metric_report_dist.py``. +""" + +from __future__ import annotations + +# Import the helpers BEFORE the slime imports so the megatron stub lands +# in sys.modules first. pytest's prepend importmode puts this file's +# directory (``tests/``) on sys.path, which is what makes the bare-name +# import work without an ``__init__.py``. +import _cp_dist_helpers # noqa: F401 +import pytest +import torch + +from vime.backends.megatron_utils.cp_utils import ( # noqa: E402 + get_logits_and_tokens_offset_with_cp, + get_sum_of_sample_mean, + reduce_train_step_metrics, + rollout_log_metric_contribution, +) + + +@pytest.fixture +def mock_dp_with_cp_group(monkeypatch): + """A sentinel "process group" object plus a no-op ``dist.all_reduce``. + + Lets the train-step report tests exercise the production call shape + (``dist.all_reduce(values, group=dp_with_cp_group)``) without standing + up a real torch.distributed runtime. The test itself simulates cross- + rank summation in pure Python and feeds the already-summed tensor in; + the no-op all_reduce leaves it untouched. + """ + import torch.distributed as dist + + monkeypatch.setattr(dist, "all_reduce", lambda tensor, group=None, op=None: None) + return object() # opaque sentinel — only used as the ``group`` argument + + +# --------------------------------------------------------------------------- +# Mirrors the actual train_one_step reporting math: +# +# per-rollout-mean path: +# reported = sum_of_reducer_per_mb / step_global_batch_size +# per-token-loss path: +# reported = sum_of_reducer_per_mb / sum_of_per_mb_num_tokens +# +# The reducer is the same callable used at train time (and inside +# log_rollout_data on the rollout side). +# --------------------------------------------------------------------------- + + +# 4 samples: rollout R0 owns indices 0,1,2 (mask sums 3+3+3=9); rollout R1 +# owns index 3 (mask sum 3). Pre-computed per-sample denom = group sum. +# Per-rollout-mean: R0 = 5, R1 = 11, sum = 16, divided by 2 rollouts → 8. +# Per-token-loss: sum of all x = 78, total clamped mask = 12, → 6.5. +_FIXED_RESPONSE_LENGTHS = [3, 3, 3, 3] +_FIXED_TOTAL_LENGTHS = [r + 4 for r in _FIXED_RESPONSE_LENGTHS] +_FIXED_LOSS_MASKS = [torch.ones(r, dtype=torch.float32) for r in _FIXED_RESPONSE_LENGTHS] +_FIXED_ROLLOUT_DENOMS = [9.0, 9.0, 9.0, 3.0] +_FIXED_X_PER_SAMPLE = [ + torch.tensor([1.0, 2.0, 3.0]), + torch.tensor([4.0, 5.0, 6.0]), + torch.tensor([7.0, 8.0, 9.0]), + torch.tensor([10.0, 11.0, 12.0]), +] +_FIXED_STEP_GBS = 2 # 2 distinct rollouts in the step +_EXPECTED_PER_ROLLOUT_MEAN_REPORT = 8.0 +_EXPECTED_PER_TOKEN_LOSS_REPORT = 78.0 / 12.0 + + +# Each entry: list of "rank"s, each rank is a list of mbs, each mb is the +# sample-index list packed into that mb. Covers: single mb, evenly split by +# rollout, split inside a rollout (R0 across mbs), uneven distribution, and +# fully singleton mbs per rank. +_PARTITION_CONFIGS = [ + [[[0, 1, 2, 3]]], # 1 rank, 1 mb + [[[0, 1, 2], [3]]], # 1 rank, 2 mbs split at rollout boundary + [[[0, 1], [2, 3]]], # 1 rank, 2 mbs splitting R0 across them — the tricky case + [[[0, 1]], [[2, 3]]], # 2 ranks, 1 mb each + [[[0, 1, 3]], [[2]]], # 2 ranks, R0 split across BOTH ranks (worst case for split-across-mb bug) + [[[0]], [[1]], [[2]], [[3]]], # 4 ranks, 1 sample per rank +] + + +def _simulate_report(partition, *, per_token_loss: bool) -> float: + """Reproduce train_one_step's reporting math for one partition config.""" + metric_sum = 0.0 + num_tokens_sum = 0 + for rank_mbs in partition: + for mb_indices in rank_mbs: + mb_total = [_FIXED_TOTAL_LENGTHS[i] for i in mb_indices] + mb_resp = [_FIXED_RESPONSE_LENGTHS[i] for i in mb_indices] + mb_masks = [_FIXED_LOSS_MASKS[i] for i in mb_indices] + mb_x = torch.cat([_FIXED_X_PER_SAMPLE[i] for i in mb_indices]) + if per_token_loss: + # Per-token-loss: caller uses ``calculate_per_token_loss=True`` + # to get ``sum_of_token`` (no per-sample denom). + reducer = get_sum_of_sample_mean(mb_total, mb_resp, mb_masks, calculate_per_token_loss=True) + num_tokens_sum += sum(max(int(m.sum().item()), 1) for m in mb_masks) + else: + mb_denoms = torch.tensor([_FIXED_ROLLOUT_DENOMS[i] for i in mb_indices], dtype=torch.float32) + reducer = get_sum_of_sample_mean(mb_total, mb_resp, mb_masks, mb_denoms) + metric_sum += reducer(mb_x).item() + if per_token_loss: + return metric_sum / num_tokens_sum + return metric_sum / _FIXED_STEP_GBS + + +@pytest.mark.unit +@pytest.mark.parametrize("partition", _PARTITION_CONFIGS) +def test_per_rollout_mean_report_invariant_to_mb_distribution(partition): + """Same samples should yield the same per-rollout-mean report regardless of + how they're spread across DP ranks / micro-batches — this is what lets us + change parallelism without changing wandb numbers.""" + assert _simulate_report(partition, per_token_loss=False) == pytest.approx(_EXPECTED_PER_ROLLOUT_MEAN_REPORT) + + +@pytest.mark.unit +@pytest.mark.parametrize("partition", _PARTITION_CONFIGS) +def test_per_token_loss_report_invariant_to_mb_distribution(partition): + """Same invariant for the per-token-loss reporting path.""" + assert _simulate_report(partition, per_token_loss=True) == pytest.approx(_EXPECTED_PER_TOKEN_LOSS_REPORT) + + +def _simulate_rollout_report(samples_per_rank): + """Reproduce log_rollout_data + gather_log_data's averaging math for the + per-token metric branch. + + Each "rank" applies the reducer once over its full sample subset, then + ``rollout_log_metric_contribution`` (the same helper data.py uses) emits + the ``(per_rank_sum, count)`` tuple. We aggregate via + ``Σsum / Σcount`` — the same shape ``gather_log_data`` uses. + """ + dp_size = len(samples_per_rank) + pairs: list[tuple[float, float]] = [] + for indices in samples_per_rank: + if not indices: + pairs.append( + rollout_log_metric_contribution( + 0.0, cp_size=1, num_rollouts_in_rollout=_FIXED_STEP_GBS, dp_size=dp_size + ) + ) + continue + tl = [_FIXED_TOTAL_LENGTHS[i] for i in indices] + rl = [_FIXED_RESPONSE_LENGTHS[i] for i in indices] + masks = [_FIXED_LOSS_MASKS[i] for i in indices] + denoms = torch.tensor([_FIXED_ROLLOUT_DENOMS[i] for i in indices], dtype=torch.float32) + x = torch.cat([_FIXED_X_PER_SAMPLE[i] for i in indices]) + reducer = get_sum_of_sample_mean(tl, rl, masks, denoms) + pairs.append( + rollout_log_metric_contribution( + reducer(x).item(), + cp_size=1, + num_rollouts_in_rollout=_FIXED_STEP_GBS, + dp_size=dp_size, + ) + ) + total_sum = sum(p[0] for p in pairs) + total_count = sum(p[1] for p in pairs) + return total_sum / total_count + + +_DP_PARTITIONS = [ + [[0, 1, 2, 3]], # 1 rank holds everything + [[0, 1, 2], [3]], # 2 ranks, balanced by rollout + [[0, 1], [2, 3]], # 2 ranks splitting R0 across mb-and-rank + [[0, 1, 3], [2]], # 2 ranks with R0 spread across BOTH (one of R0's samples is on rank 1) + [[0], [1], [2], [3]], # 4 ranks, one sample each (R0's samples spread across 3 ranks) +] + + +@pytest.mark.unit +@pytest.mark.parametrize("dp_partition", _DP_PARTITIONS) +def test_rollout_report_matches_train_report_in_single_step(dp_partition): + """In a 1-step rollout, the rollout-side report (log_rollout_data → gather) + must equal the train-side report (train_one_step ``value / step_global_batch_size``) + for the same samples — otherwise wandb numbers between phases drift. + + Both go through the same reducer with the same precomputed denominators; + the contract this test pins is that the gather count plumbing on the + rollout side sums to the same denominator the train side uses + (``step_global_batch_size``), independent of how the rollout's samples + are spread across DP ranks. + """ + rollout_report = _simulate_rollout_report(dp_partition) + assert rollout_report == pytest.approx(_EXPECTED_PER_ROLLOUT_MEAN_REPORT) + + +@pytest.mark.unit +def test_train_one_step_per_rollout_mean_report_invariant_to_cp(monkeypatch, mock_dp_with_cp_group): + """End-to-end check of train_one_step's report formula across CP sizes. + + Mirrors the actual reduction order: + 1. Each (DP, CP) rank computes per-mb reducer output. + 2. Per-rank values are summed across mbs locally. + 3. All-reduce sums across DP*CP ranks. + 4. ``reduce_train_step_metrics`` applied (the same helper + ``train_one_step`` calls, so this test stays honest if the + implementation changes). + + cp_size = 1 vs cp_size = 2 must give the same reported number — + otherwise wandb metrics would drift the moment a user enables CP. + """ + from megatron.core import mpu as _mpu + + total_lengths = [12, 12] + response_lengths = [8, 8] + loss_masks = [torch.ones(r, dtype=torch.float32) for r in response_lengths] + sample_denoms = torch.tensor([16.0, 16.0], dtype=torch.float32) + x_full = [ + torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]), + torch.tensor([10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0]), + ] + step_global_batch_size = 1 # one rollout in the step + + def simulate(cp_size: int) -> float: + monkeypatch.setattr(_mpu, "get_context_parallel_world_size", lambda: cp_size) + # values[0] is the per-token-loss path's num_tokens slot; for + # per-rollout-mean it's a zero placeholder (loss_function sets 0). + value_after_allreduce = 0.0 + for cp_rank in range(cp_size): + monkeypatch.setattr(_mpu, "get_context_parallel_rank", lambda r=cp_rank: r) + if cp_size == 1: + x_for_rank = torch.cat(x_full) + else: + x_chunks_per_sample = [] + for tl, rl, x in zip(total_lengths, response_lengths, x_full, strict=True): + prompt_length = tl - rl + _, _, _, tokens_offset = get_logits_and_tokens_offset_with_cp(tl, rl) + c0 = x[tokens_offset[0][0] - prompt_length : tokens_offset[0][1] - prompt_length] + c1 = x[tokens_offset[1][0] - prompt_length : tokens_offset[1][1] - prompt_length] + x_chunks_per_sample.append(torch.cat([c0, c1])) + x_for_rank = torch.cat(x_chunks_per_sample) + reducer = get_sum_of_sample_mean(total_lengths, response_lengths, loss_masks, sample_denoms) + value_after_allreduce += reducer(x_for_rank).item() + reduced = reduce_train_step_metrics( + [{"keys": ["metric"], "values": torch.tensor([0.0, value_after_allreduce])}], + calculate_per_token_loss=False, + step_global_batch_size=step_global_batch_size, + cp_size=cp_size, + dp_with_cp_group=mock_dp_with_cp_group, + ) + return reduced["metric"] + + assert simulate(1) == pytest.approx(simulate(2)) + + +@pytest.mark.unit +def test_train_one_step_per_token_loss_report_invariant_to_cp(monkeypatch, mock_dp_with_cp_group): + """Same end-to-end check for the per-token-loss path: divisor is + ``values[0] = num_tokens`` (computed in loss.py from FULL loss masks), + which each CP rank duplicates and all-reduce sums by ``cp_size``. The + ``cp_factor = cp_size`` multiplier inside ``reduce_train_step_metrics`` + cancels that inflation, so the report stays CP-invariant. + """ + from megatron.core import mpu as _mpu + + total_lengths = [12, 12] + response_lengths = [8, 8] + loss_masks = [torch.ones(r, dtype=torch.float32) for r in response_lengths] + num_tokens_per_mb = sum(int(m.sum().item()) for m in loss_masks) # = 16 + x_full = [ + torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]), + torch.tensor([10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0]), + ] + + def simulate(cp_size: int) -> float: + monkeypatch.setattr(_mpu, "get_context_parallel_world_size", lambda: cp_size) + value_after_allreduce = 0.0 + num_tokens_after_allreduce = 0 # each CP rank reports the same num_tokens + for cp_rank in range(cp_size): + monkeypatch.setattr(_mpu, "get_context_parallel_rank", lambda r=cp_rank: r) + if cp_size == 1: + x_for_rank = torch.cat(x_full) + else: + x_chunks_per_sample = [] + for tl, rl, x in zip(total_lengths, response_lengths, x_full, strict=True): + prompt_length = tl - rl + _, _, _, tokens_offset = get_logits_and_tokens_offset_with_cp(tl, rl) + c0 = x[tokens_offset[0][0] - prompt_length : tokens_offset[0][1] - prompt_length] + c1 = x[tokens_offset[1][0] - prompt_length : tokens_offset[1][1] - prompt_length] + x_chunks_per_sample.append(torch.cat([c0, c1])) + x_for_rank = torch.cat(x_chunks_per_sample) + reducer = get_sum_of_sample_mean( + total_lengths, response_lengths, loss_masks, calculate_per_token_loss=True + ) + value_after_allreduce += reducer(x_for_rank).item() + num_tokens_after_allreduce += num_tokens_per_mb + reduced = reduce_train_step_metrics( + [ + { + "keys": ["metric"], + "values": torch.tensor([num_tokens_after_allreduce, value_after_allreduce], dtype=torch.float32), + } + ], + calculate_per_token_loss=True, + step_global_batch_size=999, # unused in per-token-loss path + cp_size=cp_size, + dp_with_cp_group=mock_dp_with_cp_group, + ) + return reduced["metric"] + + assert simulate(1) == pytest.approx(simulate(2)) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_metric_report_dist.py b/tests/test_metric_report_dist.py new file mode 100644 index 000000000..7535027aa --- /dev/null +++ b/tests/test_metric_report_dist.py @@ -0,0 +1,303 @@ +"""Multi-process distributed tests for the cp_utils report helpers. + +Spawn ``dp_size * cp_size`` workers with real ``torch.distributed`` (gloo +backend) and exercise the actual production helpers end-to-end. The +single-process numerical contracts live in ``test_metric_report.py``; this +file pins the cross-rank plumbing. + +Mapping to the production train_one_step / log_rollout_data flows: + + - ``_train_step_distributed_worker`` mirrors ``train_one_step``: + per-rank reducer → ``reduce_train_step_metrics`` + (which calls ``dist.all_reduce`` over the dp-with-cp group and + applies the cp_size cancellation for the per-token-loss path). + - ``_rollout_log_distributed_worker`` mirrors ``log_rollout_data``: + per-rank reducer → ``rollout_log_metric_contribution`` → + ``gather_and_reduce_log_dict`` (which calls ``dist.gather_object`` + and applies per-key reductions). + +ALL (dp, cp) configurations must give the same reported number — that's +the contract a user touches when they flip any parallelism dial. +""" + +from __future__ import annotations + +# IMPORTANT: import the helpers (and the megatron stub it installs) BEFORE +# any slime import. Spawned workers re-import this module from scratch, so +# the same ordering must hold there — see ``stub_megatron_in_worker`` +# for the worker-side details. pytest's prepend importmode puts +# ``tests/`` on sys.path so the bare-name import works without an +# ``__init__.py``; mp.spawn children inherit the parent's sys.path. +import _cp_dist_helpers +import pytest +import torch +from _cp_dist_helpers import ( + FOUR_ROLLOUT_EXPECTED_REPORT, + FOUR_ROLLOUT_RESPONSE_LENGTHS, + FOUR_ROLLOUT_TOTAL_LENGTHS, + FOUR_ROLLOUT_X_VALUES, + cp_chunk_response_tensor, + free_port, + init_worker_process_group, + stub_megatron_in_worker, +) + + +def _train_step_distributed_worker( + rank: int, + world_size: int, + cp_size: int, + dp_size: int, + per_token_loss: bool, + master_port: int, + result_path: str, +) -> None: + """Per-rank entrypoint for ``mp.spawn``: init gloo pg, run one rank's + share of the train-step report, write rank-0's result to a file.""" + import torch.distributed as _dist + + cp_rank = rank % cp_size + dp_rank = rank // cp_size + stub_megatron_in_worker(cp_size, cp_rank) + + dp_with_cp_group = init_worker_process_group(rank, world_size, master_port) + try: + # Import AFTER the megatron stub override so cp_utils still binds + # against the pre-installed stub (which we've now pinned for this + # worker's CP rank). + from vime.backends.megatron_utils.cp_utils import get_sum_of_sample_mean, reduce_train_step_metrics + + all_total_lengths = FOUR_ROLLOUT_TOTAL_LENGTHS + all_response_lengths = FOUR_ROLLOUT_RESPONSE_LENGTHS + all_loss_masks = [torch.ones(r, dtype=torch.float32) for r in all_response_lengths] + all_x = [torch.tensor(v) for v in FOUR_ROLLOUT_X_VALUES] + step_global_batch_size = 4 # 4 rollouts in the step + + # Round-robin DP partition: with 4 samples / dp=N, rank i gets + # samples i, i+N, ... (matches what _split_train_data_by_dp does + # for evenly divisible cases). + my_indices = [i for i in range(4) if i % dp_size == dp_rank] + my_tl = [all_total_lengths[i] for i in my_indices] + my_rl = [all_response_lengths[i] for i in my_indices] + my_masks = [all_loss_masks[i] for i in my_indices] + my_x = [all_x[i] for i in my_indices] + my_denoms = torch.tensor([float(m.sum().item()) for m in my_masks], dtype=torch.float32) + + if cp_size == 1: + x_for_rank = torch.cat(my_x) + else: + x_for_rank = torch.cat( + [cp_chunk_response_tensor(x, tl, rl) for tl, rl, x in zip(my_tl, my_rl, my_x, strict=True)] + ) + + if per_token_loss: + reducer = get_sum_of_sample_mean(my_tl, my_rl, my_masks, calculate_per_token_loss=True) + # num_tokens is computed off the FULL mask (not the chunked + # one) in loss.py — every CP rank reports the same number, + # which is why ``reduce_train_step_metrics`` cancels by + # ``cp_factor = cp_size`` afterwards. + num_tokens = sum(int(m.sum().item()) for m in my_masks) + values_tensor = torch.tensor([float(num_tokens), reducer(x_for_rank).item()], dtype=torch.float32) + else: + reducer = get_sum_of_sample_mean(my_tl, my_rl, my_masks, my_denoms) + values_tensor = torch.tensor([0.0, reducer(x_for_rank).item()], dtype=torch.float32) + + reduced = reduce_train_step_metrics( + [{"keys": ["metric"], "values": values_tensor}], + calculate_per_token_loss=per_token_loss, + step_global_batch_size=step_global_batch_size, + cp_size=cp_size, + dp_with_cp_group=dp_with_cp_group, + ) + + if rank == 0: + with open(result_path, "w") as f: + f.write(repr(reduced["metric"])) + finally: + _dist.destroy_process_group() + + +@pytest.mark.unit +@pytest.mark.parametrize( + "dp_size,cp_size", + [(dp, cp) for dp in [1, 2, 4] for cp in [1, 2, 4]], +) +def test_train_step_per_rollout_mean_real_distributed(dp_size, cp_size, tmp_path): + """End-to-end multi-process: spawn ``dp_size * cp_size`` workers, each + runs its share with real ``torch.distributed`` (gloo); ALL parallelism + combinations must give the same reported per-rollout-mean number. + + Expected = sum of per-rollout token-means / step_gbs + = (4.5 + 45 + 450 + 4500) / 4 = 1249.875 + """ + import torch.multiprocessing as mp + + world_size = dp_size * cp_size + result_path = str(tmp_path / "result.txt") + mp.spawn( + _train_step_distributed_worker, + args=(world_size, cp_size, dp_size, False, free_port(), result_path), + nprocs=world_size, + join=True, + ) + with open(result_path) as f: + result = float(f.read()) + assert result == pytest.approx(FOUR_ROLLOUT_EXPECTED_REPORT) + + +@pytest.mark.unit +@pytest.mark.parametrize( + "dp_size,cp_size", + [(dp, cp) for dp in [1, 2, 4] for cp in [1, 2, 4]], +) +def test_train_step_per_token_loss_real_distributed(dp_size, cp_size, tmp_path): + """Same end-to-end multi-process check for the per-token-loss path. + + Expected = sum of all x / total_tokens + = (36 + 360 + 3600 + 36000) / 32 = 1249.875 + """ + import torch.multiprocessing as mp + + world_size = dp_size * cp_size + result_path = str(tmp_path / "result.txt") + mp.spawn( + _train_step_distributed_worker, + args=(world_size, cp_size, dp_size, True, free_port(), result_path), + nprocs=world_size, + join=True, + ) + with open(result_path) as f: + result = float(f.read()) + assert result == pytest.approx(FOUR_ROLLOUT_EXPECTED_REPORT) + + +def _rollout_log_distributed_worker( + rank: int, + world_size: int, + cp_size: int, + dp_size: int, + master_port: int, + result_path: str, +) -> None: + """Per-rank entrypoint for ``mp.spawn``: build a multi-key log_dict + covering all three reduction modes ``gather_and_reduce_log_dict`` + supports, run real ``dist.gather_object``, have rank 0 dump the + reduced dict via pickle for the parent to assert on. + """ + import pickle + + import torch.distributed as _dist + + cp_rank = rank % cp_size + dp_rank = rank // cp_size + stub_megatron_in_worker(cp_size, cp_rank) + + dp_group = init_worker_process_group(rank, world_size, master_port) + try: + from vime.backends.megatron_utils.cp_utils import ( + gather_and_reduce_log_dict, + get_sum_of_sample_mean, + rollout_log_metric_contribution, + ) + + all_total_lengths = FOUR_ROLLOUT_TOTAL_LENGTHS + all_response_lengths = FOUR_ROLLOUT_RESPONSE_LENGTHS + all_loss_masks = [torch.ones(r, dtype=torch.float32) for r in all_response_lengths] + all_x = [torch.tensor(v) for v in FOUR_ROLLOUT_X_VALUES] + num_rollouts_in_rollout = 4 + + my_indices = [i for i in range(4) if i % dp_size == dp_rank] + my_tl = [all_total_lengths[i] for i in my_indices] + my_rl = [all_response_lengths[i] for i in my_indices] + my_masks = [all_loss_masks[i] for i in my_indices] + my_x = [all_x[i] for i in my_indices] + my_denoms = torch.tensor([float(m.sum().item()) for m in my_masks], dtype=torch.float32) + + if cp_size == 1: + x_for_rank = torch.cat(my_x) + else: + x_for_rank = torch.cat( + [cp_chunk_response_tensor(x, tl, rl) for tl, rl, x in zip(my_tl, my_rl, my_x, strict=True)] + ) + + reducer = get_sum_of_sample_mean(my_tl, my_rl, my_masks, my_denoms) + per_rank_reducer_sum = reducer(x_for_rank).item() + + # Exercise every reduction mode the production log_rollout_data emits. + log_dict = { + # per-rollout-mean: (sum, count) via rollout_log_metric_contribution. + # gather: Σsum / Σcount = sum_DP_full / num_rollouts. + "logp_per_rollout": rollout_log_metric_contribution( + per_rank_reducer_sum, + cp_size=cp_size, + num_rollouts_in_rollout=num_rollouts_in_rollout, + dp_size=dp_size, + ), + # per-sample-mean: (Σval, num_samples) — matches the + # ``total_lengths`` style in log_rollout_data. gather: Σsum/Σcount + # = total / total_samples = per-sample mean of total_lengths. + "total_lengths_per_sample": (float(sum(my_tl)), float(len(my_tl))), + # mean-across-ranks: plain scalar — matches log_multi_turn_data + # style. gather: Σvalue / dp_world. + "rank_local_mean": float(sum(my_tl)) / len(my_tl), + } + + reduced = gather_and_reduce_log_dict(log_dict, dp_size=world_size, dp_src_rank=0, dp_group=dp_group) + + if rank == 0: + with open(result_path, "wb") as f: + pickle.dump(reduced, f) + finally: + _dist.destroy_process_group() + + +@pytest.mark.unit +@pytest.mark.parametrize( + "dp_size,cp_size", + [(dp, cp) for dp in [1, 2, 4] for cp in [1, 2, 4]], +) +def test_rollout_log_real_distributed_multi_key(dp_size, cp_size, tmp_path): + """End-to-end multi-process for ``gather_and_reduce_log_dict``. + + Covers the three key shapes ``log_rollout_data`` produces: + - per-rollout-mean ((sum, count) via ``rollout_log_metric_contribution``) + - per-sample-mean ((Σval, num_samples) tuple — e.g. ``total_lengths``) + - mean-across-ranks (plain float — e.g. multi_turn stats) + + All (dp, cp) configs must yield the same reduced numbers; matches the + expected values written in pure Python from the fixture. In particular + the per-rollout-mean number must equal what the train-step report tests + above land on (FOUR_ROLLOUT_EXPECTED_REPORT), pinning the cross-phase + contract. + """ + import pickle + + import torch.multiprocessing as mp + + world_size = dp_size * cp_size + result_path = str(tmp_path / "result.pkl") + mp.spawn( + _rollout_log_distributed_worker, + args=(world_size, cp_size, dp_size, free_port(), result_path), + nprocs=world_size, + join=True, + ) + with open(result_path, "rb") as f: + reduced = pickle.load(f) + + # per-rollout-mean: matches the train-side report — 1249.875. + assert reduced["logp_per_rollout"] == pytest.approx(FOUR_ROLLOUT_EXPECTED_REPORT) + # per-sample-mean: every sample has total_length=12, so the average is 12. + assert reduced["total_lengths_per_sample"] == pytest.approx(12.0) + # mean-across-ranks: every rank's local mean is 12, so cross-rank mean is 12. + assert reduced["rank_local_mean"] == pytest.approx(12.0) + + +# Keep an explicit reference to silence "unused import" complaints while +# documenting that importing the helpers module is load-bearing (it +# installs the megatron stub before slime is touched). +_ = _cp_dist_helpers + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/test_qwen2.5_0.5B_fanout_short.py b/tests/test_qwen2.5_0.5B_fanout_short.py new file mode 100644 index 000000000..01c5ddbb1 --- /dev/null +++ b/tests/test_qwen2.5_0.5B_fanout_short.py @@ -0,0 +1,224 @@ +"""E2E test: one prompt → random 1..3 training samples (compact / subagent fan-out). + +What this test pins +------------------- +The "compact" pattern (where one rollout execution emits a *variable* +number of training samples sharing a single ``rollout_id``) has CPU unit +coverage at the piece-level (``test_dp_schedule.py`` for the rollout- +aware step splitter, ``test_sample.py`` for ``Sample.rollout_id`` round- +trip, ``test_cp_utils.py`` for the per-rollout-mean reducer). But until +this test, **no e2e training run had ever exercised the full chain**: + + custom_generate returns list[Sample] sharing rollout_id + → _validate_rollout_id_annotated at depth ≥ 2 passes + → _split_train_data_by_dp groups by rollout_id and trims to N steps + using ``rollout_batch_size * n_samples_per_prompt / global_batch_size`` + (NOT total sample count, which would inflate steps once N>1) + → loss aggregation uses ``rollout_mask_sums`` so every sibling sample + contributes one token-weighted mean per rollout + → train_one_step's ``step_global_batch_size`` denominator equals + num_rollouts (not num_samples), keeping grad magnitude stable + independent of fan-out + +The fan-out function itself lives in +``vime/rollout/_fanout_test_helpers.py`` — it has to be at a dot-free +module path so ``importlib.import_module`` can resolve the string +``--custom-generate-function-path`` flag (this filename has dots). + +Test choices +------------ +- **Deterministic fan-out** ``N = 1 + (sample.index % MAX_FANOUT)`` for + reproducibility. Every value in {1, 2, 3} gets exercised in a single + rollout. N=1 keeps the backward-compat (no fan-out) path alive in CI. +- **Smoke + implicit step-count assertion**. ``--ci-test`` flips the + framework's built-in numerical guards (KL divergence, log_prob ≈ + ref_log_prob); a step-counting / loss-denominator regression would + trip them. Plus the helper writes one line per call to a tmp counter + file — post-train we assert the count equals + ``num_rollout * rollout_batch_size``, proving the custom path actually + drove every prompt (vs. silent fallback to default rollout). +""" + +import os +import tempfile + +import vime.utils.external_utils.command_utils as U + +TIGHT_DEVICE_MEMORY = U.get_bool_env_var("SLIME_TEST_TIGHT_DEVICE_MEMORY", "1") + +MODEL_NAME = "Qwen2.5-0.5B-Instruct" +MODEL_TYPE = "qwen2.5-0.5B" +NUM_GPUS = 4 + +# Counter file used by the compact_generate helper. We pass its path +# through to the Ray-submitted job via an env var so all worker +# processes write to the same path. +FANOUT_COUNTER_FILE = os.environ.get( + "SLIME_FANOUT_TEST_COUNTER_FILE", + os.path.join(tempfile.gettempdir(), "slime_fanout_test_counter.log"), +) + + +def prepare(): + U.exec_command("mkdir -p /root/models /root/datasets") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.hf_download_dataset("zhuzilin/dapo-math-17k") + # Clear the counter so a previous run's invocations don't bleed in. + try: + os.remove(FANOUT_COUNTER_FILE) + except FileNotFoundError: + pass + + +def execute(): + ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load /root/models/{MODEL_NAME}/ " + + # Shape: rollout_batch_size=8 prompts, n_samples_per_prompt=1 (all + # fan-out is owned by compact_generate; this knob stays at 1 so a + # regression that confuses sample count vs rollout count surfaces), + # global_batch_size=4 → 2 training steps per rollout, num_rollout=3 + # → 6 total training steps. + # + # NB no ``--group-rm``: when custom_generate returns ``list[Sample]`` + # the per-sample rm path inside ``generate_and_rm`` (vllm_rollout.py) + # handles the fan-out correctly via ``batched_async_rm`` on the flat + # sibling list. ``--group-rm`` defers rm to ``generate_and_rm_group`` + # which assumes ``group`` is already flat ``list[Sample]`` — combining + # it with a list-returning custom_generate yields a + # ``list[list[Sample]]`` and crashes ``async_rm`` + # (`'list' object has no attribute 'metadata'`). + rollout_args = ( + "--prompt-data /root/datasets/dapo-math-17k/dapo-math-17k.jsonl " + "--input-key prompt " + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type deepscaler " + "--num-rollout 3 " + "--rollout-batch-size 8 " + "--n-samples-per-prompt 1 " + "--rollout-max-response-len 8192 " + "--rollout-temperature 0.8 " + "--global-batch-size 4 " + "--balance-data " + "--custom-generate-function-path vime.rollout._fanout_test_helpers.compact_generate " + # GRPO normalization needs per-prompt grouping. The default + # ``_post_process_rewards`` (vime/ray/rollout.py) reshapes + # by ``n_samples_per_prompt`` and falls back to "one big group" + # when the per-prompt count is uneven — fan-out trips exactly + # that fallback. The helper here groups by ``Sample.group_index`` + # (the per-prompt counter the data source stamps; deepcopy in + # compact_generate preserves it across siblings) so each prompt's + # siblings normalize against each other, matching the GRPO + # semantics the default targets in the uniform case. + "--custom-reward-post-process-path vime.rollout._fanout_test_helpers.grpo_normalize_by_group_index " + ) + + perf_args = ( + "--tensor-model-parallel-size 1 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 1 " + "--expert-model-parallel-size 1 " + "--expert-tensor-parallel-size 1 " + "--use-dynamic-batch-size " + "--max-tokens-per-gpu 9216 " + ) + + grpo_args = ( + "--advantage-estimator grpo " + "--use-kl-loss " + "--kl-loss-coef 0.00 " + "--kl-loss-type low_var_kl " + "--entropy-coef 0.00 " + "--eps-clip 0.2 " + "--eps-clip-high 0.28 " + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + ) + + vllm_args = ( + "--rollout-num-gpus-per-engine 1 " + f"--vllm-gpu-memory-utilization {0.6 if TIGHT_DEVICE_MEMORY else 0.7} " + "--vllm-max-cudagraph-capture-size 32 " + ) + + ci_args = "--ci-test " + + fault_tolerance_args = ( + "--use-fault-tolerance " + "--rollout-health-check-interval 5 " + "--rollout-health-check-timeout 10 " + "--rollout-health-check-first-wait 0 " + ) + + misc_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--actor-num-nodes 1 " + "--actor-num-gpus-per-node 4 " + "--colocate " + "--megatron-to-hf-mode bridge " + ) + + train_args = ( + f"{ckpt_args} " + f"{rollout_args} " + f"{optimizer_args} " + f"{grpo_args} " + f"{U.get_default_wandb_args(__file__)} " + f"{perf_args} " + f"{vllm_args} " + f"{ci_args} " + f"{fault_tolerance_args} " + f"{misc_args} " + ) + + U.execute_train( + train_args=train_args, + num_gpus_per_node=NUM_GPUS, + megatron_model_type=MODEL_TYPE, + # Make the counter path visible inside the Ray-submitted job + # (helper picks it up via os.environ). + extra_env_vars={"SLIME_FANOUT_TEST_COUNTER_FILE": FANOUT_COUNTER_FILE}, + ) + + # Post-train assertion: compact_generate must have been called exactly + # ``num_rollout * rollout_batch_size`` = 3 * 8 = 24 times. A regression + # that bypassed the custom path (arg parser drops the flag, or the + # path is silently mis-routed) would either skip the file entirely or + # under-count. + expected_calls = 3 * 8 + try: + with open(FANOUT_COUNTER_FILE) as f: + actual_calls = sum(1 for _ in f) + except FileNotFoundError as e: + raise AssertionError( + f"compact_generate counter file {FANOUT_COUNTER_FILE} missing — the custom " + f"generate path was never invoked. Suggests --custom-generate-function-path " + f"was dropped by the arg parser or the resolved import path is wrong." + ) from e + assert actual_calls == expected_calls, ( + f"compact_generate was called {actual_calls} times, expected {expected_calls} " + f"(num_rollout=3 × rollout_batch_size=8). A mismatch points at the rollout " + f"submission loop double-submitting / under-submitting prompts." + ) + + +if __name__ == "__main__": + prepare() + os.environ.pop("http_proxy") + os.environ.pop("https_proxy") + os.environ.pop("HTTP_PROXY") + os.environ.pop("HTTPS_PROXY") + execute() diff --git a/tests/test_qwen3_4B_streaming_partial_rollout.py b/tests/test_qwen3_4B_streaming_partial_rollout.py new file mode 100644 index 000000000..898e4c2cc --- /dev/null +++ b/tests/test_qwen3_4B_streaming_partial_rollout.py @@ -0,0 +1,144 @@ +"""CI smoke test for the streaming vLLM rollout path. + +Wires ``vime.rollout.vllm_streaming_rollout.generate_streaming`` in as the +per-sample generate function, with ``--over-sampling-batch-size`` > +``--rollout-batch-size`` and ``--partial-rollout`` enabled so the rollout +loop *must* abort in-flight requests every step — exercising the streaming +abort path (partial state should already be on the sample when the SSE is +cut, then the partial groups get recycled into the data buffer). + +Uses Qwen3-4B (vs the 0.5B in other short tests) so responses on dapo-math +are long enough to actually trigger mid-stream aborts, and good enough to +produce non-zero rewards. + +vime counterpart of slime's ``test_qwen3_4B_streaming_partial_rollout`` — the +only structural change is the rollout engine: sglang args become vLLM args and +the generate function points at ``vllm_streaming_rollout``. +""" + +import os + +import vime.utils.external_utils.command_utils as U + +TIGHT_HOST_MEMORY = U.get_bool_env_var("SLIME_TEST_TIGHT_HOST_MEMORY", "1") + +MODEL_NAME = "Qwen3-4B" +MODEL_TYPE = "qwen3-4B" +NUM_GPUS = 8 + + +def prepare(): + U.exec_command("mkdir -p /root/models /root/datasets") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.hf_download_dataset("zhuzilin/dapo-math-17k") + U.convert_checkpoint(model_name=MODEL_NAME, megatron_model_type=MODEL_TYPE, num_gpus_per_node=NUM_GPUS) + + +def execute(): + ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load /root/{MODEL_NAME}_torch_dist " + + rollout_args = ( + # Streaming generate at the per-sample level — the outer rollout + # loop is still the stock vllm one (semaphore, abort orchestration). + "--custom-generate-function-path vime.rollout.vllm_streaming_rollout.generate_streaming " + "--prompt-data /root/datasets/dapo-math-17k/dapo-math-17k.jsonl " + "--input-key prompt " + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type deepscaler " + "--num-rollout 2 " + "--rollout-batch-size 4 " + # Over-sample 2x so half of every rollout's in-flight groups must + # be aborted, then partial-rollout recycles them. + "--over-sampling-batch-size 8 " + "--partial-rollout " + "--mask-offpolicy-in-partial-rollout " + "--n-samples-per-prompt 4 " + # Long enough that aborts cut samples mid-generation rather than + # everything finishing first; Qwen3-4B also generates real reasoning + # on dapo-math at this length, so the reward signal is non-trivial. + "--rollout-max-response-len 4096 " + "--rollout-temperature 0.8 " + "--global-batch-size 16 " + "--balance-data " + ) + + perf_args = ( + "--tensor-model-parallel-size 2 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 2 " + "--recompute-granularity full " + "--recompute-method uniform " + "--recompute-num-layers 1 " + "--use-dynamic-batch-size " + f"--max-tokens-per-gpu {2048 if TIGHT_HOST_MEMORY else 8192} " + ) + + grpo_args = ( + "--advantage-estimator grpo " + "--use-kl-loss " + "--kl-loss-coef 0.00 " + "--kl-loss-type low_var_kl " + "--entropy-coef 0.00 " + "--eps-clip 0.2 " + "--eps-clip-high 0.28 " + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + ) + + vllm_args = ( + "--rollout-num-gpus-per-engine 2 " + "--rollout-num-gpus 8 " + "--vllm-gpu-memory-utilization 0.8 " + "--vllm-max-num-seqs 512 " + "--vllm-max-cudagraph-capture-size 32 " + ) + + ci_args = "--ci-test " + + misc_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--actor-num-nodes 1 " + "--actor-num-gpus-per-node 8 " + "--colocate " + ) + + train_args = ( + f"{ckpt_args} " + f"{rollout_args} " + f"{optimizer_args} " + f"{grpo_args} " + f"{U.get_default_wandb_args(__file__)} " + f"{perf_args} " + f"{vllm_args} " + f"{ci_args} " + f"{misc_args} " + ) + + U.execute_train( + train_args=train_args, + num_gpus_per_node=NUM_GPUS, + megatron_model_type=MODEL_TYPE, + ) + + +if __name__ == "__main__": + prepare() + os.environ.pop("http_proxy", None) + os.environ.pop("https_proxy", None) + os.environ.pop("HTTP_PROXY", None) + os.environ.pop("HTTPS_PROXY", None) + execute() diff --git a/tests/test_sample.py b/tests/test_sample.py new file mode 100644 index 000000000..bc83414e7 --- /dev/null +++ b/tests/test_sample.py @@ -0,0 +1,272 @@ +"""CPU unit tests for ``vime.utils.types.Sample``. + +Pins two contracts that the rollout / training boundary depends on: + + 1. ``to_dict`` / ``from_dict`` round-trip — Sample crosses Ray actor + boundaries as a dict (especially in async / fully-async / partial- + rollout paths). A silent field drop or enum corruption here means a + sample loses its status / spec_info / prefix_cache_info on the way + to the trainer with no crash signal. + + 2. ``update_from_meta_info`` finish_reason → Status enum mapping + (length→TRUNCATED, abort→ABORTED, stop→COMPLETED). The match + statement at types.py:176-182 is the only place the engine's + finish_reason gets translated; a typo'd enum or removed case here + would silently mis-tag every sample. +""" + +from __future__ import annotations + +import argparse + +import pytest + +from vime.utils.types import Sample + + +# --------------------------------------------------------------------------- +# to_dict / from_dict round-trip +# --------------------------------------------------------------------------- + + +def _make_sample(**overrides) -> Sample: + """Build a Sample with one non-default value per field-category so the + round-trip test exercises every code path in to_dict/from_dict, not + just the trivial defaults case.""" + base = dict( + group_index=0, + index=42, + rollout_id=7, + prompt="hello", + tokens=[1, 2, 3], + multimodal_inputs={"images": ["fake_url"]}, + response="world", + response_length=5, + label="42", + reward=0.75, + loss_mask=[1, 1, 0, 1, 1], + weight_versions=["v1"], + rollout_log_probs=[-0.1, -0.2], + rollout_routed_experts=[[0, 1], [2, 3]], + remove_sample=False, + teacher_log_probs=[-0.3, -0.4], + status=Sample.Status.COMPLETED, + metadata={"rm_type": "math"}, + generate_function_path="some.module.fn", + train_metadata={"loss_type": "policy_loss"}, + session_id="uuid-1234", + non_generation_time=1.5, + ) + base.update(overrides) + return Sample(**base) + + +@pytest.mark.unit +def test_to_dict_serializes_status_as_string_value(): + """The ``status`` field is an enum on the dataclass; ``to_dict`` must + flatten it to its string value so it survives JSON / pickle across + Ray boundaries.""" + sample = _make_sample() + d = sample.to_dict() + assert d["status"] == "completed" # not Sample.Status.COMPLETED + assert isinstance(d["status"], str) + + +@pytest.mark.unit +def test_to_dict_flattens_spec_info_and_prefix_cache_info(): + """``spec_info`` and ``prefix_cache_info`` are nested dataclasses; + to_dict converts each via its own to_dict (types.py:133-134).""" + sample = _make_sample() + sample.spec_info.spec_accept_token_num = 10 + sample.spec_info.spec_draft_token_num = 20 + sample.prefix_cache_info.cached_tokens = 5 + sample.prefix_cache_info.total_prompt_tokens = 50 + + d = sample.to_dict() + assert d["spec_info"] == { + "spec_accept_token_num": 10, + "spec_draft_token_num": 20, + "spec_verify_ct": 0, + "completion_token_num": 0, + } + assert d["prefix_cache_info"] == {"cached_tokens": 5, "total_prompt_tokens": 50} + + +@pytest.mark.unit +def test_round_trip_preserves_every_field(): + """Serialize → deserialize → compare. If any field gets silently + dropped on either side, the new sample won't equal the old. Uses + ``__dict__`` equality (not ``__eq__`` on the dataclass, which Sample + doesn't define) so nested SpecInfo / PrefixCacheInfo also get + compared structurally.""" + original = _make_sample() + original.spec_info.spec_accept_token_num = 3 + original.prefix_cache_info.cached_tokens = 7 + + restored = Sample.from_dict(original.to_dict()) + + # Status came back as the enum, not the string value. + assert restored.status is Sample.Status.COMPLETED + # Nested infos round-tripped as the correct type. + assert isinstance(restored.spec_info, Sample.SpecInfo) + assert isinstance(restored.prefix_cache_info, Sample.PrefixCacheInfo) + assert restored.spec_info.spec_accept_token_num == 3 + assert restored.prefix_cache_info.cached_tokens == 7 + + # All non-nested fields preserved. + for field in ( + "group_index", + "index", + "rollout_id", + "prompt", + "tokens", + "multimodal_inputs", + "response", + "response_length", + "label", + "reward", + "loss_mask", + "weight_versions", + "rollout_log_probs", + "rollout_routed_experts", + "remove_sample", + "teacher_log_probs", + "metadata", + "generate_function_path", + "train_metadata", + "session_id", + "non_generation_time", + ): + assert getattr(restored, field) == getattr(original, field), f"field {field} drifted" + + +@pytest.mark.unit +def test_from_dict_preserves_unknown_fields_as_attributes(): + """``from_dict`` keeps unknown keys as attributes (types.py:148-150), + not as dataclass fields. This is what lets newer rollout code stash + extra metadata that older trainers will simply ignore — a back-compat + contract worth pinning.""" + d = _make_sample().to_dict() + d["future_extension"] = "carried through" + + restored = Sample.from_dict(d) + assert restored.future_extension == "carried through" # type: ignore[attr-defined] + + +@pytest.mark.unit +def test_round_trip_through_default_constructed_sample(): + """A bare Sample (only defaults) must also round-trip — this is the + common case for a freshly-spawned rollout. Catches regressions where + ``from_dict`` requires a key that ``to_dict`` doesn't always emit.""" + original = Sample() + restored = Sample.from_dict(original.to_dict()) + assert restored.status is Sample.Status.PENDING + assert restored.tokens == [] + assert restored.metadata == {} + + +# --------------------------------------------------------------------------- +# update_from_meta_info — finish_reason → Status mapping +# --------------------------------------------------------------------------- + + +def _make_args(speculative: bool = False) -> argparse.Namespace: + """``update_from_meta_info`` only consults ``args.vllm_speculative_config`` + — minimal stub is enough.""" + return argparse.Namespace(vllm_speculative_config=speculative) + + +@pytest.mark.unit +@pytest.mark.parametrize( + "finish_reason,expected_status", + [ + ("length", Sample.Status.TRUNCATED), + ("abort", Sample.Status.ABORTED), + ("stop", Sample.Status.COMPLETED), + ], +) +def test_status_mapping_for_each_finish_reason(finish_reason, expected_status): + """The match statement at types.py:176-182 is the one place the engine's + finish_reason ever gets translated. Each branch must hit the right + enum; a typo in the enum name would crash later in unrelated places.""" + sample = Sample() + sample.update_from_meta_info( + _make_args(), + meta_info={"finish_reason": {"type": finish_reason}}, + ) + assert sample.status is expected_status + + +@pytest.mark.unit +def test_unknown_finish_reason_leaves_status_unchanged(): + """No ``case`` matches → status stays at whatever it was. Pins the + "no default clause means no-op" behavior so a future refactor adding + a default doesn't silently break this contract.""" + sample = Sample() + sample.status = Sample.Status.PENDING + sample.update_from_meta_info( + _make_args(), + meta_info={"finish_reason": {"type": "something_new"}}, + ) + assert sample.status is Sample.Status.PENDING + + +@pytest.mark.unit +def test_weight_version_is_appended_when_present(): + """``weight_version`` in meta_info is appended to the sample's list + (types.py:173-174) — partial-rollout uses this to track which model + version produced each chunk.""" + sample = Sample() + sample.weight_versions = ["v1"] + sample.update_from_meta_info( + _make_args(), + meta_info={ + "finish_reason": {"type": "stop"}, + "weight_version": "v2", + }, + ) + assert sample.weight_versions == ["v1", "v2"] + + +@pytest.mark.unit +def test_prefix_cache_info_is_accumulated_across_calls(): + """Every call to update_from_meta_info adds to prefix_cache_info + (types.py:171). Multi-turn rollouts call this once per turn — the + counts must accumulate, not overwrite.""" + sample = Sample() + for prompt_tokens, cached_tokens in [(100, 0), (200, 50)]: + sample.update_from_meta_info( + _make_args(), + meta_info={ + "finish_reason": {"type": "stop"}, + "prompt_tokens": prompt_tokens, + "cached_tokens": cached_tokens, + }, + ) + assert sample.prefix_cache_info.cached_tokens == 50 # 0 + 50 + assert sample.prefix_cache_info.total_prompt_tokens == 300 # 100 + 200 + + +@pytest.mark.unit +def test_spec_info_only_updated_when_speculative_enabled(): + """``spec_info.add`` is gated on ``args.vllm_speculative_config`` + (types.py:166-168). Without the flag, spec stats stay at zero even + if the engine sends them.""" + meta_info = { + "finish_reason": {"type": "stop"}, + "spec_accept_token_num": 7, + "spec_draft_token_num": 10, + } + + no_spec = Sample() + no_spec.update_from_meta_info(_make_args(speculative=False), meta_info=meta_info) + assert no_spec.spec_info.spec_accept_token_num == 0 + + with_spec = Sample() + with_spec.update_from_meta_info(_make_args(speculative=True), meta_info=meta_info) + assert with_spec.spec_info.spec_accept_token_num == 7 + assert with_spec.spec_info.spec_draft_token_num == 10 + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/tests/utils/test_hf_checkpoint_saver.py b/tests/utils/test_hf_checkpoint_saver.py new file mode 100644 index 000000000..7b8851537 --- /dev/null +++ b/tests/utils/test_hf_checkpoint_saver.py @@ -0,0 +1,75 @@ +import json +from pathlib import Path + +import pytest +import torch +from safetensors.torch import load_file + +from vime.backends.megatron_utils.hf_checkpoint_saver import ( + _clear_existing_hf_weights, + _copy_hf_assets, + _SafetensorShardWriter, +) + + +NUM_GPUS = 0 + + +def test_copy_hf_assets_keeps_quantized_config_and_skips_weights(tmp_path: Path): + src = tmp_path / "src" + dst = tmp_path / "dst" + src.mkdir() + dst.mkdir() + + config = {"model_type": "tiny", "quantization_config": {"quant_method": "fp8"}} + (src / "config.json").write_text(json.dumps(config), encoding="utf-8") + (src / "tokenizer.json").write_text("{}", encoding="utf-8") + (src / "model.safetensors.index.json").write_text("{}", encoding="utf-8") + (src / "model-00001-of-00001.safetensors").write_bytes(b"weight") + (src / "pytorch_model.bin").write_bytes(b"weight") + + _copy_hf_assets(str(src), dst) + + assert json.loads((dst / "config.json").read_text(encoding="utf-8")) == config + assert (dst / "tokenizer.json").exists() + assert not (dst / "model.safetensors.index.json").exists() + assert not (dst / "model-00001-of-00001.safetensors").exists() + assert not (dst / "pytorch_model.bin").exists() + + +def test_clear_existing_hf_weights_removes_old_weight_files_only(tmp_path: Path): + (tmp_path / "config.json").write_text("{}", encoding="utf-8") + (tmp_path / "model.safetensors.index.json").write_text("{}", encoding="utf-8") + (tmp_path / "model-00001-of-00001.safetensors").write_bytes(b"weight") + (tmp_path / "pytorch_model.bin").write_bytes(b"weight") + + _clear_existing_hf_weights(tmp_path) + + assert (tmp_path / "config.json").exists() + assert not (tmp_path / "model.safetensors.index.json").exists() + assert not (tmp_path / "model-00001-of-00001.safetensors").exists() + assert not (tmp_path / "pytorch_model.bin").exists() + + +def test_safetensor_shard_writer_writes_hf_index(tmp_path: Path): + writer = _SafetensorShardWriter(tmp_path, enabled=True) + writer.write([("layers.0.weight", torch.ones(2, 2)), ("layers.0.weight_scale", torch.ones(1))]) + writer.write([("layers.1.weight", torch.zeros(2, 2))]) + writer.finalize() + + index = json.loads((tmp_path / "model.safetensors.index.json").read_text(encoding="utf-8")) + assert index["metadata"]["total_size"] == 36 + assert index["weight_map"] == { + "layers.0.weight": "model-00001-of-00002.safetensors", + "layers.0.weight_scale": "model-00001-of-00002.safetensors", + "layers.1.weight": "model-00002-of-00002.safetensors", + } + + shard0 = load_file(tmp_path / "model-00001-of-00002.safetensors") + shard1 = load_file(tmp_path / "model-00002-of-00002.safetensors") + assert torch.equal(shard0["layers.0.weight"], torch.ones(2, 2)) + assert torch.equal(shard1["layers.1.weight"], torch.zeros(2, 2)) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/vime/agent/__init__.py b/vime/agent/__init__.py new file mode 100644 index 000000000..1048d52b8 --- /dev/null +++ b/vime/agent/__init__.py @@ -0,0 +1 @@ +"""Agent rollout building blocks.""" diff --git a/vime/agent/adapters/__init__.py b/vime/agent/adapters/__init__.py new file mode 100644 index 000000000..4f19a829f --- /dev/null +++ b/vime/agent/adapters/__init__.py @@ -0,0 +1,7 @@ +"""HTTP adapters for agent rollouts.""" + +from vime.agent.adapters.anthropic import AnthropicAdapter +from vime.agent.adapters.common import BaseAdapter +from vime.agent.adapters.openai import OpenAIAdapter + +__all__ = ["AnthropicAdapter", "BaseAdapter", "OpenAIAdapter"] diff --git a/vime/agent/adapters/anthropic.py b/vime/agent/adapters/anthropic.py new file mode 100644 index 000000000..f0a85f5c0 --- /dev/null +++ b/vime/agent/adapters/anthropic.py @@ -0,0 +1,481 @@ +"""Anthropic Messages adapter for agent rollouts. + +The adapter exposes ``/v1/messages`` and ``/v1/messages/count_tokens``. It +renders each Anthropic message history with the served model's chat template, +calls vLLM's ``/inference/v1/generate`` with ``token_ids``, and +records the exact sampled token ids/logprobs as ``TurnRecord`` objects. New +code should use ``AnthropicAdapter`` and call ``finish_session()`` at trajectory +end to drain trainable ``TokenSegment`` objects. + +It also handles Claude Code sub-agent and compaction patterns by splitting one +session into ``subagent``, ``wipe``, and ``final`` segments. +""" + +from __future__ import annotations + +import asyncio +import dataclasses +import json +import logging +import secrets +from typing import Any + +from aiohttp import web + +from vime.agent.adapters.common import ADAPTER_KEY, REASONING_PARSER_KEY, TOKENIZER_KEY, TOOL_PARSER_KEY +from vime.agent.adapters.common import AdapterChain as Chain +from vime.agent.adapters.common import ( + BaseAdapter, + call_vllm_generate, + ok_response, + render_token_ids, + request_session_id, +) +from vime.agent.adapters.common import stable_hash as _hash +from vime.agent.parsing import parse_model_output +from vime.agent.trajectory import TokenSegment, TurnRecord, TurnSegment, make_turn_segment, merge_turn_segments + +logger = logging.getLogger(__name__) + + +# Tool names claude-code uses to dispatch a sub-agent. +_SUBAGENT_TOOLS = {"Task", "Agent"} + + +@dataclasses.dataclass +class Session: + main: Chain = dataclasses.field(default_factory=Chain) + active_sub: Chain | None = None # at most one sub-agent at a time + pending_dispatch_id: str = "" # tool_use_id we're waiting to close + sampling_defaults: dict = dataclasses.field(default_factory=dict) + max_context_tokens: int = 0 + lock: asyncio.Lock = dataclasses.field(default_factory=asyncio.Lock) + segments: list[TurnSegment] = dataclasses.field(default_factory=list) # frozen output + + +class AnthropicAdapter(BaseAdapter): + """Anthropic Messages-compatible HTTP adapter with session lifecycle helpers.""" + + session_cls = Session + + def __init__(self, *, tokenizer, vllm_url, tool_parser=None, reasoning_parser=None) -> None: + super().__init__( + tokenizer=tokenizer, + vllm_url=vllm_url, + tool_parser=tool_parser, + reasoning_parser=reasoning_parser, + ) + self.app.router.add_post("/v1/messages", _handle_request) + self.app.router.add_post("/v1/messages/count_tokens", _count_tokens) + self.app.router.add_get("/healthz", _ok) + self.app.router.add_get("/v1/models", _ok) + + async def finish_session(self, sid: str, *, wait_timeout: float = 5.0) -> list[TokenSegment]: + await self.shutdown_session(sid, wait_timeout=wait_timeout) + s = self.store.pop(sid, None) + if s is None: + return [] + if s.active_sub is not None and s.active_sub.turns: + s.segments.append(make_turn_segment(s.active_sub.turns, kind="subagent")) + if s.main.turns: + s.segments.append(make_turn_segment(s.main.turns, kind="final")) + + return merge_turn_segments(s.segments) + + +# ============================================================================= +# 2. Per-turn stages +# ============================================================================= + + +def _select_chain(s: Session, body: dict) -> tuple[Chain, bool, str]: + """Decide which chain this turn operates on. + + 1. fingerprint body.messages and body.system into hashes + 2. if main now contains the tool_result for a pending sub dispatch, + snapshot the sub chain into s.segments and clear s.active_sub + 3. pick main vs s.active_sub based on whether request continues main's prefix + 4. classify as 'new' | 'append' | 'wipe' against the chosen target; + a wipe also snapshots the target's current state into s.segments + + Returns (target_chain, is_sub, kind). + """ + all_msgs = body.get("messages") or [] + msg_hashes = [_hash(m) for m in all_msgs] + req_system_hash = _hash(body.get("system")) if "system" in body else s.main.system_hash + + # Close active sub-agent if its dispatch tool_result has landed on main. + if s.pending_dispatch_id and s.active_sub is not None: + tu_id = s.pending_dispatch_id + for m in all_msgs: + if not isinstance(m, dict) or m.get("role") != "user": + continue + content = m.get("content") + if not isinstance(content, list): + continue + done = any( + isinstance(b, dict) and b.get("type") == "tool_result" and b.get("tool_use_id") == tu_id + for b in content + ) + if done: + if s.active_sub.turns: + s.segments.append(make_turn_segment(s.active_sub.turns, kind="subagent")) + s.active_sub = None + s.pending_dispatch_id = "" + break + + # Route: main iff request continues main's prefix. Sub system_hash can be + # "" (armed before sub dialled in), so never route by sub equality alone. + if s.active_sub is None: + target, is_sub = s.main, False + else: + main_continues = ( + req_system_hash == s.main.system_hash + and len(msg_hashes) >= s.main.seen_msgs + and msg_hashes[: s.main.seen_msgs] == s.main.msg_hashes[: s.main.seen_msgs] + ) + target, is_sub = (s.main, False) if main_continues else (s.active_sub, True) + + # Classify; snapshot a "wipe" segment first if we're discarding work. + if target.seen_msgs == 0: + kind = "new" + else: + is_append = ( + req_system_hash == target.system_hash + and len(msg_hashes) >= target.seen_msgs + and msg_hashes[: target.seen_msgs] == target.msg_hashes[: target.seen_msgs] + ) + if is_append: + kind = "append" + else: + if target.turns: + s.segments.append(make_turn_segment(target.turns, kind="wipe")) + kind = "wipe" + + return target, is_sub, kind + + +def _flatten(c: Any) -> str: + """Recursive Anthropic content flattener: text/tool_result(content) joined + by newline, images replaced with a placeholder.""" + if c is None: + return "" + if isinstance(c, str): + return c + if not isinstance(c, list): + return str(c) + parts: list[str] = [] + for b in c: + if isinstance(b, dict): + t = b.get("type") + if t == "text": + parts.append(b.get("text", "")) + elif t == "tool_result": + parts.append(_flatten(b.get("content"))) + elif t == "image": + parts.append("[image omitted]") + elif isinstance(b, str): + parts.append(b) + return "\n".join(p for p in parts if p) + + +def _translate_anthropic(msgs: list[dict], system: Any) -> list[dict]: + """Anthropic messages + system -> chat-template messages. Pure function.""" + translated: list[dict] = [] + if system: + translated.append({"role": "system", "content": _flatten(system)}) + for m in msgs: + if not isinstance(m, dict): + continue + role, content = m.get("role"), m.get("content") + if role == "user": + blocks = content if isinstance(content, list) else [{"type": "text", "text": _flatten(content)}] + for b in blocks: + if isinstance(b, dict) and b.get("type") == "tool_result": + translated.append({"role": "tool", "content": _flatten(b.get("content"))}) + elif isinstance(b, dict) and b.get("type") == "text": + translated.append({"role": "user", "content": b.get("text", "")}) + else: + translated.append({"role": "user", "content": _flatten(b)}) + elif role == "assistant": + texts, thinkings, tcs = [], [], [] + blocks = content if isinstance(content, list) else [{"type": "text", "text": _flatten(content)}] + for b in blocks: + if not isinstance(b, dict): + continue + if b.get("type") == "text": + texts.append(b.get("text", "")) + elif b.get("type") == "thinking": + thinkings.append(b.get("thinking", "")) + elif b.get("type") == "tool_use": + tcs.append({"function": {"name": b.get("name", "tool"), "arguments": b.get("input") or {}}}) + mo: dict[str, Any] = {"role": "assistant", "content": "".join(texts)} + if thinkings: + mo["reasoning_content"] = "".join(thinkings) + if tcs: + mo["tool_calls"] = tcs + translated.append(mo) + elif role == "system": + translated.append({"role": "system", "content": _flatten(content)}) + return translated + + +def _anthropic_tools_to_chat_tools(anth_tools: list[dict] | None) -> list[dict] | None: + """Convert Anthropic tools to tokenizer chat-template tool schema.""" + if not anth_tools: + return None + ts: list[dict] = [] + for t in anth_tools: + if not isinstance(t, dict) or "name" not in t: + continue + ts.append( + { + "type": "function", + "function": { + "name": t["name"], + "description": t.get("description", ""), + "parameters": t.get("input_schema") or t.get("parameters") or {"type": "object", "properties": {}}, + }, + } + ) + return ts or None + + +def _replace_chat_messages(target: Chain, body: dict) -> None: + """new/wipe: full reset of chat state and turn log.""" + all_msgs = body.get("messages") or [] + target.chat_messages = _translate_anthropic(all_msgs, body.get("system")) + if "system" in body: + target.system_hash = _hash(body.get("system")) + target.turns.clear() + target.seen_msgs = len(all_msgs) + target.msg_hashes = [_hash(m) for m in all_msgs] + if target.tools_schema is None: + target.tools_schema = _anthropic_tools_to_chat_tools(body.get("tools")) + + +def _extend_chat_messages(target: Chain, body: dict) -> None: + """append: translate only the new tail.""" + all_msgs = body.get("messages") or [] + translated = _translate_anthropic(all_msgs[target.seen_msgs :], None) + target.chat_messages.extend(translated) + + target.seen_msgs = len(all_msgs) + target.msg_hashes = [_hash(m) for m in all_msgs] + if target.tools_schema is None: + target.tools_schema = _anthropic_tools_to_chat_tools(body.get("tools")) + + +def _build_prompt(target: Chain, body: dict, kind: str, tok) -> list[int]: + """Replace/extend chat_messages and render input ids for vLLM.""" + (_extend_chat_messages if kind == "append" else _replace_chat_messages)(target, body) + return render_token_ids(target, tok) + + +async def _generate( + prompt_ids: list[int], s: Session, body: dict, app, *, session_id: str | None = None +) -> TurnRecord: + """Call vLLM and return a TurnRecord. + + 1. build sampling_params (session defaults overlaid with body overrides) + 2. POST vLLM ``/inference/v1/generate``; on cancel/error tear down + the request (vLLM has no per-request HTTP abort endpoint) + 3. keep the exact prompt/output token ids; trajectory merge later compares + later prompt tokens with earlier outputs to build the loss mask + """ + return await call_vllm_generate( + prompt_ids, + s, + body, + app, + max_token_keys=("max_tokens",), + stop_keys=("stop_sequences",), + log_prefix="anthropic_adapter", + logger=logger, + session_id=session_id, + ) + + +def _build_reply(target: Chain, output_ids: list[int], finish: str, app) -> tuple[list[dict], str, str]: + """Turn the model's raw output ids into the reply we send back to claude-code. + + 1. parse decoded text -> (thinking, visible, tool_uses) via parsers + 2. pack into Anthropic content blocks; tag dispatch_id when a tool_use + names Task/Agent (sub-agent trigger) + 3. derive stop_reason: 'tool_use' | 'max_tokens' | 'end_turn' + + Returns (blocks, stop_reason, dispatch_id). + """ + tok = app[TOKENIZER_KEY] + + raw_output = tok.decode(output_ids, skip_special_tokens=False) if output_ids else "" + parsed = parse_model_output( + raw_output, + tools_schema=target.tools_schema, + tool_parser_name=app[TOOL_PARSER_KEY], + reasoning_parser_name=app[REASONING_PARSER_KEY], + ) + blocks, dispatch_id = _anthropic_blocks(parsed.reasoning, parsed.text, parsed.tool_uses) + return blocks, _stop_reason(parsed.tool_uses, finish), dispatch_id + + +def _anthropic_blocks(thinking: str, visible: str, tool_uses: list[dict]) -> tuple[list[dict], str]: + """Pack parsed model output into Anthropic content blocks.""" + blocks: list[dict] = [] + if thinking: + blocks.append({"type": "thinking", "thinking": thinking}) + if visible: + blocks.append({"type": "text", "text": visible}) + dispatch_id = "" + for tu in tool_uses: + tu_id = f"toolu_{secrets.token_hex(8)}" + blocks.append({"type": "tool_use", "id": tu_id, "name": tu["name"], "input": tu["input"]}) + if tu["name"] in _SUBAGENT_TOOLS: + dispatch_id = tu_id + if not blocks: + blocks.append({"type": "text", "text": ""}) + return blocks, dispatch_id + + +def _stop_reason(tool_uses: list[dict], finish: str) -> str: + if tool_uses: + return "tool_use" + if finish == "length": + return "max_tokens" + return "end_turn" + + +def _start_sub_chain(s: Session, dispatch_id: str) -> None: + """Start a fresh sub chain on this session and remember the tool_use_id + we'll watch for on main to know when this sub is done. The matching + 'sub done' step lives inside _select_chain.""" + s.pending_dispatch_id = dispatch_id + if s.active_sub is None: + s.active_sub = Chain() + + +# ============================================================================= +# 3. Request handling -- one full turn + SSE wrap +# ============================================================================= + + +def _request_session_id(request: web.Request) -> str: + return request_session_id(request, include_x_api_key=True) + + +async def _handle_request(request: web.Request) -> web.StreamResponse: + body = await request.json() + sid = _request_session_id(request) + adapter = request.app[ADAPTER_KEY] + if sid in adapter.closed: # session drained; refuse stragglers + return web.Response(status=503, text="session closed") + app = request.app + s = adapter.store.setdefault(sid, Session()) + task = asyncio.current_task() + adapter.inflight.setdefault(sid, set()).add(task) + try: + async with s.lock: # same sid -> serialized + target, is_sub, kind = _select_chain(s, body) + ideal_ids = _build_prompt(target, body, kind, app[TOKENIZER_KEY]) + turn = await _generate(ideal_ids, s, body, app, session_id=sid) + blocks, stop, did = _build_reply(target, turn.output_ids, turn.finish_reason, app) + target.turns.append(turn) + if did and not is_sub: # sub doesn't nest + _start_sub_chain(s, did) + in_tok, out_tok = len(ideal_ids), len(turn.output_ids) + if body.get("stream") is True or "text/event-stream" in request.headers.get("Accept", ""): + return await _stream_response(request, blocks, stop, in_tok, out_tok) + return web.json_response(_message_response(body, blocks, stop, in_tok, out_tok)) + finally: + adapter.inflight.get(sid, set()).discard(task) + + +def _message_response(body: dict, blocks: list[dict], stop_reason: str, in_tok: int, out_tok: int) -> dict: + return { + "id": f"msg_{secrets.token_hex(12)}", + "type": "message", + "role": "assistant", + "model": body.get("model", "vime-actor"), + "content": blocks, + "stop_reason": stop_reason, + "stop_sequence": None, + "usage": {"input_tokens": in_tok, "output_tokens": out_tok}, + } + + +async def _stream_response(request, blocks, stop_reason, in_tok, out_tok) -> web.StreamResponse: + """Stream blocks back to claude-code as an Anthropic Messages SSE + response: message_start, (content_block_start, content_block_delta, + content_block_stop)*N, message_delta, message_stop.""" + out = web.StreamResponse( + status=200, + headers={ + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + "Connection": "keep-alive", + }, + ) + await out.prepare(request) + + # message_start + ms_data = { + "type": "message_start", + "message": { + "id": f"msg_{secrets.token_hex(12)}", + "type": "message", + "role": "assistant", + "model": "vime-actor", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": in_tok, "output_tokens": 0}, + }, + } + await out.write(f"event: message_start\ndata: {json.dumps(ms_data, ensure_ascii=False)}\n\n".encode()) + + for idx, block in enumerate(blocks): + bt = block["type"] + if bt == "thinking": + start = {"type": "thinking", "thinking": ""} + delta = {"type": "thinking_delta", "thinking": block["thinking"]} + elif bt == "text": + start = {"type": "text", "text": ""} + delta = {"type": "text_delta", "text": block["text"]} + else: # tool_use + start = {"type": "tool_use", "id": block["id"], "name": block["name"], "input": {}} + delta = { + "type": "input_json_delta", + "partial_json": json.dumps(block["input"], ensure_ascii=False), + } + + cbs_data = {"type": "content_block_start", "index": idx, "content_block": start} + await out.write(f"event: content_block_start\ndata: {json.dumps(cbs_data, ensure_ascii=False)}\n\n".encode()) + + cbd_data = {"type": "content_block_delta", "index": idx, "delta": delta} + await out.write(f"event: content_block_delta\ndata: {json.dumps(cbd_data, ensure_ascii=False)}\n\n".encode()) + + cbe_data = {"type": "content_block_stop", "index": idx} + await out.write(f"event: content_block_stop\ndata: {json.dumps(cbe_data, ensure_ascii=False)}\n\n".encode()) + + md_data = { + "type": "message_delta", + "delta": {"stop_reason": stop_reason, "stop_sequence": None}, + "usage": {"input_tokens": in_tok, "output_tokens": out_tok}, + } + await out.write(f"event: message_delta\ndata: {json.dumps(md_data, ensure_ascii=False)}\n\n".encode()) + + mst_data = {"type": "message_stop"} + await out.write(f"event: message_stop\ndata: {json.dumps(mst_data, ensure_ascii=False)}\n\n".encode()) + + return out + + +# Trivial endpoints claude-code probes during a session: count_tokens runs +# every turn (return 0 -- client uses it as a hint, not a hard budget), +# healthz/v1/models are startup readiness checks. +async def _count_tokens(request: web.Request) -> web.Response: + return web.json_response({"input_tokens": 0}) + + +async def _ok(request: web.Request) -> web.Response: + return await ok_response(request) diff --git a/vime/agent/adapters/common.py b/vime/agent/adapters/common.py new file mode 100644 index 000000000..14fb2e850 --- /dev/null +++ b/vime/agent/adapters/common.py @@ -0,0 +1,308 @@ +"""Shared adapter primitives for token-capturing agent rollouts.""" + +from __future__ import annotations + +import asyncio +import dataclasses +import hashlib +import json +import logging +import uuid +from collections.abc import Callable +from typing import Any + +import aiohttp +from aiohttp import web + +from vime.agent.trajectory import TokenSegment, TurnRecord + + +ADAPTER_KEY = web.AppKey("adapter", object) +TOKENIZER_KEY = web.AppKey("tokenizer", object) +VLLM_URL_KEY = web.AppKey("vllm_url", object) +TOOL_PARSER_KEY = web.AppKey("tool_parser", object) +REASONING_PARSER_KEY = web.AppKey("reasoning_parser", object) + + +@dataclasses.dataclass +class AdapterChain: + """Protocol-neutral chat chain state used by HTTP adapters.""" + + system_hash: str = "" + chat_messages: list[dict] = dataclasses.field(default_factory=list) + tools_schema: list[dict] | None = None + seen_msgs: int = 0 + msg_hashes: list[str] = dataclasses.field(default_factory=list) + turns: list[TurnRecord] = dataclasses.field(default_factory=list) + + +class BaseAdapter: + """Base HTTP adapter with per-instance session lifecycle state.""" + + session_cls: type + + def __init__(self, *, tokenizer, vllm_url, tool_parser=None, reasoning_parser=None) -> None: + self.store: dict[str, Any] = {} + self.inflight: dict[str, set[asyncio.Task]] = {} + self.closed: set[str] = set() + self.app = web.Application(client_max_size=64 * 1024 * 1024) + self.app[ADAPTER_KEY] = self + self.app[TOKENIZER_KEY] = tokenizer + self.app[VLLM_URL_KEY] = vllm_url.rstrip("/") if isinstance(vllm_url, str) else vllm_url + self.app[TOOL_PARSER_KEY] = tool_parser + self.app[REASONING_PARSER_KEY] = reasoning_parser + + def open_session( + self, + sid: str, + *, + sampling_defaults: dict | None = None, + max_context_tokens: int = 0, + ) -> None: + register_session( + self.store, + sid, + self.session_cls, + sampling_defaults=sampling_defaults, + max_context_tokens=max_context_tokens, + ) + + async def shutdown_session(self, sid: str, *, wait_timeout: float = 5.0) -> None: + await shutdown_session_tasks(sid, self.closed, self.inflight, wait_timeout=wait_timeout) + + async def finish_session(self, sid: str, *, wait_timeout: float = 5.0) -> list[TokenSegment]: + raise NotImplementedError + + +def strip_cache_control(obj: Any) -> Any: + if isinstance(obj, dict): + return {k: strip_cache_control(v) for k, v in obj.items() if k != "cache_control"} + if isinstance(obj, list): + return [strip_cache_control(x) for x in obj] + return obj + + +def stable_hash(obj: Any) -> str: + payload = json.dumps(strip_cache_control(obj), sort_keys=True, ensure_ascii=False, default=str).encode("utf-8") + return hashlib.sha1(payload).hexdigest()[:12] + + +def json_arguments(value: Any) -> str: + if value is None: + return "{}" + if isinstance(value, str): + return value + return json.dumps(value, ensure_ascii=False) + + +def render_token_ids(chain: AdapterChain, tokenizer) -> list[int]: + enc = tokenizer.apply_chat_template( + chain.chat_messages, + tools=chain.tools_schema, + tokenize=True, + add_generation_prompt=True, + ) + ids = enc["input_ids"] if hasattr(enc, "__getitem__") and "input_ids" in enc else enc + return list(ids) + + +def request_session_id( + request: web.Request, + *, + body: dict | None = None, + include_x_api_key: bool = False, +) -> str: + auth = request.headers.get("Authorization", "") + if auth.lower().startswith("bearer "): + sid = auth[7:].strip() + if sid: + return sid + + if body is not None: + metadata = body.get("metadata") + if isinstance(metadata, dict) and metadata.get("session_id"): + return str(metadata["session_id"]) + if body.get("user"): + return str(body["user"]) + + if include_x_api_key: + api_key = request.headers.get("X-Api-Key") + if api_key: + return api_key.strip() + + return "default" + + +def register_session( + store: dict[str, Any], + sid: str, + session_factory: Callable[[], Any], + *, + sampling_defaults: dict | None = None, + max_context_tokens: int = 0, +) -> None: + if sid in store: + raise ValueError(f"session_id {sid!r} already exists; sids must be unique per agent run") + session = store[sid] = session_factory() + session.sampling_defaults = dict(sampling_defaults or {}) + session.max_context_tokens = int(max_context_tokens or 0) + + +def _sampling_params(session: Any, body: dict, *, max_token_keys: tuple[str, ...], stop_keys: tuple[str, ...]) -> dict: + sp: dict[str, Any] = { + "skip_special_tokens": False, + "spaces_between_special_tokens": False, + "no_stop_trim": True, + "max_new_tokens": 4096, + **(session.sampling_defaults or {}), + } + + for key in max_token_keys: + if body.get(key) is not None: + sp["max_new_tokens"] = min(int(sp.get("max_new_tokens", body[key])), int(body[key])) + break + + for src_k, dst_k in (("temperature", "temperature"), ("top_p", "top_p"), ("top_k", "top_k")): + if src_k in body: + sp[dst_k] = body[src_k] + + for key in stop_keys: + if body.get(key): + sp["stop"] = body[key] + break + + return sp + + +def _vllm_sampling_body(sp: dict) -> dict: + """Map the canonical (sglang-shaped) sampling dict to a vLLM ``/inference/v1/generate`` + ``sampling_params`` body. vLLM uses ``max_tokens`` (not ``max_new_tokens``) and returns + per-token logprobs when ``logprobs`` is set.""" + body: dict[str, Any] = { + "max_tokens": int(sp.get("max_new_tokens", 4096)), + "logprobs": 1, + } + if "temperature" in sp: + body["temperature"] = sp["temperature"] + if "top_p" in sp: + body["top_p"] = sp["top_p"] + tk = sp.get("top_k") + if tk is not None and (tk > 0 or tk == -1): + body["top_k"] = tk + if sp.get("stop"): + body["stop"] = sp["stop"] + if sp.get("stop_token_ids"): + body["stop_token_ids"] = sp["stop_token_ids"] + if sp.get("skip_special_tokens") is not None: + body["skip_special_tokens"] = bool(sp["skip_special_tokens"]) + return body + + +def _tokens_and_logprobs_from_choice(choice: dict) -> tuple[list[int], list[float]]: + """Parse ``token_ids`` + ``logprobs.content[i].logprob`` from a vLLM + ``/inference/v1/generate`` choice. Mirrors vime ``_inference_generate_tokens_and_logprobs``.""" + tids_raw = choice.get("token_ids") + if not (isinstance(tids_raw, list) and tids_raw and all(isinstance(x, int) for x in tids_raw)): + return [], [] + tids = [int(x) for x in tids_raw] + lp = choice.get("logprobs") + if not isinstance(lp, dict): + return tids, [0.0] * len(tids) + content = lp.get("content") + if isinstance(content, list) and content: + lps: list[float] = [] + for i in range(len(tids)): + if i < len(content) and isinstance(content[i], dict): + lps.append(float(content[i].get("logprob", 0.0))) + else: + lps.append(0.0) + return tids, lps + return tids, [0.0] * len(tids) + + +async def call_vllm_generate( + prompt_ids: list[int], + session: Any, + body: dict, + app, + *, + max_token_keys: tuple[str, ...], + stop_keys: tuple[str, ...], + log_prefix: str, + logger: logging.Logger, + session_id: str | None = None, +) -> TurnRecord: + sp = _sampling_params(session, body, max_token_keys=max_token_keys, stop_keys=stop_keys) + + if session.max_context_tokens > 0: + remaining_context = session.max_context_tokens - len(prompt_ids) + if remaining_context <= 0: + logger.warning( + "[%s] prompt exceeds max_context_tokens (%d >= %d)", + log_prefix, + len(prompt_ids), + session.max_context_tokens, + ) + return TurnRecord(prompt_ids=list(prompt_ids), output_ids=[], finish_reason="length") + sp["max_new_tokens"] = min(int(sp.get("max_new_tokens", remaining_context)), remaining_context) + + vllm_url = app[VLLM_URL_KEY] + payload: dict[str, Any] = { + "token_ids": list(prompt_ids), + "sampling_params": _vllm_sampling_body(sp), + } + # session_id routes via vllm-router's consistent_hash policy (x-session-id header); + # see vime ``vllm_rollout.py`` headers handling. + headers = {"x-session-id": session_id} if session_id and session_id != "default" else None + timeout = aiohttp.ClientTimeout(total=None, sock_read=900) + task = asyncio.current_task() + try: + async with aiohttp.ClientSession(timeout=timeout) as sess, sess.post( + f"{vllm_url}/inference/v1/generate", + json=payload, + headers=headers, + ) as r: + if r.status >= 400: + text = await r.text() + raise RuntimeError(f"vllm upstream {r.status}: {text[:400]}") + data = await r.json(content_type=None) + choice = (data.get("choices") or [{}])[0] + output_ids, output_log_probs = _tokens_and_logprobs_from_choice(choice) + fr = choice.get("finish_reason") + finish = fr if isinstance(fr, str) and fr else "stop" + except (asyncio.CancelledError, aiohttp.ClientError, asyncio.TimeoutError): + # vLLM ``/inference/v1/generate`` has no per-request HTTP abort endpoint (unlike + # sglang ``/abort_request``). Cancelling the in-flight task tears down the aiohttp + # request, which drops the streaming connection so vLLM stops generating. + if task is not None: + task.cancel() + raise + + return TurnRecord( + prompt_ids=list(prompt_ids), + output_ids=output_ids, + finish_reason=finish, + output_log_probs=output_log_probs, + ) + + +async def shutdown_session_tasks( + sid: str, + closed: set[str], + inflight: dict[str, set[asyncio.Task]], + *, + wait_timeout: float = 5.0, +) -> None: + closed.add(sid) + tasks = [t for t in inflight.pop(sid, ()) if not t.done()] + if not tasks: + return + _, pending = await asyncio.wait(tasks, timeout=wait_timeout) + for task in pending: + task.cancel() + if pending: + await asyncio.gather(*pending, return_exceptions=True) + + +async def ok_response(request: web.Request) -> web.Response: + return web.json_response({"ok": True}) diff --git a/vime/agent/adapters/openai.py b/vime/agent/adapters/openai.py new file mode 100644 index 000000000..5042d1dd6 --- /dev/null +++ b/vime/agent/adapters/openai.py @@ -0,0 +1,572 @@ +"""OpenAI-compatible adapters for agent rollouts. + +The adapter exposes ``/v1/chat/completions`` and ``/v1/responses``. Both +endpoints render incoming messages with the served model's chat template, call +vLLM's ``/inference/v1/generate`` with ``token_ids``, and record +the exact sampled token ids/logprobs as ``TurnRecord`` objects. New code should +use ``OpenAIAdapter`` and call ``finish_session()`` at trajectory end to drain +trainable ``TokenSegment`` objects. +""" + +from __future__ import annotations + +import asyncio +import dataclasses +import json +import logging +import secrets +import time +from typing import Any + +from aiohttp import web + +from vime.agent.adapters.common import ADAPTER_KEY, REASONING_PARSER_KEY, TOKENIZER_KEY, TOOL_PARSER_KEY +from vime.agent.adapters.common import AdapterChain as Chain +from vime.agent.adapters.common import BaseAdapter, call_vllm_generate +from vime.agent.adapters.common import json_arguments as _json_arguments +from vime.agent.adapters.common import ok_response, render_token_ids, request_session_id +from vime.agent.adapters.common import stable_hash as _hash +from vime.agent.parsing import ParsedModelOutput, parse_model_output +from vime.agent.trajectory import TokenSegment, TurnRecord, TurnSegment, make_turn_segment, merge_turn_segments + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class Session: + main: Chain = dataclasses.field(default_factory=Chain) + sampling_defaults: dict = dataclasses.field(default_factory=dict) + max_context_tokens: int = 0 + lock: asyncio.Lock = dataclasses.field(default_factory=asyncio.Lock) + segments: list[TurnSegment] = dataclasses.field(default_factory=list) + + +class OpenAIAdapter(BaseAdapter): + """OpenAI-compatible HTTP adapter with session lifecycle helpers.""" + + session_cls = Session + + def __init__(self, *, tokenizer, vllm_url, tool_parser=None, reasoning_parser=None) -> None: + super().__init__( + tokenizer=tokenizer, + vllm_url=vllm_url, + tool_parser=tool_parser, + reasoning_parser=reasoning_parser, + ) + self.app.router.add_post("/v1/chat/completions", _handle_chat_completions) + self.app.router.add_post("/v1/responses", _handle_responses) + self.app.router.add_get("/healthz", _ok) + self.app.router.add_get("/v1/models", _ok) + + async def finish_session(self, sid: str, *, wait_timeout: float = 5.0) -> list[TokenSegment]: + await self.shutdown_session(sid, wait_timeout=wait_timeout) + s = self.store.pop(sid, None) + if s is None: + return [] + if s.main.turns: + s.segments.append(make_turn_segment(s.main.turns, kind="final")) + return merge_turn_segments(s.segments) + + +def _flatten_content(content: Any) -> str: + """Flatten OpenAI text/content parts into a chat-template string.""" + if content is None: + return "" + if isinstance(content, str): + return content + if not isinstance(content, list): + return str(content) + + parts: list[str] = [] + for item in content: + if isinstance(item, str): + parts.append(item) + continue + if not isinstance(item, dict): + parts.append(str(item)) + continue + typ = item.get("type") + if typ in {"text", "input_text", "output_text"}: + parts.append(item.get("text", "")) + elif typ in {"image_url", "input_image"}: + parts.append("[image omitted]") + elif "content" in item: + parts.append(_flatten_content(item.get("content"))) + elif "text" in item: + parts.append(str(item.get("text") or "")) + return "\n".join(p for p in parts if p) + + +def _normalize_tool_call(call: dict[str, Any]) -> dict[str, Any]: + function = call.get("function") or {} + name = function.get("name") or call.get("name") or "tool" + arguments = function.get("arguments", call.get("arguments", {})) + out = { + "type": "function", + "function": { + "name": name, + "arguments": _json_arguments(arguments), + }, + } + if call.get("id"): + out["id"] = call["id"] + return out + + +def _translate_chat_messages(messages: list[dict]) -> list[dict]: + """OpenAI chat messages -> tokenizer chat-template messages.""" + translated: list[dict] = [] + for msg in messages: + if not isinstance(msg, dict): + continue + role = msg.get("role") + content = msg.get("content") + if role == "developer": + role = "system" + + if role in {"system", "user"}: + translated.append({"role": role, "content": _flatten_content(content)}) + elif role == "tool": + tool_msg = {"role": "tool", "content": _flatten_content(content)} + if msg.get("tool_call_id"): + tool_msg["tool_call_id"] = msg["tool_call_id"] + translated.append(tool_msg) + elif role == "assistant": + assistant: dict[str, Any] = {"role": "assistant", "content": _flatten_content(content)} + if msg.get("reasoning_content"): + assistant["reasoning_content"] = msg["reasoning_content"] + tool_calls = msg.get("tool_calls") or [] + if tool_calls: + assistant["tool_calls"] = [_normalize_tool_call(c) for c in tool_calls if isinstance(c, dict)] + translated.append(assistant) + return translated + + +def _normalize_tool(tool: dict[str, Any]) -> dict[str, Any] | None: + if not isinstance(tool, dict): + return None + if tool.get("type") != "function": + return None + if isinstance(tool.get("function"), dict): + function = tool["function"] + name = function.get("name") + if not name: + return None + return { + "type": "function", + "function": { + "name": name, + "description": function.get("description", ""), + "parameters": function.get("parameters") or {"type": "object", "properties": {}}, + }, + } + name = tool.get("name") + if not name: + return None + return { + "type": "function", + "function": { + "name": name, + "description": tool.get("description", ""), + "parameters": tool.get("parameters") or {"type": "object", "properties": {}}, + }, + } + + +def _normalize_tools(tools: list[dict] | None) -> list[dict] | None: + normalized = [_normalize_tool(t) for t in tools or []] + return [t for t in normalized if t is not None] or None + + +def _responses_input_to_messages(input_value: Any, instructions: Any = None) -> list[dict]: + """Responses API input -> OpenAI chat message list. + + This intentionally covers the common message/function-call shapes used by + agent SDKs. Unknown input items are preserved as user text where possible. + """ + messages: list[dict] = [] + if instructions: + messages.append({"role": "system", "content": _flatten_content(instructions)}) + + if isinstance(input_value, str): + messages.append({"role": "user", "content": input_value}) + return messages + + if not isinstance(input_value, list): + messages.append({"role": "user", "content": _flatten_content(input_value)}) + return messages + + for item in input_value: + if isinstance(item, str): + messages.append({"role": "user", "content": item}) + continue + if not isinstance(item, dict): + messages.append({"role": "user", "content": str(item)}) + continue + + typ = item.get("type") + if typ == "function_call_output": + messages.append( + { + "role": "tool", + "tool_call_id": item.get("call_id") or item.get("id") or "", + "content": item.get("output", ""), + } + ) + elif typ == "function_call": + messages.append( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": item.get("call_id") or item.get("id") or f"call_{secrets.token_hex(8)}", + "type": "function", + "function": { + "name": item.get("name", "tool"), + "arguments": item.get("arguments", "{}"), + }, + } + ], + } + ) + elif item.get("role"): + messages.append({"role": item.get("role"), "content": item.get("content", "")}) + elif typ == "message": + messages.append({"role": item.get("role", "user"), "content": item.get("content", "")}) + else: + messages.append({"role": "user", "content": _flatten_content(item)}) + return messages + + +def _select_kind(s: Session, messages: list[dict]) -> str: + target = s.main + msg_hashes = [_hash(m) for m in messages] + if target.seen_msgs == 0: + kind = "new" + else: + is_append = len(msg_hashes) >= target.seen_msgs and msg_hashes[: target.seen_msgs] == target.msg_hashes + if is_append: + kind = "append" + else: + if target.turns: + s.segments.append(make_turn_segment(target.turns, kind="wipe")) + kind = "wipe" + return kind + + +def _replace_chat_messages(target: Chain, messages: list[dict], tools_schema: list[dict] | None) -> None: + target.chat_messages = _translate_chat_messages(messages) + target.turns.clear() + target.seen_msgs = len(messages) + target.msg_hashes = [_hash(m) for m in messages] + if tools_schema is not None: + target.tools_schema = tools_schema + + +def _extend_chat_messages(target: Chain, messages: list[dict], tools_schema: list[dict] | None) -> None: + translated = _translate_chat_messages(messages[target.seen_msgs :]) + target.chat_messages.extend(translated) + target.seen_msgs = len(messages) + target.msg_hashes = [_hash(m) for m in messages] + if tools_schema is not None: + target.tools_schema = tools_schema + + +def _build_prompt(target: Chain, messages: list[dict], tools_schema: list[dict] | None, kind: str, tok) -> list[int]: + (_extend_chat_messages if kind == "append" else _replace_chat_messages)(target, messages, tools_schema) + return render_token_ids(target, tok) + + +async def _generate( + prompt_ids: list[int], s: Session, body: dict, app, *, session_id: str | None = None +) -> TurnRecord: + return await call_vllm_generate( + prompt_ids, + s, + body, + app, + max_token_keys=("max_output_tokens", "max_completion_tokens", "max_tokens"), + stop_keys=("stop",), + log_prefix="openai_adapter", + logger=logger, + session_id=session_id, + ) + + +def _parse_turn(target: Chain, turn: TurnRecord, app) -> ParsedModelOutput: + tok = app[TOKENIZER_KEY] + raw_output = tok.decode(turn.output_ids, skip_special_tokens=False) if turn.output_ids else "" + return parse_model_output( + raw_output, + tools_schema=target.tools_schema, + tool_parser_name=app[TOOL_PARSER_KEY], + reasoning_parser_name=app[REASONING_PARSER_KEY], + ) + + +def _openai_tool_calls(tool_uses: list[dict[str, Any]]) -> list[dict[str, Any]]: + calls: list[dict[str, Any]] = [] + for tool_use in tool_uses: + call_id = f"call_{secrets.token_hex(12)}" + calls.append( + { + "id": call_id, + "type": "function", + "function": { + "name": tool_use.get("name", "tool"), + "arguments": _json_arguments(tool_use.get("input") or {}), + }, + } + ) + return calls + + +def _finish_reason(parsed: ParsedModelOutput, finish: str) -> str: + if parsed.tool_uses: + return "tool_calls" + if finish == "length": + return "length" + return "stop" + + +def _chat_message(parsed: ParsedModelOutput) -> dict[str, Any]: + tool_calls = _openai_tool_calls(parsed.tool_uses) + message: dict[str, Any] = { + "role": "assistant", + "content": parsed.text if parsed.text else None, + } + if parsed.reasoning: + message["reasoning_content"] = parsed.reasoning + if tool_calls: + message["tool_calls"] = tool_calls + return message + + +def _usage(in_tok: int, out_tok: int) -> dict[str, int]: + return { + "prompt_tokens": in_tok, + "completion_tokens": out_tok, + "total_tokens": in_tok + out_tok, + } + + +def _responses_usage(in_tok: int, out_tok: int) -> dict[str, int]: + return { + "input_tokens": in_tok, + "output_tokens": out_tok, + "total_tokens": in_tok + out_tok, + } + + +def _request_session_id(request: web.Request, body: dict) -> str: + return request_session_id(request, body=body) + + +async def _run_turn( + request: web.Request, body: dict, messages: list[dict] +) -> tuple[TurnRecord, ParsedModelOutput, int, int]: + sid = _request_session_id(request, body) + adapter = request.app[ADAPTER_KEY] + if sid in adapter.closed: + raise web.HTTPServiceUnavailable(text="session closed") + app = request.app + s = adapter.store.setdefault(sid, Session()) + task = asyncio.current_task() + adapter.inflight.setdefault(sid, set()).add(task) + try: + async with s.lock: + target = s.main + tools_schema = _normalize_tools(body.get("tools")) + kind = _select_kind(s, messages) + prompt_ids = _build_prompt(target, messages, tools_schema, kind, app[TOKENIZER_KEY]) + turn = await _generate(prompt_ids, s, body, app, session_id=sid) + parsed = _parse_turn(target, turn, app) + target.turns.append(turn) + return turn, parsed, len(prompt_ids), len(turn.output_ids) + finally: + adapter.inflight.get(sid, set()).discard(task) + + +async def _handle_chat_completions(request: web.Request) -> web.StreamResponse: + body = await request.json() + messages = body.get("messages") or [] + if not isinstance(messages, list): + raise web.HTTPBadRequest(text="messages must be a list") + turn, parsed, in_tok, out_tok = await _run_turn(request, body, messages) + if body.get("stream"): + return await _stream_chat_completion(request, body, parsed, turn.finish_reason, in_tok, out_tok) + return web.json_response(_chat_completion_response(body, parsed, turn.finish_reason, in_tok, out_tok)) + + +def _chat_completion_response( + body: dict, + parsed: ParsedModelOutput, + finish: str, + in_tok: int, + out_tok: int, +) -> dict[str, Any]: + return { + "id": f"chatcmpl_{secrets.token_hex(12)}", + "object": "chat.completion", + "created": int(time.time()), + "model": body.get("model", "vime-actor"), + "choices": [ + { + "index": 0, + "message": _chat_message(parsed), + "finish_reason": _finish_reason(parsed, finish), + } + ], + "usage": _usage(in_tok, out_tok), + } + + +async def _stream_chat_completion( + request: web.Request, + body: dict, + parsed: ParsedModelOutput, + finish: str, + in_tok: int, + out_tok: int, +) -> web.StreamResponse: + out = web.StreamResponse( + status=200, + headers={ + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + "Connection": "keep-alive", + }, + ) + await out.prepare(request) + completion_id = f"chatcmpl_{secrets.token_hex(12)}" + created = int(time.time()) + + async def emit(choice_delta: dict[str, Any], finish_reason: str | None = None, usage: dict | None = None) -> None: + chunk = { + "id": completion_id, + "object": "chat.completion.chunk", + "created": created, + "model": body.get("model", "vime-actor"), + "choices": [{"index": 0, "delta": choice_delta, "finish_reason": finish_reason}], + } + if usage is not None: + chunk["usage"] = usage + await out.write(f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n".encode()) + + await emit({"role": "assistant"}) + if parsed.reasoning: + await emit({"reasoning_content": parsed.reasoning}) + if parsed.text: + await emit({"content": parsed.text}) + for idx, call in enumerate(_openai_tool_calls(parsed.tool_uses)): + await emit({"tool_calls": [{**call, "index": idx}]}) + await emit({}, finish_reason=_finish_reason(parsed, finish), usage=_usage(in_tok, out_tok)) + await out.write(b"data: [DONE]\n\n") + return out + + +async def _handle_responses(request: web.Request) -> web.StreamResponse: + body = await request.json() + messages = _responses_input_to_messages(body.get("input", ""), body.get("instructions")) + turn, parsed, in_tok, out_tok = await _run_turn(request, body, messages) + if body.get("stream"): + return await _stream_response(request, body, parsed, turn.finish_reason, in_tok, out_tok) + return web.json_response(_response_response(body, parsed, turn.finish_reason, in_tok, out_tok)) + + +def _response_output(parsed: ParsedModelOutput) -> list[dict[str, Any]]: + output: list[dict[str, Any]] = [] + if parsed.reasoning: + output.append( + { + "id": f"rs_{secrets.token_hex(12)}", + "type": "reasoning", + "summary": [{"type": "summary_text", "text": parsed.reasoning}], + } + ) + if parsed.text: + output.append( + { + "id": f"msg_{secrets.token_hex(12)}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": parsed.text, "annotations": []}], + } + ) + for call in _openai_tool_calls(parsed.tool_uses): + output.append( + { + "id": f"fc_{secrets.token_hex(12)}", + "type": "function_call", + "status": "completed", + "call_id": call["id"], + "name": call["function"]["name"], + "arguments": call["function"]["arguments"], + } + ) + if not output: + output.append( + { + "id": f"msg_{secrets.token_hex(12)}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "", "annotations": []}], + } + ) + return output + + +def _response_response( + body: dict, + parsed: ParsedModelOutput, + finish: str, + in_tok: int, + out_tok: int, +) -> dict[str, Any]: + status = "incomplete" if finish == "length" else "completed" + return { + "id": f"resp_{secrets.token_hex(12)}", + "object": "response", + "created_at": int(time.time()), + "status": status, + "model": body.get("model", "vime-actor"), + "output": _response_output(parsed), + "usage": _responses_usage(in_tok, out_tok), + } + + +async def _stream_response( + request: web.Request, + body: dict, + parsed: ParsedModelOutput, + finish: str, + in_tok: int, + out_tok: int, +) -> web.StreamResponse: + out = web.StreamResponse( + status=200, + headers={ + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + "Connection": "keep-alive", + }, + ) + await out.prepare(request) + response = _response_response(body, parsed, finish, in_tok, out_tok) + created = {"type": "response.created", "response": response} + await out.write(f"event: response.created\ndata: {json.dumps(created, ensure_ascii=False)}\n\n".encode()) + if parsed.text: + delta = {"type": "response.output_text.delta", "delta": parsed.text} + await out.write( + f"event: response.output_text.delta\ndata: {json.dumps(delta, ensure_ascii=False)}\n\n".encode() + ) + completed = {"type": "response.completed", "response": response} + await out.write(f"event: response.completed\ndata: {json.dumps(completed, ensure_ascii=False)}\n\n".encode()) + return out + + +async def _ok(request: web.Request) -> web.Response: + return await ok_response(request) diff --git a/vime/agent/parsing.py b/vime/agent/parsing.py new file mode 100644 index 000000000..db1913608 --- /dev/null +++ b/vime/agent/parsing.py @@ -0,0 +1,113 @@ +"""Model-output parsing helpers for agent harnesses.""" + +from __future__ import annotations + +import dataclasses +import json +import logging +import re +from typing import Any + + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass(frozen=True) +class ParsedModelOutput: + """Structured view of one decoded model output.""" + + reasoning: str + text: str + tool_uses: list[dict[str, Any]] + + +def parse_model_output( + raw_output: str, + *, + tools_schema: list[dict] | None, + tool_parser_name: str | None, + reasoning_parser_name: str | None, +) -> ParsedModelOutput: + """Parse raw model text into reasoning, visible text, and tool uses. + + When ``reasoning_parser_name`` / ``tool_parser_name`` are set the heavy + format-specific work is delegated to SGLang's reasoning and function-call + parsers (imported lazily, so they are only required when explicitly + enabled). The coding-agent example leaves both unset and relies on the XML + fallback, which covers the Anthropic-style tool-call text that some + coding-agent models still emit occasionally. + """ + reasoning, body_text = "", raw_output + if reasoning_parser_name: + from sglang.srt.parser.reasoning_parser import ReasoningParser + + r, b = ReasoningParser( + model_type=reasoning_parser_name, + stream_reasoning=False, + ).parse_non_stream(raw_output) + reasoning, body_text = r or "", b or "" + if not reasoning and "" in body_text: + reasoning, body_text = body_text.split("", 1) + + body_text, tool_uses = parse_tool_uses(body_text, tools_schema, tool_parser_name) + return ParsedModelOutput( + reasoning=reasoning, + text=(body_text or "").strip(), + tool_uses=tool_uses, + ) + + +def parse_tool_uses( + body_text: str, + tools_schema: list[dict] | None, + tool_parser_name: str | None, +) -> tuple[str, list[dict[str, Any]]]: + """Parse tool calls from body text and return visible text plus tool uses.""" + tool_uses: list[dict[str, Any]] = [] + if tool_parser_name and tools_schema: + from sglang.srt.entrypoints.openai.protocol import Function, Tool + from sglang.srt.function_call.function_call_parser import FunctionCallParser + + sg_tools = [Tool(type="function", function=Function(**d["function"])) for d in tools_schema] + parser = FunctionCallParser(tools=sg_tools, tool_call_parser=tool_parser_name) + calls = [] + if parser.has_tool_call(body_text): + try: + body_text, calls = parser.parse_non_stream(body_text) + except Exception: + logger.exception("[agent.parsing] sglang tool-call parsing failed; falling back") + for c in calls: + try: + args = json.loads(c.parameters or "{}") + except json.JSONDecodeError: + args = {"_raw_arguments": c.parameters} + tool_uses.append({"name": c.name or "tool", "input": args}) + + if not tool_uses and tools_schema: + body_text, tool_uses = parse_xml_tool_uses(body_text, tools_schema) + + return body_text, tool_uses + + +def parse_xml_tool_uses(body_text: str, tools_schema: list[dict]) -> tuple[str, list[dict[str, Any]]]: + """Fallback parser for Anthropic-style XML tool calls.""" + valid_tools = {t.get("function", {}).get("name") for t in tools_schema} + tool_uses: list[dict[str, Any]] = [] + cleaned_parts: list[str] = [] + last = 0 + for m in re.finditer( + r"\s*]+)>(.*?)\s*", + body_text, + flags=re.DOTALL, + ): + name, inner = m.group(1), m.group(2) + if name in valid_tools: + args = { + p.group(1): p.group(2).strip() + for p in re.finditer(r"]+)>(.*?)", inner, flags=re.DOTALL) + } + tool_uses.append({"name": name, "input": args}) + cleaned_parts.append(body_text[last : m.start()]) + last = m.end() + cleaned_parts.append(body_text[last:]) + return "".join(cleaned_parts), tool_uses diff --git a/vime/agent/sandbox.py b/vime/agent/sandbox.py new file mode 100644 index 000000000..6447bd2ee --- /dev/null +++ b/vime/agent/sandbox.py @@ -0,0 +1,281 @@ +"""Sandbox backends for agent rollouts. + +The public sandbox contract is intentionally small: async context management, +command execution, and file read/write. Agent examples can build task-specific +setup, runner, and evaluator logic on top of this without depending directly on +one sandbox provider. +""" + +from __future__ import annotations + +import asyncio +import io +import json +import logging +import os +from pathlib import Path +from typing import Protocol, runtime_checkable + +logger = logging.getLogger(__name__) + + +ExecResult = tuple[int, str, str] +FileContent = str | bytes | Path + + +@runtime_checkable +class Sandbox(Protocol): + """Minimal async sandbox interface used by agent rollouts. + + ``write_file`` accepts either in-memory content (``str``/``bytes``) or a + host ``Path`` to stream into the sandbox. + """ + + sandbox_id: str + + async def __aenter__(self) -> Sandbox: ... + + async def __aexit__(self, exc_type, exc, tb) -> None: ... + + async def exec( + self, + cmd: str, + *, + user: str = "root", + env: dict[str, str] | None = None, + timeout: int = 120, + check: bool = False, + ) -> ExecResult: ... + + async def write_file(self, sandbox_path: str, content: FileContent, *, user: str = "root") -> None: ... + + async def read_file(self, sandbox_path: str, *, user: str = "root") -> str: ... + + +def _getenv(*names: str, default: str = "") -> str: + for name in names: + value = os.environ.get(name) + if value is not None and value.strip(): + return value + return default + + +class E2BSandbox: + """Async context manager around e2b.AsyncSandbox.""" + + metadata_file_env = ("VIME_AGENT_SANDBOX_METADATA_FILE", "SWE_SANDBOX_METADATA_FILE") + metadata_json_env = ("VIME_AGENT_SANDBOX_METADATA_JSON", "SWE_SANDBOX_METADATA_JSON") + image_metadata_key_env = ("VIME_AGENT_SANDBOX_IMAGE_METADATA_KEY", "SWE_SANDBOX_IMAGE_METADATA_KEY") + lifetime_sec_env = ("VIME_AGENT_SANDBOX_LIFETIME_SEC", "SWE_SANDBOX_LIFETIME_SEC") + rpc_retries_env = ("VIME_AGENT_SANDBOX_RPC_RETRIES", "SWE_RPC_RETRIES") + + default_lifetime_sec = 3600 + default_rpc_retries = 3 + # With retries=3 the sleep budget is 3s, which handles common E2B h2 reset + # / SSL / pool-timeout flaps without stalling rollout steps for too long. + rpc_backoff_base_sec = 1.0 + + def __init__( + self, + image: str, + *, + timeout: int | None = None, + metadata: dict[str, str] | None = None, + image_metadata_key: str | None = None, + rpc_retries: int | None = None, + ) -> None: + self.image = image + self.timeout = timeout if timeout is not None else self._lifetime_sec_from_env() + self.metadata = dict(metadata) if metadata is not None else self._metadata_from_env() + self.image_metadata_key = image_metadata_key or self._image_metadata_key_from_env() + self.rpc_retries = rpc_retries if rpc_retries is not None else self._rpc_retries_from_env() + self._sb = None + self.sandbox_id = "" + + @classmethod + def _metadata_from_env(cls) -> dict[str, str]: + """Read E2B routing metadata from file or JSON environment values.""" + file_path = _getenv(*cls.metadata_file_env) + raw = "" + if file_path: + try: + raw = Path(file_path).read_text() + except OSError as e: + logger.warning("[agent.sandbox] metadata file %s unreadable: %s", file_path, e) + raw = "" + if not raw: + raw = _getenv(*cls.metadata_json_env) + if not raw: + return {} + try: + md = json.loads(raw) + except json.JSONDecodeError as e: + logger.warning("[agent.sandbox] metadata not valid JSON, ignoring: %s", e) + return {} + if not isinstance(md, dict): + logger.warning("[agent.sandbox] metadata must be a JSON object, got %s", type(md).__name__) + return {} + return {str(k): str(v) for k, v in md.items()} + + @classmethod + def _image_metadata_key_from_env(cls) -> str | None: + return _getenv(*cls.image_metadata_key_env) or None + + @classmethod + def _lifetime_sec_from_env(cls) -> int: + return int(_getenv(*cls.lifetime_sec_env, default=str(cls.default_lifetime_sec))) + + @classmethod + def _rpc_retries_from_env(cls) -> int: + return int(_getenv(*cls.rpc_retries_env, default=str(cls.default_rpc_retries))) + + @staticmethod + def _is_transient_rpc_error(e: BaseException) -> bool: + """True if e is a transient E2B client-side failure safe to retry.""" + name = type(e).__name__ + if name in { + "ProtocolError", + "LocalProtocolError", + "WriteError", + "ReadError", + "ConnectError", + "ConnectTimeout", + "ReadTimeout", + "WriteTimeout", + "PoolTimeout", + "RemoteProtocolError", + "SSLError", + }: + return True + msg = str(e) + if name == "SandboxException": + if "does not exist" in msg or "STOPPED state" in msg: + return False + return True + return False + + async def _rpc_retry(self, op_name: str, coro_factory): + """Run coro_factory() with retries for transient E2B RPC failures.""" + last_err = None + for attempt in range(self.rpc_retries): + try: + return await coro_factory() + except Exception as e: + if not self._is_transient_rpc_error(e): + raise + last_err = e + if attempt + 1 < self.rpc_retries: + backoff = self.rpc_backoff_base_sec * (2**attempt) + logger.debug( + "[agent.sandbox] %s transient %s, retry %d/%d in %.1fs: %s", + op_name, + type(e).__name__, + attempt + 1, + self.rpc_retries, + backoff, + str(e)[:120], + ) + await asyncio.sleep(backoff) + assert last_err is not None + raise last_err + + async def __aenter__(self) -> E2BSandbox: + if self.image_metadata_key is None: + raise RuntimeError( + "VIME_AGENT_SANDBOX_IMAGE_METADATA_KEY is not set. Export it " + "to the metadata key your E2B gateway uses for image routing. " + "The legacy SWE_SANDBOX_IMAGE_METADATA_KEY name is also " + "accepted for coding-agent examples." + ) + from e2b import AsyncSandbox # type: ignore + + md = dict(self.metadata) + md.setdefault(self.image_metadata_key, self.image) + self._sb = await AsyncSandbox.create(timeout=self.timeout, metadata=md) + self.sandbox_id = self._sb.sandbox_id + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + try: + if self._sb is not None: + await self._sb.kill() + except Exception as e: + logger.warning("[agent.sandbox] kill %s failed: %s", self.sandbox_id[:8], e) + + async def exec( + self, + cmd: str, + *, + user: str = "root", + env: dict[str, str] | None = None, + timeout: int = 120, + check: bool = False, + ) -> ExecResult: + from e2b.sandbox.commands.command_handle import CommandExitException + + try: + res = await self._rpc_retry( + f"exec({cmd[:60]!r})", + lambda: self._sb.commands.run( + cmd, + user=user, + envs=env, + timeout=timeout, + on_stdout=lambda s: None, + on_stderr=lambda s: None, + ), + ) + return res.exit_code, res.stdout or "", res.stderr or "" + except CommandExitException as e: + if check: + raise RuntimeError( + f"e2b exec failed (exit={e.exit_code}): {cmd[:120]}\n{(e.stderr or '')[:400]}" + ) from None + return e.exit_code, e.stdout or "", e.stderr or "" + + async def write_file(self, sandbox_path: str, content: FileContent, *, user: str = "root") -> None: + if isinstance(content, Path): + host_path = content + + async def _do_path(): + with open(host_path, "rb") as fp: + await self._sb.files.write( + sandbox_path, + fp, + user=user, + gzip=False, + use_octet_stream=True, + request_timeout=600, + ) + + await self._rpc_retry(f"write_file({sandbox_path} <- {host_path.name})", _do_path) + return + + if isinstance(content, bytes): + + async def _do_bytes(): + await self._sb.files.write( + sandbox_path, + io.BytesIO(content), + user=user, + gzip=False, + use_octet_stream=True, + request_timeout=600, + ) + + await self._rpc_retry(f"write_file({sandbox_path}, bytes={len(content)})", _do_bytes) + return + + await self._rpc_retry( + f"write_file({sandbox_path})", + lambda: self._sb.files.write(sandbox_path, content, user=user), + ) + + async def read_file(self, sandbox_path: str, *, user: str = "root") -> str: + try: + return await self._rpc_retry( + f"read_file({sandbox_path})", + lambda: self._sb.files.read(sandbox_path, user=user), + ) + except Exception: + return "" diff --git a/vime/agent/trajectory.py b/vime/agent/trajectory.py new file mode 100644 index 000000000..b30b69080 --- /dev/null +++ b/vime/agent/trajectory.py @@ -0,0 +1,208 @@ +"""Token-level trajectory helpers for agent rollouts.""" + +from __future__ import annotations + +import copy +import dataclasses +import logging +from typing import Any + +from vime.utils.types import Sample + + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass(frozen=True) +class TurnRecord: + """Exact token snapshot for one assistant generation. + + ``prompt_ids`` is the full tokenized prompt sent to the generator for that + turn. ``output_ids`` is the raw generated output, and + ``output_log_probs`` is aligned with it when the rollout engine returns + per-token log probabilities. + """ + + prompt_ids: list[int] + output_ids: list[int] + finish_reason: str + output_log_probs: list[float] = dataclasses.field(default_factory=list) + + +@dataclasses.dataclass(frozen=True) +class TokenSegment: + """One training segment assembled from an agent trajectory.""" + + prompt_ids: list[int] + response_ids: list[int] + loss_mask: list[int] + rollout_log_probs: list[float] = dataclasses.field(default_factory=list) + metadata: dict[str, Any] = dataclasses.field(default_factory=dict) + + +@dataclasses.dataclass(frozen=True) +class TurnSegment: + """A frozen group of turns before token-level merge.""" + + turns: list[TurnRecord] + metadata: dict[str, Any] = dataclasses.field(default_factory=dict) + + +def make_turn_segment( + turns: list[TurnRecord], + *, + kind: str = "", + metadata: dict[str, Any] | None = None, +) -> TurnSegment: + """Freeze turns and attach conventional segment metadata.""" + frozen_turns = list(turns) + segment_metadata = dict(metadata or {}) + if kind: + segment_metadata.setdefault("segment_kind", kind) + segment_metadata.setdefault("finish_reason", frozen_turns[-1].finish_reason if frozen_turns else "") + return TurnSegment(turns=frozen_turns, metadata=segment_metadata) + + +def _common_prefix_len(a: list[int], b: list[int]) -> int: + n = min(len(a), len(b)) + i = 0 + while i < n and a[i] == b[i]: + i += 1 + return i + + +def _output_log_probs(turn: TurnRecord) -> list[float]: + if len(turn.output_log_probs) == len(turn.output_ids): + return list(turn.output_log_probs) + logger.warning( + "[trajectory] turn logprob length mismatch; zeroing output logprobs (%d ids, %d logprobs)", + len(turn.output_ids), + len(turn.output_log_probs), + ) + return [0.0] * len(turn.output_ids) + + +def merge_turns(turns: list[TurnRecord], *, metadata: dict[str, Any] | None = None) -> TokenSegment | None: + """Replay turn records into one linear training segment. + + The first turn's prompt becomes the segment prompt. Later turn prompts are + stitched against ``prompt + response_so_far``. Any new prompt suffix is + non-model context and receives loss mask 0. If a later prompt diverges + inside a previous model output, the retained prefix of that whole output + turn is also masked out, because partial token matches are not a faithful + training target for that turn. + """ + if not turns: + return None + + prompt_ids = list(turns[0].prompt_ids) + response_ids: list[int] = [] + loss_mask: list[int] = [] + rollout_log_probs: list[float] = [] + output_spans: list[tuple[int, int]] = [] + + for i, turn in enumerate(turns): + if i > 0: + if turn.prompt_ids[: len(prompt_ids)] != prompt_ids: + logger.warning("[trajectory] merge prompt base changed; starting segment from drifted prompt") + prompt_ids = list(turn.prompt_ids) + response_ids = [] + loss_mask = [] + rollout_log_probs = [] + output_spans = [] + else: + prompt_suffix = turn.prompt_ids[len(prompt_ids) :] + matched_len = _common_prefix_len(response_ids, prompt_suffix) + if matched_len < len(response_ids): + logger.warning( + "[trajectory] merge prefix drift; truncating %d unstitched response tokens", + len(response_ids) - matched_len, + ) + for start, end in output_spans: + if start < matched_len < end: + loss_mask[start:matched_len] = [0] * (matched_len - start) + rollout_log_probs[start:matched_len] = [0.0] * (matched_len - start) + response_ids = response_ids[:matched_len] + loss_mask = loss_mask[:matched_len] + rollout_log_probs = rollout_log_probs[:matched_len] + output_spans = [ + (start, min(end, matched_len)) for start, end in output_spans if start < matched_len + ] + + context_tail = prompt_suffix[matched_len:] + response_ids.extend(context_tail) + loss_mask.extend([0] * len(context_tail)) + rollout_log_probs.extend([0.0] * len(context_tail)) + + output_start = len(response_ids) + response_ids.extend(turn.output_ids) + loss_mask.extend([1] * len(turn.output_ids)) + rollout_log_probs.extend(_output_log_probs(turn)) + output_spans.append((output_start, len(response_ids))) + + rollout_log_probs = [logprob if mask else 0.0 for logprob, mask in zip(rollout_log_probs, loss_mask, strict=True)] + + return TokenSegment( + prompt_ids=prompt_ids, + response_ids=response_ids, + loss_mask=loss_mask, + rollout_log_probs=rollout_log_probs, + metadata=dict(metadata or {}), + ) + + +def merge_turn_segments(segments: list[TurnSegment]) -> list[TokenSegment]: + """Merge frozen turn segments and keep every non-empty output.""" + out: list[TokenSegment] = [] + for turn_segment in segments: + token_segment = merge_turns(turn_segment.turns, metadata=turn_segment.metadata) + if token_segment is None: + continue + if token_segment.response_ids: + out.append(token_segment) + return out + + +def write_segment_to_sample(sample: Sample, segment: TokenSegment, reward: float, tokenizer) -> None: + """Populate token, mask, response, reward, and status fields from a segment.""" + sample.tokens = list(segment.prompt_ids) + list(segment.response_ids) + sample.response_length = len(segment.response_ids) + sample.loss_mask = list(segment.loss_mask) + sample.rollout_log_probs = list(segment.rollout_log_probs) + sample.response = tokenizer.decode(segment.response_ids, skip_special_tokens=False) + sample.reward = float(reward) + sample.status = Sample.Status.COMPLETED + + +def fan_out_sample_segments( + sample: Sample, + segments: list[TokenSegment], + reward: float, + tokenizer, + *, + metadata: dict[str, Any] | None = None, + rollout_id: int | None = None, +) -> list[Sample]: + """Emit one Sample per segment, splitting reward uniformly across them. + + Sibling samples share ``rollout_id`` so reducers that average by rollout do + not over-count trajectories split by compaction or sub-agent dispatch. + """ + k = len(segments) + per_segment_reward = float(reward) / max(1, k) + shared_rollout_id = getattr(sample, "index", None) if rollout_id is None else rollout_id + base_metadata = {**(sample.metadata or {}), **(metadata or {})} + + out: list[Sample] = [] + for i, segment in enumerate(segments): + sub = sample if i == 0 else copy.copy(sample) + write_segment_to_sample(sub, segment, per_segment_reward, tokenizer) + sub.rollout_id = shared_rollout_id + sub.metadata = { + **base_metadata, + **(segment.metadata or {}), + "segment_idx": i, + "num_segments": k, + } + out.append(sub) + return out diff --git a/vime/backends/megatron_utils/actor.py b/vime/backends/megatron_utils/actor.py index 8b656cb3a..5bb7c8a78 100644 --- a/vime/backends/megatron_utils/actor.py +++ b/vime/backends/megatron_utils/actor.py @@ -71,9 +71,6 @@ def init( self.tokenizer = AutoTokenizer.from_pretrained(self.args.hf_checkpoint, trust_remote_code=True) dist.barrier(group=get_gloo_group()) - self.train_parallel_config = { - "dp_size": mpu.get_data_parallel_world_size(with_context_parallel=False), - } dist.barrier(group=get_gloo_group()) if args.offload_train: @@ -85,6 +82,20 @@ def init( args, role ) + vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1 + if vpp_size > 1: + from megatron.core.utils import get_model_config + + microbatch_group_size_per_vp_stage = get_model_config(self.model[0]).microbatch_group_size_per_vp_stage + else: + microbatch_group_size_per_vp_stage = 1 + self.train_parallel_config = { + "dp_size": mpu.get_data_parallel_world_size(with_context_parallel=False), + "cp_size": mpu.get_context_parallel_world_size(), + "vpp_size": vpp_size, + "microbatch_group_size_per_vp_stage": microbatch_group_size_per_vp_stage, + } + start_rollout_id = loaded_rollout_id + 1 if role == "critic": @@ -204,6 +215,12 @@ def _get_rollout_data(self, rollout_data_ref: Box) -> RolloutBatch: rollout_data["loss_masks"] = [ torch.tensor(t, dtype=torch.int, device=torch.cuda.current_device()) for t in rollout_data["loss_masks"] ] + if "rollout_mask_sums" in rollout_data: + # Promote precomputed per-rollout mask totals to GPU tensors here + # (matching loss_masks) so the loss reducer can just divide. + rollout_data["rollout_mask_sums"] = torch.tensor( + rollout_data["rollout_mask_sums"], dtype=torch.float32, device=torch.cuda.current_device() + ) if "multimodal_train_inputs" in rollout_data: # Move multimodal training tensors to GPU in advance rollout_data["multimodal_train_inputs"] = [ @@ -387,7 +404,9 @@ def train(self, rollout_id: int, rollout_data_ref: Box, external_data=None): def train_critic(self, rollout_id: int, rollout_data: RolloutBatch): """Train critic and return CPU values (used as old-values for the next actor train).""" - data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data) + data_iterator = get_data_iterator(rollout_data) + num_microbatches = rollout_data["num_microbatches"] + global_batch_sizes = rollout_data["global_batch_sizes"] # Compute current critic values (used as old_values for value loss and for actor advantages). rollout_data.update(forward_only(get_values, self.args, self.model, data_iterator, num_microbatches)) @@ -402,6 +421,7 @@ def train_critic(self, rollout_id: int, rollout_data: RolloutBatch): self.opt_param_scheduler, data_iterator, num_microbatches, + global_batch_sizes, ) if mpu.is_pipeline_last_stage() and "values" in rollout_data: @@ -412,7 +432,9 @@ def train_critic(self, rollout_id: int, rollout_data: RolloutBatch): def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data=None) -> None: # Create data iterator for log_probs and train. - data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data) + data_iterator = get_data_iterator(rollout_data) + num_microbatches = rollout_data["num_microbatches"] + global_batch_sizes = rollout_data["global_batch_sizes"] if self.args.use_rollout_routing_replay: self.fill_routing_replay(data_iterator, num_microbatches, rollout_data) @@ -509,6 +531,7 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data self.opt_param_scheduler, data_iterator, num_microbatches, + global_batch_sizes, ) self.prof.step(rollout_id=rollout_id) diff --git a/vime/backends/megatron_utils/cp_utils.py b/vime/backends/megatron_utils/cp_utils.py index b0eab9b6a..448c154c6 100644 --- a/vime/backends/megatron_utils/cp_utils.py +++ b/vime/backends/megatron_utils/cp_utils.py @@ -54,21 +54,37 @@ def get_sum_of_sample_mean( total_lengths: list[int], response_lengths: list[int], loss_masks: list[torch.Tensor], + sample_denoms: list[torch.Tensor] | torch.Tensor | None = None, calculate_per_token_loss: bool = False, qkv_format: str = "thd", max_seq_lens: list[int] | None = None, ) -> Callable[[torch.Tensor], torch.Tensor]: """ - Calculate correct sample mean for CP + Calculate correct sample mean for CP. + + The default (``sample_denoms=None``) is the legacy per-sample mean: each + sample's denominator is its own ``loss_mask.sum()``. Callers that want a + per-rollout token-weighted mean pass pre-computed per-sample denominators + (already as GPU tensors — see actor side) where every sample in the same + rollout group carries the same value (the sum of that rollout's mask + totals across every sibling sample in the step). Pre-computing at the + step level rather than per-mb is required — otherwise a rollout whose + samples land in different micro-batches would get a partial denominator + on each side. """ + if sample_denoms is None: + sample_denoms = [m.sum() for m in loss_masks] + cp_size = mpu.get_context_parallel_world_size() if cp_size == 1: def sum_of_sample_mean(x: torch.Tensor) -> torch.Tensor: return sum( [ - (x_i * loss_mask_i).sum() / torch.clamp_min(loss_mask_i.sum(), 1) - for x_i, loss_mask_i in zip(x.split(response_lengths, dim=0), loss_masks, strict=False) + (x_i * loss_mask_i).sum() / torch.clamp_min(denom, 1) + for x_i, loss_mask_i, denom in zip( + x.split(response_lengths, dim=0), loss_masks, sample_denoms, strict=False + ) ] ) @@ -100,9 +116,9 @@ def sum_of_token(x: torch.Tensor) -> torch.Tensor: def sum_of_sample_mean(x: torch.Tensor) -> torch.Tensor: return sum( [ - (x_i * chunked_loss_mask).sum() / torch.clamp_min(loss_mask.sum(), 1) - for x_i, chunked_loss_mask, loss_mask in zip( - x.split(cp_chunk_lengths, dim=0), chunked_loss_masks, loss_masks, strict=False + (x_i * chunked_loss_mask).sum() / torch.clamp_min(denom, 1) + for x_i, chunked_loss_mask, denom in zip( + x.split(cp_chunk_lengths, dim=0), chunked_loss_masks, sample_denoms, strict=False ) ] ) @@ -120,6 +136,114 @@ def sum_of_token(x: torch.Tensor) -> torch.Tensor: return sum_of_sample_mean if not calculate_per_token_loss else sum_of_token +def reduce_train_step_metrics( + losses_reduced: list[dict], + *, + calculate_per_token_loss: bool, + step_global_batch_size: int, + cp_size: int, + dp_with_cp_group, +) -> dict[str, float]: + """Aggregate per-mb log dicts into the dict ``train_one_step`` reports. + + Pipeline (1:1 with what the train loop used to do inline): + 1. Sum each metric's per-mb ``values`` tensor locally on this rank. + 2. All-reduce across the DP*CP group (``dp_with_cp_group``). + 3. Apply the per-mode divisor / cp_factor: + - per-token-loss: divisor = ``values[0]`` = all-reduced ``num_tokens``, + CP-inflated by ``cp_size`` because every CP rank computes the same + num_tokens off the FULL (not chunked) masks; the + ``cp_factor = cp_size`` multiplier cancels that inflation, leaving + the genuine per-token average. + - per-rollout-mean: divisor = constant ``step_global_batch_size`` from + the rollout side, never all-reduced, so no CP inflation to cancel + and ``cp_factor = 1``. + + Tests pass a mock ``dp_with_cp_group`` and monkeypatch ``dist.all_reduce`` + to a no-op, then pre-aggregate virtual ranks themselves — this exercises + the same call shape as production while staying single-process. + """ + keys = losses_reduced[0]["keys"] + values = None + for x in losses_reduced: + values = x["values"] if values is None else values + x["values"] + assert len(keys) + 1 == values.numel() + dist.all_reduce(values, group=dp_with_cp_group) + values = values.tolist() + + if calculate_per_token_loss: + num_samples_or_tokens = values[0] + cp_factor = cp_size + else: + num_samples_or_tokens = step_global_batch_size + cp_factor = 1 + return {key: value * cp_factor / num_samples_or_tokens for key, value in zip(keys, values[1:], strict=False)} + + +def rollout_log_metric_contribution( + per_rank_reducer_sum: float, + *, + cp_size: int, + num_rollouts_in_rollout: int, + dp_size: int, +) -> tuple[float, float]: + """``(sum, count)`` tuple to hand the gather step for a per-rollout-mean + metric on the rollout side (``log_rollout_data``). + + Sum across DP*CP ranks of ``count`` lands on ``num_rollouts_in_rollout`` + (``dp_size`` here is the no-CP DP width; the gather covers ``dp_size * + cp_size`` ranks, and each rank emits the same ``count``, so the totals + cancel out the ``cp_size`` in the sum). Result: ``Σsum / Σcount = + sum_DP_full / num_rollouts`` — the same number ``train_one_step`` reports + for the same samples (when ``num_steps_per_rollout == 1``). + + Pair with :func:`gather_and_reduce_log_dict` to do the full end-to-end + in tests (single helper call per rank, returns the reduced number on + the source rank). + """ + sum_value = cp_size * per_rank_reducer_sum + count = num_rollouts_in_rollout / dp_size + return sum_value, count + + +def gather_and_reduce_log_dict( + log_dict: dict, + *, + dp_size: int, + dp_src_rank: int, + dp_group, +) -> dict | None: + """``dist.gather_object`` per-rank log_dicts + per-key reduction. + + Per key in the gathered dicts: + - ``(sum, count)`` tuple → ``Σsum / Σcount`` (per-rollout-mean shape; + pair with :func:`rollout_log_metric_contribution`). + - plain value → ``Σ / dp_size`` (legacy mean-across-ranks; the only + correct answer when ranks hold the same data). + + Returns the reduced dict on ``dp_src_rank``, ``None`` elsewhere. The + caller adds whatever metric-name prefix / wandb plumbing it wants — + this helper stays free of side effects so CPU multi-process unit tests + can drive it directly with real ``torch.distributed``. + """ + if dist.get_rank() == dp_src_rank: + gathered = [None] * dp_size + dist.gather_object(log_dict, gathered, dst=dp_src_rank, group=dp_group) + reduced: dict = {} + for key in log_dict: + values = [d[key] for d in gathered] + first = values[0] + if isinstance(first, tuple) and len(first) == 2: + total_sum = sum(v[0] for v in values) + total_count = sum(v[1] for v in values) + reduced[key] = total_sum / total_count if total_count else 0.0 + else: + reduced[key] = sum(values) / dp_size + return reduced + dist.gather_object(log_dict, None, dst=dp_src_rank, group=dp_group) + return None + + def all_gather_with_cp(tensor: torch.Tensor, total_length: int, response_length: int) -> torch.Tensor: """ Gather tensors across all ranks in the context parallel group. diff --git a/vime/backends/megatron_utils/data.py b/vime/backends/megatron_utils/data.py index dee37edf8..42c19e7e6 100644 --- a/vime/backends/megatron_utils/data.py +++ b/vime/backends/megatron_utils/data.py @@ -10,14 +10,17 @@ from megatron.core.packed_seq_params import PackedSeqParams from vime.utils import train_metric_utils -from vime.utils.data import get_minimum_num_micro_batch_size from vime.utils.flops_utils import calculate_fwd_flops from vime.utils.metric_utils import compute_pass_rate, compute_rollout_step -from vime.utils.seqlen_balancing import get_seqlen_balanced_partitions from vime.utils.types import RolloutBatch from ...utils import logging_utils -from .cp_utils import get_sum_of_sample_mean, slice_with_cp +from .cp_utils import ( + gather_and_reduce_log_dict, + get_sum_of_sample_mean, + rollout_log_metric_contribution, + slice_with_cp, +) logger = logging.getLogger(__name__) @@ -53,9 +56,6 @@ def get_batch( assert "tokens" in keys batch = data_iterator.get_next(keys) - if "dynamic_global_batch_size" in data_iterator.rollout_data: - batch["dynamic_global_batch_size"] = data_iterator.rollout_data["dynamic_global_batch_size"] - tokens = batch["tokens"] # use 0 as the pad token id should be fine? pad_token_id = 0 @@ -180,105 +180,69 @@ def gather_log_data( metric_name: str, args: Namespace, rollout_id: int, - log_dict: dict[str, float], + log_dict: dict[str, "float | tuple[float, float]"], ) -> dict[str, float] | None: """ - Gather per-rank metrics, reduce by mean on the DP source rank, and log. - - Expects `log_dict` to contain plain scalars. The DP source rank prints and - optionally logs to WandB/TensorBoard with a step derived from `rollout_id` and - batch sizes. Returns the reduced dict on the DP source rank; returns None on others. - """ - - if mpu.get_data_parallel_rank(with_context_parallel=True) == 0: - dp_size = mpu.get_data_parallel_world_size(with_context_parallel=True) - - gathered_log_dict = [None] * dp_size - # Not sure if this will be a performance bottleneck. - dist.gather_object( - log_dict, - gathered_log_dict, - dst=mpu.get_data_parallel_src_rank(with_context_parallel=True), - group=mpu.get_data_parallel_group_gloo(with_context_parallel=True), - ) - - reduced_log_dict = { - f"{metric_name}/{key}": sum([d[key] for d in gathered_log_dict]) / dp_size for key in log_dict - } - logger.info(f"{metric_name} {rollout_id}: {reduced_log_dict}") + Gather per-rank metrics, reduce on the DP source rank, and log to W&B / TB. - # Calculate step once to avoid duplication - step = compute_rollout_step(args, rollout_id) - reduced_log_dict["rollout/step"] = step - logging_utils.log(args, reduced_log_dict, step_key="rollout/step") + Each value in ``log_dict`` is either: + * a ``(sum, count)`` tuple → reduced as ``Σsum / Σcount``; + * a plain scalar → reduced as ``Σ / dp_size`` (mean across ranks). - return reduced_log_dict - else: - dist.gather_object( - log_dict, - None, - dst=mpu.get_data_parallel_src_rank(with_context_parallel=True), - group=mpu.get_data_parallel_group_gloo(with_context_parallel=True), - ) + The gather + reduce step is delegated to + :func:`cp_utils.gather_and_reduce_log_dict` so it can be exercised by + CPU multi-process unit tests directly. This function adds the + ``metric_name`` prefix and the W&B / TB logging side effects. + """ + reduced = gather_and_reduce_log_dict( + log_dict, + dp_size=mpu.get_data_parallel_world_size(with_context_parallel=True), + dp_src_rank=mpu.get_data_parallel_src_rank(with_context_parallel=True), + dp_group=mpu.get_data_parallel_group_gloo(with_context_parallel=True), + ) + if reduced is None: return None + reduced_log_dict = {f"{metric_name}/{k}": v for k, v in reduced.items()} + logger.info(f"{metric_name} {rollout_id}: {reduced_log_dict}") + # Calculate step once to avoid duplication + step = compute_rollout_step(args, rollout_id) + reduced_log_dict["rollout/step"] = step + logging_utils.log(args, reduced_log_dict, step_key="rollout/step") + return reduced_log_dict class DataIterator: - """Micro-batch iterator over rollout dicts. - - Supports either fixed contiguous micro-batches or an explicit per-step - index schedule (for dynamic batch sizing / sequence-length balancing). - """ + """Iterator over a rollout dict following an explicit micro-batch index schedule.""" def __init__( self, rollout_data: RolloutBatch, - micro_batch_size: int | None = None, - micro_batch_indices: list[list[int]] | None = None, + micro_batch_indices: list[list[int]], ) -> None: - """Initialize an iterator over `rollout_data`. + """Initialize an iterator over ``rollout_data``. Args: - rollout_data: Dict of per-sample fields for the local step. - micro_batch_size: Fixed contiguous slice size when not using dynamic scheduling. - micro_batch_indices: Explicit indices per micro-batch when using dynamic balancing. - Must be mutually exclusive with `micro_batch_size`. + rollout_data: Dict of per-sample fields for this DP rank. + micro_batch_indices: List of mbs, each mbs being the local sample indices to select. """ self.rollout_data = rollout_data - self.micro_batch_size = micro_batch_size self.micro_batch_indices = micro_batch_indices - assert micro_batch_size is None or micro_batch_indices is None self.offset = 0 def get_next(self, keys: Sequence[str]) -> dict[str, list[object] | None]: """Return the next micro-batch for the requested keys. - - If `micro_batch_indices` is provided, selects rows according to the current - index list for each requested key. - - Otherwise, slices a contiguous window of size `micro_batch_size` starting - at the current offset. - Returns a dict mapping each key to a list subset (or None if absent). """ batch = {} + indices = self.micro_batch_indices[self.offset] for key in keys: vals = self.rollout_data.get(key, None) if vals is None: batch[key] = None else: - if self.micro_batch_indices is not None: - indices = self.micro_batch_indices[self.offset] - batch[key] = [vals[i] for i in indices] - else: - assert self.offset + self.micro_batch_size <= len( - vals - ), f"offset: {self.offset}, micro_batch_size: {self.micro_batch_size}, len(vals): {len(vals)}" - batch[key] = vals[self.offset : self.offset + self.micro_batch_size] - - if self.micro_batch_indices is not None: - self.offset += 1 - else: - self.offset += self.micro_batch_size + batch[key] = [vals[i] for i in indices] + self.offset += 1 return batch def reset(self) -> "DataIterator": @@ -287,102 +251,11 @@ def reset(self) -> "DataIterator": return self -def get_data_iterator( - args: Namespace, - model: torch.nn.Module | Sequence[torch.nn.Module], - rollout_data: RolloutBatch, -) -> tuple[list[DataIterator], list[int]]: - """ - Create iterators and a micro-batch schedule for a rollout step. - - - If `use_dynamic_batch_size` is False, splits into fixed-size contiguous - micro-batches of `micro_batch_size`. - - If True, computes the number of micro-batches per local step based on - `max_tokens_per_gpu` and per-sample lengths, all-reduces to a DP-wide - maximum, optionally enforces divisibility for Virtual Pipeline Parallelism (VPP), and builds a balanced - index schedule to equalize token counts across micro-batches. - - Returns `(data_iterators, num_microbatches)` where: - - `data_iterators`: list of `DataIterator`, one per VPP stage (size 1 if VPP disabled) - - `num_microbatches`: list[int], one per local step in the rollout (length = steps) - """ - dp_size = mpu.get_data_parallel_world_size(with_context_parallel=False) - dp_group = mpu.get_data_parallel_group() - vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size() - if vpp_size is None: - vpp_size = 1 - if vpp_size > 1: - from megatron.core.utils import get_model_config - - config = get_model_config(model[0]) - microbatch_group_size_per_vp_stage = config.microbatch_group_size_per_vp_stage - cp_size = mpu.get_context_parallel_world_size() - - num_local_samples = len(rollout_data["total_lengths"]) - global_batch_size = rollout_data.get("dynamic_global_batch_size", args.global_batch_size) - num_local_gbs = global_batch_size // dp_size - num_steps_per_rollout = num_local_samples // num_local_gbs - - if global_batch_size != args.global_batch_size: - logger.info( - f"Using dynamic global_batch_size={global_batch_size} (original={args.global_batch_size}), " - f"num_local_samples={num_local_samples}, num_steps_per_rollout={num_steps_per_rollout}" - ) - - def _generate_data_iterator(rollout_data, micro_batch_size, micro_batch_indices=None): - data_iterator = [] - for _ in range(vpp_size): - data_iterator.append(DataIterator(rollout_data, micro_batch_size, micro_batch_indices)) - return data_iterator - - if not args.use_dynamic_batch_size: - num_microbatches = [num_local_gbs // args.micro_batch_size for _ in range(num_steps_per_rollout)] - data_iterator = _generate_data_iterator(rollout_data, args.micro_batch_size) - else: - assert args.max_tokens_per_gpu is not None - # calculate the number of mirobatches for each step - samples = rollout_data["total_lengths"] - assert len(samples) == num_local_samples - num_microbatches = [] - for i in range(num_steps_per_rollout): - start, end = i * num_local_gbs, (i + 1) * num_local_gbs - num_microbatches.append( - get_minimum_num_micro_batch_size(samples[start:end], args.max_tokens_per_gpu * cp_size) - ) - - num_microbatches = torch.tensor(num_microbatches, dtype=torch.int, device=torch.cuda.current_device()) - dist.all_reduce(num_microbatches, op=dist.ReduceOp.MAX, group=dp_group) - - if vpp_size > 1: - # vpp requies the number of microbatches to be divisible by vpp_size - num_microbatches = torch.clamp( - num_microbatches // microbatch_group_size_per_vp_stage * microbatch_group_size_per_vp_stage, - min=1, - ) - - num_microbatches = num_microbatches.tolist() - - # balance the each micro batch - samples = rollout_data["total_lengths"] - # balance the number of mirobatches across steps - micro_batch_indices = [] - for i, num_mbs in enumerate(num_microbatches): - start, end = i * num_local_gbs, (i + 1) * num_local_gbs - samples = rollout_data["total_lengths"][start:end] - partitions = get_seqlen_balanced_partitions(samples, num_mbs, equal_size=False) - for j in range(num_mbs): - for k in range(len(partitions[j])): - partitions[j][k] += start - micro_batch_indices.extend(partitions) - - assert len(set(sum(micro_batch_indices, []))) == num_local_samples - - data_iterator = _generate_data_iterator(rollout_data, None, micro_batch_indices) - - return ( - data_iterator, - num_microbatches, - ) +def get_data_iterator(rollout_data: RolloutBatch) -> list[DataIterator]: + """Build one ``DataIterator`` per VPP stage from the pre-computed schedule in ``rollout_data``.""" + vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1 + micro_batch_indices = rollout_data["micro_batch_indices"] + return [DataIterator(rollout_data, micro_batch_indices) for _ in range(vpp_size)] def log_rollout_data( @@ -406,6 +279,16 @@ def log_rollout_data( loss_masks = rollout_data["loss_masks"] total_lengths = rollout_data["total_lengths"] max_seq_lens = rollout_data.get("max_seq_lens", None) + # Same per-rollout denominators the training loss uses, so reported + # log_probs / returns / advantages / etc. live in the same per-rollout + # mean space (rather than per-sample) as the gradient signal. + rollout_mask_sums = rollout_data.get("rollout_mask_sums", None) + # For per-rollout-mean metrics: ``rollout_log_metric_contribution`` + # produces the ``(sum, count)`` tuple so gather_log_data's + # ``Σsum / Σcount`` lands on ``sum_DP_full / num_rollouts`` — the + # same number train_one_step reports for the same samples. + dp_world = mpu.get_data_parallel_world_size(with_context_parallel=False) + num_rollouts_in_rollout = sum(rollout_data["global_batch_sizes"]) for key, val in rollout_data.items(): if key in [ @@ -413,15 +296,20 @@ def log_rollout_data( "multimodal_train_inputs", "loss_masks", "sample_indices", + "rollout_ids", + "rollout_mask_sums", "rollout_routed_experts", "max_seq_lens", - "dynamic_global_batch_size", + "global_batch_sizes", + "num_microbatches", + "micro_batch_indices", ]: continue - # Upload per sample mean for each rollout value - # There are the following assumptions: - # - Each dp rank has the same number of samples + # Emit (sum, count) so gather_log_data can do a weighted average across + # DP ranks. This stops the legacy "every rank has the same N samples" + # assumption from biasing means once uneven-DP partitioning lands. if isinstance(val, (list, tuple)): + count = len(val) if isinstance(val[0], torch.Tensor): # NOTE: Here we have to do the clone().detach(), otherwise the tensor will be # modified in place and will cause problem for the next rollout. @@ -435,25 +323,38 @@ def log_rollout_data( "teacher_log_probs", "opd_reverse_kl", ]: - val = torch.cat(val).clone().detach() + tensor = torch.cat(val).clone().detach() sum_of_sample_mean = get_sum_of_sample_mean( total_lengths, response_lengths, loss_masks, + rollout_mask_sums, qkv_format=args.qkv_format, max_seq_lens=max_seq_lens, ) - val = cp_size * sum_of_sample_mean(val) / len(loss_masks) - else: - val = torch.cat(val).clone().detach() - val = val.mean() * cp_size + # Compute (sum, count) via the shared helper so this + # path and the unit tests stay in sync. + sum_value, count = rollout_log_metric_contribution( + sum_of_sample_mean(tensor).item(), + cp_size=cp_size, + num_rollouts_in_rollout=num_rollouts_in_rollout, + dp_size=dp_world, + ) + log_dict[key] = (sum_value, count) + continue + tensor = torch.cat(val).clone().detach() + # val.mean() * cp_size is the per-sample mean for one rank; + # multiply by count to get the per-rank sum. + per_rank_sum = tensor.mean() * cp_size * count + sum_value = per_rank_sum.item() else: - val = sum(val) / len(val) + sum_value = sum(val) + log_dict[key] = (sum_value, count) elif isinstance(val, torch.Tensor): - val = val.float().mean() + # Scalar tensor (one per rank): treat as count=1. + log_dict[key] = (val.float().mean().item(), 1) else: raise ValueError(f"Unsupported type: {type(val)} for key: {key}") - log_dict[key] = val.item() if isinstance(val, torch.Tensor) else val reduced_log_dict = gather_log_data("rollout", args, rollout_id, log_dict) if args.ci_test and reduced_log_dict is not None: @@ -531,8 +432,14 @@ def quantile(total_value, n_quantiles, data) -> dict: for p, val in correct_response_length_percentile.items(): rollout_data[f"correct_length/{p}"] = [val] * num_correct_responses if len(correct_entropy) > 0: + # NOTE: per-sample-mean over the correct subset, not per-rollout. + # A rollout's siblings may not all be correct, and slicing + # ``rollout_mask_sums`` here would leave a denom that still + # includes incorrect siblings — meaningless for a "correct-only" + # entropy report. Per-sample-mean over the filtered subset is + # the cleanest semantic. sum_of_sample_mean = get_sum_of_sample_mean( - correct_total_lengths, correct_response_lengths, correct_loss_masks + correct_total_lengths, correct_response_lengths, correct_loss_masks, sample_denoms=None ) correct_entropy = sum_of_sample_mean(torch.cat(correct_entropy, dim=0)) rollout_data["correct_entropy"] = [correct_entropy.item()] * num_correct_responses diff --git a/vime/backends/megatron_utils/hf_checkpoint_saver.py b/vime/backends/megatron_utils/hf_checkpoint_saver.py new file mode 100644 index 000000000..76f0a6ef6 --- /dev/null +++ b/vime/backends/megatron_utils/hf_checkpoint_saver.py @@ -0,0 +1,199 @@ +import json +import logging +import os +import shutil +from pathlib import Path +from typing import Any + +import torch + +logger = logging.getLogger(__name__) + +_HF_WEIGHT_FILE_NAMES = { + "model.safetensors.index.json", + "pytorch_model.bin.index.json", + "tf_model.h5", + "flax_model.msgpack", +} +_HF_WEIGHT_FILE_SUFFIXES = (".safetensors", ".bin", ".pt", ".pth", ".ckpt", ".msgpack") + + +def save_hf_model_direct(args, rollout_id: int, model) -> None: + """Save a Megatron model as an HF safetensors checkpoint without Megatron Bridge.""" + import torch.distributed as dist + from transformers import AutoConfig + + from .update_weight.common import named_params_and_buffers + from .update_weight.hf_weight_iterator_direct import HfWeightIteratorDirect + + path = Path(args.save_hf.format(rollout_id=rollout_id)) + is_save_rank = _is_global_rank_zero() + hf_checkpoint = Path(args.hf_checkpoint).resolve() + save_path = path.resolve() + if hf_checkpoint == save_path: + raise ValueError("--save-hf must not point to the same directory as --hf-checkpoint") + if not hf_checkpoint.is_dir(): + raise ValueError(f"--hf-checkpoint must be a local directory when using raw --save-hf: {args.hf_checkpoint}") + + setup_error = None + if is_save_rank: + try: + logger.info("Saving model in HuggingFace format to %s with raw Megatron-to-HF conversion", path) + path.mkdir(parents=True, exist_ok=True) + _clear_existing_hf_weights(path) + _copy_hf_assets(args.hf_checkpoint, path) + except Exception as e: + setup_error = repr(e) + + _raise_if_rank_zero_failed("prepare raw HuggingFace save directory", setup_error) + + metadata_error = None + payload: list[Any] = [None] + if is_save_rank: + try: + hf_config = AutoConfig.from_pretrained(args.hf_checkpoint, trust_remote_code=True) + payload = [ + ( + type(hf_config).__name__.lower() if args.model_name is None else args.model_name, + getattr(hf_config, "quantization_config", None), + ) + ] + except Exception as e: + metadata_error = repr(e) + _raise_if_rank_zero_failed("load HuggingFace conversion metadata", metadata_error) + + if dist.is_available() and dist.is_initialized(): + dist.broadcast_object_list(payload, src=0) + model_name, quantization_config = payload[0] + + hf_weight_iterator = HfWeightIteratorDirect( + args=args, + model=model, + model_name=model_name, + quantization_config=quantization_config, + ) + megatron_local_weights = dict(named_params_and_buffers(args, model, convert_to_global_name=True)) + writer = _SafetensorShardWriter(path, enabled=is_save_rank) + + for hf_named_tensors in hf_weight_iterator.get_hf_weight_chunks( + megatron_local_weights, progress_desc="Save HF checkpoint" + ): + write_error = None + try: + writer.write(hf_named_tensors) + except Exception as e: + write_error = repr(e) + _raise_if_rank_zero_failed("write raw HuggingFace weight shard", write_error) + del hf_named_tensors + if torch.cuda.is_available(): + torch.cuda.ipc_collect() + + finalize_error = None + if is_save_rank: + try: + writer.finalize() + except Exception as e: + finalize_error = repr(e) + _raise_if_rank_zero_failed("finalize raw HuggingFace checkpoint", finalize_error) + + if is_save_rank: + logger.info("Successfully saved HuggingFace model to %s", path) + + +class _SafetensorShardWriter: + def __init__(self, path: Path, *, enabled: bool) -> None: + self.path = path + self.enabled = enabled + self.total_size = 0 + self.weight_map: dict[str, str] = {} + self.shard_files: list[str] = [] + + def write(self, named_tensors) -> None: + if not self.enabled: + return + + from safetensors.torch import save_file + + state_dict = {} + for name, tensor in named_tensors: + if name in self.weight_map or name in state_dict: + raise ValueError(f"Duplicate HF tensor while saving: {name}") + self.total_size += tensor.numel() * tensor.element_size() + state_dict[name] = _tensor_for_safetensors(tensor) + + if not state_dict: + return + + filename = f"model-{len(self.shard_files) + 1:05d}.safetensors" + save_file(state_dict, self.path / filename, metadata={"format": "pt"}) + self.shard_files.append(filename) + for name in state_dict: + self.weight_map[name] = filename + + def finalize(self) -> None: + if not self.enabled: + return + if not self.shard_files: + raise ValueError("No HF tensors were produced while saving") + + total_files = len(self.shard_files) + rename_map = {} + for idx, old_name in enumerate(self.shard_files, start=1): + new_name = f"model-{idx:05d}-of-{total_files:05d}.safetensors" + os.replace(self.path / old_name, self.path / new_name) + rename_map[old_name] = new_name + + final_weight_map = {name: rename_map[filename] for name, filename in self.weight_map.items()} + index_data = {"metadata": {"total_size": self.total_size}, "weight_map": final_weight_map} + with open(self.path / "model.safetensors.index.json", "w", encoding="utf-8") as f: + json.dump(index_data, f, indent=2) + + +def _tensor_for_safetensors(tensor: torch.Tensor) -> torch.Tensor: + tensor = tensor.detach() + if not tensor.is_contiguous(): + tensor = tensor.contiguous() + if tensor.device.type != "cpu": + tensor = tensor.cpu() + return tensor + + +def _clear_existing_hf_weights(path: Path) -> None: + for item in path.iterdir(): + if item.is_file() and _is_hf_weight_file(item): + item.unlink() + + +def _copy_hf_assets(origin_hf_dir: str, output_dir: Path) -> None: + origin = Path(origin_hf_dir) + if not origin.is_dir(): + raise ValueError(f"--hf-checkpoint must be a local directory when using raw --save-hf: {origin_hf_dir}") + + for item in origin.iterdir(): + if item.is_file(): + if _is_hf_weight_file(item): + continue + shutil.copy2(item, output_dir / item.name) + + +def _is_hf_weight_file(path: Path) -> bool: + name = path.name + return name in _HF_WEIGHT_FILE_NAMES or name.endswith(_HF_WEIGHT_FILE_SUFFIXES) + + +def _is_global_rank_zero() -> bool: + import torch.distributed as dist + + return not (dist.is_available() and dist.is_initialized()) or dist.get_rank() == 0 + + +def _raise_if_rank_zero_failed(context: str, error: str | None) -> None: + import torch.distributed as dist + + if dist.is_available() and dist.is_initialized(): + payload = [error] + dist.broadcast_object_list(payload, src=0) + error = payload[0] + + if error is not None: + raise RuntimeError(f"Failed to {context}: {error}") diff --git a/vime/backends/megatron_utils/loss.py b/vime/backends/megatron_utils/loss.py index 6f95550b2..3f6ab29c5 100644 --- a/vime/backends/megatron_utils/loss.py +++ b/vime/backends/megatron_utils/loss.py @@ -931,12 +931,17 @@ def policy_loss_function( tis_func = vanilla_tis_function pg_loss, modified_response_masks, tis_metrics = tis_func(**tis_kwargs) - # [decouple IS and rejection] Rebuild sum_of_sample_mean with modified_response_masks for denominator correction - # modified_response_masks will be sliced with cp in get_sum_of_sample_mean + # [decouple IS and rejection] Rebuild sum_of_sample_mean with + # modified_response_masks for numerator correction (rejected tokens + # zeroed in pg_loss). Denominators stay the precomputed per-rollout + # totals from ``rollout_mask_sums`` (based on original loss_masks) — + # same normalizer as the outer reducer, so pg_loss and the rest of the + # reported metrics live in the same per-rollout-mean space. sum_of_sample_mean = get_sum_of_sample_mean( total_lengths, response_lengths, modified_response_masks, + batch["rollout_mask_sums"], args.calculate_per_token_loss, args.qkv_format, max_seq_lens, @@ -1136,6 +1141,7 @@ def loss_function( args: Namespace, batch: RolloutBatch, num_microbatches: int, + step_global_batch_size: int, logits: torch.Tensor, ) -> tuple[torch.Tensor, int | torch.Tensor, dict[str, list[str] | torch.Tensor]]: """Dispatch to the configured loss and rescale for Megatron integration. @@ -1147,10 +1153,14 @@ def loss_function( Args: args: Configuration specifying `loss_type`, `calculate_per_token_loss`, - `global_batch_size`, and optionally `custom_loss_function_path`. + and optionally `custom_loss_function_path`. batch: Mini-batch with "loss_masks", "response_lengths", and other keys required by the selected loss function. num_microbatches: Number of gradient accumulation steps. + step_global_batch_size: Sample count for the current training step + (total across DP). Replaces the legacy ``args.global_batch_size`` + fallback so the train side stops depending on "every DP rank holds + the same N samples". logits: Model outputs (policy or value head). Returns: @@ -1162,12 +1172,12 @@ def loss_function( "values" (1D tensor: [count, metric1, metric2, ...]). """ num_tokens = sum([torch.clamp_min(loss_mask.sum(), 1) for loss_mask in batch["loss_masks"]]) - num_samples = len(batch["response_lengths"]) sum_of_sample_mean = get_sum_of_sample_mean( batch["total_lengths"], batch["response_lengths"], batch["loss_masks"], + batch["rollout_mask_sums"], args.calculate_per_token_loss, args.qkv_format, batch.get("max_seq_lens", None), @@ -1199,10 +1209,12 @@ def loss_function( loss = loss + 0 * logits.sum() # Here we need to divide by cp_size because to cancel the multiply in Megatron. - global_batch_size = batch.get("dynamic_global_batch_size", args.global_batch_size) if not args.calculate_per_token_loss: loss = ( - loss * num_microbatches / global_batch_size * mpu.get_data_parallel_world_size(with_context_parallel=True) + loss + * num_microbatches + / step_global_batch_size + * mpu.get_data_parallel_world_size(with_context_parallel=True) ) else: loss = loss * mpu.get_context_parallel_world_size() @@ -1212,9 +1224,16 @@ def loss_function( (num_tokens if args.calculate_per_token_loss else torch.tensor(1, device=logits.device)), { "keys": list(log.keys()), + # values[0] is the consumer's reporting denominator after + # all-reduce. For per-token-loss it must equal step total tokens + # (only known by summing per-mb num_tokens across mbs / DP). For + # per-rollout-mean it is a constant — ``step_global_batch_size`` — + # so we leave a 0 placeholder here and let ``train_one_step`` + # substitute the constant directly, instead of routing it through + # per-mb fractions. "values": torch.tensor( [ - num_samples if not args.calculate_per_token_loss else num_tokens, + num_tokens if args.calculate_per_token_loss else 0, ] + list(log.values()), device=logits.device, diff --git a/vime/backends/megatron_utils/model.py b/vime/backends/megatron_utils/model.py index c8ee435a5..550d0b036 100644 --- a/vime/backends/megatron_utils/model.py +++ b/vime/backends/megatron_utils/model.py @@ -31,6 +31,7 @@ from vime.utils.memory_utils import clear_memory from .checkpoint import load_checkpoint, save_checkpoint +from .cp_utils import reduce_train_step_metrics from .data import DataIterator, get_batch from .loss import loss_function from .model_provider import get_model_provider_func @@ -147,7 +148,15 @@ def get_optimizer_param_scheduler(args: Namespace, optimizer: MegatronOptimizer) Returns: OptimizerParamScheduler: Initialized scheduler bound to ``optimizer``. """ - # Iteration-based training. + # Iteration-based training. ``train_iters`` is an estimate of the total + # number of training steps — it's only used to size Megatron's LR decay + # schedule (and ``lr_decay_iters`` defaults to it). With variable per-rollout + # sample counts (dynamic sampling / filtering / custom step splitter) the + # *actual* total can drift; the schedule still tracks the true progress via + # ``opt_param_scheduler.num_steps`` (samples consumed, also persisted across + # resume), so the worst case is the cosine/linear schedule reaches its + # plateau slightly early or late. Pass ``--lr-decay-iters`` explicitly if you + # need exact decay control. args.train_iters = args.num_rollout * args.rollout_batch_size * args.n_samples_per_prompt // args.global_batch_size if args.lr_decay_iters is None: args.lr_decay_iters = args.train_iters @@ -416,6 +425,7 @@ def train_one_step( optimizer: MegatronOptimizer, opt_param_scheduler: OptimizerParamScheduler, num_microbatches: int, + step_global_batch_size: int, microbatch_pbar=None, ) -> tuple[dict[str, float], float]: """Execute a single pipeline-parallel training step. @@ -432,6 +442,13 @@ def train_one_step( optimizer (MegatronOptimizer): Optimizer instance. opt_param_scheduler (OptimizerParamScheduler): LR/WD scheduler. num_microbatches (int): Number of microbatches to process. + step_global_batch_size (int): Rollout count for this training step + (total across DP; one "rollout" = one execution of one of the + ``n_samples_per_prompt`` rollouts, which may emit >1 training + sample under compact / subagent). Used both as the loss + normalizer inside the closure and as the LR scheduler + ``increment``. In the common case (1 rollout = 1 sample) this + equals the per-step sample count, so behavior is unchanged. Returns: tuple[dict[str, float], float]: Reduced loss dictionary (last stage only) @@ -484,6 +501,7 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p "rollout_log_probs", "max_seq_lens", "teacher_log_probs", + "rollout_mask_sums", ], args.data_pad_size_multiplier, args.qkv_format, @@ -526,7 +544,7 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p if os.environ.get("ENABLE_ROUTING_REPLAY", "0") == "1": os.environ["ROUTING_REPLAY_STAGE"] = old_stage - return output_tensor, partial(loss_function, args, batch, num_microbatches) + return output_tensor, partial(loss_function, args, batch, num_microbatches, step_global_batch_size) # Forward pass. forward_backward_func = get_forward_backward_func() @@ -567,7 +585,7 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p # Update learning rate. assert update_successful - opt_param_scheduler.step(increment=args.global_batch_size) + opt_param_scheduler.step(increment=step_global_batch_size) # release grad for model_chunk in model: @@ -575,22 +593,13 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p optimizer.zero_grad() if mpu.is_pipeline_last_stage(ignore_virtual=True): - # Average loss across microbatches. - keys = losses_reduced[0]["keys"] - values = None - for x in losses_reduced: - if values is None: - values = x["values"] - else: - values += x["values"] - assert len(keys) + 1 == values.numel() - torch.distributed.all_reduce(values, group=mpu.get_data_parallel_group(with_context_parallel=True)) - - loss_reduced = {} - values = values.tolist() - num_samples_or_tokens = values[0] - for key, value in zip(keys, values[1:], strict=False): - loss_reduced[key] = value * mpu.get_context_parallel_world_size() / num_samples_or_tokens + loss_reduced = reduce_train_step_metrics( + losses_reduced, + calculate_per_token_loss=args.calculate_per_token_loss, + step_global_batch_size=step_global_batch_size, + cp_size=mpu.get_context_parallel_world_size(), + dp_with_cp_group=mpu.get_data_parallel_group(with_context_parallel=True), + ) return loss_reduced, grad_norm return {}, grad_norm @@ -607,6 +616,7 @@ def train( opt_param_scheduler: OptimizerParamScheduler, data_iterator: Sequence[DataIterator], num_microbatches: Sequence[int], + global_batch_sizes: Sequence[int], ) -> None: """Run training over a rollout consisting of multiple steps. @@ -620,9 +630,20 @@ def train( opt_param_scheduler (OptimizerParamScheduler): LR/WD scheduler. data_iterator (Sequence[DataIterator]): Iterable(s) yielding training batches. num_microbatches (Sequence[int]): Microbatches per step in the rollout. + global_batch_sizes (Sequence[int]): Rollout count per step (total + across DP; one "rollout" = one execution of one of the + ``n_samples_per_prompt`` rollouts of a prompt). Same length as + ``num_microbatches``; consumed by ``train_one_step`` for loss + scaling and LR scheduler increments. Equals per-step sample count + in the common case (1 rollout = 1 sample). """ args = get_args() + assert len(num_microbatches) == len(global_batch_sizes), ( + f"num_microbatches and global_batch_sizes must have the same length, " + f"got {len(num_microbatches)} vs {len(global_batch_sizes)}" + ) + for iterator in data_iterator: iterator.reset() @@ -717,6 +738,7 @@ def train( optimizer, opt_param_scheduler, num_microbatches[step_id], + global_batch_sizes[step_id], microbatch_pbar=microbatch_pbar, ) @@ -770,6 +792,8 @@ def train( for param_group_id, param_group in enumerate(optimizer.param_groups): log_dict[f"train/{role_tag}lr-pg_{param_group_id}"] = opt_param_scheduler.get_lr(param_group) + # Per-step gbs — uneven step sizes are easy to miss without this. + log_dict[f"train/{role_tag}global_batch_size"] = global_batch_sizes[step_id] log_dict["train/step"] = accumulated_step_id logging_utils.log(args, log_dict, step_key="train/step") @@ -853,6 +877,19 @@ def save_hf_model(args, rollout_id: int, model: Sequence[DDP]) -> None: model (Sequence[DDP]): Sequence of DDP-wrapped model chunks. rollout_id (int): Rollout ID for path formatting. """ + if args.megatron_to_hf_mode != "bridge": + try: + from vime.backends.megatron_utils.hf_checkpoint_saver import save_hf_model_direct + + save_hf_model_direct(args, rollout_id, model) + except Exception as e: + if ( + mpu.get_data_parallel_rank(with_context_parallel=True) == 0 + and mpu.get_tensor_model_parallel_rank() == 0 + ): + logger.error(f"Failed to save HuggingFace format: {e}") + return + should_log = ( mpu.get_data_parallel_rank(with_context_parallel=True) == 0 and mpu.get_tensor_model_parallel_rank() == 0 ) diff --git a/vime/backends/megatron_utils/update_weight/hf_weight_iterator_base.py b/vime/backends/megatron_utils/update_weight/hf_weight_iterator_base.py index ef7d62e8a..369f8c2d4 100644 --- a/vime/backends/megatron_utils/update_weight/hf_weight_iterator_base.py +++ b/vime/backends/megatron_utils/update_weight/hf_weight_iterator_base.py @@ -21,7 +21,7 @@ def __init__(self, args, model, model_name, quantization_config): self.quantization_config = quantization_config @abstractmethod - def get_hf_weight_chunks(self, megatron_local_weights): + def get_hf_weight_chunks(self, megatron_local_weights, progress_desc: str = "Update weights"): """ Mental model of the API: megatron_model.to_hf_magically().named_parameters() diff --git a/vime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py b/vime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py index d905e84da..d9aa7338f 100644 --- a/vime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py +++ b/vime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py @@ -49,7 +49,7 @@ def __init__(self, *args, **kwargs): ) _patch_bridge_expert_cache_to_cpu() - def get_hf_weight_chunks(self, megatron_local_weights): + def get_hf_weight_chunks(self, megatron_local_weights, progress_desc: str = "Update weights"): # TODO support quantization (e.g. modify megatron-bridge to provide megatron param name) renamed_megatron_local_weights = {strip_param_name_prefix(k): v for k, v in megatron_local_weights.items()} with megatron_bridge_utils.patch_megatron_model(self.model): diff --git a/vime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py b/vime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py index b84743ae9..d345adde8 100644 --- a/vime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py +++ b/vime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py @@ -20,11 +20,11 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.megatron_local_param_info_buckets = _get_megatron_local_param_info_buckets(self.args, self.model) - def get_hf_weight_chunks(self, megatron_local_weights): + def get_hf_weight_chunks(self, megatron_local_weights, progress_desc: str = "Update weights"): rank = dist.get_rank() for megatron_local_param_infos in tqdm( - self.megatron_local_param_info_buckets, disable=rank != 0, desc="Update weights" + self.megatron_local_param_info_buckets, disable=rank != 0, desc=progress_desc ): megatron_full_params = _get_megatron_full_params(megatron_local_param_infos, megatron_local_weights) hf_named_tensors = self._convert_to_hf_named_tensors(megatron_full_params, megatron_local_param_infos) diff --git a/vime/ray/rollout.py b/vime/ray/rollout.py index 71f4f6e62..572f95514 100644 --- a/vime/ray/rollout.py +++ b/vime/ray/rollout.py @@ -22,12 +22,12 @@ GPU_MEMORY_TYPE_CUDA_GRAPH = "cuda_graph" from vime.rollout.base_types import call_rollout_fn from vime.utils import logging_utils +from vime.utils.dp_schedule import build_dp_schedule from vime.utils.health_monitor import RolloutHealthMonitor from vime.utils.http_utils import _wrap_ipv6, find_available_port, get_host_info, init_http_client from vime.utils.logging_utils import configure_logger, init_tracking from vime.utils.metric_utils import compute_pass_rate, compute_rollout_step, compute_statistics, dict_add_prefix from vime.utils.misc import Box, group_by, load_function -from vime.utils.seqlen_balancing import get_seqlen_balanced_partitions from vime.utils.types import Sample from ..utils.metric_utils import has_repetition @@ -498,7 +498,7 @@ def generate(self, rollout_id): # if debug rollout only, we don't convert samples to train data and directly return return data = self._convert_samples_to_train_data(data) - return self._split_train_data_by_dp(data, self.train_parallel_config["dp_size"]) + return self._split_train_data_by_dp(data) def eval(self, rollout_id): if self.args.debug_train_only: @@ -593,56 +593,19 @@ def _get_rollout_data(self, rollout_id): data = call_rollout_fn(self.generate_rollout, self.args, rollout_id, self.data_source, evaluation=False) metrics = data.metrics data = data.samples + # Enforce the rollout_id contract before flattening: any list[Sample] + # encountered in the nested output must have rollout_id set on every + # element. Default rollouts inherit it from the data source; compact / + # subagent paths that split one rollout into N training samples must + # set the same rollout_id on every sibling so the loss reducer counts + # the rollout once instead of N times. + _validate_rollout_id_annotated(data) # flatten the data if it is a list of lists while isinstance(data[0], list): data = list(itertools.chain.from_iterable(data)) - if not self.args.disable_rollout_trim_samples and not self.args.debug_rollout_only: - global_batch_size = self.args.global_batch_size - if self.args.use_dynamic_global_batch_size: - logger.info(f"Collected {len(data)} samples from rollout to train with dynamic global batch size") - # TODO: this is a temporary solution, we should directly save dynamic_global_batch_size to rollout data - self._dynamic_global_batch_size = self._compute_dynamic_global_batch_size(len(data)) - global_batch_size = self._dynamic_global_batch_size - - if len(data) % global_batch_size != 0: - trim_len = (len(data) // global_batch_size) * global_batch_size - if trim_len == 0: - raise ValueError(f"Not enough samples {len(data)} for global_batch_size {global_batch_size}") - origin_data_length = len(data) - data = data[:trim_len] - logger.info(f"trim number of samples from {origin_data_length} to {trim_len}") - logger.info(f"Final collected {len(data)} samples from rollout to train") - return data, metrics - def _compute_dynamic_global_batch_size(self, num_samples: int) -> int: - """Calculate dynamic global_batch_size to ensure only one training step. - - Strategy: global_batch_size = num_samples rounded down to a multiple of dp_size - This ensures num_steps_per_rollout = num_samples // global_batch_size = 1 - """ - dp_size = self.train_parallel_config["dp_size"] - original_gbs = self.args.global_batch_size - - # Round down to a multiple of dp_size to ensure only one training step - dynamic_gbs = (num_samples // dp_size) * dp_size - - if dynamic_gbs == 0: - # Too few samples, use at least dp_size - dynamic_gbs = dp_size - logger.warning(f"num_samples={num_samples} < dp_size={dp_size}, using dp_size as global_batch_size") - - # Calculate how many samples will be discarded - wasted = num_samples - dynamic_gbs - - if dynamic_gbs != original_gbs or wasted > 0: - logger.info( - f"Dynamic global_batch_size: {original_gbs} -> {dynamic_gbs} (num_samples={num_samples}, dp_size={dp_size}, num_steps=1, wasted={wasted})" - ) - - return dynamic_gbs - def _save_debug_rollout_data(self, data, rollout_id, evaluation: bool): # TODO to be refactored (originally Buffer._set_data) if (path_template := self.args.save_debug_rollout_data) is not None: @@ -701,6 +664,16 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl assert len(raw_rewards) == len(samples) assert len(rewards) == len(samples) + # Rollout id (one per rollout execution). Default rollouts emit one + # sample per rollout, so we fall back to ``sample.index`` (unique). + # Compact / subagent paths that emit multiple training samples per + # rollout set ``rollout_id`` explicitly so all siblings share a + # value; the loss reducer then aggregates them as one rollout. + if samples[0].rollout_id is None: + rollout_ids = list(range(len(samples))) + else: + rollout_ids = [sample.rollout_id for sample in samples] + train_data = { "tokens": [sample.tokens for sample in samples], "response_lengths": [sample.response_length for sample in samples], @@ -710,6 +683,7 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl "raw_reward": raw_rewards, "truncated": [1 if sample.status == Sample.Status.TRUNCATED else 0 for sample in samples], "sample_indices": [sample.index for sample in samples], + "rollout_ids": rollout_ids, } # loss mask @@ -728,6 +702,23 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl loss_masks.append(sample.loss_mask) train_data["loss_masks"] = loss_masks + # Per-rollout aggregate, precomputed at the step level (where we can + # see every sample of every rollout) and broadcast per-sample so the + # per-mb loss reducer uses the correct whole-rollout denominator even + # when a rollout's samples land in different micro-batches (first-fit + # packing can split a rollout across mbs): + # + # ``rollout_mask_sums[i]`` — sum of loss-mask totals over every + # sample in sample i's rollout. Used as the reducer's denominator + # so summing partial contributions across mbs yields one + # token-weighted mean per rollout. + rollout_id_list = train_data["rollout_ids"] + mask_sums_per_sample = [sum(m) for m in loss_masks] + rollout_total_mask: dict[int, int] = {} + for rid, ms in zip(rollout_id_list, mask_sums_per_sample, strict=True): + rollout_total_mask[rid] = rollout_total_mask.get(rid, 0) + ms + train_data["rollout_mask_sums"] = [rollout_total_mask[rid] for rid in rollout_id_list] + # Overwrite raw_reward when available. Mixed-source batches may only # populate this field for a subset of samples (e.g. SWE but not code). if any(sample.metadata and "raw_reward" in sample.metadata for sample in samples): @@ -761,27 +752,35 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl def set_train_parallel_config(self, config: dict): self.train_parallel_config = config - def _split_train_data_by_dp(self, data, dp_size): - """Split the train data by data parallel size.""" - rollout_data = {} - - if "prompt" in data: - rollout_data["prompt"] = data["prompt"] - + def _split_train_data_by_dp(self, data): + """Compute the DP/mbs schedule and package each rank's rollout_data + into a Ray Box. The schedule itself is computed by + :func:`build_dp_schedule` so it stays unit-testable without Ray/vllm. + + Step split is by rollout id (``samples[i].rollout_id``, falling back + to ``samples[i].index``); each step holds exactly + ``args.global_batch_size`` rollouts so the training-step count per + rollout is fixed at ``rollout_batch_size * n_samples_per_prompt // + global_batch_size`` regardless of how many training samples each + rollout produced. + """ + dp_size = self.train_parallel_config["dp_size"] total_lengths = [len(t) for t in data["tokens"]] data["total_lengths"] = total_lengths - if self.args.balance_data: - partitions = get_seqlen_balanced_partitions(total_lengths, dp_size, equal_size=True) - else: - partitions = [range(i, len(total_lengths), dp_size) for i in range(dp_size)] + partitions, micro_batch_indices, num_microbatches, global_batch_sizes = build_dp_schedule( + self.args, + self.train_parallel_config, + total_lengths, + global_batch_size=self.args.global_batch_size, + rollout_indices=data["rollout_ids"], + ) + # Package per-rank rollout_data rollout_data_refs = [] - - for i in range(dp_size): - rollout_data = {} - partition = partitions[i] - rollout_data["partition"] = partition + for r in range(dp_size): + partition = partitions[r] + rollout_data = {"partition": partition} for key in [ "tokens", "multimodal_train_inputs", @@ -791,6 +790,8 @@ def _split_train_data_by_dp(self, data, dp_size): "loss_masks", "round_number", "sample_indices", + "rollout_ids", + "rollout_mask_sums", "rollout_log_probs", "rollout_routed_experts", "prompt", @@ -798,23 +799,51 @@ def _split_train_data_by_dp(self, data, dp_size): ]: if key not in data: continue - val = [data[key][j] for j in partition] - rollout_data[key] = val + rollout_data[key] = [data[key][j] for j in partition] # keys that need to be splited at train side - for key in [ - "raw_reward", - "total_lengths", - ]: + for key in ["raw_reward", "total_lengths"]: if key not in data: continue rollout_data[key] = data[key] - # Pass dynamic global_batch_size to training side - if hasattr(self, "_dynamic_global_batch_size"): - rollout_data["dynamic_global_batch_size"] = self._dynamic_global_batch_size + rollout_data["global_batch_sizes"] = global_batch_sizes + rollout_data["num_microbatches"] = num_microbatches + rollout_data["micro_batch_indices"] = micro_batch_indices[r] rollout_data_refs.append(Box(ray.put(rollout_data))) return rollout_data_refs +def _validate_rollout_id_annotated(node, depth=0): + """Walk the rollout function's nested output and validate ``rollout_id`` only + when a compact / subagent pattern is detected. + + "Compact" = the rollout function wraps multiple training samples from one + rollout execution into a ``list[Sample]``. In slime's convention the + default rollout shape is ``list[list[Sample]]`` (depth-2: prompt × rollout) + so its leaf ``list[Sample]`` lands at depth 1 and we skip validation, + preserving backward compatibility. A compact rollout adds a third level: + ``list[list[list[Sample]]]`` (prompt × rollout × samples-from-one-rollout), + so the leaf ``list[Sample]`` lands at depth ≥ 2. At that point we require + every sibling to carry a non-None ``rollout_id`` and to share the same + value, so the loss reducer counts the rollout once instead of N times. + """ + if isinstance(node, Sample): + return + assert isinstance(node, list), f"unexpected rollout output node type: {type(node).__name__}" + if node and isinstance(node[0], Sample): + if depth >= 2 and len(node) > 1: + rids = [s.rollout_id for s in node] + missing = [i for i, r in enumerate(rids) if r is None] + assert not missing, ( + f"Compact rollout returned {len(node)} samples but rollout_id is unset on " + f"positions {missing}. Set Sample.rollout_id on every sibling so the loss " + "reducer can aggregate them as one rollout instead of N." + ) + assert len(set(rids)) == 1, f"Sibling samples from one compact rollout must share rollout_id; got {rids}." + return + for item in node: + _validate_rollout_id_annotated(item, depth + 1) + + def _allocate_rollout_engine_addr_and_ports_external(args, rollout_engines): addr_and_ports = {} for rank, _ in rollout_engines: diff --git a/vime/rollout/_fanout_test_helpers.py b/vime/rollout/_fanout_test_helpers.py new file mode 100644 index 000000000..3590ca472 --- /dev/null +++ b/vime/rollout/_fanout_test_helpers.py @@ -0,0 +1,114 @@ +"""Test-internal compact-rollout helpers used by ``test_qwen2.5_0.5B_fanout_short.py``. + +The underscore prefix marks this as test infrastructure — it is not part +of the user-facing vime API and is not re-exported anywhere. It lives +in ``vime/`` only so the test can reference it by a dotted module path +(``--custom-generate-function-path`` / ``--custom-reward-post-process-path`` +resolve a string via ``importlib.import_module``, which can't handle the +dots in the e2e test's filename). + +Two helpers: + + - ``compact_generate``: fans one input sample out to N siblings + sharing the same ``rollout_id``. That's the contract the rest of the + framework (per-rollout step splitter, per-rollout-mean reducer, + ``_validate_rollout_id_annotated`` validator) is built around. + + - ``grpo_normalize_by_group_index``: replaces the default + ``_post_process_rewards`` reshape-by-shape logic with a proper + ``group_index``-keyed grouping. The default at + ``vime/ray/rollout.py:_post_process_rewards`` assumes every prompt + produced exactly ``n_samples_per_prompt`` samples and reshapes by + that constant; when compact/fanout makes the per-prompt count uneven, + the reshape fails and the fallback ``view(-1, total)`` collapses + everything into ONE group, destroying per-prompt centering. + ``group_index`` (set by the data source per-prompt, preserved through + ``deepcopy``) is the right key here. +""" + +import copy +import os +from collections import defaultdict + +MAX_FANOUT = 3 + +# Each invocation appends one line. The test file reads this after train +# completes to assert the framework actually drove the custom path for +# every prompt (no silent bypass / no double-submission). +COUNTER_FILE_ENV = "SLIME_FANOUT_TEST_COUNTER_FILE" + + +async def compact_generate(args, sample, sampling_params): + """One prompt → N siblings, deterministic N = 1 + (index % MAX_FANOUT). + + Strategy: call vLLM once, deepcopy N-1 times. Bounded GPU cost — + we're pinning the framework's per-rollout handling, not generation + diversity. + """ + from vime.rollout.vllm_rollout import generate + + counter_path = os.environ.get(COUNTER_FILE_ENV) + if counter_path: + try: + with open(counter_path, "a") as f: + f.write(f"{sample.index}\n") + except OSError: + # Counter file is best-effort — never fail training because of it. + pass + + base_sample = await generate(args, sample, sampling_params) + + n = 1 + (sample.index % MAX_FANOUT) + siblings = [] + for _ in range(n): + s = copy.deepcopy(base_sample) + # Critical invariant: all siblings share ``rollout_id`` so the + # per-rollout reducer aggregates them as ONE rollout (not N) and + # the rollout-aware step splitter keeps them in the same step. + # ``group_index`` is inherited via ``deepcopy`` and is what the + # post-process reward hook below groups on for GRPO normalize. + s.rollout_id = sample.index + siblings.append(s) + return siblings + + +def grpo_normalize_by_group_index(args, samples): + """Drop-in ``--custom-reward-post-process-path`` for compact/fanout. + + The default ``_post_process_rewards`` (``vime/ray/rollout.py``) + reshapes the flat reward tensor as ``(-1, n_samples_per_prompt)`` + when ``total == n_samples_per_prompt * rollout_batch_size``, falling + back to ``view(-1, total)`` (= one giant group) otherwise. With + fanout the count per prompt is uneven, so the fallback fires and + centering is computed across ALL samples in the batch instead of + per-prompt — that's silently wrong for GRPO. + + This helper groups by ``Sample.group_index`` (the data-source-set + per-prompt counter, preserved through deepcopy in + ``compact_generate``) and applies the same mean-center + optional + std-normalize the default does, just with the correct grouping. + + Returns ``(raw_rewards, normalized_rewards)`` matching the input + ``samples`` order — same shape as the default's return contract. + """ + import torch + + raw_rewards = [s.get_reward_value(args) for s in samples] + + # group_index → list of (original_position, raw_reward) + groups: dict[int, list[tuple[int, float]]] = defaultdict(list) + for i, s in enumerate(samples): + groups[s.group_index].append((i, raw_rewards[i])) + + out = [0.0] * len(samples) + use_std = getattr(args, "grpo_std_normalization", True) + for indexed in groups.values(): + positions = [p for p, _ in indexed] + rewards = torch.tensor([r for _, r in indexed], dtype=torch.float) + rewards = rewards - rewards.mean() + if use_std: + rewards = rewards / (rewards.std() + 1e-6) + for pos, r in zip(positions, rewards.tolist(), strict=True): + out[pos] = r + + return raw_rewards, out diff --git a/vime/rollout/forge_load.py b/vime/rollout/forge_load.py new file mode 100644 index 000000000..b533e291d --- /dev/null +++ b/vime/rollout/forge_load.py @@ -0,0 +1,114 @@ +"""Load a forged rollout dump from disk so memory-test runs can keep +vLLM alive while bypassing real generation. + +Plug in by setting: + --rollout-function-path vime.rollout.forge_load.generate_rollout + --load-forge-rollout-data + +The path follows the same {rollout_id} format convention as +--load-debug-rollout-data: + - Literal path (recommended for memory tests): + --load-forge-rollout-data /path/to/forged_dump/rollout_data/0.pt + Every rollout reuses the same file (rollout_id is left untouched so + the framework's per-rollout bookkeeping still works). + - Template path (matches --save-debug-rollout-data layout): + --load-forge-rollout-data /path/to/dumps/{rollout_id}.pt + Each rollout loads its own file. If a rollout_id has no file we fall + back to 0.pt for the training path; eval has no equivalent fallback. + +Unlike --load-debug-rollout-data, this path does NOT set +skip_vllm=True / debug_train_only=True (see +vime/utils/arguments.py: skip_vllm computation in _pre_parse_mode and +the debug_train_only flip when load_debug_rollout_data is set), so +vLLM servers, router, weight_update and the full colocate +offload/onload dance still run. That is exactly what we want when +measuring real GPU memory. +""" + +import logging +import os +from pathlib import Path + +import torch + +from vime.rollout.base_types import RolloutFnEvalOutput, RolloutFnTrainOutput +from vime.utils.types import Sample + +logger = logging.getLogger(__name__) + + +def _resolve_path(args, rollout_id: int, evaluation: bool) -> str | None: + tpl = getattr(args, "load_forge_rollout_data", None) + if not tpl: + raise RuntimeError( + "--load-forge-rollout-data not set. Pass the dump path, " + "e.g. /path/to/rollout_data/0.pt (literal) or " + "/path/to/rollout_data/{rollout_id}.pt (template)." + ) + # In literal-path mode (no {rollout_id} placeholder) we can't distinguish + # train vs eval files, so eval is a no-op. Use template mode if you want + # to also replay an eval dump (--load-forge-rollout-data .../{rollout_id}.pt + # with eval_.pt files alongside the train ones). + if evaluation and "{rollout_id}" not in tpl: + return None + rid_str = ("eval_" if evaluation else "") + str(rollout_id) + path = tpl.format(rollout_id=rid_str) + if os.path.exists(path): + return path + # Fallback only for the training path: many memory tests have just 0.pt + # but want --num-rollout > 1. Eval has no equivalent fallback (we don't + # want to silently feed training samples to the eval pipeline). + if not evaluation: + fallback = tpl.format(rollout_id="0") + if os.path.exists(fallback): + logger.info("forge_load: %s missing, falling back to %s", path, fallback) + return fallback + return None + + +def generate_rollout(args, rollout_id, data_source, evaluation: bool = False): + path = _resolve_path(args, rollout_id, evaluation) + + if evaluation: + # Eval is optional for a memory-test run. If no eval dump, no-op. + if path is None: + logger.info("forge_load: no eval dump found; returning empty eval result") + return RolloutFnEvalOutput(data={}) + logger.info("forge_load: loading eval samples from %s", path) + blob = torch.load(path, weights_only=False) + samples = [Sample.from_dict(s) for s in blob["samples"]] + # See train-path note: don't overwrite rollout_id. + reward_key = args.eval_reward_key or args.reward_key + rewards = [s.reward if (not reward_key or s.reward is None) else s.reward[reward_key] for s in samples] + return RolloutFnEvalOutput( + data={ + "forge_eval": { + "rewards": [r if r is not None else 0.0 for r in rewards], + "truncated": [s.status == Sample.Status.TRUNCATED for s in samples], + "samples": samples, + } + } + ) + + if path is None: + raise RuntimeError( + f"forge_load: no dump found for rollout_id={rollout_id} " + f"(--load-forge-rollout-data={args.load_forge_rollout_data!r})" + ) + + logger.info("forge_load: loading samples from %s", path) + blob = torch.load(path, weights_only=False) + samples = [Sample.from_dict(s) for s in blob["samples"]] + # IMPORTANT: do NOT overwrite sample.rollout_id with the current rollout_id. + # Default-shape rollouts leave rollout_id=None and vime falls back to + # sample.index in vime/ray/rollout.py (the dp-schedule grouping key). + # Forcing all samples to share one rollout_id collapses them into a single + # "rollout", which trips the num_rollouts >= global_batch_size assert in + # vime/utils/dp_schedule.py. + logger.info( + "forge_load: loaded %d samples for rollout_id=%d from %s", + len(samples), + rollout_id, + Path(path).name, + ) + return RolloutFnTrainOutput(samples=samples) diff --git a/vime/rollout/vllm_streaming_rollout.py b/vime/rollout/vllm_streaming_rollout.py new file mode 100644 index 000000000..138cbfd73 --- /dev/null +++ b/vime/rollout/vllm_streaming_rollout.py @@ -0,0 +1,250 @@ +"""Streaming vLLM rollout (example). + +Drop-in alternative to :func:`vime.rollout.vllm_rollout.generate` that consumes +vLLM's ``/inference/v1/generate`` SSE stream incrementally instead of awaiting +one final JSON response. The win is on **abort**: every chunk we receive lands +directly on ``sample`` (tokens, response text, log-probs), so when a +partial-rollout recycling or weight-update abort fires mid-generation, the +partial state is already on the sample — we don't depend on the engine +returning the collected text. + +Wire it in as the per-sample generate function:: + + --rollout-function-path vime.rollout.vllm_rollout.generate_rollout \\ + --custom-generate-function-path vime.rollout.vllm_streaming_rollout.generate_streaming + +The outer rollout loop (semaphore, dp_rank balancing, abort orchestration, +partial-rollout buffer hand-off) is still owned by ``vllm_rollout``; this file +only replaces the inner HTTP call. + +vime/vLLM counterpart of ``slime.rollout.sglang_streaming_rollout``. The +behavioural difference from sglang matters here: sglang's streamed +``meta_info.output_token_logprobs`` is **cumulative** (every chunk references +the full list-so-far), whereas vLLM's ``/inference/v1/generate`` SSE chunks +carry **delta** ``token_ids`` + ``logprobs`` per +``GenerateResponseStreamChoice`` — so we *accumulate* the per-chunk deltas +(``+=``) rather than overwriting from each chunk. Each delta choice has the +same shape as the non-streaming choice, so :func:`_inference_generate_tokens_and_logprobs` +parses it unchanged. +""" + +import json +import logging +from argparse import Namespace +from typing import Any + +from vime.rollout.vllm_rollout import ( + GenerateState, + _align_engine_tokens_and_logprobs, + _align_mm_feature_placeholders_to_tokens, + _apply_vllm_routed_experts, + _base_dataset_prompt_ids, + _build_inference_sampling_params, + _coerce_flat_int_token_ids, + _inference_generate_tokens_and_logprobs, + _mm_render_response_to_generate_body, + _prepare_prompt_ids, + _vllm_meta_from_generate_choice, +) +from vime.utils import http_utils +from vime.utils.processing_utils import encode_image_for_rollout_engine +from vime.utils.trace_utils import build_vllm_meta_trace_attrs, trace_span +from vime.utils.types import Sample + +__all__ = ["generate_streaming"] + +logger = logging.getLogger(__name__) + + +async def generate_streaming(args: Namespace, sample: Sample, sampling_params: dict[str, Any]) -> Sample: + """Streaming counterpart to :func:`vime.rollout.vllm_rollout.generate`. + + Writes the accumulated state from each SSE chunk onto ``sample`` so an abort + that cuts the stream still leaves a coherent partial sample behind. + """ + if args.ci_test: + assert isinstance(sample.prompt, str) + + state = GenerateState(args) + base = f"http://{args.vllm_router_ip}:{args.vllm_router_port}" + url = f"{base}/inference/v1/generate" + + assert ( + sample.status == Sample.Status.PENDING or sample.status == Sample.Status.ABORTED + ), f"Sample status is {sample.status}" + + prompt_ids = _prepare_prompt_ids(sample, state.tokenizer, state.processor) + base_prompt_ids = _base_dataset_prompt_ids(sample, state.tokenizer, state.processor) + + # Multimodal samples use the same render-dance as the non-streaming text + # path (/v1/chat/completions/render → features), then stream the generate + # call. Streaming only changes how output is returned (SSE deltas vs one + # JSON); the image render (input prep) is identical. Built below once + # sampling params + token_ids are resolved. + images = sample.multimodal_inputs.get("images") if sample.multimodal_inputs else None + + params = dict(sampling_params) + if len(sample.response) > 0: + params["max_new_tokens"] -= len(sample.tokens) - len(base_prompt_ids) + + assert params["max_new_tokens"] >= 0, ( + f"max_new_tokens: {params['max_new_tokens']} should not be less than 0 " + f"(after partial continuation adjustment; tokens={len(sample.tokens)}, base_prompt={len(base_prompt_ids)})" + ) + if params["max_new_tokens"] == 0: + sample.status = Sample.Status.TRUNCATED + return sample + inference_sampling_params = _build_inference_sampling_params(params) + + if not sample.tokens: + sample.tokens = prompt_ids + + # vLLM ``/inference/v1/generate`` is token-only. On partial continuation, + # send the full prompt+response prefix so the engine continues from the + # current sample state (mirrors the non-streaming text path). + if len(sample.response) > 0: + token_ids = _coerce_flat_int_token_ids(sample.tokens) + else: + token_ids = prompt_ids + + # Use session_id for consistent_hash routing (vime convention: x-session-id + # header + policy "consistent_hash"). See vllm_rollout.generate. + headers = None + if sample.session_id and getattr(args, "router_policy", None) == "consistent_hash": + headers = {"x-session-id": sample.session_id} + + payload: dict[str, Any] + if images: + # Same render-dance as vllm_rollout.generate's MM path, then stream. + # mm placeholders live in the (stable) prompt prefix, so re-rendering and + # re-aligning to the current token_ids holds across partial continuations. + content: list[dict[str, Any]] = [{"type": "text", "text": sample.prompt}] + for image in images: + content.append({"type": "image_url", "image_url": {"url": encode_image_for_rollout_engine(image)}}) + render_payload = {"model": args.hf_checkpoint, "messages": [{"role": "user", "content": content}]} + with trace_span(sample, "vllm_mm_render", attrs={"model": args.hf_checkpoint}): + render_data = await http_utils.post( + f"{base}/v1/chat/completions/render", render_payload, headers=headers + ) + payload = _mm_render_response_to_generate_body(render_data, args.hf_checkpoint) + if token_ids: + _align_mm_feature_placeholders_to_tokens(payload, token_ids) + payload["token_ids"] = token_ids + payload["sampling_params"] = inference_sampling_params + payload["stream"] = True + else: + payload = { + "model": args.hf_checkpoint, + "token_ids": token_ids, + "sampling_params": inference_sampling_params, + "stream": True, + } + + # Snapshot pre-call sample state. vLLM's SSE chunks are *deltas* within this + # call; on each chunk we append the delta and rebuild the post-call view of + # the sample = prior state + accumulated deltas. A mid-stream break leaves + # the sample exactly at the boundary of the last chunk we observed. + base_tokens = list(sample.tokens) + base_response = sample.response or "" + base_response_length = sample.response_length + base_log_probs = list(sample.rollout_log_probs or []) + base_loss_mask = list(sample.loss_mask) if sample.loss_mask is not None else None + + skip_sp = params.get("skip_special_tokens") + skip_decode = True if skip_sp is None else bool(skip_sp) + + call_tokens: list[int] = [] + call_log_probs: list[float] = [] + last_choice: dict[str, Any] | None = None + last_usage: dict[str, Any] | None = None + finish_reason: Any = None + + client = http_utils._http_client + assert client is not None, "http client not initialized; call init_http_client first" + + with trace_span( + sample, "vllm_inference_generate_stream", attrs={"max_new_tokens": params["max_new_tokens"]} + ) as span: + async with client.stream("POST", url, json=payload, headers=headers) as response: + response.raise_for_status() + async for raw_line in response.aiter_lines(): + if not raw_line or not raw_line.startswith("data:"): + continue + data_str = raw_line[len("data:") :].strip() + if not data_str or data_str == "[DONE]": + continue + try: + chunk = json.loads(data_str) + except json.JSONDecodeError: + logger.warning("vllm_streaming: skipping non-JSON chunk: %r", data_str[:120]) + continue + + choices = chunk.get("choices") or [] + if not choices: + # usage-only / keepalive chunk + if chunk.get("usage"): + last_usage = chunk["usage"] + continue + choice = choices[0] + last_choice = choice + if chunk.get("usage"): + last_usage = chunk["usage"] + if choice.get("finish_reason"): + finish_reason = choice["finish_reason"] + + # Each streamed choice carries only this chunk's *delta* tokens + # (GenerateResponseStreamChoice), so accumulate. + delta_tokens, delta_log_probs = _inference_generate_tokens_and_logprobs(choice) + if delta_tokens: + call_tokens += delta_tokens + call_log_probs += delta_log_probs + + # Surface partial state on the sample immediately. If the outer + # abort path cuts us, whatever we've written so far is what + # survives. Decode the *accumulated* tokens (not the per-chunk + # delta) so multi-token characters straddling a chunk boundary + # decode correctly. + sample.tokens = base_tokens + call_tokens + sample.response = base_response + ( + state.tokenizer.decode(call_tokens, skip_special_tokens=skip_decode) if call_tokens else "" + ) + sample.response_length = base_response_length + len(call_tokens) + sample.rollout_log_probs = base_log_probs + call_log_probs + if base_loss_mask is not None: + assert args.partial_rollout and args.mask_offpolicy_in_partial_rollout + sample.loss_mask = base_loss_mask + [1] * len(call_tokens) + + if state.aborted: + break + + if finish_reason and last_choice is not None: + span.update(build_vllm_meta_trace_attrs({"choices": [last_choice], "usage": last_usage})) + + if finish_reason and last_choice is not None: + # Finalize exactly like the non-streaming path: align logprobs to tokens, + # rebuild meta + output_token_logprobs, then let Sample own status. + new_response_tokens, new_response_log_probs = _align_engine_tokens_and_logprobs(call_tokens, call_log_probs) + + meta = _vllm_meta_from_generate_choice(args, last_choice, last_usage) + if new_response_tokens: + meta["output_token_logprobs"] = [ + [float(lp), int(tid)] for lp, tid in zip(new_response_log_probs, new_response_tokens, strict=True) + ] + + sample.tokens = base_tokens + new_response_tokens + sample.response = base_response + ( + state.tokenizer.decode(new_response_tokens, skip_special_tokens=skip_decode) if new_response_tokens else "" + ) + sample.response_length = base_response_length + len(new_response_tokens) + sample.rollout_log_probs = base_log_probs + new_response_log_probs + if base_loss_mask is not None: + assert args.partial_rollout and args.mask_offpolicy_in_partial_rollout + sample.loss_mask = base_loss_mask + [1] * len(new_response_tokens) + + sample.update_from_meta_info(args, meta) + # MoE routing replay (when requested) ships on the terminal choice. + _apply_vllm_routed_experts(args, sample, last_choice) + elif state.aborted: + sample.status = Sample.Status.ABORTED + + return sample diff --git a/vime/utils/arguments.py b/vime/utils/arguments.py index dc4efbd41..9c600d3d5 100644 --- a/vime/utils/arguments.py +++ b/vime/utils/arguments.py @@ -739,6 +739,8 @@ def add_algo_arguments(parser): help=( "Path to save the model in HuggingFace format when using Megatron backend. " "The model will be saved to `save_hf.format(rollout_id)`. " + "In raw Megatron-to-HF mode, weights are saved with the same quantization config " + "as `--hf-checkpoint`. " ), ) reset_arg(parser, "--seed", type=int, default=1234) @@ -1132,6 +1134,21 @@ def add_debug_arguments(parser): ) # --load-debug-rollout-data, --debug-rollout-only, --debug-train-only # are parsed early in _pre_parse_mode() and merged later. + parser.add_argument( + "--load-forge-rollout-data", + type=str, + default=None, + help=( + "Path (or {rollout_id} template) to a dumped rollout .pt file replayed by " + "vime.rollout.forge_load.generate_rollout. Mirrors --load-debug-rollout-data's " + "format(rollout_id=...) convention: a path without the placeholder is treated as " + "a literal file and reused across every rollout_id; a path containing {rollout_id} " + "loads a per-rollout file (with eval_.pt for the eval pipeline). Unlike " + "--load-debug-rollout-data, this does NOT force debug_train_only / skip_vllm -- " + "vLLM servers, router, weight_update and the colocate offload/onload dance all " + "stay live, which is the point (memory measurement at long context)." + ), + ) parser.add_argument( "--load-debug-rollout-data-subsample", type=float, @@ -1306,18 +1323,6 @@ def add_rollout_buffer_arguments(parser): "can process all samples including filtered ones." ), ) - parser.add_argument( - "--disable-rollout-trim-samples", - action="store_true", - default=False, - help="disable trim samples in rollout buffer when converting samples to train data", - ) - parser.add_argument( - "--use-dynamic-global-batch-size", - action="store_true", - default=False, - help="enable dynamic global batch size, disable trim samples in rollout buffer when converting samples to train data", - ) return parser def add_custom_megatron_plugins_arguments(parser): diff --git a/vime/utils/data.py b/vime/utils/data.py index 3f1417437..d158ea627 100644 --- a/vime/utils/data.py +++ b/vime/utils/data.py @@ -282,20 +282,6 @@ def __len__(self): return len(self.samples) -def get_minimum_num_micro_batch_size(total_lengths, max_tokens_per_gpu): - # use first fit to get the number of micro batches - batches = [] - for length in total_lengths: - for i in range(len(batches)): - if batches[i] + length <= max_tokens_per_gpu: - batches[i] += length - break - else: - batches.append(length) - - return len(batches) - - def process_rollout_data(args, rollout_data_ref, dp_rank, dp_size): assert len(rollout_data_ref) == dp_size rollout_data = ray.get(rollout_data_ref[dp_rank].inner) diff --git a/vime/utils/dp_schedule.py b/vime/utils/dp_schedule.py new file mode 100644 index 000000000..e30621ca5 --- /dev/null +++ b/vime/utils/dp_schedule.py @@ -0,0 +1,191 @@ +"""Per-rollout DP/microbatch scheduling. + +Pure-Python logic that decides, for one rollout's worth of sample lengths, +how to group samples into micro-batches and which DP rank owns each mbs. +Lives outside the ray/sglang-importing modules so it can be unit-tested +under CPU-only CI. + +The scheduling philosophy is **pack first, distribute second**: + + 1. Group samples by rollout id (``rollout_indices[i]`` = + ``samples[i].index``) and split rollouts into steps of + ``global_batch_size`` rollouts each. In the common case one rollout + emits one training sample so this is the same as a contiguous chunk; + under compact / subagent one rollout may emit multiple training + samples, in which case all of those samples stay in the same step. + 2. For each step, pack its samples into ``K`` micro-batches with a + single first-fit pass (dynamic batch) or fixed-size chunking + (static batch). + 3. Adjust ``K`` to a multiple of ``dp_size * (mb_group if vpp>1 else 1)`` + by splitting the largest multi-sample bins (dynamic only). + 4. Distribute the ``K`` mbs across ``dp_size`` ranks, ``K / dp_size`` + each, with either a strided round-robin or a Karmarkar-Karp pass on + mbs token sums. + +Invariants guaranteed by :func:`build_dp_schedule` (asserted by the tests): + - every DP rank runs the **same** ``num_microbatches`` per training step + (required for PP sync); + - every mbs (dynamic path) holds ``<= max_tokens_per_gpu * cp_size`` + tokens, with one exception — an individual sample larger than that cap + lands alone in its own mbs (and that mbs is the only one allowed to + exceed the cap); + - the union of per-rank sample indices equals the set of samples kept + after trimming trailing rollouts (every kept sample placed exactly + once); + - flattening ``micro_batch_indices`` for a rank yields + ``range(num_samples_rank)`` (each rank's samples are tiled exactly + once by its mbs schedule). +""" + +from __future__ import annotations + +import logging +from typing import Any + +from vime.utils.seqlen_balancing import expand_bins_by_splitting, first_fit_pack, get_seqlen_balanced_partitions + +logger = logging.getLogger(__name__) + + +def _pack_step_into_mbs( + step_lengths: list[int], + *, + use_dynamic_batch_size: bool, + max_per_bin: int | None, + micro_batch_size: int | None, +) -> list[list[int]]: + """Group a step's samples into mbs. Returns ``mbs[k]`` = local indices into ``step_lengths``.""" + if use_dynamic_batch_size: + assert max_per_bin is not None + return first_fit_pack(step_lengths, max_per_bin) + assert micro_batch_size is not None + n = len(step_lengths) + return [list(range(i, min(i + micro_batch_size, n))) for i in range(0, n, micro_batch_size)] + + +def build_dp_schedule( + args: Any, + train_parallel_config: dict, + total_lengths: list[int], + *, + global_batch_size: int, + rollout_indices: list[int], +) -> tuple[list[list[int]], list[list[list[int]]], list[int], list[int]]: + """Compute the per-rank DP partition and micro-batch schedule. + + See module docstring for the pack-first-distribute-second strategy. + + Args: + args: Namespace with ``micro_batch_size``, ``use_dynamic_batch_size``, + ``max_tokens_per_gpu``, ``balance_data``. + train_parallel_config: ``{"dp_size", "cp_size", "vpp_size", + "microbatch_group_size_per_vp_stage"}``. + total_lengths: token count per sample, indexed globally. + global_batch_size: number of rollouts (NOT training samples) per + training step. Number of training steps = + ``num_rollouts // global_batch_size``; trailing rollouts whose + samples don't fit are dropped. + rollout_indices: rollout id for each sample (``samples[i].index``). + Samples sharing the same id are kept together in one step. + + Returns: + ``(partitions, micro_batch_indices, num_microbatches, global_batch_sizes)``. + ``global_batch_sizes[s]`` = rollout count for step s (constant + ``global_batch_size`` for every step). + """ + dp_size = train_parallel_config["dp_size"] + cp_size = train_parallel_config["cp_size"] + vpp_size = train_parallel_config["vpp_size"] + mb_group = train_parallel_config["microbatch_group_size_per_vp_stage"] + + max_per_bin = None + if args.use_dynamic_batch_size: + assert args.max_tokens_per_gpu is not None + max_per_bin = args.max_tokens_per_gpu * cp_size + + # mbs count per step must be divisible by (dp_size * mb_group_for_vpp) so + # every rank ends up with the same num_mbs and (for VPP) the per-rank mbs + # count is a multiple of mb_group. + align_to = dp_size * (mb_group if vpp_size > 1 else 1) + + # Group samples by rollout id (preserve first-occurrence order). All + # samples from one rollout stay in a single step so the per-rollout loss + # reducer is well-defined. + rollout_id_to_samples: dict[int, list[int]] = {} + for sample_pos, rid in enumerate(rollout_indices): + rollout_id_to_samples.setdefault(rid, []).append(sample_pos) + rollout_ids = list(rollout_id_to_samples.keys()) + + num_steps = len(rollout_ids) // global_batch_size + assert num_steps >= 1, ( + f"num_rollouts ({len(rollout_ids)}) < global_batch_size ({global_batch_size}); " + f"need at least one rollout per step." + ) + + partitions: list[list[int]] = [[] for _ in range(dp_size)] + micro_batch_indices: list[list[list[int]]] = [[] for _ in range(dp_size)] + num_microbatches: list[int] = [] + global_batch_sizes: list[int] = [] + + for step_i in range(num_steps): + step_rollouts = rollout_ids[step_i * global_batch_size : (step_i + 1) * global_batch_size] + sample_indices = [pos for rid in step_rollouts for pos in rollout_id_to_samples[rid]] + step_lengths = [total_lengths[i] for i in sample_indices] + global_batch_sizes.append(global_batch_size) + assert len(sample_indices) >= dp_size, ( + f"step {step_i}: {len(sample_indices)} samples < dp_size {dp_size}; " + f"each step needs at least one sample per rank." + ) + + # 1. Pack samples in this step into mbs with one global pass. + # ``step_mbs`` indices are LOCAL into ``sample_indices``. + step_mbs = _pack_step_into_mbs( + step_lengths, + use_dynamic_batch_size=args.use_dynamic_batch_size, + max_per_bin=max_per_bin, + micro_batch_size=getattr(args, "micro_batch_size", None), + ) + + # 2. Align mbs count to a multiple of ``align_to``. + target_K = max(((len(step_mbs) + align_to - 1) // align_to) * align_to, align_to) + if target_K != len(step_mbs): + if args.use_dynamic_batch_size: + expand_bins_by_splitting(step_mbs, target_K, step_lengths) + assert len(step_mbs) == target_K, ( + f"dynamic path: could only produce {len(step_mbs)} mbs after maximal splitting; " + f"need {target_K}. step {step_i} has {len(sample_indices)} samples, below the " + f"alignment threshold ({align_to})." + ) + else: + raise AssertionError( + f"static path: num_mbs ({len(step_mbs)}) is not a multiple of " + f"dp_size * mb_group ({align_to}); got " + f"step_size={len(sample_indices)}, micro_batch_size={args.micro_batch_size}, " + f"dp_size={dp_size}, mb_group={mb_group if vpp_size > 1 else 1}. " + f"Splitting static mbs would break the fixed-size invariant; adjust the config " + f"so step_size % (dp_size * micro_batch_size * mb_group) == 0." + ) + + K = len(step_mbs) + num_mbs_per_rank = K // dp_size + num_microbatches.append(num_mbs_per_rank) + + # 3. Distribute mbs across ranks: KK on mbs token sums when balance_data is on, + # otherwise a strided round-robin. Both produce ``num_mbs_per_rank`` mbs per + # rank (equal_size=True is what KK needs for PP to stay synced). + if args.balance_data: + mbs_token_sums = [sum(step_lengths[i] for i in bin_) for bin_ in step_mbs] + rank_mbs_idx = get_seqlen_balanced_partitions(mbs_token_sums, dp_size, equal_size=True) + else: + rank_mbs_idx = [list(range(r, K, dp_size)) for r in range(dp_size)] + + # 4. Build per-rank partitions (global sample indices) and micro_batch_indices + # (local indices into partitions[r]). + for r in range(dp_size): + for mbs_idx in rank_mbs_idx[r]: + mbs_locals = step_mbs[mbs_idx] # local indices into sample_indices + local_start = len(partitions[r]) + partitions[r].extend(sample_indices[i] for i in mbs_locals) + micro_batch_indices[r].append(list(range(local_start, local_start + len(mbs_locals)))) + + return partitions, micro_batch_indices, num_microbatches, global_batch_sizes diff --git a/vime/utils/seqlen_balancing.py b/vime/utils/seqlen_balancing.py index a5dd71f94..5736d8850 100644 --- a/vime/utils/seqlen_balancing.py +++ b/vime/utils/seqlen_balancing.py @@ -177,6 +177,58 @@ def _check_and_sort_partitions(partitions): return _check_and_sort_partitions(partitions) +def first_fit_pack(total_lengths, max_tokens_per_bin): + """First-fit bin packing. + + Returns ``list[list[int]]`` — each bin is a list of indices into ``total_lengths``. + Bin sums are ``<= max_tokens_per_bin`` whenever every individual ``length`` fits; + an oversized sample lands alone in its own bin with sum equal to its length. + """ + bins: list[list[int]] = [] + bin_sums: list[int] = [] + for idx, length in enumerate(total_lengths): + for j in range(len(bins)): + if bin_sums[j] + length <= max_tokens_per_bin: + bins[j].append(idx) + bin_sums[j] += length + break + else: + bins.append([idx]) + bin_sums.append(length) + return bins + + +def _split_bin_by_tokens(bin_indices: list[int], lengths) -> list[list[int]]: + """Split a bin's indices into two halves balanced by total tokens (LPT heuristic). + + Returns ``[left, right]`` where both lists together cover ``bin_indices``. Because + each half is a strict subset of ``bin_indices``, both have token sums ``<=`` the + original bin's sum — useful when you need to grow a bin packing without ever + creating a bin larger than the originals. + """ + halves: list[list[int]] = [[], []] + sums = [0, 0] + for idx in sorted(bin_indices, key=lambda i: -lengths[i]): + h = 0 if sums[0] <= sums[1] else 1 + halves[h].append(idx) + sums[h] += lengths[idx] + return halves + + +def expand_bins_by_splitting(bins: list[list[int]], target_count: int, lengths) -> None: + """Grow ``bins`` in place to ``target_count`` by repeatedly splitting the largest + multi-sample bin via :func:`_split_bin_by_tokens`. Stops early if every remaining + bin is a singleton (no bin can be split further).""" + while len(bins) < target_count: + candidates = [(sum(lengths[i] for i in b), idx) for idx, b in enumerate(bins) if len(b) > 1] + if not candidates: + break + _, idx = max(candidates) + left, right = _split_bin_by_tokens(bins[idx], lengths) + bins[idx] = left + bins.append(right) + + def get_reverse_idx(idx_map): reverse_idx_map = copy.deepcopy(idx_map) diff --git a/vime/utils/types.py b/vime/utils/types.py index 9ac3916c2..a05a2d3bd 100644 --- a/vime/utils/types.py +++ b/vime/utils/types.py @@ -11,6 +11,14 @@ class Sample: group_index: int | None = None index: int | None = None + # Id of the rollout this sample came from. Defaults to ``None`` and the + # downstream pipeline falls back to ``index`` (so the default rollout + # path, where one execution = one training sample, sees rollout_id == + # index). Compact / subagent paths that split one rollout execution into + # multiple training samples should set the same ``rollout_id`` on every + # sibling, so loss aggregation averages within the rollout instead of + # over-counting it. + rollout_id: int | None = None # prompt prompt: str | list[dict[str, str]] = "" tokens: list[int] = field(default_factory=list)