SaylorTwift's picture
SaylorTwift HF Staff
RECIPE: document per-checkpoint eval worker
b6f8fe5 verified
|
Raw History Blame Contribute Delete
14.7 kB

Recipe: OpenCode-only agentic RL on LFM2.5-2.6B, shipped as a LoRA adapter

Train a small open model inside the OpenCode agent harness β€” unmodified, running in cloud sandboxes β€” with async GRPO, and distribute the result as a ~170 MB LoRA adapter on top of LiquidAI/LFM2.5-2.6B.

Everything runs on Hugging Face infrastructure: training on HF Jobs, sandboxes on HF Sandboxes (Harbor's hf-sandbox backend). The only credential needed is an HF token.

Based on the stack from The ultimate guide to multi-harness RL (OpenEnv capture proxy + Harbor + TRL), narrowed to one harness and adapted for LoRA and HF-only infrastructure.


Architecture

                 β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ HF Job (h200x2) ───────────────────────────┐
                 β”‚                                                                       β”‚
 GRPO groups ──► β”‚  TRL AsyncGRPOTrainer (GPU 0)                                         β”‚
 of 8 rollouts   β”‚    Β· LoRA r=32 on q/k/v/out_proj + w1/w2/w3                            β”‚
                 β”‚    Β· merges adapter β†’ streams FULL weights over NCCL (136 ms/step)    β”‚
                 β”‚            β”‚                           β–²                              β”‚
                 β”‚            β–Ό weight sync               β”‚ tokens + logprobs             β”‚
                 β”‚  vLLM 0.31 (GPU 1) ◄── capture proxy β—„β”€β”˜ (mounted at /capture)        β”‚
                 β”‚    Β· --return-tokens-as-token-ids                                     β”‚
                 β”‚    Β· --logprobs-mode processed_logprobs                               β”‚
                 β”‚    Β· --tool-call-parser lfm2 --enable-auto-tool-choice                β”‚
                 β”‚            β–²                           β”‚                              β”‚
                 β”‚  OpenEnv Harbor env server (CPU) β”€β”€β”€β”€β”€β”€β”˜                              β”‚
                 β”‚    Β· task API Β· capture proxy Β· per-rollout sandbox lifecycle         β”‚
                 β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
                                 β”‚ https://<job-id>--8000.hf.jobs/capture
                                 β”‚ (session id as x-api-key β€” Jobs proxy strips Authorization)
                 β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β–Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
                 β”‚  HF Sandboxes (1 per rollout)  β”‚  cpu-basic Jobs running the task image
                 β”‚  OpenCode β†’ runs the task β†’    β”‚  pulls Kaggle data from a public bucket,
                 β”‚  verifier grades the answer    β”‚  agent loop owned entirely by OpenCode
                 β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜

Data flow per rollout: the trainer's HarnessRolloutWorker asks the env server for a rollout β†’ env server boots an HF Sandbox with the task β†’ OpenCode runs inside, calling the model through the capture proxy β†’ proxy records exact prompt/completion token IDs + per-token logprobs from vLLM β†’ sandbox verifier scores the answer β†’ reward (correctness + tool-efficiency bonus) and the token-faithful trace go back to the trainer β†’ GRPO updates the LoRA β†’ merged weights stream to vLLM for the next group.


Key design decisions (and why)

Decision Why
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.
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.
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.
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.
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.
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.
vllm/vllm-openai:v0.31.0 image Top of trl 1.15's supported band; ships the lfm2 tool parser (for LFM2.5's `<
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.
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.

Exact recipe

Versions

Component Version Notes
trl 1.15.0 (trl[peft]) HarnessRolloutWorker in trl.experimental.async_grpo.openenv_harness
vLLM 0.31.0 (docker image) needs kernels==0.16.0 pip-installed for trl's kernels-community/flash-attn3
openenv 0.10.0 (openenv[harbor]) harbor_env comes from the OpenEnv repo (envs/), not the wheel β€” clone + .pth
harbor β‰₯0.24.0 (harbor[hf-sandbox]) the version with EnvironmentType.HF_SANDBOX
transformers β‰₯5.2.0 (--no-deps next to vLLM's stack) v5 removed PushToHubCallback β€” see callback below
Base model LiquidAI/LFM2.5-2.6B Lfm2ForCausalLM, 30 layers (20 conv + 10 attention), tied embeddings

Data

  • Train: 1,000 SmolDataEnvs Harbor tasks (400 medium / 600 hard, fixed order), pinned to dataset revision 4719635555de… with per-instruction SHA-256 checks.
  • Test: 250 held-out tasks (no notebook/question overlap), revision b130595b57….
  • Task lists + prep script: prepare_tasks.py + data/{train,test}.json in the code repo. It also injects docker_image = "docker.io/savatar101/env-data-agent-train:base" into every task.toml (required by the hf-sandbox backend) and sets 1 CPU / 4 GB.
  • Task data files are pulled into each sandbox from the public bucket AdithyaSK/jupyter-agent-kaggle-all via the task's own healthcheck hook.

Training hyperparameters

Knob Value
Algorithm async GRPO (AsyncGRPOTrainer + HarnessRolloutWorker, loop owned by the harness)
LoRA r=32, Ξ±=64, dropout 0, targets q_proj,k_proj,v_proj,out_proj,w1,w2,w3
Steps 1,000 (1 GRPO group of 8 rollouts β‰ˆ 1 task per step)
Learning rate 1e-6
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)
Output cap 4,096 tokens, same in training and eval (a mismatch collapsed one of the article's Qwen runs)
Agent budget 32 model calls/rollout (max_tool_calling_iterations β€” the loop-owning default is 8!)
Staleness max 4 optimizer steps (default)
Reward env_reward (1/0) + 0.1 Γ— max(0, (32 βˆ’ tool_calls) / 32) on correct answers
Hardware HF Jobs h200x2 (~90 s/step early, ~35–50 s/step once the rollout queue saturates)

Launch

# 0. Upload the code bundle (one time)
hf repos create YOU/lfm-opencode-lora-rl-code --type model --private
hf upload YOU/lfm-opencode-lora-rl-code . \
  --include "*.py" --include "*.sh" --include "*.md" --include "data/*.json"

# 1. Launch (Python API β€” the CLI lacks expose_public)
uv run --with 'huggingface_hub>=1.29' python - <<'EOF'
import os
from huggingface_hub import run_job
token = os.popen("hf auth token").read().strip()
run_job(
    image="vllm/vllm-openai:v0.31.0",
    command=["bash", "-c",
             "pip -q install huggingface_hub && "
             "hf download YOU/lfm-opencode-lora-rl-code --local-dir /app && "
             "bash /app/run_job.sh"],
    flavor="h200x2", timeout="72h",
    env={"ADAPTER_REPO": "YOU/LFM2.5-2.6B-opencode-lora-RL", "MAX_STEPS": "1000"},
    secrets={"HF_TOKEN": token},
    expose_public=[8000],
    token=token,
)
EOF

# 2. Watch
hf jobs logs -f YOU/<job-id>

run_job.sh does, in order: install git+packages β†’ clone OpenEnv & .pth its envs/ β†’ patch the OpenCode seam to send the capture session as x-api-key β†’ prepare the pinned tasks β†’ start vLLM (GPU 1) β†’ start the env server with SPACE_HOST pointed at the job's own public URL β†’ start the 30-min watchdog β†’ train (GPU 0) β†’ final eval.

MAX_STEPS=20 gives a ~40-min smoke run β€” do that first. (SANDBOX=daytona/e2b

  • the matching API key also works if you prefer those backends; then the Jobs-proxy plumbing is unnecessary.)

Evaluation (two layers)

Per-checkpoint, live β€” a separate eval-worker job (run_eval_watch.sh, 1Γ—H200) polls the adapter repo; every time a new checkpoint-N appears it hot-loads it as a LoRA into its running vLLM (runtime LoRA updates, no restart) and evals it on a fixed 100-task test subset. Base-model cells are computed once and reused. Results: eval/<checkpoint-N>_results.json + a rolling eval/summary.json β€” the live learning curve. Launch it alongside training:

uv run --with 'huggingface_hub>=1.29' python - <<'EOF'
import os
from huggingface_hub import run_job
token = os.popen("hf auth token").read().strip()
run_job(
    image="vllm/vllm-openai:v0.31.0",
    command=["bash", "-c",
             "pip -q install huggingface_hub && "
             "hf download YOU/lfm-opencode-lora-rl-code --local-dir /app && "
             "bash /app/run_eval_watch.sh"],
    flavor="h200", timeout="48h",
    env={"ADAPTER_REPO": "YOU/LFM2.5-2.6B-opencode-lora-RL", "N_TASKS": "100"},
    secrets={"HF_TOKEN": token},
    expose_public=[8000],
    token=token,
)
EOF

Final, full β€” after training, the training job itself restarts vLLM with the final adapter and runs all 250 test tasks twice (base vs final) β†’ eval/final_results.json. (purpose="eval" is required when passing eval_sampling.) A standalone version (run_eval.sh) runs the same comparison on any checkpoint.


Pitfalls (each one cost us a run β€” skip them all)

  1. trl 1.16 rejects LFM2.5. ValueError: Lfm2ForCausalLM has ['conv'] layers β€” tree packing is attention-only. Pin trl==1.15.0.
  2. ImportError: kernels β€” trl 1.15 loads kernels-community/flash-attn3; pip install kernels==0.16.0.
  3. train_dataset is required by the AsyncGRPOTrainer ctor even with an injected rollout worker. Build it from HarborSessionFactory(...).prompt_rows() β€” create() rejects any prompt the env server doesn't serve, and there is no Hub-dataset shortcut.
  4. eos_token_ids is a required worker arg the article's snippet omits.
  5. The rollout loop caps the agent at 8 model calls by default β€” set max_tool_calling_iterations explicitly.
  6. LFM2.5's chat template isn't recognized by TRL's add_response_schema; set tokenizer.response_template = trl.chat_template_utils.lfm2_2_5_template yourself.
  7. OpenCode always sends tool_choice="auto" β†’ vLLM needs --enable-auto-tool-choice --tool-call-parser lfm2, or every rollout 400s.
  8. The HF Jobs proxy strips Authorization. Session keys must travel as x-api-key; patch the OpenCode seam (options.headers) β€” run_job.sh does it.
  9. transformers v5 has no PushToHubCallback β€” hub_strategy="every_save" is vestigial; push checkpoints in your own on_save callback or lose them on crash.
  10. A tunnel is a single point of failure. Gradio share died at hour 11. Prefer the Jobs proxy; add a stall watchdog regardless (30 min without an optimizer step β†’ fail).
  11. hf jobs run --env takes one variable per flag, and the vLLM image has no python (only python3) and no git.
  12. run_rollout(eval_sampling=...) requires purpose="eval" or every cell fails validation.

Outputs

  • Adapter: YOU/LFM2.5-2.6B-opencode-lora-RL β€” adapter_model.safetensors + adapter_config.json (+ checkpoint-{100,200,...}/ every 100 steps). Load with PeftModel.from_pretrained("LiquidAI/LFM2.5-2.6B", YOU/...) or serve with vllm serve LiquidAI/LFM2.5-2.6B --enable-lora --max-lora-rank 32 --lora-modules opencode-lora=YOU/....
  • Eval: eval/final_results.json β€” pass@1 base vs trained on 250 held-out tasks, plus both/only-base/only-adapter breakdown.
  • Cost ballpark: h200x2 for ~10–25h (task-difficulty dependent) + thousands of cpu-basic sandbox-minutes. Smoke run first: MAX_STEPS=20, ~40 min.

Repo layout

File Role
train_lora.py trainer + worker wiring, reward, LoRA config, upload & heartbeat callbacks
run_job.sh job entrypoint: setup, seam patch, 3 processes, watchdog, final eval
run_eval.sh standalone base-vs-adapter eval job
run_eval_watch.sh + eval_watch.py per-checkpoint eval worker (live learning curve)
eval_adapter.py held-out eval loop (pass@1, agreement breakdown)
prepare_tasks.py + data/ pinned, hash-checked task materialization (+ docker_image patch)
serve_policy.sh / serve_env.sh the two servers, for multi-machine setups
export_adapter.py publish a checkpoint with a model card (+ optional merged model)