RomanNum3ral
/
soup
mirror of https://github.com/razor-ai/soup.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

refactor(prm): python-review fixes (TrainingConfig typing, PRM input cap, env dedup, type tightening) (v0.71.30)

This commit is contained in:
Alpamys 2026-07-05 17:19:42 +05:00
parent df5011cb7b
commit db2dc7ed65
7 changed files with 150 additions and 98 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/soup_cli/config/schema.py
View File

@ -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

29
src/soup_cli/envs/_common.py Normal file
View File

@ -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)]

35
src/soup_cli/envs/calculator.py
View File

@ -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)

31
src/soup_cli/envs/guess_number.py
View File

@ -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)

34
src/soup_cli/envs/retrieval_qa.py
View File

@ -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)

58
src/soup_cli/trainer/grpo.py
View File

@ -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

59
src/soup_cli/utils/prm_reward.py
View File

@ -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