"""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 import threading from functools import partial from datasets import Dataset from huggingface_hub import HfApi from peft import LoraConfig from transformers import AutoTokenizer, TrainerCallback 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 ( CaptureContractError, HarnessRolloutWorker, _HarnessRolloutLoop, ) 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 = os.environ.get("TRAIN_SPLIT", "/app/prepared/datasets/train") SANDBOX = os.environ.get("SANDBOX", "hf-sandbox") # Harbor backend 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 class UploadCheckpointsCallback(TrainerCallback): """Push every checkpoint to the Hub, in the background. transformers v5 removed PushToHubCallback: periodic _save_checkpoint writes only to local disk (save_model is called with _internal_call=True, which skips the push), and `hub_strategy="every_save"` is vestigial. Without this callback a crashed job loses the whole run. """ def __init__(self, repo_id): self.repo_id = repo_id self._threads = [] def on_save(self, args, state, control, **kwargs): if not (self.repo_id and state.is_world_process_zero): return control ckpt = os.path.join(args.output_dir, f"checkpoint-{state.global_step}") t = threading.Thread(target=self._upload, args=(ckpt,), daemon=True) t.start() self._threads.append(t) return control def _upload(self, ckpt): try: HfApi().upload_folder( repo_id=self.repo_id, folder_path=ckpt, path_in_repo=os.path.basename(ckpt), # Keep adapter + trainer state, skip multi-GB optimizer shards. ignore_patterns=["optimizer.pt", "optimizer/**", "*.pth"], ) print(f"[hub] uploaded {ckpt}", flush=True) except Exception as exc: # never let an upload hiccup kill training print(f"[hub] checkpoint upload failed: {exc}", flush=True) def on_train_end(self, args, state, control, **kwargs): for t in self._threads: t.join(timeout=600) return control class ResilientHarnessRolloutLoop(_HarnessRolloutLoop): """Score a contract-violating capture as unscorable instead of killing the worker. TRL makes CaptureContractError fatal: never train on a corrupted capture. Right in principle — but in practice the corruption source is transient infra (a dropped websocket between worker and env server), and one bad capture used to kill a 25h run at step 551. Dropping the rollout keeps bad data out of training just the same. To keep the safety property for SYSTEMATIC corruption (e.g. every capture broken, as in the 401 incident), re-raise after a full two groups' worth of consecutive contract errors. """ MAX_CONSECUTIVE_CONTRACT_ERRORS = 16 def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._contract_errors = 0 def _run_session(self, prompt, group_id=0): import logging try: result = super()._run_session(prompt, group_id) except CaptureContractError as exc: self._contract_errors += 1 if self._contract_errors >= self.MAX_CONSECUTIVE_CONTRACT_ERRORS: raise logging.warning("capture contract violation; scoring unscorable (%s)", exc) return self._EMPTY_ROLLOUT, None self._contract_errors = 0 return result class ResilientHarnessRolloutWorker(HarnessRolloutWorker): _loop_cls = ResilientHarnessRolloutLoop class HeartbeatCallback(TrainerCallback): """Touch a file at every optimizer step; run_job.sh's watchdog kills the job if it goes stale. The first full run stalled silently for 13h when the rollout path died — steps just stopped, and nothing noticed.""" PATH = "/tmp/rl-heartbeat" def on_step_end(self, args, state, control, **kwargs): with open(self.PATH, "w") as f: f.write(str(state.global_step)) return control 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_factory(): return 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", ) def make_rollout_worker(dataset, tokenizer): return ResilientHarnessRolloutWorker( harness_session_factory=make_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) # LFM2.5-2.6B's chat_template.jinja doesn't byte-match any template TRL's # add_response_schema recognizes, and the rollout-loop ctor requires a # response parser. TRL ships one purpose-built for LFM2.5's Python-style # tool calls; set it directly. (Only the ctor needs it: in loop-owning mode # the harness, not TRL, parses responses.) from trl.chat_template_utils import lfm2_2_5_template tokenizer.response_template = lfm2_2_5_template # The trainer dataset is discovered from the env server, not loaded from the # Hub: HarborSessionFactory.create() only accepts instructions the server # itself serves, and prompt_rows() is its blessed export (one row per task, # shared by all num_generations of a group). discovery = make_factory()() dataset = Dataset.from_list(discovery.prompt_rows()) print(f"discovered {len(dataset)} tasks from {ENV_URL} split={TRAIN_SPLIT}") 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: # push_to_hub covers the END-of-training push; periodic checkpoints are # pushed by UploadCheckpointsCallback (transformers v5 dropped that path). config_kwargs.update(push_to_hub=True, hub_model_id=ADAPTER_REPO) 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, train_dataset=dataset, # required by the ctor even with an injected worker processing_class=tokenizer, rollout_worker=worker, peft_config=peft_config, callbacks=[UploadCheckpointsCallback(ADAPTER_REPO), HeartbeatCallback()], ) resume = os.environ.get("RESUME_FROM", "") or None if resume: print(f"resuming from {resume}") trainer.train(resume_from_checkpoint=resume) trainer.save_model("./checkpoints/final") # adapter_model.safetensors + adapter_config.json if ADAPTER_REPO: trainer.push_to_hub() if __name__ == "__main__": main()