SaylorTwift HF Staff

Resilient worker (transient CaptureContractError -> unscorable, systematic still fatal), resume-from-repo, eval worker base-cells download fix

e7168a9 verified 11 days ago

Raw History Blame Contribute Delete

10.6 kB

"""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()

	"""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()