From 31398eb552b741c9bcfa054f9eb71a451ca0adba Mon Sep 17 00:00:00 2001 From: smirnovlad Date: Tue, 30 Jun 2026 10:56:40 +0200 Subject: [PATCH 1/2] docs: add guide for running the service with the PRM scorer --- docs/service/running_with_prm.md | 281 +++++++++++++++++++++++++++++++ 1 file changed, 281 insertions(+) create mode 100644 docs/service/running_with_prm.md diff --git a/docs/service/running_with_prm.md b/docs/service/running_with_prm.md new file mode 100644 index 00000000..83a23f4c --- /dev/null +++ b/docs/service/running_with_prm.md @@ -0,0 +1,281 @@ +# Running the ThinkBooster Service with the PRM Scorer + +How to launch the ThinkBooster service so you can hit it over an OpenAI-compatible +API and run **offline best-of-N with a PRM scorer**, where: + +- **generation** happens **remotely** on OpenRouter (e.g. `openai/gpt-oss-20b`), and +- the **PRM reward model** runs **locally on one GPU** and re-ranks the sampled trajectories. + +This is the recommended single-GPU setup ("PRM-only mode"). You do **not** need a local +generation model — no `VLLM_MODEL_PATH`. The only thing the GPU does is run the PRM. + +``` + extra_body={"tts_api_key": ""} + your client ──HTTP──► ThinkBooster service ──── generates N trajectories ────► OpenRouter + (openai SDK / curl) /v1/offline_bon/prm (gpt-oss-20b) + │ + │ scores each trajectory with a local PRM (GPU) + ▼ + Qwen2.5-Math-PRM-7B ── picks the best trajectory ──► response +``` + +--- + +## 1. Prerequisites + +| Need | Detail | +|---|---| +| GPU | 1× NVIDIA GPU, **≥24 GB** (PRM-7B in bf16 ≈ 15 GB + KV cache). 3090/4090/A6000/A100/H100 all work. | +| NVIDIA driver | Required. For the Docker path you also need `nvidia-container-toolkit` installed on the host. | +| OpenRouter key | Generation runs on OpenRouter — get a key at https://openrouter.ai/keys | +| Tooling | `git` + **either** Docker **or** a Python 3.10+ environment. | + +--- + +## 2. Get the code + +```bash +git clone https://github.com/IINemo/thinkbooster.git +cd thinkbooster + +# lm-polygraph is a separate repo (gitignored). It is required both for the +# Docker build and for the scorers. Clone it INTO the repo root: +git clone https://github.com/IINemo/lm-polygraph.git + +# .env is copied for service config. The OpenRouter key here is OPTIONAL because each +# request carries its own key via tts_api_key (see §4). Set it only as a fallback. +cp service_app/.env.example .env +# (optional) edit .env and set: OPENROUTER_API_KEY=sk-or-... # used only if a request omits tts_api_key +``` + +--- + +## 3. Launch the service — pick ONE option + +> On a single-GPU box, **Option B (bare metal) is the quickest and is our deployment path.** +> Use **Option A (Docker)** if you want an isolated/reproducible container. + +### Option A — Docker (GPU-enabled) + +The committed `docker-compose.yml` is API-only (no GPU). Add a small override file that +attaches the GPU and the PRM env. Create **`docker-compose.gpu.yml`** in the repo root: + +```yaml +services: + thinkbooster: + environment: + # --- PRM scorer (local, on GPU) --- + - PRM_MODEL_PATH=Qwen/Qwen2.5-Math-PRM-7B + - PRM_DEVICE=cuda:0 + - PRM_USE_VLLM=true + - PRM_GPU_MEMORY_UTILIZATION=0.9 + # generation is on OpenRouter, so DO NOT set VLLM_MODEL_PATH + - VLLM_WORKER_MULTIPROC_METHOD=spawn # avoids CUDA init errors + - CUDA_VISIBLE_DEVICES=0 + - HF_HOME=/app/.hf # cache model weights on the mounted volume + volumes: + - ~/.cache/huggingface:/app/.hf # reuse HF cache → PRM weights downloaded once + deploy: + resources: + reservations: + devices: + - driver: nvidia + count: 1 + capabilities: [gpu] +``` + +Build and start (PRM weights download on first run — a few minutes): + +```bash +docker compose -f docker-compose.yml -f docker-compose.gpu.yml up --build -d + +# health + logs +curl http://localhost:8001/health # {"status":"healthy", ...} +docker compose logs -f thinkbooster +``` + +Service is now on **`http://:8001`** (docs at `/docs`). + +> **Build gotchas** +> - The build fails if `lm-polygraph/` is missing — clone it first (step 2). +> - If the build errors on `README.md` not found, add `COPY README.md .` to +> `service_app/Dockerfile` just before the `pip install` line, then rebuild. + +### Option B — Bare metal in tmux (single GPU, simplest) + +```bash +# one-time: install deps + clone lm-polygraph + scorers +./setup.sh + +# keep the service alive after you disconnect +apt update && apt install -y tmux +tmux new -t thinkbooster + +# inside tmux — PRM-only mode (generation on OpenRouter, PRM on GPU) +export CUDA_VISIBLE_DEVICES=0 +export PRM_MODEL_PATH=Qwen/Qwen2.5-Math-PRM-7B +export PRM_DEVICE=cuda:0 +export PRM_USE_VLLM=true +export PRM_GPU_MEMORY_UTILIZATION=0.9 +export VLLM_WORKER_MULTIPROC_METHOD=spawn # avoids CUDA init errors +export PORT=8001 # default bare-metal port is 8080 +python service_app/main.py +``` + +Detach from tmux with `Ctrl-b d`. Logs are written to `logs//