OpenCode-only LoRA RL on LFM2.5-2.6B (verified against trl 1.16.0 / openenv 0.10.0)

README.md

@@ -0,0 +1,104 @@
+# OpenCode-only RL on LFM2.5-2.6B with LoRA
+
+Reproduction of the OpenCode-only run from
+[The ultimate guide to multi-harness RL](https://huggingface.co/spaces/FineEnvs/multi-harness-rl),
+but training a **LoRA adapter** instead of full weights, so the result is a small,
+distributable adapter on top of `LiquidAI/LFM2.5-2.6B`.
+
+Stack: OpenEnv (capture proxy + Harbor integration) → Harbor (tasks + sandboxes) →
+TRL `AsyncGRPOTrainer` with `HarnessRolloutWorker` → vLLM policy server.
+**Every API in this repo was verified against the released wheels (trl 1.16.0,
+openenv 0.10.0) and the FineEnvs production env Space — not just the article.**
+
+## Why merged-weight sync (source-verified)
+
+Generation here goes harness → capture proxy → vLLM, addressed by a **fixed model
+name**. TRL's adapter-only sync publishes versioned adapters (`trl-policy-v{N}`) and
+only the trainer-built default worker learns the versioned name; an injected
+`HarnessRolloutWorker` never gets one (`async_grpo_trainer.py` builds the name only
+in its own worker branch). So the server runs *without* `--enable-lora` and the
+trainer merges the LoRA into the base and NCCL-streams full weights each step
+(~5.2 GB for 2.6B — trivial on one node). vLLM also generates 1.33–1.39x slower per
+token when serving an adapter, so merged sync is the better default anyway.
+**Checkpoints still contain only the LoRA adapter** — that is what you distribute.
+
+## Files
+
+| File | Purpose |
+|---|---|
+| `train_lora.py` | Worker wiring + `AsyncGRPOConfig` + `LoraConfig` (r=32 on q/k/v/out_proj, w1/2/3) |
+| `serve_policy.sh` | vLLM with capture-proxy flags + NCCL weight transfer |
+| `serve_env.sh` | `openenv harbor serve` with production env vars (`MAX_CONCURRENT_ENVS`, pinned OpenCode version, 4096 output cap) |
+| `run_job.sh` | Single-container entrypoint for HF Jobs (setup + all 3 processes) |
+| `export_adapter.py` | Publish any adapter checkpoint to the Hub (+ optional merged model) |
+
+## Launch on HF Jobs
+
+Two runs: a smoke run to validate the stack, then the full run.
+
+```bash
+# 0. One-time: put this folder on the Hub so the job can pull it
+hf repos create YOU/lfm-opencode-lora-rl-code --type model
+hf upload YOU/lfm-opencode-lora-rl-code . --exclude "__pycache__/**" "checkpoints/**"
+
+# 1. Smoke run (~1h): validates weight sync, capture contract, importance ratio
+hf jobs run vllm/vllm-openai:latest --flavor h200x2 --timeout 6h --detach \
+  --secrets HF_TOKEN="$HF_TOKEN" DAYTONA_API_KEY="..." \
+  --env ADAPTER_REPO=YOU/LFM2.5-2.6B-opencode-lora-RL MAX_STEPS=20 \
+  -- bash -c 'pip -q install huggingface_hub && hf download YOU/lfm-opencode-lora-rl-code --local-dir /app && bash /app/run_job.sh'
+
+# 2. Watch it: weight sync should log merged transfers (no .vllm_lora/),
+#    and capture-proxy importance ratios should sit at ~0.998–0.9999.
+hf jobs ps
+hf jobs logs <job-id> --follow
+
+# 3. Full run (the article's OpenCode-only run took ~32h on 2x H100)
+hf jobs run vllm/vllm-openai:latest --flavor h200x2 --timeout 72h --detach \
+  --secrets HF_TOKEN="$HF_TOKEN" DAYTONA_API_KEY="..." \
+  --env ADAPTER_REPO=YOU/LFM2.5-2.6B-opencode-lora-RL MAX_STEPS=1000 \
+  -- bash -c 'pip -q install huggingface_hub && hf download YOU/lfm-opencode-lora-rl-code --local-dir /app && bash /app/run_job.sh'
+```
+
+Checkpoints (adapter-only) push to `ADAPTER_REPO` at every save, so the job can die
+without losing training. `E2B_API_KEY` + `SANDBOX=e2b` works instead of Daytona.
+
+## What you distribute
+
+Every checkpoint contains `adapter_model.safetensors` + `adapter_config.json`
+(~50 MB at r=32). `export_adapter.py` adds a model card and pushes:
+
+```bash
+python export_adapter.py checkpoints/checkpoint-1000 YOU/LFM2.5-2.6B-opencode-lora-RL [--also-merged]
+```
+
+Consumers load it with `PeftModel.from_pretrained("LiquidAI/LFM2.5-2.6B", YOU/...)` or
+`vllm serve LiquidAI/LFM2.5-2.6B --enable-lora --max-lora-rank 32 --lora-modules opencode-lora=YOU/...`
+(`Lfm2ForCausalLM` implements `SupportsLoRA`).
+
+## Recipe choices carried over from the article (don't regress these)
+
+- `top_p=1.0`, `top_k=0` — truncation biases the policy and breaks logprob agreement
+  (importance ratio dropped to ~0.99 with top_p<1).
+- Same 4096-token output cap in training and evaluation (`OPENENV_MAX_OUTPUT_TOKENS`
+  + worker `max_tokens`). A Qwen run with 16k training answers collapsed at 4k eval.
+- Correctness + <=0.1 tool-efficiency bonus on correct answers. Without it,
+  all-correct GRPO groups carry no signal and tool calls creep up (13 → 41).
+- `num_generations=8`, one harness per group, so advantage compares actions.
+- `max_tool_calling_iterations=32` on the worker — the loop-owning default would
+  silently cap the agent at 8 model calls.
+- `max_staleness=4` (config default) matches the article's "never more than four
+  optimizer steps old".
+
+## Honest caveats
+
+- `HarnessRolloutWorker` + `peft_config` is not a combination the article tested
+  (their runs were full fine-tunes). The pieces are all released and compatible on
+  paper — the smoke run exists to prove it before you spend 32+ GPU-hours.
+- Expect the article's single-harness shape: gains concentrate in OpenCode
+  (34 → 58% there in their run) and transfer unevenly elsewhere.
+- LoRA r=32 may cap peak accuracy vs full FT; no published comparison on this stack.
+- Costs: h200x2 for ~32–46h is real money, plus ~1000s of Daytona sandbox-hours.
+  The smoke run is the cheap insurance.
+- Task data: SmolDataEnvs-harbor-train (1,000 tasks) / -test (250), public,
+  non-gated. For the pinned, hash-checked variant see the `prepare.py` flow in the
+  [FineEnvs env Space](https://huggingface.co/spaces/FineEnvs/smoldataenv-multi-harness-harbor).

export_adapter.py

@@ -0,0 +1,95 @@
+"""Publish a trained LoRA checkpoint to the Hub.
+
+Usage:
+    python export_adapter.py checkpoints/checkpoint-1000 you/LFM2.5-2.6B-opencode-lora-RL
+    python export_adapter.py checkpoints/final you/LFM2.5-2.6B-opencode-lora-RL --also-merged
+"""
+
+import argparse
+import shutil
+import tempfile
+from pathlib import Path
+
+from huggingface_hub import HfApi
+
+BASE_MODEL = "LiquidAI/LFM2.5-2.6B"
+
+MODEL_CARD = """---
+license: other
+base_model: {base}
+tags:
+  - lora
+  - rl
+  - grpo
+  - agentic-rl
+  - opencode
+---
+
+# {repo} — LoRA adapter
+
+LoRA adapter (rank {rank}) trained with async GRPO **inside the OpenCode agent harness**
+on [SmolDataEnvs](https://huggingface.co/collections/FineEnvs/smoldataenvs) tasks, using
+the OpenEnv capture proxy + Harbor + TRL stack from
+[The ultimate guide to multi-harness RL](https://huggingface.co/spaces/FineEnvs/multi-harness-rl).
+
+Fits `{base}` only.
+
+## Load it
+
+```python
+from transformers import AutoModelForCausalLM
+from peft import PeftModel
+
+base = AutoModelForCausalLM.from_pretrained("{base}")
+model = PeftModel.from_pretrained(base, "{repo}")
+```
+
+```bash
+# vLLM
+vllm serve {base} --enable-lora --max-lora-rank {rank} \\
+    --lora-modules opencode-lora={repo}
+```
+"""
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("checkpoint", type=Path, help="checkpoint dir containing adapter_model.safetensors")
+    p.add_argument("repo_id", help="Hub repo for the adapter, e.g. you/LFM2.5-2.6B-opencode-lora-RL")
+    p.add_argument("--also-merged", action="store_true",
+                   help="also push a merged full model to <repo_id>-merged")
+    args = p.parse_args()
+
+    assert (args.checkpoint / "adapter_model.safetensors").exists(), \
+        f"no adapter_model.safetensors in {args.checkpoint} — is this a PEFT checkpoint?"
+
+    import json
+    rank = json.loads((args.checkpoint / "adapter_config.json").read_text())["r"]
+
+    api = HfApi()
+    with tempfile.TemporaryDirectory() as tmp:
+        tmp = Path(tmp)
+        shutil.copy(args.checkpoint / "adapter_model.safetensors", tmp)
+        shutil.copy(args.checkpoint / "adapter_config.json", tmp)
+        (tmp / "README.md").write_text(MODEL_CARD.format(base=BASE_MODEL, repo=args.repo_id, rank=rank))
+        api.create_repo(args.repo_id, exist_ok=True)
+        api.upload_folder(repo_id=args.repo_id, folder_path=tmp)
+    print(f"adapter -> https://huggingface.co/{args.repo_id}")
+
+    if args.also_merged:
+        import torch
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+        from peft import PeftModel
+
+        base = AutoModelForCausalLM.from_pretrained(BASE_MODEL, torch_dtype=torch.bfloat16)
+        model = PeftModel.from_pretrained(base, args.checkpoint)
+        merged = model.merge_and_unload()
+        merged_id = f"{args.repo_id}-merged"
+        api.create_repo(merged_id, exist_ok=True)
+        merged.push_to_hub(merged_id)
+        AutoTokenizer.from_pretrained(BASE_MODEL).push_to_hub(merged_id)
+        print(f"merged  -> https://huggingface.co/{merged_id}")
+
+
+if __name__ == "__main__":
+    main()

run_job.sh

@@ -0,0 +1,75 @@
+#!/usr/bin/env bash
+# Single-job entrypoint for HF Jobs (2-GPU flavor, e.g. h200x2).
+# Runs all three processes in one container:
+#   GPU 1: vLLM policy server          (serve_policy.sh)
+#   CPU:   OpenEnv Harbor env server   (serve_env.sh)
+#   GPU 0: TRL AsyncGRPOTrainer        (train_lora.py)
+#
+# Expected env (via `hf jobs run --secrets/--env`):
+#   HF_TOKEN          (secret)  hub downloads + checkpoint pushes
+#   DAYTONA_API_KEY   (secret)  sandboxes — or E2B_API_KEY with SANDBOX=e2b
+#   ADAPTER_REPO      (env)     Hub repo that receives adapter checkpoints
+#   MAX_STEPS         (env)     20 for the smoke run, 1000 for the full run
+#   SANDBOX           (env)     daytona (default) | e2b
+set -euo pipefail
+cd "$(dirname "$0")"
+
+echo "=== [setup] packages ==="
+# vllm + torch come from the base image (vllm/vllm-openai). Per TRL docs,
+# transformers>=5.2.0 may need --no-deps next to vllm's pinned stack.
+python -m pip install --quiet 'transformers>=5.2.0' --no-deps
+python -m pip install --quiet \
+    'trl[peft]>=1.16.0' \
+    'openenv[harbor]' \
+    'harbor[daytona]>=0.22.0' \
+    math-verify \
+    aiohttp
+
+echo "=== [setup] OpenEnv repo (harbor_env lives in envs/, not the wheel) ==="
+git clone --depth 1 https://github.com/huggingface/OpenEnv.git /opt/OpenEnv
+SITE=$(python -c "import site; print(site.getsitepackages()[0])")
+echo "/opt/OpenEnv/envs" > "$SITE/openenv_examples.pth"
+
+echo "=== [check] imports ==="
+python - <<'EOF'
+import trl, openenv, harbor_env  # noqa: F401
+from trl.experimental.async_grpo.openenv_harness import HarnessRolloutWorker  # noqa: F401
+print("trl", trl.__version__, "| openenv", openenv.__version__)
+EOF
+
+echo "=== [start] vLLM policy server (GPU 1) ==="
+CUDA_VISIBLE_DEVICES=1 VLLM_SERVER_DEV_MODE=1 \
+vllm serve "${MODEL:-LiquidAI/LFM2.5-2.6B}" \
+    --port 8001 \
+    --max-model-len 32768 \
+    --return-tokens-as-token-ids \
+    --logprobs-mode processed_logprobs \
+    --weight-transfer-config '{"backend":"nccl"}' &
+VLLM_PID=$!
+
+echo "=== [start] env server (CPU) ==="
+bash serve_env.sh &
+ENV_PID=$!
+
+echo "=== [wait] servers ==="
+python - <<'EOF'
+import time, urllib.request
+for name, url in [("vllm", "http://127.0.0.1:8001/v1/models"),
+                  ("env",  "http://127.0.0.1:8000/health")]:
+    for i in range(360):  # up to 30 min for model load
+        try:
+            urllib.request.urlopen(url, timeout=5)
+            print(name, "ready")
+            break
+        except Exception:
+            time.sleep(5)
+    else:
+        raise SystemExit(f"{name} server never became ready")
+EOF
+
+echo "=== [train] GPU 0, MAX_STEPS=${MAX_STEPS:-1000} ==="
+CUDA_VISIBLE_DEVICES=0 VLLM_URL=http://127.0.0.1:8001 ENV_URL=http://127.0.0.1:8000 \
+    python train_lora.py
+
+kill $VLLM_PID $ENV_PID 2>/dev/null || true
+echo "=== done ==="

serve_env.sh

@@ -0,0 +1,20 @@
+#!/usr/bin/env bash
+# OpenEnv Harbor environment server: Task API + capture proxy + web UI.
+# Sandboxes run on Daytona (default) or E2B — needs DAYTONA_API_KEY / E2B_API_KEY.
+#
+# Concurrency and capture settings are env vars, not CLI flags (verified against
+# openenv 0.10.0 and the FineEnvs production env Space). MAX_CONCURRENT_ENVS must
+# exceed the GRPO group size (8) or rollouts are rejected with CAPACITY_REACHED.
+set -euo pipefail
+
+export MAX_CONCURRENT_ENVS="${MAX_CONCURRENT_ENVS:-16}"
+export OPENENV_MAX_OUTPUT_TOKENS=4096           # same cap as training + eval
+export OPENENV_HARBOR_REWARD_KEY="correctness,reward"
+# Pin the harness build so rollouts are reproducible (versions from the FineEnvs Space).
+export OPENENV_HARBOR_AGENT_VERSIONS='{"opencode":"1.18.31"}'
+
+openenv harbor serve \
+    --dataset FineEnvs/SmolDataEnvs-harbor-train,FineEnvs/SmolDataEnvs-harbor-test \
+    --port "${ENV_PORT:-8000}" \
+    --capture-port "${CAPTURE_PORT:-8100}" \
+    --expose gradio

serve_policy.sh

@@ -0,0 +1,25 @@
+#!/usr/bin/env bash
+# Policy server: vLLM serving LiquidAI/LFM2.5-2.6B on GPU 1.
+#
+# Flag provenance:
+#   --return-tokens-as-token-ids     capture proxy needs the exact sampled token ids
+#   --logprobs-mode processed_logprobs  proxy needs per-token logprobs that match the policy
+#   --weight-transfer-config         TRL AsyncGRPOTrainer pushes merged weights over NCCL
+#   VLLM_SERVER_DEV_MODE=1           exposes the weight-update endpoint TRL posts to
+#
+# Deliberately NOT passed:
+#   --enable-lora                    adapter sync rotates adapter names, which the
+#                                    OpenEnv capture proxy cannot address; the trainer
+#                                    merges the LoRA into the base and syncs full weights
+set -euo pipefail
+
+MODEL="${MODEL:-LiquidAI/LFM2.5-2.6B}"
+PORT="${PORT:-8001}"
+
+CUDA_VISIBLE_DEVICES=1 VLLM_SERVER_DEV_MODE=1 \
+vllm serve "$MODEL" \
+    --port "$PORT" \
+    --max-model-len 32768 \
+    --return-tokens-as-token-ids \
+    --logprobs-mode processed_logprobs \
+    --weight-transfer-config '{"backend":"nccl"}'

train_lora.py

@@ -0,0 +1,143 @@
+"""OpenCode-only async GRPO on LFM2.5-2.6B with a LoRA adapter.
+
+All APIs verified against the released wheels (trl 1.16.0, openenv 0.10.0):
+- HarnessRolloutOutcome fields: env_reward, tool_call_count, tool_failure_count, ...
+- _AsyncRolloutLoop requires eos_token_ids; max_tool_calling_iterations defaults to
+  max_turns=8 in loop-owning mode, so it is set explicitly.
+- AsyncGRPOConfig is TrainingArguments-derived (no use_vllm/vllm_mode); the vLLM
+  server is given as vllm_server_base_url; dtype defaults to float32, so set it.
+- With an injected rollout_worker, the trainer never assigns a versioned lora_name,
+  so adapter-only sync cannot address the policy through the capture proxy: the
+  vLLM server runs WITHOUT --enable-lora and the trainer streams merged weights
+  (its default when the server lacks --enable-lora). Checkpoints still contain
+  only the LoRA adapter.
+
+Reward: correctness (1/0) + tool-efficiency bonus of at most 0.1 on correct
+answers, the recipe the article's LFM runs used. outcome.tool_call_count comes
+from the harness's own ATIF trace.
+"""
+
+import os
+from functools import partial
+
+from datasets import load_dataset
+from peft import LoraConfig
+from transformers import AutoTokenizer
+from trl.experimental.async_grpo.async_grpo_config import AsyncGRPOConfig
+from trl.experimental.async_grpo.async_grpo_trainer import AsyncGRPOTrainer
+from trl.experimental.async_grpo.openenv_harness import HarnessRolloutWorker
+
+from harbor_env.harness import HarborSessionFactory  # via .pth to OpenEnv/envs
+
+MODEL = "LiquidAI/LFM2.5-2.6B"
+VLLM_URL = os.environ.get("VLLM_URL", "http://127.0.0.1:8001")   # serve_policy.sh
+ENV_URL = os.environ.get("ENV_URL", "http://127.0.0.1:8000")     # serve_env.sh
+TRAIN_SPLIT = "FineEnvs/SmolDataEnvs-harbor-train"
+SANDBOX = os.environ.get("SANDBOX", "daytona")                   # or "e2b"
+MAX_STEPS = int(os.environ.get("MAX_STEPS", "1000"))
+ADAPTER_REPO = os.environ.get("ADAPTER_REPO", "")                # push checkpoints here
+
+TOOL_CALL_BUDGET = 32   # reference for the efficiency bonus; tune to the task pool
+MAX_BONUS = 0.1
+
+
+def correctness_plus_efficiency(outcome):
+    """1/0 correctness, plus <= 0.1 for a correct answer in fewer tool calls.
+
+    A wrong answer stays 0, so the bonus can never be earned by giving up early.
+    Returning None keeps an unverified rollout (dead sandbox) unscorable.
+    """
+    reward = outcome.env_reward
+    if reward is None or reward <= 0:
+        return reward
+    bonus = MAX_BONUS * max(0.0, (TOOL_CALL_BUDGET - outcome.tool_call_count) / TOOL_CALL_BUDGET)
+    return reward + min(bonus, MAX_BONUS)
+
+
+def make_rollout_worker(dataset, tokenizer):
+    factory = partial(
+        HarborSessionFactory,
+        server_url=ENV_URL,
+        split=TRAIN_SPLIT,
+        llm_url=VLLM_URL,
+        model=MODEL,
+        harness="opencode",       # OpenCode-only run
+        sandbox=SANDBOX,
+        reward_key="correctness,reward",
+    )
+    return HarnessRolloutWorker(
+        harness_session_factory=factory,
+        harness_adapter=None,     # let the harness own the loop (black-box capture)
+        rollout_reward_fn=correctness_plus_efficiency,
+        model_name=MODEL,
+        dataset=dataset,
+        processing_class=tokenizer,
+        eos_token_ids=[tokenizer.eos_token_id],   # required loop arg
+        reward_funcs=[],
+        num_generations=8,        # one GRPO group = 8 rollouts of one task, one harness
+        max_inflight_tasks=8,
+        vllm_server_url=VLLM_URL,
+        max_tokens=4096,          # same cap in training and eval — do not regress this
+        max_tool_calling_iterations=32,   # loop default would cap the agent at 8 turns
+        temperature=1.0,
+        top_p=1.0,                # no truncation: keeps sampled dist. == policy dist.
+        top_k=0,                  # 0 = disabled in current vLLM
+        chat_template_kwargs={"enable_thinking": False},
+    )
+
+
+def main():
+    tokenizer = AutoTokenizer.from_pretrained(MODEL)
+    dataset = load_dataset(TRAIN_SPLIT, split="train")
+    worker = make_rollout_worker(dataset, tokenizer)
+
+    config_kwargs = {
+        "output_dir": "./checkpoints",
+        "vllm_server_base_url": VLLM_URL,
+        "dtype": "bfloat16",     # config default is float32
+        "bf16": True,
+        "learning_rate": 1e-6,
+        "max_steps": MAX_STEPS,
+        "logging_steps": 1,
+        # Checkpoints ARE your LoRA adapters. Never set save_strategy="no".
+        "save_strategy": "steps",
+        "save_steps": 100,
+        "save_total_limit": 12,
+        "report_to": [],
+        "seed": 42,
+    }
+    if ADAPTER_REPO:
+        # Jobs are ephemeral: push every checkpoint (adapter-only) to the Hub.
+        config_kwargs.update(
+            push_to_hub=True,
+            hub_model_id=ADAPTER_REPO,
+            hub_strategy="every_save",
+        )
+    config = AsyncGRPOConfig(**config_kwargs)
+
+    # LoRA on attention + MLP projections. Module names verified against the
+    # LFM2.5-2.6B safetensors index (self_attn.{q,k,v,out}_proj, feed_forward.w{1,2,3}).
+    # Short-conv projections (conv.in_proj/out_proj) are left out for portability.
+    peft_config = LoraConfig(
+        r=32,
+        lora_alpha=64,
+        lora_dropout=0.0,
+        target_modules=["q_proj", "k_proj", "v_proj", "out_proj", "w1", "w2", "w3"],
+        task_type="CAUSAL_LM",
+    )
+
+    trainer = AsyncGRPOTrainer(
+        model=MODEL,
+        args=config,
+        processing_class=tokenizer,
+        rollout_worker=worker,
+        peft_config=peft_config,
+    )
+    trainer.train()
+    trainer.save_model("./checkpoints/final")   # adapter_model.safetensors + adapter_config.json
+    if ADAPTER_REPO:
+        trainer.push_to_hub()
+
+
+if __name__ == "__main__":
+    main()