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:

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).

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):

.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).