SaylorTwift HF Staff commited on
Commit
af7d04e
Β·
verified Β·
1 Parent(s): 231d3f3

Add RECIPE.md: shareable setup + run recipe with pitfalls

Browse files
Files changed (1) hide show
  1. RECIPE.md +205 -0
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) |