Add RECIPE.md: shareable setup + run recipe with pitfalls
Browse files
RECIPE.md
ADDED
|
@@ -0,0 +1,205 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Recipe: OpenCode-only agentic RL on LFM2.5-2.6B, shipped as a LoRA adapter
|
| 2 |
+
|
| 3 |
+
Train a small open model *inside the OpenCode agent harness* β unmodified, running in
|
| 4 |
+
cloud sandboxes β with async GRPO, and distribute the result as a ~170 MB LoRA
|
| 5 |
+
adapter on top of [`LiquidAI/LFM2.5-2.6B`](https://huggingface.co/LiquidAI/LFM2.5-2.6B).
|
| 6 |
+
|
| 7 |
+
Everything runs on Hugging Face infrastructure: training on HF Jobs, sandboxes on HF
|
| 8 |
+
Sandboxes (Harbor's `hf-sandbox` backend). The only credential needed is an HF token.
|
| 9 |
+
|
| 10 |
+
Based on the stack from [The ultimate guide to multi-harness RL](https://huggingface.co/spaces/FineEnvs/multi-harness-rl)
|
| 11 |
+
(OpenEnv capture proxy + Harbor + TRL), narrowed to one harness and adapted for LoRA
|
| 12 |
+
and HF-only infrastructure.
|
| 13 |
+
|
| 14 |
+
---
|
| 15 |
+
|
| 16 |
+
## Architecture
|
| 17 |
+
|
| 18 |
+
```
|
| 19 |
+
ββββββββββββββββββββββββββββ HF Job (h200x2) ββββββββββββββββββββββββββββ
|
| 20 |
+
β β
|
| 21 |
+
GRPO groups βββΊ β TRL AsyncGRPOTrainer (GPU 0) β
|
| 22 |
+
of 8 rollouts β Β· LoRA r=32 on q/k/v/out_proj + w1/w2/w3 β
|
| 23 |
+
β Β· merges adapter β streams FULL weights over NCCL (136 ms/step) β
|
| 24 |
+
β β β² β
|
| 25 |
+
β βΌ weight sync β tokens + logprobs β
|
| 26 |
+
β vLLM 0.31 (GPU 1) βββ capture proxy βββ (mounted at /capture) β
|
| 27 |
+
β Β· --return-tokens-as-token-ids β
|
| 28 |
+
β Β· --logprobs-mode processed_logprobs β
|
| 29 |
+
β Β· --tool-call-parser lfm2 --enable-auto-tool-choice β
|
| 30 |
+
β β² β β
|
| 31 |
+
β OpenEnv Harbor env server (CPU) βββββββ β
|
| 32 |
+
β Β· task API Β· capture proxy Β· per-rollout sandbox lifecycle β
|
| 33 |
+
βββββββββββββββββ¬ββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 34 |
+
β https://<job-id>--8000.hf.jobs/capture
|
| 35 |
+
β (session id as x-api-key β Jobs proxy strips Authorization)
|
| 36 |
+
βββββββββββββββββΌβββββββββββββββββ
|
| 37 |
+
β HF Sandboxes (1 per rollout) β cpu-basic Jobs running the task image
|
| 38 |
+
β OpenCode β runs the task β β pulls Kaggle data from a public bucket,
|
| 39 |
+
β verifier grades the answer β agent loop owned entirely by OpenCode
|
| 40 |
+
ββββββββββββββββββββββββββββββββββ
|
| 41 |
+
```
|
| 42 |
+
|
| 43 |
+
**Data flow per rollout:** the trainer's `HarnessRolloutWorker` asks the env server
|
| 44 |
+
for a rollout β env server boots an HF Sandbox with the task β OpenCode runs inside,
|
| 45 |
+
calling the model through the capture proxy β proxy records exact prompt/completion
|
| 46 |
+
token IDs + per-token logprobs from vLLM β sandbox verifier scores the answer β
|
| 47 |
+
reward (correctness + tool-efficiency bonus) and the token-faithful trace go back to
|
| 48 |
+
the trainer β GRPO updates the LoRA β merged weights stream to vLLM for the next group.
|
| 49 |
+
|
| 50 |
+
---
|
| 51 |
+
|
| 52 |
+
## Key design decisions (and why)
|
| 53 |
+
|
| 54 |
+
| Decision | Why |
|
| 55 |
+
|---|---|
|
| 56 |
+
| **LoRA r=32, Ξ±=64, on `q,k,v,out_proj,w1,w2,w3`** | Distributable artifact (~170 MB). Module names verified against the model's safetensors index. Short-conv projections left out for portability. |
|
| 57 |
+
| **Merged-weight sync, not vLLM `--enable-lora` sync** | Generation flows through the capture proxy addressed by a *fixed model name*; TRL's adapter sync rotates versioned adapter names that only its own built-in worker knows. Merged sync is 136 ms/step over NCCL for 2.6B β and vLLM serves adapters ~1.35Γ slower per token anyway. Checkpoints still contain only the adapter. |
|
| 58 |
+
| **HF Sandboxes (`hf-sandbox` backend, harbor β₯0.24)** | No Daytona/E2B account; each rollout is an HF Jobs cpu-basic machine under the same token. Requires `docker_image` in task.toml β SmolDataEnvs tasks all `FROM` one prebuilt image, so it's a one-line patch. |
|
| 59 |
+
| **Jobs-proxy exposure instead of a Gradio tunnel** | A Gradio tunnel died 11h into run #1 and stalled it silently. The Jobs proxy (`https://<job-id>--8000.hf.jobs/capture`) is first-party and stable. |
|
| 60 |
+
| **Session id as `x-api-key`** | The Jobs proxy **strips the `Authorization` header** (verified with a header-echo probe). The capture proxy also accepts `x-api-key`; OpenCode's `@ai-sdk/openai-compatible` provider supports custom headers. One-line seam patch. |
|
| 61 |
+
| **trl==1.15.0, not 1.16** | 1.16 made tree packing unconditional and *refuses hybrid models* β LFM2.5 has conv layers. 1.15.0 has `HarnessRolloutWorker` (#6947) but predates tree packing. Bonus: it matches the article's actual training era. |
|
| 62 |
+
| **vllm/vllm-openai:v0.31.0 image** | Top of trl 1.15's supported band; ships the `lfm2` tool parser (for LFM2.5's `<|tool_call_start|>[fn(args)]<|tool_call_end|>` format) and predates vLLM's `mcp>=2` pin, which conflicts with openenv's fastmcp 3.x. |
|
| 63 |
+
| **Reward = correctness + β€0.1 tool-efficiency bonus** | From the article: without it, all-correct GRPO groups carry no contrast and tool calls creep up. A wrong answer scores 0 regardless, so the bonus can't reward giving up. |
|
| 64 |
+
| **Watchdog + checkpoint pushes every 100 steps** | transformers v5 silently dropped `PushToHubCallback` β periodic saves stay local unless you push them yourself. A heartbeat-touched file + 30-min-stall killer turns silent stalls into loud, resumable failures. |
|
| 65 |
+
|
| 66 |
+
---
|
| 67 |
+
|
| 68 |
+
## Exact recipe
|
| 69 |
+
|
| 70 |
+
### Versions
|
| 71 |
+
|
| 72 |
+
| Component | Version | Notes |
|
| 73 |
+
|---|---|---|
|
| 74 |
+
| trl | **1.15.0** (`trl[peft]`) | `HarnessRolloutWorker` in `trl.experimental.async_grpo.openenv_harness` |
|
| 75 |
+
| vLLM | **0.31.0** (docker image) | needs `kernels==0.16.0` pip-installed for trl's `kernels-community/flash-attn3` |
|
| 76 |
+
| openenv | **0.10.0** (`openenv[harbor]`) | `harbor_env` comes from the OpenEnv **repo** (`envs/`), not the wheel β clone + `.pth` |
|
| 77 |
+
| harbor | **β₯0.24.0** (`harbor[hf-sandbox]`) | the version with `EnvironmentType.HF_SANDBOX` |
|
| 78 |
+
| transformers | **β₯5.2.0** (`--no-deps` next to vLLM's stack) | v5 removed PushToHubCallback β see callback below |
|
| 79 |
+
| Base model | `LiquidAI/LFM2.5-2.6B` | `Lfm2ForCausalLM`, 30 layers (20 conv + 10 attention), tied embeddings |
|
| 80 |
+
|
| 81 |
+
### Data
|
| 82 |
+
|
| 83 |
+
- **Train:** 1,000 SmolDataEnvs Harbor tasks (400 medium / 600 hard, fixed order),
|
| 84 |
+
pinned to dataset revision `4719635555deβ¦` with per-instruction SHA-256 checks.
|
| 85 |
+
- **Test:** 250 held-out tasks (no notebook/question overlap), revision `b130595b57β¦`.
|
| 86 |
+
- Task lists + prep script: `prepare_tasks.py` + `data/{train,test}.json` in the code repo.
|
| 87 |
+
It also injects `docker_image = "docker.io/savatar101/env-data-agent-train:base"` into
|
| 88 |
+
every task.toml (required by the hf-sandbox backend) and sets 1 CPU / 4 GB.
|
| 89 |
+
- Task data files are pulled into each sandbox from the public bucket
|
| 90 |
+
`AdithyaSK/jupyter-agent-kaggle-all` via the task's own healthcheck hook.
|
| 91 |
+
|
| 92 |
+
### Training hyperparameters
|
| 93 |
+
|
| 94 |
+
| Knob | Value |
|
| 95 |
+
|---|---|
|
| 96 |
+
| Algorithm | async GRPO (`AsyncGRPOTrainer` + `HarnessRolloutWorker`, loop owned by the harness) |
|
| 97 |
+
| LoRA | r=32, Ξ±=64, dropout 0, targets `q_proj,k_proj,v_proj,out_proj,w1,w2,w3` |
|
| 98 |
+
| Steps | 1,000 (1 GRPO group of 8 rollouts β 1 task per step) |
|
| 99 |
+
| Learning rate | 1e-6 |
|
| 100 |
+
| Sampling | temperature 1.0, top_p 1.0, top_k 0 β **no truncation** (it breaks logprob agreement; importance ratio should sit at 0.998β1.0) |
|
| 101 |
+
| Output cap | 4,096 tokens, **same in training and eval** (a mismatch collapsed one of the article's Qwen runs) |
|
| 102 |
+
| Agent budget | 32 model calls/rollout (`max_tool_calling_iterations` β the loop-owning default is 8!) |
|
| 103 |
+
| Staleness | max 4 optimizer steps (default) |
|
| 104 |
+
| Reward | `env_reward` (1/0) + `0.1 Γ max(0, (32 β tool_calls) / 32)` on correct answers |
|
| 105 |
+
| Hardware | HF Jobs `h200x2` (~90 s/step early, ~35β50 s/step once the rollout queue saturates) |
|
| 106 |
+
|
| 107 |
+
### Launch
|
| 108 |
+
|
| 109 |
+
```bash
|
| 110 |
+
# 0. Upload the code bundle (one time)
|
| 111 |
+
hf repos create YOU/lfm-opencode-lora-rl-code --type model --private
|
| 112 |
+
hf upload YOU/lfm-opencode-lora-rl-code . \
|
| 113 |
+
--include "*.py" --include "*.sh" --include "*.md" --include "data/*.json"
|
| 114 |
+
|
| 115 |
+
# 1. Launch (Python API β the CLI lacks expose_public)
|
| 116 |
+
uv run --with 'huggingface_hub>=1.29' python - <<'EOF'
|
| 117 |
+
import os
|
| 118 |
+
from huggingface_hub import run_job
|
| 119 |
+
token = os.popen("hf auth token").read().strip()
|
| 120 |
+
run_job(
|
| 121 |
+
image="vllm/vllm-openai:v0.31.0",
|
| 122 |
+
command=["bash", "-c",
|
| 123 |
+
"pip -q install huggingface_hub && "
|
| 124 |
+
"hf download YOU/lfm-opencode-lora-rl-code --local-dir /app && "
|
| 125 |
+
"bash /app/run_job.sh"],
|
| 126 |
+
flavor="h200x2", timeout="72h",
|
| 127 |
+
env={"ADAPTER_REPO": "YOU/LFM2.5-2.6B-opencode-lora-RL", "MAX_STEPS": "1000"},
|
| 128 |
+
secrets={"HF_TOKEN": token},
|
| 129 |
+
expose_public=[8000],
|
| 130 |
+
token=token,
|
| 131 |
+
)
|
| 132 |
+
EOF
|
| 133 |
+
|
| 134 |
+
# 2. Watch
|
| 135 |
+
hf jobs logs -f YOU/<job-id>
|
| 136 |
+
```
|
| 137 |
+
|
| 138 |
+
`run_job.sh` does, in order: install git+packages β clone OpenEnv & `.pth` its `envs/` β
|
| 139 |
+
**patch the OpenCode seam to send the capture session as `x-api-key`** β prepare the
|
| 140 |
+
pinned tasks β start vLLM (GPU 1) β start the env server with `SPACE_HOST` pointed at
|
| 141 |
+
the job's own public URL β start the 30-min watchdog β train (GPU 0) β final eval.
|
| 142 |
+
|
| 143 |
+
`MAX_STEPS=20` gives a ~40-min smoke run β do that first. (`SANDBOX=daytona`/`e2b`
|
| 144 |
+
+ the matching API key also works if you prefer those backends; then the Jobs-proxy
|
| 145 |
+
plumbing is unnecessary.)
|
| 146 |
+
|
| 147 |
+
### Evaluation (built into the job)
|
| 148 |
+
|
| 149 |
+
After training, the job restarts vLLM with `--enable-lora --lora-modules trained=$ADAPTER_REPO`
|
| 150 |
+
and runs all 250 test tasks twice β base model vs trained adapter β through the same
|
| 151 |
+
harness/sandbox path (`purpose="eval"`, required when passing `eval_sampling`).
|
| 152 |
+
Pass@1 for each lands in the adapter repo at `eval/final_results.json`.
|
| 153 |
+
A standalone version (`run_eval.sh`) runs the same comparison on any checkpoint.
|
| 154 |
+
|
| 155 |
+
---
|
| 156 |
+
|
| 157 |
+
## Pitfalls (each one cost us a run β skip them all)
|
| 158 |
+
|
| 159 |
+
1. **trl 1.16 rejects LFM2.5.** `ValueError: Lfm2ForCausalLM has ['conv'] layers` β tree
|
| 160 |
+
packing is attention-only. Pin `trl==1.15.0`.
|
| 161 |
+
2. **`ImportError: kernels`** β trl 1.15 loads `kernels-community/flash-attn3`;
|
| 162 |
+
`pip install kernels==0.16.0`.
|
| 163 |
+
3. **`train_dataset` is required** by the AsyncGRPOTrainer ctor even with an injected
|
| 164 |
+
rollout worker. Build it from `HarborSessionFactory(...).prompt_rows()` β `create()`
|
| 165 |
+
rejects any prompt the env server doesn't serve, and there is no Hub-dataset shortcut.
|
| 166 |
+
4. **`eos_token_ids` is a required worker arg** the article's snippet omits.
|
| 167 |
+
5. **The rollout loop caps the agent at 8 model calls by default** β set
|
| 168 |
+
`max_tool_calling_iterations` explicitly.
|
| 169 |
+
6. **LFM2.5's chat template isn't recognized** by TRL's `add_response_schema`; set
|
| 170 |
+
`tokenizer.response_template = trl.chat_template_utils.lfm2_2_5_template` yourself.
|
| 171 |
+
7. **OpenCode always sends `tool_choice="auto"`** β vLLM needs
|
| 172 |
+
`--enable-auto-tool-choice --tool-call-parser lfm2`, or every rollout 400s.
|
| 173 |
+
8. **The HF Jobs proxy strips `Authorization`.** Session keys must travel as
|
| 174 |
+
`x-api-key`; patch the OpenCode seam (options.headers) β `run_job.sh` does it.
|
| 175 |
+
9. **transformers v5 has no `PushToHubCallback`** β `hub_strategy="every_save"` is
|
| 176 |
+
vestigial; push checkpoints in your own `on_save` callback or lose them on crash.
|
| 177 |
+
10. **A tunnel is a single point of failure.** Gradio share died at hour 11. Prefer the
|
| 178 |
+
Jobs proxy; add a stall watchdog regardless (30 min without an optimizer step β fail).
|
| 179 |
+
11. **`hf jobs run --env` takes one variable per flag**, and the vLLM image has no
|
| 180 |
+
`python` (only `python3`) and no `git`.
|
| 181 |
+
12. **`run_rollout(eval_sampling=...)` requires `purpose="eval"`** or every cell fails
|
| 182 |
+
validation.
|
| 183 |
+
|
| 184 |
+
## Outputs
|
| 185 |
+
|
| 186 |
+
- **Adapter:** `YOU/LFM2.5-2.6B-opencode-lora-RL` β `adapter_model.safetensors` +
|
| 187 |
+
`adapter_config.json` (+ `checkpoint-{100,200,...}/` every 100 steps).
|
| 188 |
+
Load with `PeftModel.from_pretrained("LiquidAI/LFM2.5-2.6B", YOU/...)` or serve with
|
| 189 |
+
`vllm serve LiquidAI/LFM2.5-2.6B --enable-lora --max-lora-rank 32 --lora-modules opencode-lora=YOU/...`.
|
| 190 |
+
- **Eval:** `eval/final_results.json` β pass@1 base vs trained on 250 held-out tasks,
|
| 191 |
+
plus both/only-base/only-adapter breakdown.
|
| 192 |
+
- **Cost ballpark:** h200x2 for ~10β25h (task-difficulty dependent) + thousands of
|
| 193 |
+
cpu-basic sandbox-minutes. Smoke run first: `MAX_STEPS=20`, ~40 min.
|
| 194 |
+
|
| 195 |
+
## Repo layout
|
| 196 |
+
|
| 197 |
+
| File | Role |
|
| 198 |
+
|---|---|
|
| 199 |
+
| `train_lora.py` | trainer + worker wiring, reward, LoRA config, upload & heartbeat callbacks |
|
| 200 |
+
| `run_job.sh` | job entrypoint: setup, seam patch, 3 processes, watchdog, final eval |
|
| 201 |
+
| `run_eval.sh` | standalone base-vs-adapter eval job |
|
| 202 |
+
| `eval_adapter.py` | held-out eval loop (pass@1, agreement breakdown) |
|
| 203 |
+
| `prepare_tasks.py` + `data/` | pinned, hash-checked task materialization (+ `docker_image` patch) |
|
| 204 |
+
| `serve_policy.sh` / `serve_env.sh` | the two servers, for multi-machine setups |
|
| 205 |
+
| `export_adapter.py` | publish a checkpoint with a model card (+ optional merged model) |
|