| # vLLM rollout (optional extra) |
|
|
| This directory is **not** part of the `align` package and nothing under `align/` |
| imports it (`tests/test_dependency_firewall.py::test_extras_are_not_reachable_from_core`). |
| The default GRPO rollout is the dependency-free native sampler, |
| `align/rollout/native.py`. Use this only when generation throughput is the |
| bottleneck and a dedicated engine GPU is available. |
|
|
| ## Runs in its own venv |
|
|
| `rollout_worker.py` and `rollout_pilot.py` run under `.venv-rollout` |
| (vllm 0.26.0, ray, pyarrow, torch cu130), **never** the `align` venv, whose |
| only dependencies are torch, triton, numpy and safetensors. Build it with |
| `uv venv .venv-rollout` plus `torch==2.11.0 triton==3.6.0 vllm==0.26.0 ray pyarrow`. |
|
|
| ## How it talks to the trainer |
|
|
| Only through `align/rollout/protocol.py` (stdlib-only by contract): one JSON |
| file per message in a spool directory, atomically written, with a heartbeat |
| file for liveness. The worker puts `rollout_worker.py`'s repo root on |
| `sys.path` and imports that single module; `import align` and |
| `import align.rollout` are lazy / stdlib-only so no `align` dependency reaches |
| this venv. |
|
|
| Trainer side: |
|
|
| ```python |
| from align.rollout.client import make_rollout_client, spawn_worker |
| |
| proc = spawn_worker(spool, gpu=3, log_path=".../rollout_worker.log") |
| client = make_rollout_client("spool", spool=spool, timeout_s=3600) |
| ``` |
|
|
| `spawn_worker` launches `extras/vllm_rollout/rollout_worker.py` with the |
| required environment (flashinfer sampler off, caches under `/apps`, torchrun |
| variables scrubbed, own process group so teardown reaps the engine core). |
|
|
| ## Logprob basis caveat |
|
|
| The worker requests `SamplingParams(logprobs=0)`, which in vLLM are the |
| engine's **raw** logprobs (before temperature). The native backend returns |
| `log_softmax(logits / T)` untruncated, which is what the GRPO module recomputes |
| teacher-forced. The two coincide at `temperature == 1.0` (the stage-b setting). |
| For `T != 1` the spool backend's logprobs differ from the training-side basis |
| by the temperature rescaling, so rho will not be ~1 on-policy there; use |
| `T = 1.0` with this worker, or set vLLM's `logprobs_mode` to a processed mode |
| and accept that top-p/top-k then enter the logprob too. |
|
|
| Weight sync is filesystem-based: a changed `weights_dir` tears the engine down |
| and rebuilds it (~60 s), so sync every `sync_every` optimizer steps and let |
| `use_rho_correction=True` absorb the off-policyness. |
|
|
| ## Lean proving generation driver (experiment P2): `lean_gen.py` + `lean_prompts.py` |
|
|
| Generates Goedel-Prover-V2 proof attempts with vLLM and writes the shared Generation |
| record JSONL (one line per sample). Runs in `.venv-rollout`; `lean_prompts.py` and the |
| pure half of `lean_gen.py` are stdlib-only at import and are unit-tested from the align |
| venv (`tests/test_lean_gen_contract.py`, 36 tests, 15/15 hand mutants killed). |
|
|
| ```bash |
| cd /apps/development/rlhf/align && export TMPDIR=$PWD/.tmp/p2/<you> |
| # round 0 |
| .venv-rollout/bin/python extras/vllm_rollout/lean_gen.py --problems problems.jsonl \ |
| --out gen_r0.jsonl --n 4 --cap-tokens 16384 --tp 4 --seed 0 --wait-lock |
| # correction round 1 (parents = failed round-0 samples; needs a verification JSONL, see below) |
| .venv-rollout/bin/python extras/vllm_rollout/lean_gen.py --problems problems.jsonl \ |
| --out gen_r1.jsonl --n 1 --cap-tokens 16384 --round 1 \ |
| --corrections-from verif_r0.jsonl gen_r0.jsonl --wait-lock |
| # plan + prompts only, no GPU / lock |
| ... --dry-run |
| ``` |
|
|
| Rerunning the same command resumes. `<out>.manifest.json` has one `sessions[]` entry per |
| launch (model rev, vllm/torch versions, args, load/generate wall, tokens, finish counts, KV |
| blocks, peak GPU memory); `<out>.skipped.jsonl` lists correction prompts that did not fit. |
|
|
| **Fidelity.** `python extras/vllm_rollout/lean_prompts.py fidelity --model <snapshot> --out r.json` |
| (rollout venv) imports the reference `src/utils.py` and asserts identical prompt strings and |
| token ids on all 244 minif2f problems, identical `extrac_code` + `problem_check` on 63 synthetic |
| outputs, `get_error_str` / `generate_correction_prompt` on 60 random error sets, and |
| `compile.py:handle` vs `compile_view`. |
|
|
| **What a record means.** |
| * `cap_tokens` is the label every reported number carries. `sampling.max_tokens` is what vLLM |
| got: `min(cap, max_model_len - len(prompt))`. They differ only when the context window ends the |
| sample (`context_clipped: true`); at `--cap-tokens 40960 --max-model-len 40960` every sample is |
| clipped by its prompt length (~250-600 tokens). `--max-model-len` above the model's 40960 |
| position limit raises. |
| * `completion_ids` are vLLM's `token_ids`: they end with the stop token (`<|im_end|>`) when |
| `finish_reason == "stop"`. Train on them as they are; do not re-tokenize `completion_text`. |
| * `messages` are the PROMPT messages (what the chat template rendered; no completion), so a |
| correction round's history is `parent.messages + [assistant: parent.completion_text, user: errors]` |
| exactly as Goedel builds it. `lean_code` is Goedel's `full_code` or null, with the reason in |
| `lean_code_error` (extra field). |
| * Sampling: temperature 1.0, top_p 0.95, no top_k (Goedel's `SamplingParams`; the model's |
| `generation_config.json` says top_k=20 / T=0.6 but explicit params override it). One seed per |
| sample, derived from (seed, problem, sample_idx, round); not the cap. |
|
|
| **Correction rounds need two extra fields in the VERIFICATION record** (the shared contract does |
| not carry them): `repl_errors` (list of Lean REPL messages with `severity == "error"`, each |
| `{pos:{line,column}, endPos:{line,column}|null, data}`) and `repl_code` (the exact text those |
| positions index into; for the Goedel flow that is `lean_prompts.compile_view(lean_code)`, i.e. the |
| code with `import` / `set_option` / `open` lines removed). Without them `<error>` markers land on |
| the wrong lines, so a parent lacking them is skipped and counted (`no_repl_errors`), never guessed. |
|
|
| **Throughput/memory choices.** `gpu_memory_utilization=0.90` (20.7 GB of the A10G's 22.5 GiB; |
| measured peak `nvidia-smi` use 21.7-22.5 GB, i.e. 0.5-1.3 GB of slack for NCCL/allocator; 0.92 was |
| not tried), `max_num_seqs=128`, prefix caching on, CUDA graphs on, FLASH_ATTN backend (works on |
| sm_86), `VLLM_USE_FLASHINFER_SAMPLER=0`. KV pool: 28,533 blocks x 16 = 456,528 tokens (TP=4: each |
| GPU holds 2 of 8 KV heads, 36 KB/token). The driver steps the engine itself (`add_request`/`step`) |
| with a refill window instead of `LLM.generate` chunks, so a 16k-token straggler does not idle the |
| GPUs. Measured on Goedel-Prover-V2-8B, TP=4, seeded sampling (minif2f in file order, T=1.0, top_p=0.95): |
| |
| | run | cap | samples | mean / median / p90 completion tokens | finish=length | end-to-end tok/s | tok/s while >=64 running | |
| |---|---|---|---|---|---|---| |
| | 32 problems x n=4, max_num_seqs 128 | 16384 | 128 | 5066 / 3242 / 13993 | 10/128 (7.8%) | 918 | 1602 | |
| | 64 problems x n=4, max_num_seqs 256 | 16384 | 256 | 4996 / 3033 / 13641 | 21/256 (8.2%) | 1043 | 1363 (KV pool 100% full at peak) | |
| | 8 problems x n=2, max_num_seqs 128 | 40960 (context-clipped, max_tokens ~40.6-40.8k) | 16 | 4450 / 2662 / 7339 | 0/16 | 256 (tail of 1 sequence) | n/a (max 16 running) | |
|
|
| Per-sequence decode speed is ~25 tok/s at 85-128 concurrent sequences and ~74 tok/s alone: TP=4 |
| over PCIe (no NVLink) is allreduce-bound, so aggregate throughput saturates near 1.4-2.2k tok/s |
| beyond ~100 concurrent sequences, and 256 sequences only fill the KV pool. Keep `--max-num-seqs 128`. |
|
|
| ### Class-B trajectories, exit status, un-merged adapter (lane V2) |
|
|
| * `--trajectory-cap-tokens 40000` enforces the class-B cumulative cap; every record carries a |
| `trajectory` ledger (root id, round, cumulative completion tokens, truncation, budget clipping); |
| failed parents with no compiler error are `skipped_no_feedback` in `<out>.skipped.jsonl` |
| (or re-sampled with `--no-feedback-policy resample`). Goedel's rule (repair only problems |
| without an accepted sample) is the default; `--repair-solved` lifts it. Pass every earlier |
| round's `VERIF GEN` pair in `--corrections-from` so round-2+ lineage is complete (otherwise the |
| entry is `partial` and a cap refuses). `build_trajectory_ledger()` aggregates a final state per root. |
| * Exit status: 0 only for a complete sample set, 143 on SIGTERM/SIGINT. `pipeline_gen_verify.py` |
| also checks the manifest (exit 4 if not `complete`) and reaps the generator's process group, |
| sweeps `/dev/shm` and waits for the GPUs to drain after every generator exit. |
| * Un-merged adapter: `--lora-adapter <peft dir> [--lora-name adapter]` serves it via vLLM |
| `enable_lora`/`LoRARequest` (no bf16 merge). Adapter sha256 goes in the manifest, each record and |
| `model_rev` (`<base>+lora-<sha12>`). GPU smoke (run after the chains exit, takes the heavy lock): |
|
|
| ```bash |
| .venv-rollout/bin/python extras/vllm_rollout/lean_gen.py --problems <5 problems>.jsonl \ |
| --out .tmp/v2_smoke/gen_lora.jsonl --n 2 --cap-tokens 2048 --tp 1 --max-model-len 8192 \ |
| --lora-adapter <adapter_dir> --wait-lock |
| # expect: manifest sessions[-1].adapter.sha256 set, vllm engine log shows LoRA enabled, every |
| # record has "adapter", exit 0; then rerun the same command (must resume, 0 new samples). |
| ``` |
|
|
