diff --git a/src/soup_cli/config/schema.py b/src/soup_cli/config/schema.py index 07c77e9..122030f 100644 --- a/src/soup_cli/config/schema.py +++ b/src/soup_cli/config/schema.py @@ -1077,7 +1077,7 @@ class TrainingConfig(BaseModel): @field_validator("prm_reward", mode="before") @classmethod - def _validate_prm_reward_field(cls, value): + def _validate_prm_reward_field(cls, value: Any) -> Optional[str]: """v0.71.30 — shape-only validation (containment enforced at load).""" if value is None: return None diff --git a/src/soup_cli/envs/_common.py b/src/soup_cli/envs/_common.py new file mode 100644 index 0000000..2fdefa2 --- /dev/null +++ b/src/soup_cli/envs/_common.py @@ -0,0 +1,29 @@ +"""Shared scaffold for the bundled rollout envs — v0.71.30. + +Each env module supplies a per-row generator over a seeded ``random.Random``; +this helper owns the deterministic seeding + fixed row count so the three env +modules stay tiny and cannot drift on the boilerplate. +""" + +from __future__ import annotations + +import random +from collections.abc import Callable + +# Fixed curriculum size per env (deterministic; overridable by callers). +DEFAULT_ROWS = 64 + + +def seeded_rows( + seed: int, + make_row: Callable[[random.Random], dict[str, str]], + count: int = DEFAULT_ROWS, +) -> list[dict[str, str]]: + """Build ``count`` deterministic ``{"prompt","answer"}`` rows. + + A fresh ``random.Random(seed)`` is created on every call, so the output is + identical across calls (determinism the tests assert). ``make_row`` receives + that RNG and returns one ``{"prompt","answer"}`` row. + """ + rng = random.Random(seed) + return [make_row(rng) for _ in range(count)] diff --git a/src/soup_cli/envs/calculator.py b/src/soup_cli/envs/calculator.py index 43688fa..888b2c3 100644 --- a/src/soup_cli/envs/calculator.py +++ b/src/soup_cli/envs/calculator.py @@ -12,31 +12,30 @@ from __future__ import annotations import random from typing import Any +from soup_cli.envs._common import seeded_rows + _SEED = 20730 -_DEFAULT_ROWS = 64 _OPS = ("+", "-", "*") -def rollout(prompts: Any = None) -> list[dict]: +def _make_row(rng: random.Random) -> dict[str, str]: + op = _OPS[rng.randrange(len(_OPS))] + if op == "*": + a = rng.randint(2, 12) + b = rng.randint(2, 12) + else: + a = rng.randint(0, 99) + b = rng.randint(0, 99) + result = {"+": a + b, "-": a - b, "*": a * b}[op] + prompt = f"What is {a} {op} {b}? Reply with just the number." + return {"prompt": prompt, "answer": str(result)} + + +def rollout(prompts: Any = None) -> list[dict[str, str]]: """Return a deterministic list of ``{"prompt", "answer"}`` arithmetic rows. ``prompts`` (the seed prompts from the GRPO dataset) is accepted for the openenv contract but does not change the generated curriculum — the env is a self-contained deterministic seeder. """ - rng = random.Random(_SEED) - rows: list[dict] = [] - for _ in range(_DEFAULT_ROWS): - op = _OPS[rng.randrange(len(_OPS))] - if op == "*": - a = rng.randint(2, 12) - b = rng.randint(2, 12) - else: - a = rng.randint(0, 99) - b = rng.randint(0, 99) - result = {"+": a + b, "-": a - b, "*": a * b}[op] - prompt = ( - f"What is {a} {op} {b}? Reply with just the number." - ) - rows.append({"prompt": prompt, "answer": str(result)}) - return rows + return seeded_rows(_SEED, _make_row) diff --git a/src/soup_cli/envs/guess_number.py b/src/soup_cli/envs/guess_number.py index 165565e..c44d129 100644 --- a/src/soup_cli/envs/guess_number.py +++ b/src/soup_cli/envs/guess_number.py @@ -17,27 +17,28 @@ from __future__ import annotations import random from typing import Any +from soup_cli.envs._common import seeded_rows + _SEED = 20731 -_DEFAULT_ROWS = 64 -def rollout(prompts: Any = None) -> list[dict]: +def _make_row(rng: random.Random) -> dict[str, str]: + a = rng.randint(2, 9) + b = rng.randint(2, 9) + answer = a * b + prompt = ( + f"I'm thinking of a number between 1 and 100. " + f"It equals {a} times {b}. What is the number? " + "Reply with just the number." + ) + return {"prompt": prompt, "answer": str(answer)} + + +def rollout(prompts: Any = None) -> list[dict[str, str]]: """Return a deterministic list of ``{"prompt", "answer"}`` deduction rows. Each puzzle states two factors whose product is the answer, plus a range, so the answer is uniquely deducible. ``prompts`` is accepted for the openenv contract but does not change the curriculum. """ - rng = random.Random(_SEED) - rows: list[dict] = [] - for _ in range(_DEFAULT_ROWS): - a = rng.randint(2, 9) - b = rng.randint(2, 9) - answer = a * b - prompt = ( - f"I'm thinking of a number between 1 and 100. " - f"It equals {a} times {b}. What is the number? " - "Reply with just the number." - ) - rows.append({"prompt": prompt, "answer": str(answer)}) - return rows + return seeded_rows(_SEED, _make_row) diff --git a/src/soup_cli/envs/retrieval_qa.py b/src/soup_cli/envs/retrieval_qa.py index 75d35fa..ab24167 100644 --- a/src/soup_cli/envs/retrieval_qa.py +++ b/src/soup_cli/envs/retrieval_qa.py @@ -13,8 +13,9 @@ from __future__ import annotations import random from typing import Any +from soup_cli.envs._common import seeded_rows + _SEED = 20732 -_DEFAULT_ROWS = 64 # (entity, attribute, value) fact templates — the answer is always ``value``, # which is embedded verbatim in the document so it is a retrievable span. @@ -30,25 +31,24 @@ _FACTS = ( ) -def rollout(prompts: Any = None) -> list[dict]: +def _make_row(rng: random.Random) -> dict[str, str]: + # Pick 3 distinct facts as the document; ask about one of them. + facts = rng.sample(_FACTS, 3) + entity, attribute, value = facts[rng.randrange(len(facts))] + doc = " ".join(f"{e} {a} {v}." for (e, a, v) in facts) + prompt = ( + f"Document: {doc}\n" + f"Question: What does {entity} {attribute.rstrip()}? " + "Reply with just the value." + ) + return {"prompt": prompt, "answer": value} + + +def rollout(prompts: Any = None) -> list[dict[str, str]]: """Return a deterministic list of ``{"prompt", "answer"}`` retrieval rows. Each row embeds a short document containing several facts and asks about one of them; the answer is a span present in the document. ``prompts`` is accepted for the openenv contract but does not change the curriculum. """ - rng = random.Random(_SEED) - rows: list[dict] = [] - for _ in range(_DEFAULT_ROWS): - # Pick 3 distinct facts as the document; ask about one of them. - facts = rng.sample(_FACTS, 3) - target = facts[rng.randrange(len(facts))] - entity, attribute, value = target - doc = " ".join(f"{e} {a} {v}." for (e, a, v) in facts) - prompt = ( - f"Document: {doc}\n" - f"Question: What does {entity} {attribute.rstrip()}? " - "Reply with just the value." - ) - rows.append({"prompt": prompt, "answer": value}) - return rows + return seeded_rows(_SEED, _make_row) diff --git a/src/soup_cli/trainer/grpo.py b/src/soup_cli/trainer/grpo.py index 67dbc1f..db2f163 100644 --- a/src/soup_cli/trainer/grpo.py +++ b/src/soup_cli/trainer/grpo.py @@ -10,7 +10,7 @@ from typing import Any, Optional from rich.console import Console -from soup_cli.config.schema import SoupConfig +from soup_cli.config.schema import SoupConfig, TrainingConfig from soup_cli.utils.gpu import estimate_batch_size, model_size_from_name console = Console() @@ -91,9 +91,7 @@ def _make_grpo_trainer_variant_cached(base_cls: type, variant: str) -> type: # Fall back to the original loss — defence-in-depth so a # TRL internal rename does not crash the training loop. self._warn_fallback("missing per-token log-prob inputs") - return super().compute_loss( - model, inputs, return_outputs=return_outputs, **kwargs - ) + return super().compute_loss(model, inputs, return_outputs=return_outputs, **kwargs) beta_attr = getattr(self.args, "beta", None) beta = float(beta_attr) if beta_attr is not None else 0.0 @@ -110,14 +108,10 @@ def _make_grpo_trainer_variant_cached(base_cls: type, variant: str) -> type: ) except (TypeError, ValueError) as exc: self._warn_fallback(f"kernel error: {exc}") - return super().compute_loss( - model, inputs, return_outputs=return_outputs, **kwargs - ) + return super().compute_loss(model, inputs, return_outputs=return_outputs, **kwargs) if variant_loss is None: self._warn_fallback("kernel returned None") - return super().compute_loss( - model, inputs, return_outputs=return_outputs, **kwargs - ) + return super().compute_loss(model, inputs, return_outputs=return_outputs, **kwargs) if return_outputs: return variant_loss, None return variant_loss @@ -135,7 +129,7 @@ def _read_attr(obj: Any, name: str) -> Any: return getattr(obj, name, None) -def _select_reward_fn(tcfg: Any, device: str, trust_remote_code: bool) -> Any: +def _select_reward_fn(tcfg: TrainingConfig, device: str, trust_remote_code: bool) -> Any: """Choose the GRPO reward function (v0.71.30). When ``tcfg.prm_reward`` is set, a trained Soup PRM scores each completion's @@ -152,6 +146,7 @@ def _select_reward_fn(tcfg: Any, device: str, trust_remote_code: bool) -> Any: return load_reward_fn(tcfg.reward_fn, verifiable_domain=tcfg.verifiable_domain) + class GRPOTrainerWrapper: """High-level wrapper for GRPO training from SoupConfig. @@ -276,16 +271,13 @@ class GRPOTrainerWrapper: # lists) and will raise ValueError if the template is missing. if not getattr(self.tokenizer, "chat_template", None): self.tokenizer.chat_template = ( - "{% for msg in messages %}" - "{{ msg['content'] }}\n" - "{% endfor %}" + "{% for msg in messages %}{{ msg['content'] }}\n{% endfor %}" ) trainable, total = self.model.get_nb_trainable_parameters() pct = 100 * trainable / total console.print( - f"[green]LoRA applied:[/] {trainable:,} trainable" - f" / {total:,} total ({pct:.2f}%)" + f"[green]LoRA applied:[/] {trainable:,} trainable / {total:,} total ({pct:.2f}%)" ) # --- Batch size --- @@ -330,9 +322,7 @@ class GRPOTrainerWrapper: tokenizer=self.tokenizer, reward_fn=reward_fn, ) - train_data = _prepare_grpo_dataset( - [dict(row) for row in rollout_result.rows] - ) + train_data = _prepare_grpo_dataset([dict(row) for row in rollout_result.rows]) console.print( f"[green]Rollout backend '{tcfg.rollout_backend}':[/] " f"{len(train_data)} prompts collected " @@ -355,8 +345,7 @@ class GRPOTrainerWrapper: import math total_steps = ( - math.ceil(len(train_ds) / batch_size / tcfg.gradient_accumulation_steps) - * tcfg.epochs + math.ceil(len(train_ds) / batch_size / tcfg.gradient_accumulation_steps) * tcfg.epochs ) warmup_steps = int(total_steps * tcfg.warmup_ratio) @@ -448,11 +437,13 @@ class GRPOTrainerWrapper: self.trainer._soup_grpo_delta = float(tcfg.grpo_delta) # v0.53.11 #127 — wire the live stability callback. from soup_cli.utils.peft_wiring import attach_grpo_stability_callback + attach_grpo_stability_callback(self.trainer, tcfg) # v0.71.11 #235/#238/#240 — wire the live RL callbacks (reward-hack, # echo-trap, mid-epoch RL checkpoint). from soup_cli.utils.peft_wiring import attach_rl_callbacks + attach_rl_callbacks( self.trainer, tcfg, @@ -468,6 +459,7 @@ class GRPOTrainerWrapper: attach_plugin_callback, attach_relora_callback, ) + attach_relora_callback(self.trainer, tcfg) # v0.53.5 #114/#115 — dynamic curriculum live callback. attach_curriculum_callback(self.trainer, tcfg, str(output_dir), console) @@ -492,14 +484,17 @@ class GRPOTrainerWrapper: from soup_cli.utils.quant_menu import build_quantization_config_for_loader quant_config_obj = build_quantization_config_for_loader( - tcfg=tcfg, base=cfg.base, console=console, + tcfg=tcfg, + base=cfg.base, + console=console, ) console.print(f"[dim]Loading model: {cfg.base}[/]") # On CPU, use device_map="cpu" to avoid meta tensors from "auto" dev_map = "cpu" if self.device == "cpu" else "auto" model_kwargs = { - "trust_remote_code": self._trust_remote_code, "device_map": dev_map, + "trust_remote_code": self._trust_remote_code, + "device_map": dev_map, } if quant_config_obj is not None: model_kwargs["quantization_config"] = quant_config_obj @@ -535,6 +530,7 @@ class GRPOTrainerWrapper: apply_post_lora_patches, apply_pre_lora_patches, ) + apply_pre_lora_patches(self.model, cfg.base) self.model = get_peft_model(self.model, lora_config) apply_post_lora_patches(self.model) @@ -548,9 +544,14 @@ class GRPOTrainerWrapper: # v0.35.0 #60 — multi-trainer wiring of v0.28.0 speed/memory features. from soup_cli.utils.v028_features import apply_v028_speed_memory + apply_v028_speed_memory( - model=self.model, tcfg=tcfg, base_model=cfg.base, - console=console, device=self.device, backend=cfg.backend, + model=self.model, + tcfg=tcfg, + base_model=cfg.base, + console=console, + device=self.device, + backend=cfg.backend, ) def _setup_unsloth(self, cfg, tcfg): @@ -586,7 +587,9 @@ class GRPOTrainerWrapper: self.trainer.add_callback( SoupTrainerCallback( - display, tracker=tracker, run_id=run_id, + display, + tracker=tracker, + run_id=run_id, loss_watchdog=self.config.training.loss_watchdog, loss_watchdog_threshold=self.config.training.loss_watchdog_threshold, loss_watchdog_patience=self.config.training.loss_watchdog_patience, @@ -597,7 +600,8 @@ class GRPOTrainerWrapper: from soup_cli.utils.v028_features import activation_offloading_context with activation_offloading_context( - self.config.training, self._output_dir, + self.config.training, + self._output_dir, ): self.trainer.train(resume_from_checkpoint=resume_from_checkpoint) duration = time.time() - start diff --git a/src/soup_cli/utils/prm_reward.py b/src/soup_cli/utils/prm_reward.py index 73b7e48..f3a2e4f 100644 --- a/src/soup_cli/utils/prm_reward.py +++ b/src/soup_cli/utils/prm_reward.py @@ -29,7 +29,13 @@ Security: from __future__ import annotations import math -from typing import Any +from typing import TYPE_CHECKING, Any, Union + +if TYPE_CHECKING: # pragma: no cover — type-only, keeps the module torch-free + from soup_cli.config.schema import TrainingConfig + +# A GRPO completion is either a plain string or a chat message list. +Completion = Union[str, list] # Aggregation modes for folding per-step scores into a scalar reward. AGGREGATE_MODES: tuple[str, ...] = ("min", "prod", "last") @@ -37,6 +43,14 @@ AGGREGATE_MODES: tuple[str, ...] = ("min", "prod", "last") # Bounds — a pathological completion cannot blow up the forward pass. _MAX_STEPS = 64 _MAX_STEP_CHARS = 2_000 +# Hard ceiling on the assembled (prompt + steps) token length fed to the PRM +# forward pass. The PRM is a second, independently-loaded model whose context +# window may differ from the policy model's, and the reward fires every GRPO +# step, so cap defensively regardless of the policy-side max_length. +_MAX_INPUT_TOKENS = 8_192 +# Cap the rendered prompt CONTEXT chars before tokenising (a huge prompt is +# expensive to tokenise even before the token cap applies). +_MAX_PROMPT_CHARS = 8_000 def split_steps(text: Any) -> list[str]: @@ -75,9 +89,7 @@ def aggregate_step_scores(scores: list[float], mode: Any) -> float: not poison the whole reward). """ if isinstance(mode, bool) or not isinstance(mode, str) or mode not in AGGREGATE_MODES: - raise ValueError( - f"prm_aggregate must be one of {AGGREGATE_MODES}; got {mode!r}" - ) + raise ValueError(f"prm_aggregate must be one of {AGGREGATE_MODES}; got {mode!r}") if not scores: return 0.0 clean = [_finite(s) for s in scores] @@ -108,9 +120,7 @@ class PRMScorer: trust_remote_code: bool = False, ) -> None: if isinstance(aggregate, bool) or aggregate not in AGGREGATE_MODES: - raise ValueError( - f"aggregate must be one of {AGGREGATE_MODES}; got {aggregate!r}" - ) + raise ValueError(f"aggregate must be one of {AGGREGATE_MODES}; got {aggregate!r}") self.prm_path = prm_path self.aggregate = aggregate self.device = device @@ -158,7 +168,7 @@ class PRMScorer: def _render_prompt(self, prompt: Any) -> str: """Best-effort render of a GRPO prompt (str or message list) to text.""" if isinstance(prompt, str): - return prompt + return prompt[:_MAX_PROMPT_CHARS] if isinstance(prompt, (list, tuple)): parts: list[str] = [] for msg in prompt: @@ -166,18 +176,17 @@ class PRMScorer: parts.append(str(msg.get("content", ""))) else: parts.append(str(msg)) - return "\n".join(p for p in parts if p) + return "\n".join(p for p in parts if p)[:_MAX_PROMPT_CHARS] return "" - def _completion_text(self, completion: Any) -> str: + def _completion_text(self, completion: Completion) -> str: if isinstance(completion, str): return completion if isinstance(completion, dict): return str(completion.get("content", "")) if isinstance(completion, (list, tuple)): parts = [ - str(m.get("content", "")) if isinstance(m, dict) else str(m) - for m in completion + str(m.get("content", "")) if isinstance(m, dict) else str(m) for m in completion ] return "".join(parts) return str(completion) @@ -190,9 +199,7 @@ class PRMScorer: tokenizer = self._tokenizer # Prompt context (trained distribution) then step boundaries. prefix_ids = ( - tokenizer(prompt_text, add_special_tokens=False)["input_ids"] - if prompt_text - else [] + tokenizer(prompt_text, add_special_tokens=False)["input_ids"] if prompt_text else [] ) input_ids = list(prefix_ids) step_positions: list[int] = [] @@ -204,6 +211,19 @@ class PRMScorer: step_positions.append(len(input_ids) - 1) if not step_positions: return 0.0 + # Cap the assembled length to the PRM's own context window (bounded by a + # hard ceiling) — the PRM may have a smaller window than the policy + # model, and this reward fires every GRPO step. Keep the early steps. + config = getattr(self._model, "config", None) + max_pos = getattr(config, "max_position_embeddings", _MAX_INPUT_TOKENS) + cap = _MAX_INPUT_TOKENS + if isinstance(max_pos, int) and not isinstance(max_pos, bool): + cap = min(max_pos, _MAX_INPUT_TOKENS) + if len(input_ids) > cap: + input_ids = input_ids[:cap] + step_positions = [p for p in step_positions if p < cap] + if not step_positions: + return 0.0 ids = torch.tensor([input_ids], dtype=torch.long, device=self.device) with torch.no_grad(): outputs = self._model(input_ids=ids, output_hidden_states=True) @@ -227,7 +247,7 @@ class PRMScorer: return rewards -def load_reward_head_weights(prm_path: str) -> dict: +def load_reward_head_weights(prm_path: str) -> dict[str, Any]: """Load ``reward_head.{weight,bias}`` tensors from a Soup-trained PRM dir. Scans every ``*.safetensors`` shard in ``prm_path`` via @@ -253,7 +273,7 @@ def load_reward_head_weights(prm_path: str) -> dict: with safe_open(shard, framework="pt") as handle: for key in handle.keys(): # noqa: SIM118 — safe_open handle API if key.startswith("reward_head."): - collected[key[len("reward_head."):]] = handle.get_tensor(key) + collected[key[len("reward_head.") :]] = handle.get_tensor(key) if "weight" not in collected or "bias" not in collected: raise ValueError( f"No reward_head weights found in {prm_path!r} — this is not a " @@ -263,7 +283,7 @@ def load_reward_head_weights(prm_path: str) -> dict: def build_prm_reward_fn( - tcfg: Any, + tcfg: "TrainingConfig", device: str, trust_remote_code: bool, ) -> PRMScorer: @@ -290,8 +310,7 @@ def build_prm_reward_fn( cwd = os.path.realpath(os.getcwd()) if os.path.commonpath([real, cwd]) != cwd: raise ValueError( - "prm_reward path must stay under the current working " - f"directory; got {prm_path!r}" + f"prm_reward path must stay under the current working directory; got {prm_path!r}" ) prm_path = real