mirror of https://github.com/razor-ai/soup.git
feat(rl): live GRPO/RL callbacks — reward-hack, echo-trap, RL ckpt, ULD, MiniLLM, iterative-DPO (v0.71.11)
Lifts the v0.70.0 schema-only build_*_callback / build_uld_projection / run_iterative_dpo stubs. Validated end-to-end on SmolLM2-135M. Closes #235, #236, #237, #238, #239, #240, #159, #160 - #235 RewardHackCallback: info_rm cluster-sep / rm_ensemble divergence, OK/WARN/HACK, halt on HACK. Shared thread-safe RLSignalBuffer captures per-step rewards by wrapping the reward fns (no TRL monkeypatching). - #236 ULD: Wasserstein-1 / top-k aligned distill loss in DistillTrainer. - #237 MiniLLM: teacher-mixed length-normalised reverse-KL + pretrain anchor. - #238 RLCheckpointCallback: adapter + optimizer.pt + manifest + keep_last prune. - #239 run_iterative_dpo: sample -> RM-score -> build-pairs -> DPO-train per round. - #240 EchoTrapCallback: n-gram repetition OK/WARN/TRAP, halt on TRAP. - #159 one-shot WARNING when a GRPO variant compute_loss falls back to super(). - #160 in-place ref-model EMA (no state_dict round-trip) + 0-overlap warning. Tests 13142 -> 13203 (+62 in tests/test_v07111.py).
This commit is contained in:
Alpamys
parent
6df553cb4e
commit
f316d334bc
52
CHANGELOG.md
52
CHANGELOG.md
|
|
@ -12,6 +12,58 @@ reproducing 70+ versions of notes.
|
|||
|
||||
## [Unreleased]
|
||||
|
||||
## [0.71.11] - 2026-06-04
|
||||
|
||||
### Added
|
||||
- **GRPO / RL callbacks — live wiring** (closes #235, #236, #237, #238, #239,
|
||||
#240, #159, #160). The reward-hacking, cross-tokenizer distillation, MiniLLM,
|
||||
mid-epoch RL checkpoint, iterative-DPO and echo-trap surfaces that shipped
|
||||
schema-only in v0.70.0 are now real, validated end-to-end on SmolLM2-135M.
|
||||
- **Reward-hacking detector is live** (#235). `--reward-hack-detector
|
||||
info_rm|rm_ensemble` now installs a GRPO `TrainerCallback` that reads the
|
||||
per-step rewards (via a shared, thread-safe reward-fn capture buffer),
|
||||
computes an InfoRM cluster-separation drop (`info_rm`) or RM-ensemble
|
||||
divergence (`rm_ensemble`), classifies OK/WARN/HACK, logs the verdict to
|
||||
`state.log_history`, and halts training on HACK when `--reward-hack-halt` is
|
||||
set. `rm_ensemble` requires ≥2 reward functions.
|
||||
- **Cross-tokenizer ULD distillation is live** (#236). `task: distill` with
|
||||
`--uld-strategy wasserstein|topk_align` now computes a real Wasserstein-1
|
||||
(sorted-CDF) or top-k-aligned distillation loss inside the distill trainer,
|
||||
handling student/teacher vocab-size mismatch by clamping teacher ids to the
|
||||
teacher vocab.
|
||||
- **MiniLLM reverse-KL distillation is live** (#237). `--minillm-enabled` adds
|
||||
a teacher-mixed, length-normalised reverse-KL term plus an optional
|
||||
pretrain-anchor SFT term (`--minillm-pretrain-anchor-path` /
|
||||
`--minillm-pretrain-anchor-weight`) that keeps the student near coherent
|
||||
language. The anchor corpus reader is cwd-contained + symlink-rejecting with
|
||||
a per-line byte cap.
|
||||
- **Mid-epoch RL checkpoint is live** (#238). `--rl-checkpoint-save-every-steps
|
||||
N` writes a real adapter + optimizer state + JSON manifest every N steps
|
||||
during PPO/GRPO and prunes to `--rl-checkpoint-keep-last`, so a long RL run
|
||||
survives a crash without losing the optimizer momentum.
|
||||
- **`soup iterative-dpo` orchestrator is live** (#239). Runs the full
|
||||
sample → reward-score → build-pairs → DPO-train loop across rounds: each
|
||||
round samples completions from the previous round's adapter, the next round
|
||||
trains a fresh LoRA from the base on that round's harvested pairs.
|
||||
`--plan-only` still renders the plan without running.
|
||||
- **Echo-trap detector is live** (#240). `--echo-trap-enabled` installs a GRPO
|
||||
callback that scores per-trajectory n-gram repetition, classifies
|
||||
OK/WARN/TRAP against `--echo-trap-threshold`, logs the verdict, and halts on
|
||||
TRAP when `--echo-trap-halt` is set (catches RAGEN-style degenerate
|
||||
repetition in multi-turn agent RL).
|
||||
- **GRPO variant fallback now warns once** (#159). When a `--grpo-variant`
|
||||
custom `compute_loss` falls back to the base trainer (because the installed
|
||||
TRL renamed the loss inputs), the trainer logs a one-shot WARNING instead of
|
||||
silently degrading to the default objective.
|
||||
|
||||
### Changed
|
||||
- **GRPO reference-model EMA no longer materialises full state dicts** (#160).
|
||||
`--ref-model-ema-alpha` now updates the reference model in place by iterating
|
||||
`named_parameters()` (`ref = (1-α)·ref + α·policy`), eliminating the three
|
||||
model-sized allocations per step the v0.53.11 path made. A total
|
||||
name/shape-mismatch (0 shared parameters) logs a one-shot WARNING so a
|
||||
misconfigured EMA can't silently no-op.
|
||||
|
||||
## [0.71.10] - 2026-06-03
|
||||
|
||||
### Added
|
||||
|
|
|
|||
|
|
@ -120,7 +120,7 @@ src/soup_cli/
|
|||
templates/ - 17 built-in soup.yaml templates (YAML + manifest.json) with load_template loader (v0.39.0, +bco v0.40.0)
|
||||
ui/ - Web UI (FastAPI + HTML/JS SPA)
|
||||
|
||||
tests/ - Test suite (280 files, 13142 tests)
|
||||
tests/ - Test suite (281 files, 13203 tests)
|
||||
examples/ - Real-world config examples and datasets
|
||||
```
|
||||
|
||||
|
|
|
|||
28
README.md
28
README.md
|
|
@ -49,21 +49,21 @@ infrastructure instead of improving models. Soup fixes that.
|
|||
|
||||
## What's New
|
||||
|
||||
**v0.71.10 — RAG family (live).** The retrieval-augmented fine-tuning, steering, and citation
|
||||
surfaces are now real, validated on SmolLM2-135M:
|
||||
**v0.71.11 — GRPO / RL callbacks (live).** The reward-safety, distillation, and RL-loop
|
||||
surfaces from v0.70.0 are now real, validated end-to-end on SmolLM2-135M:
|
||||
|
||||
- **`data.format: raft`** — RAFT (retrieval-augmented fine-tuning): train on a query + golden
|
||||
document mixed with distractors, answer-only loss, each doc labelled `[doc-N]` so the model
|
||||
learns to cite the supporting source and ignore noise.
|
||||
- **`soup ra-dit`** — one-shot two-stage orchestrator: trains the retriever then the generator
|
||||
and records the trained retriever as the generator's paired retriever. A `soup train` of a
|
||||
generator stage auto-links the latest RA-DIT retriever from the Registry.
|
||||
- **`soup steer train --method caa|iti|repe` + `soup serve --steer <name>`** — fit a
|
||||
contrastive-activation / inference-time-intervention / representation-engineering control
|
||||
vector from `{positive, negative}` pairs and apply it at decode time via a forward hook.
|
||||
- **`soup eval citation`** — score citation precision / recall / F1 over predictions or RAFT
|
||||
rows; with `citation_faithful: true`, `[doc-id]` spans get a boosted per-token loss weight.
|
||||
A new `citation` failure mode joins `soup diagnose`.
|
||||
- **`soup train --reward-hack-detector info_rm|rm_ensemble`** — a live GRPO callback that watches
|
||||
the per-step rewards, computes an InfoRM cluster-separation drop or RM-ensemble divergence,
|
||||
classifies OK/WARN/HACK, and halts on HACK with `--reward-hack-halt`.
|
||||
- **`soup train --echo-trap-enabled`** — catches RAGEN-style degenerate repetition in multi-turn
|
||||
agent RL: scores per-trajectory n-gram repetition, OK/WARN/TRAP, halts with `--echo-trap-halt`.
|
||||
- **`soup train --rl-checkpoint-save-every-steps N`** — mid-epoch RL checkpoints (adapter +
|
||||
optimizer state + manifest) so a long GRPO/PPO run survives a crash; pruned to `--rl-checkpoint-keep-last`.
|
||||
- **`task: distill` with `--uld-strategy` or `--minillm-enabled`** — live cross-tokenizer ULD
|
||||
(Wasserstein-1 / top-k aligned) and MiniLLM teacher-mixed length-normalised reverse-KL with an
|
||||
optional pretrain anchor.
|
||||
- **`soup iterative-dpo`** — runs the full sample → reward-score → build-pairs → DPO-train loop
|
||||
across rounds; each round samples from the previous adapter and trains a fresh LoRA on the harvest.
|
||||
|
||||
Full history: [CHANGELOG.md](CHANGELOG.md) · [GitHub Releases](https://github.com/MakazhanAlpamys/Soup/releases).
|
||||
|
||||
|
|
|
|||
|
|
@ -229,12 +229,12 @@ soup expect <data.jsonl> <suite.yaml> Expectations suite: PII / token-le
|
|||
soup data gen-magpie --base <m> --provider ollama|vllm --target N --output <jsonl> [--base-url <url>] [--quality-filter] Magpie synthetic generator — live (v0.69.0; live v0.71.6)
|
||||
soup data persona-mix --prompts <jsonl> --n N --output <jsonl> Persona-Hub diversity sampler (v0.69.0)
|
||||
soup data brain-rot <data.jsonl> [--strict] Brain-rot detector — arXiv 2510.13928 (v0.69.0)
|
||||
soup iterative-dpo --base-model <m> --reward-model <rm> --prompts <p.jsonl> --output-dir <out> --rounds N --pairs-per-round N [--plan-only] Iterative DPO loop driver (v0.70.0; live runner v0.70.1)
|
||||
soup train --reward-hack-detector info_rm|rm_ensemble [--reward-hack-halt] Reward-hacking detector for GRPO/PPO (v0.70.0; live callback v0.70.1)
|
||||
soup train --uld-strategy wasserstein|topk_align [--uld-top-k N] Cross-tokenizer ULD on task='distill' (v0.70.0; live projection v0.70.1)
|
||||
soup train --minillm-enabled [--minillm-teacher-mix-ratio 0.3] MiniLLM reverse-KL on-policy distillation (v0.70.0; live v0.70.1)
|
||||
soup train --rl-checkpoint-save-every-steps N [--rl-checkpoint-keep-last N] Mid-epoch checkpoint for PPO/GRPO (v0.70.0; live save_state v0.70.1)
|
||||
soup train --echo-trap-enabled [--echo-trap-threshold 0.6 --echo-trap-halt] RAGEN echo-trap detector for multi-turn agent RL (v0.70.0; live callback v0.70.1)
|
||||
soup iterative-dpo --base-model <m> --reward-model <rm> --prompts <p.jsonl> --output-dir <out> --rounds N --pairs-per-round N [--plan-only] Iterative DPO loop driver — LIVE sample→score→pair→train (v0.70.0; live v0.71.11)
|
||||
soup train --reward-hack-detector info_rm|rm_ensemble [--reward-hack-halt] Reward-hacking detector for GRPO — LIVE callback (v0.70.0; live v0.71.11)
|
||||
soup train --uld-strategy wasserstein|topk_align [--uld-top-k N] Cross-tokenizer ULD on task='distill' — LIVE W1/topk loss (v0.70.0; live v0.71.11)
|
||||
soup train --minillm-enabled [--minillm-teacher-mix-ratio 0.3] MiniLLM reverse-KL distillation — LIVE (v0.70.0; live v0.71.11)
|
||||
soup train --rl-checkpoint-save-every-steps N [--rl-checkpoint-keep-last N] Mid-epoch checkpoint for GRPO/PPO — LIVE (v0.70.0; live v0.71.11)
|
||||
soup train --echo-trap-enabled [--echo-trap-threshold 0.6 --echo-trap-halt] RAGEN echo-trap detector for GRPO — LIVE callback (v0.70.0; live v0.71.11)
|
||||
soup version [--full] [--json] Show version (--full: system info, --json: JSON output)
|
||||
soup --verbose <command> Full traceback on errors
|
||||
```
|
||||
|
|
|
|||
|
|
@ -42,7 +42,7 @@
|
|||
|
||||
## Loop Hardening
|
||||
|
||||
The v0.70.0 release ships 6 surfaces that protect the training loop from the failure modes that cost a real GPU-hour. Schema + math kernels live now; live trainer-callback wiring lands in v0.70.1 (matches the project's stub-then-live cadence).
|
||||
Six surfaces protect the training loop from the failure modes that cost a real GPU-hour. The schema + math kernels shipped in v0.70.0; the live trainer-callback wiring shipped in **v0.71.11**, validated end-to-end on SmolLM2-135M.
|
||||
|
||||
```bash
|
||||
# Reward-hacking detector — auto-halt when the policy starts gaming the RM
|
||||
|
|
@ -68,14 +68,14 @@ soup train --config grpo.yaml \
|
|||
--rl-checkpoint-include-optimizer
|
||||
|
||||
# Iterative DPO loop driver — sample -> RM-score -> re-pair -> retrain
|
||||
# (drop --plan-only to run the loop; --plan-only just renders the per-round plan)
|
||||
soup iterative-dpo \
|
||||
--base-model meta-llama/Llama-3.1-8B \
|
||||
--reward-model ./output_rm \
|
||||
--prompts ./prompts.jsonl \
|
||||
--output-dir ./iterative_dpo_out \
|
||||
--rounds 5 \
|
||||
--pairs-per-round 1000 \
|
||||
--plan-only
|
||||
--pairs-per-round 1000
|
||||
|
||||
# RAGEN echo-trap detector — auto-halt when trajectories collapse to self-repetition
|
||||
# (Zhu et al. 2025 arXiv:2504.14437)
|
||||
|
|
@ -88,7 +88,7 @@ soup train --config grpo.yaml \
|
|||
|
||||
`--echo-trap-tokenizer-aware` switches echo-trap n-grams from whitespace tokens to the active tokenizer's integer ids. This catches subword repetition that punctuation-heavy decoded text can hide, but the score becomes tokenizer-specific rather than vocabulary-agnostic.
|
||||
|
||||
Every detector composes with v0.34 `soup why` (anomaly explainer), v0.32 spike recovery, and v0.53.11 #127 `GRPOStabilityCallback` so a single training run can have InfoRM + echo-trap + spike-recovery + ref-model regen all active simultaneously without duplicating trajectory / state collection. Live trainer-callback wiring for all 6 Parts lands in v0.70.1 (`build_reward_hack_callback`, `build_uld_projection`, `build_minillm_callback`, `build_rl_checkpoint_callback`, `run_iterative_dpo`, `build_echo_trap_callback`); today every CLI / config flag is validated at schema-load so misconfigured runs fail loudly at config-load time.
|
||||
Every detector composes with v0.34 `soup why` (anomaly explainer), v0.32 spike recovery, and the v0.53.11 #127 `GRPOStabilityCallback` so a single GRPO run can have InfoRM + echo-trap + spike-recovery + in-place ref-model EMA all active simultaneously without duplicating trajectory / state collection. The reward-hack and echo-trap callbacks read the per-step rewards through a shared, thread-safe capture buffer (Soup wraps your reward functions so it never has to monkeypatch TRL); `rm_ensemble` needs ≥2 reward functions to compute a divergence. The MiniLLM teacher-mix is an offline distribution-blend analog of the paper's on-policy teacher-mixed *sampling*, and ULD compares the distributions after clamping teacher ids to the teacher vocab (correct for same-family / extended-vocab pairs; a genuinely different tokenization needs a sequence-alignment step). The reference-model EMA (`--ref-model-ema-alpha`) updates in place — no full `state_dict` round-trip — so it is cheap at 70B+ scale.
|
||||
|
||||
|
||||
## Unlearning (`task='unlearn'`, NPO / SimNPO / RMU)
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
|||
|
||||
[project]
|
||||
name = "soup-cli"
|
||||
version = "0.71.10"
|
||||
version = "0.71.11"
|
||||
description = "Fine-tune LLMs in one command. No SSH, no config hell."
|
||||
readme = "README.md"
|
||||
license = "Apache-2.0"
|
||||
|
|
|
|||
|
|
@ -1,3 +1,3 @@
|
|||
"""Soup CLI — Fine-tune LLMs in one command."""
|
||||
|
||||
__version__ = "0.71.10"
|
||||
__version__ = "0.71.11"
|
||||
|
|
|
|||
|
|
@ -75,21 +75,29 @@ def main(
|
|||
if plan_only:
|
||||
console.print(
|
||||
Panel(
|
||||
"[green]Plan rendered[/green] (--plan-only). To execute, "
|
||||
"drop the flag once v0.70.1 ships.",
|
||||
"[green]Plan rendered[/green] (--plan-only). Drop the flag "
|
||||
"to execute the sample → score → pair → train loop.",
|
||||
title="Iterative DPO",
|
||||
)
|
||||
)
|
||||
raise typer.Exit(code=0)
|
||||
|
||||
try:
|
||||
run_iterative_dpo(plan)
|
||||
except NotImplementedError as exc:
|
||||
result = run_iterative_dpo(plan)
|
||||
except Exception as exc: # noqa: BLE001 — CLI boundary: friendly exit-1
|
||||
console.print(
|
||||
Panel(
|
||||
"[yellow]Live runner deferred to v0.70.1.[/yellow] "
|
||||
f"{escape(str(exc))}",
|
||||
f"[red]Iterative DPO failed:[/red] {escape(str(exc))}",
|
||||
title="Iterative DPO",
|
||||
)
|
||||
)
|
||||
raise typer.Exit(code=3) from exc
|
||||
raise typer.Exit(code=1) from exc
|
||||
|
||||
console.print(
|
||||
Panel(
|
||||
f"[green]Done[/green] — {result.rounds_completed} round(s); "
|
||||
f"final adapter: {escape(result.final_adapter)}; "
|
||||
f"pairs/round: {list(result.per_round_pairs)}",
|
||||
title="Iterative DPO",
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -61,6 +61,61 @@ def update_ema(ref_state: dict, policy_state: dict, alpha: float) -> dict:
|
|||
return ref_state
|
||||
|
||||
|
||||
def _validate_alpha(alpha) -> float:
|
||||
"""Shared alpha guard: non-bool float in (0, 1], finite."""
|
||||
import math
|
||||
|
||||
if not isinstance(alpha, (int, float)) or isinstance(alpha, bool):
|
||||
raise TypeError("alpha must be a non-bool float")
|
||||
alpha_f = float(alpha)
|
||||
if not math.isfinite(alpha_f):
|
||||
raise ValueError("alpha must be finite (no NaN/Inf)")
|
||||
if not (0.0 < alpha_f <= 1.0):
|
||||
raise ValueError(f"alpha must be in (0, 1], got {alpha}")
|
||||
return alpha_f
|
||||
|
||||
|
||||
def update_ema_in_place(ref_model, policy_model, alpha: float) -> int:
|
||||
"""In-place EMA of the reference model from the policy (v0.71.11 #160).
|
||||
|
||||
``ref.param = (1-α)·ref.param + α·policy.param`` for every shared
|
||||
parameter, mutating the reference tensors directly. This replaces the
|
||||
v0.53.11 path that materialised BOTH full ``state_dict()`` copies AND
|
||||
a ``load_state_dict`` round-trip — three full model-sized allocations
|
||||
per step. Iterating ``named_parameters()`` and updating in place keeps
|
||||
the per-step memory overhead at zero (no extra model-sized buffers),
|
||||
which matters at 70B+ scale.
|
||||
|
||||
Only parameters present (by name) and shape-matching in both models
|
||||
are updated; mismatches are skipped (defensive against PEFT vs base
|
||||
naming). Returns the number of parameters actually updated so callers
|
||||
can detect a total-name-mismatch no-op (code-review LOW fix).
|
||||
"""
|
||||
import torch
|
||||
|
||||
alpha_f = _validate_alpha(alpha)
|
||||
ref_params = dict(ref_model.named_parameters())
|
||||
updated = 0
|
||||
with torch.no_grad():
|
||||
for name, p_tensor in policy_model.named_parameters():
|
||||
r_tensor = ref_params.get(name)
|
||||
if r_tensor is None:
|
||||
continue
|
||||
if (
|
||||
hasattr(r_tensor, "shape")
|
||||
and hasattr(p_tensor, "shape")
|
||||
and r_tensor.shape != p_tensor.shape
|
||||
):
|
||||
continue
|
||||
src = p_tensor.data
|
||||
if hasattr(src, "to") and hasattr(r_tensor, "device"):
|
||||
src = src.to(r_tensor.device)
|
||||
# r = (1-α)·r + α·p, fully in place (no model-sized temporaries).
|
||||
r_tensor.data.mul_(1.0 - alpha_f).add_(src, alpha=alpha_f)
|
||||
updated += 1
|
||||
return updated
|
||||
|
||||
|
||||
def check_tis_threshold(log_ratio, threshold: float) -> bool:
|
||||
"""Return True iff the max absolute log-ratio exceeds the TIS threshold.
|
||||
|
||||
|
|
@ -171,6 +226,10 @@ class GRPOStabilityCallback(_TrainerCallbackBase): # type: ignore[misc, valid-t
|
|||
if self.replay_buffer_size is not None:
|
||||
self._replay = deque(maxlen=int(self.replay_buffer_size))
|
||||
self._tis_alerts = 0
|
||||
# One-shot guard so a total name-mismatch EMA no-op warns exactly
|
||||
# once per run (v0.71.11 code-review LOW — mirrors the #159
|
||||
# fallback-warn pattern in trainer/grpo.py).
|
||||
self._ema_noop_warned = False
|
||||
# Set during on_train_begin (lazy — model is constructed by Trainer
|
||||
# before the first event fires).
|
||||
self._policy_model: Any = None
|
||||
|
|
@ -229,6 +288,9 @@ class GRPOStabilityCallback(_TrainerCallbackBase): # type: ignore[misc, valid-t
|
|||
instability.
|
||||
"""
|
||||
# Live EMA update of reference model from current policy.
|
||||
# v0.71.11 #160 — in-place update (no full state_dict / load round
|
||||
# trip). Mutates the ref parameters directly so a 70B+ run pays
|
||||
# zero extra model-sized allocations per step.
|
||||
if (
|
||||
self.ref_model_ema_alpha is not None
|
||||
and self._ref_model is not None
|
||||
|
|
@ -236,23 +298,18 @@ class GRPOStabilityCallback(_TrainerCallbackBase): # type: ignore[misc, valid-t
|
|||
):
|
||||
try:
|
||||
policy = model if model is not None else self._policy_model
|
||||
ref_sd = self._ref_model.state_dict()
|
||||
pol_sd = policy.state_dict()
|
||||
update_ema(ref_sd, pol_sd, self.ref_model_ema_alpha)
|
||||
# v0.53.11 review fix (security HIGH) — strict=True with
|
||||
# try/except for key mismatch. strict=False silently
|
||||
# dropped unknown keys, masking corruption from a crafted
|
||||
# checkpoint. We catch the RuntimeError and downgrade to
|
||||
# strict=False with a WARNING so operators see the drift.
|
||||
try:
|
||||
self._ref_model.load_state_dict(ref_sd, strict=True)
|
||||
except RuntimeError as key_err:
|
||||
updated = update_ema_in_place(
|
||||
self._ref_model, policy, self.ref_model_ema_alpha
|
||||
)
|
||||
if updated == 0 and not self._ema_noop_warned:
|
||||
self._ema_noop_warned = True
|
||||
logger.warning(
|
||||
"EMA load_state_dict key mismatch; falling back to "
|
||||
"strict=False (potential silent corruption): %s",
|
||||
key_err,
|
||||
"ref_model_ema_alpha is set but the EMA update matched "
|
||||
"0 shared parameters between the reference and policy "
|
||||
"models (name/shape mismatch) — the reference model is "
|
||||
"NOT being updated. Check that both models share the "
|
||||
"same architecture."
|
||||
)
|
||||
self._ref_model.load_state_dict(ref_sd, strict=False)
|
||||
except Exception as exc: # noqa: BLE001 — never crash training
|
||||
logger.debug("EMA update skipped: %s", exc)
|
||||
# Surface counters to log_history.
|
||||
|
|
|
|||
|
|
@ -238,20 +238,62 @@ class DistillTrainerWrapper:
|
|||
for p in self.teacher.parameters():
|
||||
p.requires_grad_(False)
|
||||
|
||||
# Cross-tokenizer distillation is not supported — vocab sizes must
|
||||
# match so KL between logit distributions is well-defined.
|
||||
# v0.71.11 #236 — cross-tokenizer ULD projection. When uld_strategy
|
||||
# is set, a vocab-size mismatch is EXPECTED (that's the whole point)
|
||||
# and the ULD loss handles it; otherwise vocab sizes must match so
|
||||
# the column-wise KL is well-defined.
|
||||
teacher_vocab = getattr(self.teacher.config, "vocab_size", None)
|
||||
student_vocab = getattr(self.model.config, "vocab_size", None)
|
||||
if (
|
||||
uld_projection = None
|
||||
if tcfg.uld_strategy is not None:
|
||||
from soup_cli.utils.uld import ULDConfig, build_uld_projection
|
||||
|
||||
uld_projection = build_uld_projection(
|
||||
ULDConfig(
|
||||
strategy=tcfg.uld_strategy,
|
||||
student_vocab_size=int(student_vocab or 32000),
|
||||
teacher_vocab_size=int(teacher_vocab or 32000),
|
||||
top_k=tcfg.uld_top_k,
|
||||
)
|
||||
)
|
||||
console.print(
|
||||
f"[green]Cross-tokenizer ULD enabled[/] "
|
||||
f"(strategy={tcfg.uld_strategy})"
|
||||
)
|
||||
elif (
|
||||
teacher_vocab is not None
|
||||
and student_vocab is not None
|
||||
and teacher_vocab != student_vocab
|
||||
):
|
||||
raise ValueError(
|
||||
f"Teacher vocab size ({teacher_vocab}) != student vocab "
|
||||
f"size ({student_vocab}). Cross-tokenizer distillation is "
|
||||
"not supported in v0.53.2; use a teacher that shares the "
|
||||
"student tokenizer family."
|
||||
f"size ({student_vocab}). Cross-tokenizer distillation needs "
|
||||
"training.uld_strategy (wasserstein / topk_align); use a "
|
||||
"teacher that shares the student tokenizer family, or set "
|
||||
"uld_strategy."
|
||||
)
|
||||
|
||||
# v0.71.11 #237 — MiniLLM on-policy distillation modifier.
|
||||
minillm_cb = None
|
||||
if tcfg.minillm_enabled:
|
||||
from soup_cli.utils.minillm import MiniLLMConfig, build_minillm_callback
|
||||
|
||||
minillm_cb = build_minillm_callback(
|
||||
MiniLLMConfig(
|
||||
teacher_mix_ratio=float(tcfg.minillm_teacher_mix_ratio),
|
||||
length_normalize=bool(tcfg.minillm_length_normalize),
|
||||
pretrain_anchor_weight=float(
|
||||
tcfg.minillm_pretrain_anchor_weight
|
||||
),
|
||||
pretrain_anchor_path=tcfg.minillm_pretrain_anchor_path,
|
||||
),
|
||||
tokenizer=self.tokenizer,
|
||||
temperature=temperature,
|
||||
)
|
||||
console.print(
|
||||
f"[green]MiniLLM distillation enabled[/] "
|
||||
f"(teacher_mix={tcfg.minillm_teacher_mix_ratio}, "
|
||||
f"anchor={tcfg.minillm_pretrain_anchor_weight})"
|
||||
)
|
||||
|
||||
# Dataset prep — reuse the SFT formatter so distill sees
|
||||
|
|
@ -308,6 +350,9 @@ class DistillTrainerWrapper:
|
|||
)
|
||||
|
||||
teacher_ref = self.teacher
|
||||
_teacher_vocab = teacher_vocab
|
||||
_uld_projection = uld_projection
|
||||
_minillm_cb = minillm_cb
|
||||
|
||||
class _DistillTrainer(Trainer):
|
||||
def compute_loss(
|
||||
|
|
@ -347,14 +392,43 @@ class DistillTrainerWrapper:
|
|||
for k, v in inputs.items()
|
||||
if k != "labels"
|
||||
}
|
||||
# v0.71.11 #236 — when ULD bridges different vocabs, clamp the
|
||||
# student token ids to the teacher's range so the (possibly
|
||||
# smaller) teacher embedding never index-errors.
|
||||
if (
|
||||
_uld_projection is not None
|
||||
and _teacher_vocab is not None
|
||||
and "input_ids" in teacher_inputs
|
||||
):
|
||||
teacher_inputs["input_ids"] = teacher_inputs[
|
||||
"input_ids"
|
||||
].clamp(max=int(_teacher_vocab) - 1)
|
||||
with torch.no_grad():
|
||||
teacher_out = teacher_ref(**teacher_inputs)
|
||||
teacher_logits = teacher_out.logits.to(student_logits.device)
|
||||
|
||||
distill_loss = _compute_distill_term(
|
||||
student_logits, teacher_logits, divergence, temperature
|
||||
)
|
||||
anchor = None
|
||||
if _uld_projection is not None:
|
||||
# v0.71.11 #236 — cross-tokenizer ULD distillation loss.
|
||||
distill_loss = _uld_projection(
|
||||
student_logits,
|
||||
teacher_logits,
|
||||
attention_mask=inputs.get("attention_mask"),
|
||||
)
|
||||
elif _minillm_cb is not None:
|
||||
# v0.71.11 #237 — MiniLLM teacher-mixed reverse-KL +
|
||||
# pretrain anchor.
|
||||
distill_loss = _minillm_cb.distill_term(
|
||||
student_logits, teacher_logits, labels
|
||||
)
|
||||
anchor = _minillm_cb.anchor_term(model)
|
||||
else:
|
||||
distill_loss = _compute_distill_term(
|
||||
student_logits, teacher_logits, divergence, temperature
|
||||
)
|
||||
total = _CE_WEIGHT * ce_loss + _DISTILL_WEIGHT * distill_loss
|
||||
if anchor is not None:
|
||||
total = total + anchor
|
||||
return (total, outputs) if return_outputs else total
|
||||
|
||||
# ``DataCollatorForSeq2Seq`` pads ``input_ids`` and ``attention_mask``
|
||||
|
|
@ -376,6 +450,11 @@ class DistillTrainerWrapper:
|
|||
),
|
||||
)
|
||||
|
||||
# v0.71.11 #237 — attach the MiniLLM callback for lifecycle (the
|
||||
# loss terms are applied directly in compute_loss above).
|
||||
if minillm_cb is not None:
|
||||
self.trainer.add_callback(minillm_cb)
|
||||
|
||||
# v0.40.6 #67 — ReLoRA callback.
|
||||
from soup_cli.utils.peft_wiring import (
|
||||
attach_curriculum_callback,
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
|
|
@ -13,6 +14,7 @@ from soup_cli.config.schema import SoupConfig
|
|||
from soup_cli.utils.gpu import estimate_batch_size, model_size_from_name
|
||||
|
||||
console = Console()
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def make_grpo_trainer_variant(base_cls: type, variant: str) -> type:
|
||||
|
|
@ -42,6 +44,29 @@ def _make_grpo_trainer_variant_cached(base_cls: type, variant: str) -> type:
|
|||
"""GRPOTrainer subclass that routes compute_loss through Soup's variants."""
|
||||
|
||||
_soup_grpo_variant: str = variant
|
||||
# v0.71.11 #159 — one-shot WARNING flag so a silent fallback to the
|
||||
# stock TRL loss surfaces exactly once (not on every step).
|
||||
_soup_fallback_warned: bool = False
|
||||
|
||||
def _warn_fallback(self, reason: str) -> None:
|
||||
"""Emit a one-shot WARNING when the variant kernel falls back.
|
||||
|
||||
v0.71.11 #159 — when a TRL internal rename or a kernel error
|
||||
makes ``compute_loss`` delegate to the stock GRPO loss, the
|
||||
operator's selected variant silently stops applying. Warn once
|
||||
so the run isn't quietly training the wrong objective.
|
||||
"""
|
||||
if self._soup_fallback_warned:
|
||||
return
|
||||
self._soup_fallback_warned = True
|
||||
logger.warning(
|
||||
"GRPO variant %r compute_loss fell back to the stock TRL "
|
||||
"loss (%s); the selected objective is NOT being applied. "
|
||||
"This usually means a TRL version renamed the per-token "
|
||||
"log-prob inputs.",
|
||||
self._soup_grpo_variant,
|
||||
reason,
|
||||
)
|
||||
|
||||
def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
|
||||
# v0.53.11 review fix (code-review HIGH) — read kernel inputs
|
||||
|
|
@ -65,6 +90,7 @@ def _make_grpo_trainer_variant_cached(base_cls: type, variant: str) -> type:
|
|||
if logp_new is None or logp_old is None or advantages is None:
|
||||
# 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
|
||||
)
|
||||
|
|
@ -82,11 +108,13 @@ def _make_grpo_trainer_variant_cached(base_cls: type, variant: str) -> type:
|
|||
delta=delta,
|
||||
completion_mask=completion_mask,
|
||||
)
|
||||
except (TypeError, ValueError):
|
||||
except (TypeError, ValueError) as exc:
|
||||
self._warn_fallback(f"kernel error: {exc}")
|
||||
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
|
||||
)
|
||||
|
|
@ -198,6 +226,23 @@ class GRPOTrainerWrapper:
|
|||
|
||||
reward_fn = load_reward_fn(tcfg.reward_fn)
|
||||
|
||||
# v0.71.11 #235/#240 — when the reward-hack or echo-trap detector is
|
||||
# enabled, wrap the reward function(s) with a capture shim so the
|
||||
# callbacks can observe the step's rewards + completions. The buffer
|
||||
# is created here (before GRPOTrainer construction) and handed to
|
||||
# the callbacks after the trainer is built.
|
||||
from soup_cli.utils.peft_wiring import rl_callbacks_need_buffer
|
||||
|
||||
self._rl_buffer = None
|
||||
if rl_callbacks_need_buffer(tcfg):
|
||||
from soup_cli.utils.rl_signal_buffer import (
|
||||
RLSignalBuffer,
|
||||
wrap_reward_funcs,
|
||||
)
|
||||
|
||||
self._rl_buffer = RLSignalBuffer()
|
||||
reward_fn = wrap_reward_funcs(reward_fn, self._rl_buffer)
|
||||
|
||||
if use_unsloth:
|
||||
self._setup_unsloth(cfg, tcfg)
|
||||
else:
|
||||
|
|
@ -342,6 +387,18 @@ class GRPOTrainerWrapper:
|
|||
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,
|
||||
buffer=self._rl_buffer,
|
||||
tokenizer=self.tokenizer,
|
||||
output_dir=str(output_dir),
|
||||
task="grpo",
|
||||
)
|
||||
|
||||
# v0.40.6 #67 — ReLoRA callback (magnitude-prune LoRA every N steps).
|
||||
from soup_cli.utils.peft_wiring import (
|
||||
attach_curriculum_callback,
|
||||
|
|
|
|||
|
|
@ -25,7 +25,7 @@ from __future__ import annotations
|
|||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Iterable, Sequence
|
||||
from typing import Iterable, Optional, Sequence
|
||||
|
||||
VERDICTS: tuple[str, ...] = ("OK", "WARN", "TRAP")
|
||||
_VALID_VERDICTS: frozenset[str] = frozenset(VERDICTS)
|
||||
|
|
@ -294,43 +294,177 @@ class EchoTrapReport:
|
|||
)
|
||||
|
||||
|
||||
def _split_whitespace(text: str) -> list[str]:
|
||||
"""Whitespace tokenisation for the string echo path."""
|
||||
return text.split()
|
||||
|
||||
|
||||
def _get_trainer_callback_base():
|
||||
"""Lazy-resolve ``transformers.TrainerCallback`` (mirror v0.53.11)."""
|
||||
try:
|
||||
from transformers import TrainerCallback
|
||||
|
||||
return TrainerCallback
|
||||
except ImportError:
|
||||
return object
|
||||
|
||||
|
||||
_TrainerCallbackBase = _get_trainer_callback_base()
|
||||
|
||||
|
||||
class EchoTrapCallback(_TrainerCallbackBase): # type: ignore[misc, valid-type]
|
||||
"""Live HF TrainerCallback for echo-trap detection (v0.71.11 #240).
|
||||
|
||||
Reads the GRPO step's generated completions (via the shared
|
||||
:class:`~soup_cli.utils.rl_signal_buffer.RLSignalBuffer`), scores
|
||||
trajectory repetition, and classifies OK / WARN / TRAP. When
|
||||
``tokenizer_aware`` and a tokenizer are supplied, scores over
|
||||
tokenizer ids (subword-repetition sensitive); otherwise whitespace
|
||||
tokens.
|
||||
|
||||
The aggregate echo signal is surfaced to ``state.log_history``;
|
||||
``halt_on_trap`` sets ``control.should_training_stop`` on TRAP.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
threshold: float,
|
||||
halt_on_trap: bool = True,
|
||||
ngram_n: int = 2,
|
||||
tokenizer_aware: bool = False,
|
||||
buffer: object = None,
|
||||
tokenizer: object = None,
|
||||
) -> None:
|
||||
if isinstance(threshold, bool):
|
||||
raise ValueError("threshold must not be bool")
|
||||
if not isinstance(threshold, (int, float)):
|
||||
raise ValueError(
|
||||
f"threshold must be a number, got {type(threshold).__name__}"
|
||||
)
|
||||
fv = float(threshold)
|
||||
if not math.isfinite(fv) or not (0.0 <= fv <= 1.0):
|
||||
raise ValueError(f"threshold must be in [0.0, 1.0], got {threshold}")
|
||||
if not isinstance(halt_on_trap, bool):
|
||||
raise TypeError(
|
||||
f"halt_on_trap must be bool, got {type(halt_on_trap).__name__}"
|
||||
)
|
||||
if not isinstance(tokenizer_aware, bool):
|
||||
raise TypeError(
|
||||
"tokenizer_aware must be bool, got "
|
||||
f"{type(tokenizer_aware).__name__}"
|
||||
)
|
||||
self.threshold = fv
|
||||
self.halt_on_trap = halt_on_trap
|
||||
self.ngram_n = _check_ngram_n(ngram_n)
|
||||
self.tokenizer_aware = tokenizer_aware
|
||||
self.buffer = buffer
|
||||
self.tokenizer = tokenizer
|
||||
self._last_report: Optional[EchoTrapReport] = None
|
||||
self._traps_seen = 0
|
||||
|
||||
def compute_signal(self, snapshot: dict) -> Optional[float]:
|
||||
"""Compute the aggregate echo signal from a buffer snapshot.
|
||||
|
||||
Returns ``None`` when no completions are available.
|
||||
"""
|
||||
completions = snapshot.get("completions", []) or []
|
||||
if not completions:
|
||||
return None
|
||||
if self.tokenizer_aware and self.tokenizer is not None:
|
||||
id_trajectories: list[list[int]] = []
|
||||
for text in completions:
|
||||
try:
|
||||
ids = self.tokenizer.encode(text, add_special_tokens=False)
|
||||
except (TypeError, ValueError):
|
||||
ids = []
|
||||
id_trajectories.append([int(i) for i in ids])
|
||||
return score_echo_signal_tokenized(id_trajectories, ngram_n=self.ngram_n)
|
||||
trajectories = [_split_whitespace(text) for text in completions]
|
||||
return score_echo_signal(trajectories, ngram_n=self.ngram_n)
|
||||
|
||||
def observe_signal(
|
||||
self, signal: float, step: int, n_trajectories: int
|
||||
) -> EchoTrapReport:
|
||||
"""Classify a signal and build the :class:`EchoTrapReport`."""
|
||||
clamped = max(0.0, min(1.0, float(signal)))
|
||||
verdict = classify_echo_signal(clamped)
|
||||
if verdict == "TRAP":
|
||||
self._traps_seen += 1
|
||||
report = EchoTrapReport(
|
||||
signal=clamped,
|
||||
verdict=verdict,
|
||||
step=max(0, int(step)),
|
||||
trajectories_seen=max(0, int(n_trajectories)),
|
||||
details=(
|
||||
f"ngram_n={self.ngram_n}",
|
||||
f"tokenizer_aware={self.tokenizer_aware}",
|
||||
f"signal={clamped:.4f} verdict={verdict} thr={self.threshold}",
|
||||
),
|
||||
)
|
||||
self._last_report = report
|
||||
return report
|
||||
|
||||
def last_report(self) -> Optional[EchoTrapReport]:
|
||||
"""Return the most recent :class:`EchoTrapReport` (or None)."""
|
||||
return self._last_report
|
||||
|
||||
def on_step_end(self, args, state, control, **kwargs):
|
||||
"""Per-step probe — read completions from the capture buffer."""
|
||||
if self.buffer is None:
|
||||
return control
|
||||
try:
|
||||
snapshot = self.buffer.snapshot()
|
||||
signal = self.compute_signal(snapshot)
|
||||
if signal is None:
|
||||
return control
|
||||
n_traj = len(snapshot.get("completions", []) or [])
|
||||
step = int(getattr(state, "global_step", 0) or 0)
|
||||
report = self.observe_signal(signal, step, n_traj)
|
||||
log_history = getattr(state, "log_history", None)
|
||||
if log_history is not None:
|
||||
log_history.append({
|
||||
"echo_trap_signal": report.signal,
|
||||
"echo_trap_verdict": report.verdict,
|
||||
})
|
||||
# TRAP verdict + over-threshold → optional halt.
|
||||
if (
|
||||
report.verdict == "TRAP"
|
||||
and report.signal >= self.threshold
|
||||
and self.halt_on_trap
|
||||
and control is not None
|
||||
):
|
||||
try:
|
||||
control.should_training_stop = True
|
||||
except Exception: # noqa: BLE001 — never crash training
|
||||
pass
|
||||
return control
|
||||
except Exception: # noqa: BLE001 — instrumentation must never crash
|
||||
return control
|
||||
|
||||
|
||||
def build_echo_trap_callback(
|
||||
*,
|
||||
threshold: float,
|
||||
halt_on_trap: bool = True,
|
||||
ngram_n: int = 2,
|
||||
tokenizer_aware: bool = False,
|
||||
):
|
||||
"""Live HF Trainer callback for echo-trap detection.
|
||||
buffer: object = None,
|
||||
tokenizer: object = None,
|
||||
) -> "EchoTrapCallback":
|
||||
"""Build the live echo-trap HF Trainer callback (v0.71.11 #240).
|
||||
|
||||
Deferred to v0.70.1. Validates inputs at the public boundary so
|
||||
misconfigured callers fail fast (mirrors v0.50.0 / v0.62.0 /
|
||||
v0.67.0 / v0.69.0 / Part A/B/C/D/E deferred-live policy).
|
||||
Lifts the v0.70.0 ``NotImplementedError`` stub. Validates every input
|
||||
at the public boundary (mirrors v0.50.0 / v0.61.0 fail-fast policy),
|
||||
then returns an :class:`EchoTrapCallback`.
|
||||
"""
|
||||
# Threshold validation runs FIRST so a bad numeric value fires the
|
||||
# actionable error before we check the bool flags.
|
||||
if isinstance(threshold, bool):
|
||||
raise ValueError("threshold must not be bool")
|
||||
if not isinstance(threshold, (int, float)):
|
||||
raise ValueError(
|
||||
f"threshold must be a number, got {type(threshold).__name__}"
|
||||
)
|
||||
fv = float(threshold)
|
||||
if not math.isfinite(fv) or not (0.0 <= fv <= 1.0):
|
||||
raise ValueError(f"threshold must be in [0.0, 1.0], got {threshold}")
|
||||
if not isinstance(halt_on_trap, bool):
|
||||
raise TypeError(
|
||||
f"halt_on_trap must be bool, got {type(halt_on_trap).__name__}"
|
||||
)
|
||||
if not isinstance(tokenizer_aware, bool):
|
||||
raise TypeError(
|
||||
"tokenizer_aware must be bool, got "
|
||||
f"{type(tokenizer_aware).__name__}"
|
||||
)
|
||||
_check_ngram_n(ngram_n)
|
||||
raise NotImplementedError(
|
||||
f"Live echo-trap HF Trainer callback (threshold={fv}) is deferred "
|
||||
"to v0.70.1. v0.70.0 ships the schema + math kernels only."
|
||||
return EchoTrapCallback(
|
||||
threshold=threshold,
|
||||
halt_on_trap=halt_on_trap,
|
||||
ngram_n=ngram_n,
|
||||
tokenizer_aware=tokenizer_aware,
|
||||
buffer=buffer,
|
||||
tokenizer=tokenizer,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -339,6 +473,7 @@ def build_echo_trap_callback(
|
|||
# without circular dependencies.
|
||||
__all__ = [
|
||||
"VERDICTS",
|
||||
"EchoTrapCallback",
|
||||
"EchoTrapReport",
|
||||
"build_echo_trap_callback",
|
||||
"classify_echo_signal",
|
||||
|
|
|
|||
|
|
@ -23,14 +23,18 @@ Security:
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from typing import Tuple
|
||||
from typing import Any, Callable, Optional, Tuple
|
||||
|
||||
_MIN_ROUNDS = 1
|
||||
_MAX_ROUNDS = 100
|
||||
_MIN_PAIRS_PER_ROUND = 10
|
||||
_MAX_PAIRS_PER_ROUND = 1_000_000
|
||||
_MAX_PATH_LEN = 4096
|
||||
_MAX_PROMPT_ROWS = 1_000_000
|
||||
_MAX_ROW_BYTES = 1_000_000
|
||||
|
||||
|
||||
def validate_rounds(value: object) -> int:
|
||||
|
|
@ -80,6 +84,10 @@ def _check_path(value: object, field: str) -> str:
|
|||
raise ValueError(f"{field} must be non-empty")
|
||||
if "\x00" in value:
|
||||
raise ValueError(f"{field} must not contain null bytes")
|
||||
# Reject newlines / CR / tab so a crafted base_model / path cannot inject
|
||||
# extra keys into the round YAML the runner renders (security review HIGH).
|
||||
if any(c in value for c in ("\n", "\r", "\t")):
|
||||
raise ValueError(f"{field} must not contain newline / tab characters")
|
||||
if len(value) > _MAX_PATH_LEN:
|
||||
raise ValueError(f"{field} exceeds {_MAX_PATH_LEN} chars")
|
||||
return value
|
||||
|
|
@ -198,18 +206,299 @@ def build_iterative_dpo_plan(
|
|||
)
|
||||
|
||||
|
||||
def run_iterative_dpo(plan):
|
||||
"""Execute the iterative-DPO loop. Deferred to v0.70.1.
|
||||
@dataclass(frozen=True)
|
||||
class IterativeDPOResult:
|
||||
"""Frozen result of a completed iterative-DPO run.
|
||||
|
||||
Validates plan type at the public boundary so misconfigured callers
|
||||
fail fast (mirrors v0.50.0 / v0.62.0 / v0.67.0 / v0.69.0 deferred-live
|
||||
policy).
|
||||
- ``rounds_completed``: number of rounds that ran end-to-end.
|
||||
- ``final_adapter``: adapter path of the last completed round.
|
||||
- ``per_round_pairs``: tuple of pair counts actually written per round.
|
||||
"""
|
||||
|
||||
rounds_completed: int
|
||||
final_adapter: str
|
||||
per_round_pairs: Tuple[int, ...]
|
||||
|
||||
|
||||
def _load_prompts(path: str) -> list[str]:
|
||||
"""Read prompt strings from a cwd-contained JSONL file.
|
||||
|
||||
Accepts ``{"prompt": "..."}`` / ``{"prompt": [messages]}`` /
|
||||
``{"messages": [...]}`` shapes; falls back to the raw line. Cwd
|
||||
containment + symlink rejection + DoS caps (security review MEDIUM).
|
||||
"""
|
||||
from soup_cli.utils.paths import enforce_under_cwd_and_no_symlink
|
||||
|
||||
real = enforce_under_cwd_and_no_symlink(path, "prompts_path")
|
||||
out: list[str] = []
|
||||
seen = 0
|
||||
with open(real, encoding="utf-8") as fh:
|
||||
for line in fh:
|
||||
seen += 1
|
||||
if seen > _MAX_PROMPT_ROWS:
|
||||
break
|
||||
line = line.strip()
|
||||
if not line or len(line) > _MAX_ROW_BYTES:
|
||||
continue
|
||||
try:
|
||||
obj = json.loads(line)
|
||||
except ValueError:
|
||||
out.append(line)
|
||||
continue
|
||||
if isinstance(obj, dict):
|
||||
prompt = obj.get("prompt")
|
||||
if isinstance(prompt, str):
|
||||
out.append(prompt)
|
||||
elif isinstance(prompt, list):
|
||||
out.append(_messages_to_text(prompt))
|
||||
elif "messages" in obj:
|
||||
out.append(_messages_to_text(obj["messages"]))
|
||||
elif "instruction" in obj:
|
||||
out.append(str(obj["instruction"]))
|
||||
return out
|
||||
|
||||
|
||||
def _messages_to_text(messages: Any) -> str:
|
||||
parts: list[str] = []
|
||||
if isinstance(messages, list):
|
||||
for m in messages:
|
||||
if isinstance(m, dict) and m.get("role") != "assistant":
|
||||
parts.append(str(m.get("content", "")))
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
def build_pairs_from_scored(
|
||||
scored: list[tuple[str, float]],
|
||||
) -> Optional[tuple[str, str]]:
|
||||
"""Pick (chosen, rejected) from a list of (completion, score).
|
||||
|
||||
Chosen = highest score, rejected = lowest. Returns ``None`` when there
|
||||
are fewer than 2 distinct-scored completions (no usable pair).
|
||||
"""
|
||||
if len(scored) < 2:
|
||||
return None
|
||||
ordered = sorted(scored, key=lambda t: t[1])
|
||||
rejected, r_score = ordered[0]
|
||||
chosen, c_score = ordered[-1]
|
||||
if c_score <= r_score:
|
||||
return None
|
||||
return chosen, rejected
|
||||
|
||||
|
||||
def _write_pairs_jsonl(pairs: list[tuple[str, str, str]], path: str) -> int:
|
||||
"""Atomically write (prompt, chosen, rejected) rows; returns the count."""
|
||||
from soup_cli.utils.paths import atomic_write_text
|
||||
|
||||
os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
|
||||
lines = []
|
||||
for prompt, chosen, rejected in pairs:
|
||||
lines.append(
|
||||
json.dumps(
|
||||
{"prompt": prompt, "chosen": chosen, "rejected": rejected}
|
||||
)
|
||||
)
|
||||
atomic_write_text("\n".join(lines) + ("\n" if lines else ""), path)
|
||||
return len(pairs)
|
||||
|
||||
|
||||
def _default_sample_fn(
|
||||
*,
|
||||
base_model: str,
|
||||
adapter_path: Optional[str],
|
||||
prompts: list[str],
|
||||
num_samples: int,
|
||||
max_new_tokens: int,
|
||||
device: Optional[str],
|
||||
) -> list[list[str]]:
|
||||
"""Generate ``num_samples`` completions per prompt (default seam).
|
||||
|
||||
Round 0 samples from ``base_model``; later rounds load the previous
|
||||
round's LoRA adapter on top of the base via PEFT (the adapter dir is
|
||||
NOT a standalone model — code-review HIGH fix).
|
||||
"""
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
tok = AutoTokenizer.from_pretrained(base_model)
|
||||
if tok.pad_token is None:
|
||||
tok.pad_token = tok.eos_token
|
||||
dev = device or ("cuda" if torch.cuda.is_available() else "cpu")
|
||||
model = AutoModelForCausalLM.from_pretrained(base_model).to(dev)
|
||||
if adapter_path is not None:
|
||||
from peft import PeftModel
|
||||
|
||||
model = PeftModel.from_pretrained(model, adapter_path).to(dev)
|
||||
model.eval()
|
||||
out: list[list[str]] = []
|
||||
for prompt in prompts:
|
||||
enc = tok(prompt, return_tensors="pt").to(dev)
|
||||
completions: list[str] = []
|
||||
with torch.no_grad():
|
||||
gen = model.generate(
|
||||
**enc,
|
||||
do_sample=True,
|
||||
num_return_sequences=num_samples,
|
||||
max_new_tokens=max_new_tokens,
|
||||
pad_token_id=tok.pad_token_id,
|
||||
)
|
||||
prompt_len = enc["input_ids"].shape[1]
|
||||
for seq in gen:
|
||||
completions.append(
|
||||
tok.decode(seq[prompt_len:], skip_special_tokens=True)
|
||||
)
|
||||
out.append(completions)
|
||||
return out
|
||||
|
||||
|
||||
def _default_score_fn(
|
||||
*,
|
||||
reward_model: str,
|
||||
prompt: str,
|
||||
completions: list[str],
|
||||
device: Optional[str],
|
||||
) -> list[float]:
|
||||
"""Score completions with a sequence-classification reward model."""
|
||||
import torch
|
||||
from transformers import AutoModelForSequenceClassification, AutoTokenizer
|
||||
|
||||
tok = AutoTokenizer.from_pretrained(reward_model)
|
||||
if tok.pad_token is None:
|
||||
tok.pad_token = tok.eos_token
|
||||
dev = device or ("cuda" if torch.cuda.is_available() else "cpu")
|
||||
rm = AutoModelForSequenceClassification.from_pretrained(reward_model).to(dev)
|
||||
rm.eval()
|
||||
scores: list[float] = []
|
||||
with torch.no_grad():
|
||||
for completion in completions:
|
||||
enc = tok(
|
||||
prompt, completion, return_tensors="pt", truncation=True
|
||||
).to(dev)
|
||||
logits = rm(**enc).logits
|
||||
scores.append(float(logits.reshape(-1)[0]))
|
||||
return scores
|
||||
|
||||
|
||||
def _default_train_fn(
|
||||
*,
|
||||
base_model: str,
|
||||
pairs_path: str,
|
||||
adapter_path: str,
|
||||
) -> None:
|
||||
"""Run a DPO round via a ``soup train`` subprocess (no shell).
|
||||
|
||||
Each round trains a fresh LoRA from ``base_model`` (always the plan's
|
||||
base, never a prior adapter dir — code-review HIGH fix) on the round's
|
||||
pairs. The YAML is rendered via ``yaml.safe_dump`` so no value can
|
||||
inject extra keys (security review HIGH).
|
||||
"""
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
|
||||
import yaml
|
||||
|
||||
yaml_text = yaml.safe_dump(
|
||||
{
|
||||
"base": base_model,
|
||||
"task": "dpo",
|
||||
"data": {"train": pairs_path, "format": "dpo", "max_length": 256},
|
||||
"training": {"epochs": 1, "batch_size": 1},
|
||||
"output": {"dir": adapter_path},
|
||||
},
|
||||
default_flow_style=False,
|
||||
sort_keys=False,
|
||||
)
|
||||
fd, tmp_yaml = tempfile.mkstemp(suffix=".yaml", prefix=".soup_idpo_", dir=os.getcwd())
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as fh:
|
||||
fh.write(yaml_text)
|
||||
subprocess.run( # noqa: S603 — argv list, no shell
|
||||
[sys.executable, "-m", "soup_cli.cli", "train", "--config", tmp_yaml, "--yes"],
|
||||
check=True,
|
||||
)
|
||||
finally:
|
||||
try:
|
||||
os.remove(tmp_yaml)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def run_iterative_dpo(
|
||||
plan,
|
||||
*,
|
||||
sample_fn: Optional[Callable] = None,
|
||||
score_fn: Optional[Callable] = None,
|
||||
train_fn: Optional[Callable] = None,
|
||||
num_samples: int = 4,
|
||||
max_new_tokens: int = 64,
|
||||
device: Optional[str] = None,
|
||||
) -> IterativeDPOResult:
|
||||
"""Execute the iterative-DPO loop (v0.71.11 #239).
|
||||
|
||||
For each round: sample completions from the current model, RM-score
|
||||
them, build (chosen, rejected) pairs, write them, then run a DPO
|
||||
round to produce the round's adapter. The next round samples from the
|
||||
previous adapter.
|
||||
|
||||
The ``sample_fn`` / ``score_fn`` / ``train_fn`` seams default to real
|
||||
implementations (load model + generate / load RM + score / subprocess
|
||||
``soup train``); tests inject fast fakes.
|
||||
"""
|
||||
if not isinstance(plan, IterativeDPOPlan):
|
||||
raise TypeError(
|
||||
f"plan must be IterativeDPOPlan, got {type(plan).__name__}"
|
||||
)
|
||||
raise NotImplementedError(
|
||||
"Live iterative-DPO loop runner is deferred to v0.70.1. "
|
||||
"v0.70.0 ships the schema + plan-only renderer only."
|
||||
sample_fn = sample_fn or _default_sample_fn
|
||||
score_fn = score_fn or _default_score_fn
|
||||
train_fn = train_fn or _default_train_fn
|
||||
|
||||
# ``sample_adapter`` is the LoRA the current policy is sampling from:
|
||||
# None on round 0 (sample from base), then the previous round's adapter.
|
||||
# Training always starts from ``plan.base_model`` (a fresh LoRA per
|
||||
# round) — the round's pairs carry the improvement signal. This keeps
|
||||
# the default seams from ever treating an adapter dir as a full model.
|
||||
sample_adapter: Optional[str] = None
|
||||
per_round: list[int] = []
|
||||
final_adapter = plan.base_model
|
||||
rounds_completed = 0
|
||||
|
||||
for rnd in plan.rounds:
|
||||
prompts = _load_prompts(rnd.prompts_path)
|
||||
sampled = sample_fn(
|
||||
base_model=plan.base_model,
|
||||
adapter_path=sample_adapter,
|
||||
prompts=prompts,
|
||||
num_samples=num_samples,
|
||||
max_new_tokens=max_new_tokens,
|
||||
device=device,
|
||||
)
|
||||
pairs: list[tuple[str, str, str]] = []
|
||||
for prompt, completions in zip(prompts, sampled):
|
||||
scores = score_fn(
|
||||
reward_model=plan.reward_model,
|
||||
prompt=prompt,
|
||||
completions=list(completions),
|
||||
device=device,
|
||||
)
|
||||
scored = list(zip(completions, scores))
|
||||
picked = build_pairs_from_scored(scored)
|
||||
if picked is not None:
|
||||
pairs.append((prompt, picked[0], picked[1]))
|
||||
if len(pairs) >= rnd.pairs_count:
|
||||
break
|
||||
written = _write_pairs_jsonl(pairs, rnd.pairs_path)
|
||||
per_round.append(written)
|
||||
train_fn(
|
||||
base_model=plan.base_model,
|
||||
pairs_path=rnd.pairs_path,
|
||||
adapter_path=rnd.adapter_path,
|
||||
)
|
||||
sample_adapter = rnd.adapter_path
|
||||
final_adapter = rnd.adapter_path
|
||||
rounds_completed += 1
|
||||
|
||||
return IterativeDPOResult(
|
||||
rounds_completed=rounds_completed,
|
||||
final_adapter=final_adapter,
|
||||
per_round_pairs=tuple(per_round),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ from dataclasses import dataclass
|
|||
from typing import Optional
|
||||
|
||||
_MAX_ANCHOR_PATH_LEN = 4096
|
||||
_MAX_ANCHOR_ROW_BYTES = 1_000_000
|
||||
|
||||
|
||||
def _check_unit_float(value: object, field: str) -> float:
|
||||
|
|
@ -134,18 +135,223 @@ class MiniLLMConfig:
|
|||
)
|
||||
|
||||
|
||||
def build_minillm_callback(config):
|
||||
"""Build the MiniLLM HF Trainer callback. Deferred to v0.70.1.
|
||||
def minillm_distill_term(
|
||||
student_logits,
|
||||
teacher_logits,
|
||||
labels,
|
||||
*,
|
||||
config: "MiniLLMConfig",
|
||||
temperature: float = 1.0,
|
||||
):
|
||||
"""MiniLLM reverse-KL distillation term (v0.71.11 #237).
|
||||
|
||||
Validates the config type at the public boundary so misconfigured
|
||||
callers fail fast (mirrors v0.50.0 / v0.62.0 / v0.67.0 / v0.69.0
|
||||
deferred-live policy).
|
||||
Computes ``KL(student || target)`` where ``target`` is the
|
||||
teacher-mixed rollout distribution
|
||||
``ratio * teacher + (1 - ratio) * student_detached`` — the offline
|
||||
analog of MiniLLM's teacher-mixed sampling (Gu et al. 2024 §3.1). At
|
||||
``ratio = 1`` this reduces to standard reverse-KL distillation; at
|
||||
``ratio < 1`` the target stays closer to the student's current
|
||||
distribution (keeps the student near a known-good policy).
|
||||
|
||||
When ``config.length_normalize`` is set, the per-token reverse-KL is
|
||||
averaged over the *valid* (``labels != -100``) tokens per sequence so
|
||||
long completions don't dominate the gradient.
|
||||
|
||||
Differentiable w.r.t. the student logits. Returns a scalar.
|
||||
"""
|
||||
import torch
|
||||
|
||||
if not isinstance(config, MiniLLMConfig):
|
||||
raise TypeError(
|
||||
f"config must be MiniLLMConfig, got {type(config).__name__}"
|
||||
)
|
||||
if isinstance(temperature, bool) or not isinstance(temperature, (int, float)):
|
||||
raise TypeError("temperature must be a non-bool number")
|
||||
t = float(temperature)
|
||||
if not math.isfinite(t) or t <= 0.0:
|
||||
raise ValueError("temperature must be finite and positive")
|
||||
|
||||
s = student_logits / t
|
||||
teach = teacher_logits / t
|
||||
log_p_s = torch.log_softmax(s, dim=-1)
|
||||
p_s = log_p_s.exp()
|
||||
p_t = torch.softmax(teach, dim=-1)
|
||||
ratio = float(config.teacher_mix_ratio)
|
||||
target = ratio * p_t + (1.0 - ratio) * p_s.detach()
|
||||
log_target = target.clamp(min=1e-12).log()
|
||||
# reverse KL(student || target) per token.
|
||||
rkl = (p_s * (log_p_s - log_target)).sum(dim=-1) # [B, T]
|
||||
|
||||
if config.length_normalize and labels is not None:
|
||||
valid = (labels != -100).to(rkl.dtype)
|
||||
per_seq = (rkl * valid).sum(dim=-1) / valid.sum(dim=-1).clamp(min=1.0)
|
||||
loss = per_seq.mean()
|
||||
else:
|
||||
loss = rkl.mean()
|
||||
return loss * (t * t)
|
||||
|
||||
|
||||
def _get_trainer_callback_base():
|
||||
"""Lazy-resolve ``transformers.TrainerCallback`` (mirror v0.53.11)."""
|
||||
try:
|
||||
from transformers import TrainerCallback
|
||||
|
||||
return TrainerCallback
|
||||
except ImportError:
|
||||
return object
|
||||
|
||||
|
||||
_TrainerCallbackBase = _get_trainer_callback_base()
|
||||
|
||||
|
||||
class MiniLLMCallback(_TrainerCallbackBase): # type: ignore[misc, valid-type]
|
||||
"""Live MiniLLM helper + HF TrainerCallback (v0.71.11 #237).
|
||||
|
||||
Carries the :class:`MiniLLMConfig` and provides the loss terms the
|
||||
distill trainer applies:
|
||||
|
||||
- :meth:`distill_term` — teacher-mixed length-normalised reverse-KL.
|
||||
- :meth:`anchor_term` — pretrain-anchor SFT CE on a small batch read
|
||||
lazily from ``pretrain_anchor_path``, scaled by
|
||||
``pretrain_anchor_weight`` (prevents the student drifting away from
|
||||
coherent language).
|
||||
|
||||
Honest scope: the teacher-mix is the offline distribution-blend
|
||||
analog of MiniLLM's on-policy teacher-mixed *sampling*; a full
|
||||
autoregressive rollout loop (sample → teacher-score) is a larger
|
||||
follow-up documented as a known limitation.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: MiniLLMConfig,
|
||||
*,
|
||||
tokenizer: Optional[object] = None,
|
||||
temperature: float = 1.0,
|
||||
max_anchor_rows: int = 8,
|
||||
anchor_max_length: int = 128,
|
||||
) -> None:
|
||||
if not isinstance(config, MiniLLMConfig):
|
||||
raise TypeError(
|
||||
f"config must be MiniLLMConfig, got {type(config).__name__}"
|
||||
)
|
||||
self.config = config
|
||||
self.tokenizer = tokenizer
|
||||
self.temperature = float(temperature)
|
||||
self.max_anchor_rows = int(max_anchor_rows)
|
||||
self.anchor_max_length = int(anchor_max_length)
|
||||
self._anchor_inputs: Optional[dict] = None
|
||||
self._anchor_loaded = False
|
||||
|
||||
def distill_term(self, student_logits, teacher_logits, labels):
|
||||
"""Compute the teacher-mixed length-normalised reverse-KL term."""
|
||||
return minillm_distill_term(
|
||||
student_logits,
|
||||
teacher_logits,
|
||||
labels,
|
||||
config=self.config,
|
||||
temperature=self.temperature,
|
||||
)
|
||||
|
||||
def _load_anchor(self) -> Optional[dict]:
|
||||
"""Lazily tokenise a small pretrain-anchor batch (cwd-contained)."""
|
||||
if self._anchor_loaded:
|
||||
return self._anchor_inputs
|
||||
self._anchor_loaded = True
|
||||
path = self.config.pretrain_anchor_path
|
||||
if (
|
||||
path is None
|
||||
or self.config.pretrain_anchor_weight <= 0.0
|
||||
or self.tokenizer is None
|
||||
):
|
||||
return None
|
||||
from soup_cli.utils.paths import enforce_under_cwd_and_no_symlink
|
||||
|
||||
# cwd-containment + TOCTOU symlink rejection (security review fix —
|
||||
# mirrors v0.53.7 / v0.65 reader policy; was a bare is_under_cwd).
|
||||
enforce_under_cwd_and_no_symlink(path, "minillm_pretrain_anchor_path")
|
||||
import json
|
||||
|
||||
texts: list[str] = []
|
||||
with open(path, encoding="utf-8") as fh:
|
||||
for line in fh:
|
||||
# Per-line byte cap — defends against a single multi-MB line
|
||||
# blowing up tokenisation memory (matches v0.53.7 #106 caps).
|
||||
if len(line) > _MAX_ANCHOR_ROW_BYTES:
|
||||
continue
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
obj = json.loads(line)
|
||||
text = obj.get("text") or obj.get("content") or ""
|
||||
except (ValueError, AttributeError):
|
||||
text = line
|
||||
if text:
|
||||
texts.append(str(text))
|
||||
if len(texts) >= self.max_anchor_rows:
|
||||
break
|
||||
if not texts:
|
||||
return None
|
||||
enc = self.tokenizer(
|
||||
texts,
|
||||
return_tensors="pt",
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=self.anchor_max_length,
|
||||
)
|
||||
self._anchor_inputs = {
|
||||
"input_ids": enc["input_ids"],
|
||||
"attention_mask": enc["attention_mask"],
|
||||
}
|
||||
return self._anchor_inputs
|
||||
|
||||
def anchor_term(self, model):
|
||||
"""Pretrain-anchor SFT cross-entropy scaled by the anchor weight.
|
||||
|
||||
Returns ``None`` when the anchor is disabled / unavailable so the
|
||||
caller can skip the term cleanly.
|
||||
"""
|
||||
weight = float(self.config.pretrain_anchor_weight)
|
||||
if weight <= 0.0:
|
||||
return None
|
||||
anchor = self._load_anchor()
|
||||
if anchor is None:
|
||||
return None
|
||||
import torch
|
||||
|
||||
device = next(model.parameters()).device
|
||||
input_ids = anchor["input_ids"].to(device)
|
||||
attention_mask = anchor["attention_mask"].to(device)
|
||||
out = model(input_ids=input_ids, attention_mask=attention_mask)
|
||||
logits = out.logits
|
||||
shift_logits = logits[:, :-1, :].contiguous()
|
||||
shift_labels = input_ids[:, 1:].contiguous()
|
||||
# Mask padding so the anchor CE only counts real tokens.
|
||||
pad_mask = attention_mask[:, 1:].contiguous().bool()
|
||||
shift_labels = shift_labels.masked_fill(~pad_mask, -100)
|
||||
ce = torch.nn.functional.cross_entropy(
|
||||
shift_logits.view(-1, shift_logits.size(-1)),
|
||||
shift_labels.view(-1),
|
||||
ignore_index=-100,
|
||||
)
|
||||
return weight * ce
|
||||
|
||||
|
||||
def build_minillm_callback(
|
||||
config,
|
||||
*,
|
||||
tokenizer: Optional[object] = None,
|
||||
temperature: float = 1.0,
|
||||
) -> "MiniLLMCallback":
|
||||
"""Build the live MiniLLM callback / loss helper (v0.71.11 #237).
|
||||
|
||||
Lifts the v0.70.0 ``NotImplementedError`` stub. Validates the config
|
||||
type at the public boundary (fail-fast policy), then returns a
|
||||
:class:`MiniLLMCallback`.
|
||||
"""
|
||||
if not isinstance(config, MiniLLMConfig):
|
||||
raise TypeError(
|
||||
f"config must be MiniLLMConfig, got {type(config).__name__}"
|
||||
)
|
||||
raise NotImplementedError(
|
||||
"Live MiniLLM HF Trainer callback is deferred to v0.70.1. "
|
||||
"v0.70.0 ships the schema + validators only."
|
||||
)
|
||||
return MiniLLMCallback(config, tokenizer=tokenizer, temperature=temperature)
|
||||
|
|
|
|||
|
|
@ -196,6 +196,109 @@ def attach_grpo_stability_callback(trainer: Any, tcfg: Any) -> bool:
|
|||
return True
|
||||
|
||||
|
||||
def rl_callbacks_need_buffer(tcfg: Any) -> bool:
|
||||
"""True when a reward-fn capture buffer is needed (v0.71.11 #235/#240).
|
||||
|
||||
The reward-hack + echo-trap callbacks observe the GRPO step's rewards
|
||||
+ completions through the shared
|
||||
:class:`~soup_cli.utils.rl_signal_buffer.RLSignalBuffer`. The
|
||||
RL-checkpoint callback does not.
|
||||
"""
|
||||
return (
|
||||
getattr(tcfg, "reward_hack_detector", None) is not None
|
||||
or bool(getattr(tcfg, "echo_trap_enabled", False))
|
||||
)
|
||||
|
||||
|
||||
def attach_rl_callbacks(
|
||||
trainer: Any,
|
||||
tcfg: Any,
|
||||
*,
|
||||
buffer: Any = None,
|
||||
tokenizer: Any = None,
|
||||
output_dir: str = ".",
|
||||
task: str = "grpo",
|
||||
) -> int:
|
||||
"""Attach the v0.71.11 live RL callbacks; return how many were attached.
|
||||
|
||||
Wires (when their schema fields are set):
|
||||
- reward-hacking detector (#235) — reads ``buffer``.
|
||||
- echo-trap detector (#240) — reads ``buffer`` + ``tokenizer``.
|
||||
- mid-epoch RL checkpoint (#238) — saves under ``output_dir``.
|
||||
|
||||
The schema cross-validators already gate these fields to RL tasks on
|
||||
non-mlx backends, so this helper trusts the caller's config.
|
||||
"""
|
||||
attached = 0
|
||||
|
||||
detector = getattr(tcfg, "reward_hack_detector", None)
|
||||
if detector is not None:
|
||||
from soup_cli.utils.reward_hacking import build_reward_hack_callback
|
||||
|
||||
try:
|
||||
trainer.add_callback(
|
||||
build_reward_hack_callback(
|
||||
detector=detector,
|
||||
halt_on_hack=bool(getattr(tcfg, "reward_hack_halt", False)),
|
||||
buffer=buffer,
|
||||
)
|
||||
)
|
||||
attached += 1
|
||||
except (TypeError, ValueError) as exc:
|
||||
logger.debug("attach reward-hack callback rejected: %s", exc)
|
||||
|
||||
if bool(getattr(tcfg, "echo_trap_enabled", False)):
|
||||
from soup_cli.utils.echo_trap import build_echo_trap_callback
|
||||
|
||||
try:
|
||||
trainer.add_callback(
|
||||
build_echo_trap_callback(
|
||||
threshold=float(getattr(tcfg, "echo_trap_threshold", 0.6)),
|
||||
halt_on_trap=bool(getattr(tcfg, "echo_trap_halt", False)),
|
||||
tokenizer_aware=bool(
|
||||
getattr(tcfg, "echo_trap_tokenizer_aware", False)
|
||||
),
|
||||
buffer=buffer,
|
||||
tokenizer=tokenizer,
|
||||
)
|
||||
)
|
||||
attached += 1
|
||||
except (TypeError, ValueError) as exc:
|
||||
logger.debug("attach echo-trap callback rejected: %s", exc)
|
||||
|
||||
save_every = getattr(tcfg, "rl_checkpoint_save_every_steps", None)
|
||||
if save_every is not None:
|
||||
from soup_cli.utils.rl_checkpoint import (
|
||||
RLCheckpointConfig,
|
||||
build_rl_checkpoint_callback,
|
||||
)
|
||||
|
||||
try:
|
||||
ckpt_cfg = RLCheckpointConfig(
|
||||
save_every_steps=int(save_every),
|
||||
include_optimizer_state=bool(
|
||||
getattr(tcfg, "rl_checkpoint_include_optimizer", True)
|
||||
),
|
||||
include_ref_model=bool(
|
||||
getattr(tcfg, "rl_checkpoint_include_ref_model", False)
|
||||
),
|
||||
include_rollout_buffer=bool(
|
||||
getattr(tcfg, "rl_checkpoint_include_rollout_buffer", False)
|
||||
),
|
||||
keep_last=int(getattr(tcfg, "rl_checkpoint_keep_last", 3)),
|
||||
)
|
||||
trainer.add_callback(
|
||||
build_rl_checkpoint_callback(
|
||||
ckpt_cfg, output_dir=output_dir, task=task
|
||||
)
|
||||
)
|
||||
attached += 1
|
||||
except (TypeError, ValueError) as exc:
|
||||
logger.debug("attach RL-checkpoint callback rejected: %s", exc)
|
||||
|
||||
return attached
|
||||
|
||||
|
||||
def attach_plugin_callback(trainer: Any, console: Any = None) -> bool:
|
||||
"""Attach :class:`SoupPluginCallback` when any enabled plugin implements a hook.
|
||||
|
||||
|
|
|
|||
|
|
@ -35,7 +35,7 @@ from __future__ import annotations
|
|||
import math
|
||||
import types
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Sequence
|
||||
from typing import Any, Optional, Sequence
|
||||
|
||||
_MAX_DETECTOR_NAME_LEN = 32
|
||||
_MAX_RM_ENSEMBLE_SIZE = 32
|
||||
|
|
@ -316,34 +316,280 @@ class RewardHackReport:
|
|||
)
|
||||
|
||||
|
||||
def compute_separation_from_stats(reward_mean: object, reward_std: object) -> float:
|
||||
"""Cluster-separation proxy from logged reward mean + std.
|
||||
|
||||
Used by the live callback's ``on_log`` fallback path (e.g. PPO, or
|
||||
GRPO when the reward-fn capture buffer is unavailable). Returns the
|
||||
signal-to-noise ratio ``mean / sqrt(std^2 + eps)`` — high when rewards
|
||||
are well-separated relative to their noise, drops when rewards bunch
|
||||
up (the InfoRM reward-hacking signal).
|
||||
|
||||
Both inputs must be finite numbers (bool rejected). Non-finite or
|
||||
negative std raises ``ValueError``.
|
||||
"""
|
||||
if isinstance(reward_mean, bool) or isinstance(reward_std, bool):
|
||||
raise ValueError("reward_mean / reward_std must not be bool")
|
||||
if not isinstance(reward_mean, (int, float)) or not isinstance(
|
||||
reward_std, (int, float)
|
||||
):
|
||||
raise ValueError("reward_mean / reward_std must be numbers")
|
||||
m = float(reward_mean)
|
||||
s = float(reward_std)
|
||||
if not math.isfinite(m) or not math.isfinite(s):
|
||||
raise ValueError("reward_mean / reward_std must be finite")
|
||||
if s < 0.0:
|
||||
raise ValueError("reward_std must be non-negative")
|
||||
return m / math.sqrt(s * s + _EPS)
|
||||
|
||||
|
||||
def _health_from_signal(detector: str, raw_signal: float) -> float:
|
||||
"""Map a raw detector signal to a 'health' value (higher = healthier).
|
||||
|
||||
- ``info_rm``: separation IS health (higher separation = healthier).
|
||||
- ``rm_ensemble``: divergence rises when hacking, so health is
|
||||
``1 / (1 + divergence)`` ∈ (0, 1].
|
||||
"""
|
||||
if detector == "rm_ensemble":
|
||||
return 1.0 / (1.0 + max(0.0, raw_signal))
|
||||
return max(0.0, raw_signal)
|
||||
|
||||
|
||||
def _get_trainer_callback_base():
|
||||
"""Lazy-resolve ``transformers.TrainerCallback`` (mirror v0.53.11).
|
||||
|
||||
Resolved at module-import-of-class time so a torch-less environment can
|
||||
still import this utility module (falls back to ``object``).
|
||||
"""
|
||||
try:
|
||||
from transformers import TrainerCallback
|
||||
|
||||
return TrainerCallback
|
||||
except ImportError:
|
||||
return object
|
||||
|
||||
|
||||
_TrainerCallbackBase = _get_trainer_callback_base()
|
||||
|
||||
|
||||
class RewardHackCallback(_TrainerCallbackBase): # type: ignore[misc, valid-type]
|
||||
"""Live HF TrainerCallback for the reward-hacking detector (v0.71.11 #235).
|
||||
|
||||
Reads the per-completion rewards a GRPO step produced (via the shared
|
||||
:class:`~soup_cli.utils.rl_signal_buffer.RLSignalBuffer`) and tracks
|
||||
whether the reward signal is degrading:
|
||||
|
||||
- ``info_rm``: splits the rewards top-half (good) / bottom-half (bad)
|
||||
and computes the InfoRM cluster-separation. A drop from the
|
||||
training-start baseline = the reward model losing its grip.
|
||||
- ``rm_ensemble``: needs ≥2 reward functions; computes pairwise
|
||||
variance across them. Rising disagreement = unreliable signal.
|
||||
|
||||
When no buffer is wired (e.g. PPO), it falls back to ``on_log`` reading
|
||||
the ``reward`` / ``reward_std`` stats TRL logs.
|
||||
|
||||
Per-step health is surfaced to ``state.log_history`` so the v0.34.0
|
||||
anomaly explainer can flag it. ``halt_on_hack`` sets
|
||||
``control.should_training_stop`` on a HACK verdict.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
detector: str,
|
||||
halt_on_hack: bool = True,
|
||||
baseline_signal: Optional[float] = None,
|
||||
buffer: Any = None,
|
||||
) -> None:
|
||||
self.detector = validate_hack_detector(detector)
|
||||
if not isinstance(halt_on_hack, bool):
|
||||
raise TypeError(
|
||||
f"halt_on_hack must be bool, got {type(halt_on_hack).__name__}"
|
||||
)
|
||||
self.halt_on_hack = halt_on_hack
|
||||
if baseline_signal is not None:
|
||||
if isinstance(baseline_signal, bool):
|
||||
raise TypeError("baseline_signal must not be bool")
|
||||
if not isinstance(baseline_signal, (int, float)):
|
||||
raise TypeError("baseline_signal must be a number or None")
|
||||
if (
|
||||
not math.isfinite(float(baseline_signal))
|
||||
or float(baseline_signal) < 0.0
|
||||
):
|
||||
raise ValueError("baseline_signal must be finite and non-negative")
|
||||
baseline_signal = float(baseline_signal)
|
||||
self.buffer = buffer
|
||||
# Health baselines are recorded on the first observed signal.
|
||||
self._baseline_raw: Optional[float] = baseline_signal
|
||||
self._baseline_health: Optional[float] = (
|
||||
None
|
||||
if baseline_signal is None
|
||||
else _health_from_signal(self.detector, baseline_signal)
|
||||
)
|
||||
self._last_report: Optional[RewardHackReport] = None
|
||||
self._hacks_seen = 0
|
||||
|
||||
# --- pure signal computation (testable without transformers) ---
|
||||
|
||||
def compute_signal(self, snapshot: dict) -> Optional[float]:
|
||||
"""Compute the raw detector signal from a buffer snapshot.
|
||||
|
||||
Returns ``None`` when there isn't enough data (e.g. fewer than 4
|
||||
rewards for info_rm, or fewer than 2 reward functions for the
|
||||
ensemble detector).
|
||||
"""
|
||||
if self.detector == "rm_ensemble":
|
||||
per_func = snapshot.get("per_func", {}) or {}
|
||||
# Keep ONLY positions finite across EVERY RM so the per-prompt
|
||||
# variance compares the same prompt across RMs (code-review LOW
|
||||
# fix — per-column None-filtering would misalign prompts).
|
||||
cols = [list(v) for v in per_func.values() if v]
|
||||
if len(cols) < 2:
|
||||
return None
|
||||
n = min(len(c) for c in cols)
|
||||
aligned: list[list[float]] = [[] for _ in cols]
|
||||
for i in range(n):
|
||||
row = [c[i] for c in cols]
|
||||
if all(
|
||||
isinstance(v, (int, float))
|
||||
and not isinstance(v, bool)
|
||||
and math.isfinite(float(v))
|
||||
for v in row
|
||||
):
|
||||
for k, v in enumerate(row):
|
||||
aligned[k].append(float(v))
|
||||
if not aligned[0]:
|
||||
return None
|
||||
return compute_rm_ensemble_divergence(aligned)
|
||||
# info_rm: split into good / bad halves by sorted reward.
|
||||
rewards = snapshot.get("rewards", []) or []
|
||||
finite_rewards = [
|
||||
float(v)
|
||||
for v in rewards
|
||||
if isinstance(v, (int, float))
|
||||
and not isinstance(v, bool)
|
||||
and math.isfinite(float(v))
|
||||
]
|
||||
if len(finite_rewards) < 4:
|
||||
return None
|
||||
ordered = sorted(finite_rewards)
|
||||
half = len(ordered) // 2
|
||||
bad = ordered[:half]
|
||||
good = ordered[half:]
|
||||
if not bad or not good:
|
||||
return None
|
||||
return compute_cluster_separation(good, bad)
|
||||
|
||||
def observe_signal(self, raw_signal: float, step: int) -> RewardHackReport:
|
||||
"""Fold a raw signal into the running baseline + classify.
|
||||
|
||||
Records a baseline on the first positive health value, then
|
||||
computes the relative drop in health and classifies it. Returns
|
||||
the :class:`RewardHackReport`; the caller decides on halting.
|
||||
"""
|
||||
health = _health_from_signal(self.detector, raw_signal)
|
||||
if self._baseline_health is None and health > 0.0:
|
||||
self._baseline_health = health
|
||||
self._baseline_raw = raw_signal
|
||||
base_health = self._baseline_health
|
||||
if base_health is None or base_health <= 0.0:
|
||||
drop_pct = 0.0
|
||||
else:
|
||||
drop_pct = max(0.0, (base_health - health) / base_health)
|
||||
verdict = classify_hack_signal(drop_pct)
|
||||
if verdict == "HACK":
|
||||
self._hacks_seen += 1
|
||||
report = RewardHackReport(
|
||||
detector=self.detector,
|
||||
signal=max(0.0, float(raw_signal)),
|
||||
verdict=verdict,
|
||||
step=max(0, int(step)),
|
||||
baseline_signal=max(0.0, float(self._baseline_raw or raw_signal)),
|
||||
details=(
|
||||
f"detector={self.detector}",
|
||||
f"health={health:.4f} baseline={base_health}",
|
||||
f"drop_pct={drop_pct:.4f} verdict={verdict}",
|
||||
),
|
||||
)
|
||||
self._last_report = report
|
||||
return report
|
||||
|
||||
def last_report(self) -> Optional[RewardHackReport]:
|
||||
"""Return the most recent :class:`RewardHackReport` (or None)."""
|
||||
return self._last_report
|
||||
|
||||
# --- HF TrainerCallback surface ---
|
||||
|
||||
def _record(self, state, report: RewardHackReport, control):
|
||||
log_history = getattr(state, "log_history", None)
|
||||
if log_history is not None:
|
||||
log_history.append({
|
||||
"reward_hack_signal": report.signal,
|
||||
"reward_hack_verdict": report.verdict,
|
||||
})
|
||||
if report.verdict == "HACK" and self.halt_on_hack and control is not None:
|
||||
try:
|
||||
control.should_training_stop = True
|
||||
except Exception: # noqa: BLE001 — never crash training
|
||||
pass
|
||||
return control
|
||||
|
||||
def on_step_end(self, args, state, control, **kwargs):
|
||||
"""Per-step probe — read the reward-fn capture buffer if present."""
|
||||
if self.buffer is None:
|
||||
return control
|
||||
try:
|
||||
snapshot = self.buffer.snapshot()
|
||||
raw = self.compute_signal(snapshot)
|
||||
if raw is None:
|
||||
return control
|
||||
step = int(getattr(state, "global_step", 0) or 0)
|
||||
report = self.observe_signal(raw, step)
|
||||
return self._record(state, report, control)
|
||||
except Exception: # noqa: BLE001 — instrumentation must never crash
|
||||
return control
|
||||
|
||||
def on_log(self, args, state, control, logs=None, **kwargs):
|
||||
"""Fallback probe — when no buffer is wired, use logged reward stats.
|
||||
|
||||
info_rm only (rm_ensemble needs per-function rewards). Reads
|
||||
``reward`` / ``reward_std`` (TRL GRPO + TRL PPO both log these).
|
||||
"""
|
||||
if self.buffer is not None or self.detector != "info_rm":
|
||||
return control
|
||||
if not isinstance(logs, dict):
|
||||
return control
|
||||
mean = logs.get("reward")
|
||||
std = logs.get("reward_std", 0.0)
|
||||
if not isinstance(mean, (int, float)) or isinstance(mean, bool):
|
||||
return control
|
||||
if not isinstance(std, (int, float)) or isinstance(std, bool):
|
||||
std = 0.0
|
||||
try:
|
||||
raw = compute_separation_from_stats(mean, std)
|
||||
except ValueError:
|
||||
return control
|
||||
step = int(getattr(state, "global_step", 0) or 0)
|
||||
report = self.observe_signal(raw, step)
|
||||
return self._record(state, report, control)
|
||||
|
||||
|
||||
def build_reward_hack_callback(
|
||||
*,
|
||||
detector: str,
|
||||
halt_on_hack: bool = True,
|
||||
baseline_signal: Optional[float] = None,
|
||||
):
|
||||
"""Live HF Trainer callback factory — deferred to v0.70.1.
|
||||
buffer: Any = None,
|
||||
) -> "RewardHackCallback":
|
||||
"""Build the live reward-hacking HF Trainer callback (v0.71.11 #235).
|
||||
|
||||
Validates inputs at construction time so misconfigured runs fail
|
||||
fast even though the live callback is not yet wired. Raises
|
||||
``NotImplementedError`` with explicit v0.70.1 marker after
|
||||
validation (mirrors v0.50.0 ``apply_variant_loss`` / v0.61.0
|
||||
``apply_unlearn_loss`` policy).
|
||||
Lifts the v0.70.0 ``NotImplementedError`` stub. Validates every input
|
||||
at the public boundary (mirrors v0.50.0 / v0.61.0 fail-fast policy),
|
||||
then returns a :class:`RewardHackCallback`.
|
||||
"""
|
||||
# Validation order: name first (cheap, allowlist) then bool guard.
|
||||
validate_hack_detector(detector)
|
||||
if not isinstance(halt_on_hack, bool):
|
||||
raise TypeError(
|
||||
f"halt_on_hack must be bool, got {type(halt_on_hack).__name__}"
|
||||
)
|
||||
if baseline_signal is not None:
|
||||
if isinstance(baseline_signal, bool):
|
||||
raise TypeError("baseline_signal must not be bool")
|
||||
if not isinstance(baseline_signal, (int, float)):
|
||||
raise TypeError("baseline_signal must be a number or None")
|
||||
if not math.isfinite(float(baseline_signal)) or float(baseline_signal) < 0.0:
|
||||
raise ValueError("baseline_signal must be finite and non-negative")
|
||||
raise NotImplementedError(
|
||||
f"Live reward-hack callback for detector={detector!r} is deferred to "
|
||||
"v0.70.1. v0.70.0 ships the schema + math kernels only."
|
||||
return RewardHackCallback(
|
||||
detector=detector,
|
||||
halt_on_hack=halt_on_hack,
|
||||
baseline_signal=baseline_signal,
|
||||
buffer=buffer,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -27,8 +27,9 @@ Security:
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
from typing import Any, Optional
|
||||
|
||||
_MAX_SAVE_EVERY_STEPS = 10_000_000
|
||||
_MIN_KEEP_LAST = 1
|
||||
|
|
@ -192,18 +193,184 @@ class RLCheckpointState:
|
|||
}
|
||||
|
||||
|
||||
def build_rl_checkpoint_callback(config):
|
||||
"""Live HF Trainer callback for mid-epoch RL checkpoints.
|
||||
def _get_trainer_callback_base():
|
||||
"""Lazy-resolve ``transformers.TrainerCallback`` (mirror v0.53.11)."""
|
||||
try:
|
||||
from transformers import TrainerCallback
|
||||
|
||||
Deferred to v0.70.1. Validates the config type at the public
|
||||
boundary so misconfigured callers fail fast (mirrors v0.50.0 /
|
||||
v0.62.0 / v0.67.0 / v0.69.0 deferred-live policy).
|
||||
return TrainerCallback
|
||||
except ImportError:
|
||||
return object
|
||||
|
||||
|
||||
_TrainerCallbackBase = _get_trainer_callback_base()
|
||||
|
||||
|
||||
def _step_number(name: str) -> int:
|
||||
"""Extract the integer step from a ``step-<N>`` checkpoint dir name."""
|
||||
try:
|
||||
return int(name.split("-", 1)[1])
|
||||
except (IndexError, ValueError):
|
||||
return -1
|
||||
|
||||
|
||||
class RLCheckpointCallback(_TrainerCallbackBase): # type: ignore[misc, valid-type]
|
||||
"""Live HF TrainerCallback for mid-epoch RL checkpoints (v0.71.11 #238).
|
||||
|
||||
Saves an RL-aware checkpoint every ``save_every_steps`` steps under
|
||||
``<output_dir>/rl-checkpoints/step-<N>/``:
|
||||
|
||||
- the policy adapter (``model.save_pretrained``),
|
||||
- the optimizer state (when ``include_optimizer_state``), and
|
||||
- a ``manifest.json`` (:class:`RLCheckpointState`).
|
||||
|
||||
Older checkpoints beyond ``keep_last`` are pruned at write time. The
|
||||
optimizer is read from the ``optimizer`` kwarg HF Trainer passes to
|
||||
callbacks, so this works for any HF-Trainer-based RL loop (GRPO/PPO).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: RLCheckpointConfig,
|
||||
*,
|
||||
output_dir: str,
|
||||
task: str = "grpo",
|
||||
soup_version: Optional[str] = None,
|
||||
) -> None:
|
||||
if not isinstance(config, RLCheckpointConfig):
|
||||
raise TypeError(
|
||||
f"config must be RLCheckpointConfig, got {type(config).__name__}"
|
||||
)
|
||||
self.config = config
|
||||
self.output_dir = _validate_dir_shape(output_dir, "output_dir")
|
||||
# Containment: the run dir (and everything we write under it) must
|
||||
# stay under cwd (security review LOW). is_under_cwd is realpath-based
|
||||
# so it works before the dir exists; the per-file atomic_write_text on
|
||||
# the manifest adds the symlink-rejection at write time.
|
||||
from soup_cli.utils.paths import is_under_cwd
|
||||
|
||||
if not is_under_cwd(self.output_dir):
|
||||
raise ValueError("output_dir must stay under the current directory")
|
||||
if task not in _RL_TASKS:
|
||||
raise ValueError(f"task={task!r} must be one of {sorted(_RL_TASKS)}")
|
||||
self.task = task
|
||||
if soup_version is None:
|
||||
from soup_cli import __version__ as _v
|
||||
|
||||
soup_version = _v
|
||||
self.soup_version = soup_version
|
||||
self._saved: list[int] = []
|
||||
|
||||
def _ckpt_root(self) -> str:
|
||||
import os
|
||||
|
||||
return os.path.join(self.output_dir, "rl-checkpoints")
|
||||
|
||||
def save_checkpoint(self, *, step: int, model, optimizer) -> str:
|
||||
"""Write a checkpoint for ``step``; returns the directory path.
|
||||
|
||||
Pure of HF Trainer state — exercised directly by tests with a tiny
|
||||
peft model + a real torch optimizer.
|
||||
"""
|
||||
import os
|
||||
|
||||
from soup_cli.utils.paths import atomic_write_text
|
||||
|
||||
root = self._ckpt_root()
|
||||
ckpt_dir = os.path.join(root, f"step-{int(step)}")
|
||||
os.makedirs(ckpt_dir, exist_ok=True)
|
||||
|
||||
if model is not None and hasattr(model, "save_pretrained"):
|
||||
model.save_pretrained(ckpt_dir)
|
||||
|
||||
has_optimizer = False
|
||||
if self.config.include_optimizer_state and optimizer is not None:
|
||||
try:
|
||||
import torch
|
||||
|
||||
torch.save(
|
||||
optimizer.state_dict(),
|
||||
os.path.join(ckpt_dir, "optimizer.pt"),
|
||||
)
|
||||
has_optimizer = True
|
||||
except Exception: # noqa: BLE001 — best-effort, manifest reflects it
|
||||
has_optimizer = False
|
||||
|
||||
manifest = RLCheckpointState(
|
||||
step=int(step),
|
||||
checkpoint_dir=ckpt_dir,
|
||||
task=self.task,
|
||||
has_optimizer=has_optimizer,
|
||||
has_ref_model=bool(self.config.include_ref_model),
|
||||
has_rollout_buffer=bool(self.config.include_rollout_buffer),
|
||||
soup_version=self.soup_version,
|
||||
)
|
||||
atomic_write_text(
|
||||
json.dumps(manifest.to_dict(), indent=2),
|
||||
os.path.join(ckpt_dir, "manifest.json"),
|
||||
)
|
||||
self._saved.append(int(step))
|
||||
self._prune()
|
||||
return ckpt_dir
|
||||
|
||||
def _prune(self) -> None:
|
||||
"""Keep only the ``keep_last`` most-recent step checkpoints."""
|
||||
import os
|
||||
import shutil
|
||||
|
||||
root = self._ckpt_root()
|
||||
if not os.path.isdir(root):
|
||||
return
|
||||
entries = []
|
||||
for name in os.listdir(root):
|
||||
if not name.startswith("step-"):
|
||||
continue
|
||||
full = os.path.join(root, name)
|
||||
if os.path.islink(full) or not os.path.isdir(full):
|
||||
continue
|
||||
entries.append((_step_number(name), full))
|
||||
entries.sort(key=lambda t: t[0], reverse=True)
|
||||
for _, path in entries[self.config.keep_last:]:
|
||||
try:
|
||||
shutil.rmtree(path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def on_step_end(self, args, state, control, model=None, **kwargs):
|
||||
"""Per-step hook — save a checkpoint on the configured cadence."""
|
||||
try:
|
||||
step = int(getattr(state, "global_step", 0) or 0)
|
||||
if step <= 0 or step % self.config.save_every_steps != 0:
|
||||
return control
|
||||
optimizer = kwargs.get("optimizer")
|
||||
self.save_checkpoint(step=step, model=model, optimizer=optimizer)
|
||||
except Exception: # noqa: BLE001 — checkpoint failure must not crash run
|
||||
pass
|
||||
return control
|
||||
|
||||
|
||||
def build_rl_checkpoint_callback(
|
||||
config,
|
||||
*,
|
||||
output_dir: Optional[str] = None,
|
||||
task: str = "grpo",
|
||||
soup_version: Optional[str] = None,
|
||||
) -> "RLCheckpointCallback":
|
||||
"""Build the live mid-epoch RL checkpoint callback (v0.71.11 #238).
|
||||
|
||||
Lifts the v0.70.0 ``NotImplementedError`` stub. ``output_dir`` is the
|
||||
trainer's run directory under which ``rl-checkpoints/`` is created.
|
||||
Validates config type at the public boundary (fail-fast policy).
|
||||
"""
|
||||
if not isinstance(config, RLCheckpointConfig):
|
||||
raise TypeError(
|
||||
f"config must be RLCheckpointConfig, got {type(config).__name__}"
|
||||
)
|
||||
raise NotImplementedError(
|
||||
"Live mid-epoch RL checkpoint callback is deferred to v0.70.1. "
|
||||
"v0.70.0 ships the schema + state manifest only."
|
||||
if output_dir is None:
|
||||
raise ValueError("output_dir is required to build the RL checkpoint callback")
|
||||
return RLCheckpointCallback(
|
||||
config,
|
||||
output_dir=output_dir,
|
||||
task=task,
|
||||
soup_version=soup_version,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,221 @@
|
|||
"""Shared RL signal buffer — v0.71.11 (#235 / #240).
|
||||
|
||||
Reward-fn capture mechanism that lets the live reward-hacking detector
|
||||
(#235) and echo-trap detector (#240) callbacks observe the real
|
||||
per-completion rewards + the generated completions of a GRPO step
|
||||
without monkey-patching TRL internals.
|
||||
|
||||
How it works: when ``reward_hack_detector`` or ``echo_trap_enabled`` is
|
||||
set, the GRPO wrapper wraps every reward function with a small capturing
|
||||
shim (:func:`wrap_reward_funcs`) that records the completions + returned
|
||||
rewards into a single shared :class:`RLSignalBuffer` before forwarding
|
||||
the verbatim result. The callbacks then read a snapshot in
|
||||
``on_step_end``.
|
||||
|
||||
Design notes:
|
||||
- The reward-function signature ``reward_func(prompts, completions,
|
||||
**kwargs) -> list[float]`` is stable across TRL versions, so wrapping
|
||||
it is far more robust than hooking ``_generate_and_score_completions``.
|
||||
- The capture shim is exception-safe — a buffer error MUST NEVER break
|
||||
the reward computation (training would crash).
|
||||
- ``__name__`` is preserved so TRL's per-function logging keys stay
|
||||
correct (TRL logs ``rewards/<func_name>``).
|
||||
- No torch import at module top (pure Python; the buffer stores plain
|
||||
floats + strings).
|
||||
|
||||
Security:
|
||||
- Bounded buffers (``_MAX_COMPLETIONS`` / ``_MAX_COMPLETION_CHARS``) so a
|
||||
pathological run can't blow up RAM via the captured completions.
|
||||
- Reward values coerced to floats; non-finite / non-numeric dropped.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import threading
|
||||
from typing import Any, Optional, Sequence
|
||||
|
||||
_MAX_COMPLETIONS = 1024
|
||||
_MAX_COMPLETION_CHARS = 100_000
|
||||
|
||||
|
||||
def _completion_to_text(completion: Any) -> str:
|
||||
"""Best-effort extraction of the assistant text from a completion.
|
||||
|
||||
TRL completions come in two shapes:
|
||||
- plain string (non-conversational reward), OR
|
||||
- list of message dicts ``[{"role": "assistant", "content": "..."}]``
|
||||
(conversational). We concatenate the ``content`` of every dict.
|
||||
"""
|
||||
if isinstance(completion, str):
|
||||
text = completion
|
||||
elif isinstance(completion, dict):
|
||||
text = str(completion.get("content", ""))
|
||||
elif isinstance(completion, (list, tuple)):
|
||||
parts: list[str] = []
|
||||
for item in completion:
|
||||
if isinstance(item, dict):
|
||||
parts.append(str(item.get("content", "")))
|
||||
else:
|
||||
parts.append(str(item))
|
||||
text = "".join(parts)
|
||||
else:
|
||||
text = str(completion)
|
||||
if len(text) > _MAX_COMPLETION_CHARS:
|
||||
text = text[:_MAX_COMPLETION_CHARS]
|
||||
return text
|
||||
|
||||
|
||||
def _coerce_rewards(rewards: Any) -> list[Optional[float]]:
|
||||
"""Coerce a reward list to floats; non-numeric / non-finite → None."""
|
||||
out: list[Optional[float]] = []
|
||||
try:
|
||||
iterator = list(rewards)
|
||||
except TypeError:
|
||||
return out
|
||||
for r in iterator:
|
||||
if isinstance(r, bool) or not isinstance(r, (int, float)):
|
||||
out.append(None)
|
||||
continue
|
||||
fv = float(r)
|
||||
out.append(fv if math.isfinite(fv) else None)
|
||||
return out
|
||||
|
||||
|
||||
class RLSignalBuffer:
|
||||
"""Thread-safe ring of the most-recently observed GRPO step signal.
|
||||
|
||||
Stores the latest completions (as strings) and the per-reward-function
|
||||
reward lists. ``snapshot`` returns a consistent copy under the lock
|
||||
plus an element-wise aggregate reward (sum across functions).
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._lock = threading.Lock()
|
||||
self._completions: list[str] = []
|
||||
self._per_func: dict[str, list[Optional[float]]] = {}
|
||||
|
||||
def record(
|
||||
self,
|
||||
*,
|
||||
func_name: str,
|
||||
completions: Any,
|
||||
rewards: Any,
|
||||
) -> None:
|
||||
"""Record one reward-function call. Exception-safe by contract.
|
||||
|
||||
Completions are normalised to strings and capped at
|
||||
``_MAX_COMPLETIONS``. Within a single GRPO step every reward
|
||||
function sees the *same* completions, so overwriting on each call
|
||||
is correct.
|
||||
"""
|
||||
texts: list[str] = []
|
||||
if completions is not None:
|
||||
try:
|
||||
seq = list(completions)
|
||||
except TypeError:
|
||||
seq = []
|
||||
for c in seq[:_MAX_COMPLETIONS]:
|
||||
texts.append(_completion_to_text(c))
|
||||
coerced = _coerce_rewards(rewards)
|
||||
name = func_name if isinstance(func_name, str) and func_name else "reward"
|
||||
with self._lock:
|
||||
if texts:
|
||||
self._completions = texts
|
||||
self._per_func[name] = coerced
|
||||
|
||||
def snapshot(self) -> dict[str, Any]:
|
||||
"""Return a consistent copy of the latest step signal.
|
||||
|
||||
Keys:
|
||||
- ``completions``: list[str] (latest observed).
|
||||
- ``per_func``: dict name -> list[float] (None-filtered per entry).
|
||||
- ``rewards``: element-wise sum across functions (aggregate reward
|
||||
GRPO uses for advantages). None values are treated as 0 in the
|
||||
sum but tracked so a fully-None column drops out.
|
||||
"""
|
||||
with self._lock:
|
||||
completions = list(self._completions)
|
||||
per_func = {k: list(v) for k, v in self._per_func.items()}
|
||||
aggregate = _aggregate_rewards(per_func.values())
|
||||
return {
|
||||
"completions": completions,
|
||||
"per_func": per_func,
|
||||
"rewards": aggregate,
|
||||
}
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Reset the buffer (used by tests / between runs)."""
|
||||
with self._lock:
|
||||
self._completions = []
|
||||
self._per_func = {}
|
||||
|
||||
|
||||
def _aggregate_rewards(
|
||||
columns: "Sequence[list[Optional[float]]] | Any",
|
||||
) -> list[float]:
|
||||
"""Element-wise sum across per-function reward lists.
|
||||
|
||||
Aligns on the shortest list length; ``None`` entries contribute 0.
|
||||
A position that is ``None`` in every function is dropped (returns the
|
||||
truncated prefix up to the first all-None column-position only when no
|
||||
finite value appears — simpler: we keep finite floats, treating None
|
||||
as 0).
|
||||
"""
|
||||
col_lists = [c for c in columns if c]
|
||||
if not col_lists:
|
||||
return []
|
||||
n = min(len(c) for c in col_lists)
|
||||
out: list[float] = []
|
||||
for i in range(n):
|
||||
total = 0.0
|
||||
seen = False
|
||||
for c in col_lists:
|
||||
v = c[i]
|
||||
if v is not None:
|
||||
total += v
|
||||
seen = True
|
||||
if seen:
|
||||
out.append(total)
|
||||
else:
|
||||
out.append(0.0)
|
||||
return out
|
||||
|
||||
|
||||
def _make_capturing(fn: Any, buffer: RLSignalBuffer) -> Any:
|
||||
"""Wrap a single reward function with a capturing shim.
|
||||
|
||||
Forwards ``*args, **kwargs`` verbatim, records into ``buffer``, and
|
||||
returns the inner result unchanged. The capture is wrapped in a broad
|
||||
``except`` so a buffer bug can never break the reward computation.
|
||||
"""
|
||||
name = getattr(fn, "__name__", "reward")
|
||||
|
||||
def _wrapped(*args: Any, **kwargs: Any) -> Any:
|
||||
result = fn(*args, **kwargs)
|
||||
try:
|
||||
completions = kwargs.get("completions")
|
||||
if completions is None:
|
||||
# Positional: TRL calls reward_func(prompts, completions, ...)
|
||||
if len(args) >= 2:
|
||||
completions = args[1]
|
||||
elif len(args) == 1:
|
||||
completions = args[0]
|
||||
buffer.record(func_name=name, completions=completions, rewards=result)
|
||||
except Exception: # noqa: BLE001 — capture MUST NOT break the reward
|
||||
pass
|
||||
return result
|
||||
|
||||
_wrapped.__name__ = name
|
||||
return _wrapped
|
||||
|
||||
|
||||
def wrap_reward_funcs(reward_funcs: Any, buffer: RLSignalBuffer) -> Any:
|
||||
"""Wrap one reward function (or a list) for capture into ``buffer``.
|
||||
|
||||
Preserves the single-callable-vs-list shape so the caller can hand the
|
||||
result straight back to ``GRPOTrainer(reward_funcs=...)``.
|
||||
"""
|
||||
if isinstance(reward_funcs, (list, tuple)):
|
||||
return [_make_capturing(fn, buffer) for fn in reward_funcs]
|
||||
return _make_capturing(reward_funcs, buffer)
|
||||
|
|
@ -183,18 +183,120 @@ class ULDConfig:
|
|||
)
|
||||
|
||||
|
||||
def build_uld_projection(config):
|
||||
"""Build the projection module that bridges teacher / student vocabs.
|
||||
def _sorted_w1(p_s_sorted, p_t_sorted):
|
||||
"""1D Wasserstein-1 between two descending-sorted prob tensors.
|
||||
|
||||
Deferred to v0.70.1. Validates the config type at the public boundary
|
||||
so misconfigured callers fail fast (mirrors v0.50.0 ``apply_variant_loss``
|
||||
/ v0.62.0 ``apply_steering`` / v0.67.0 ``apply_bank_to_serve`` policy).
|
||||
Pads the shorter vocab to the longer with zeros, then returns the L1
|
||||
norm of the CDF difference (the closed-form W1 for distributions on a
|
||||
discrete line). Both inputs are ``[..., V*]`` (last dim = sorted
|
||||
probabilities summing to ~1). Returns ``[...]`` (per-position W1).
|
||||
"""
|
||||
import torch
|
||||
|
||||
v_s = p_s_sorted.shape[-1]
|
||||
v_t = p_t_sorted.shape[-1]
|
||||
common = max(v_s, v_t)
|
||||
if v_s < common:
|
||||
p_s_sorted = torch.nn.functional.pad(p_s_sorted, (0, common - v_s))
|
||||
if v_t < common:
|
||||
p_t_sorted = torch.nn.functional.pad(p_t_sorted, (0, common - v_t))
|
||||
cdf_s = torch.cumsum(p_s_sorted, dim=-1)
|
||||
cdf_t = torch.cumsum(p_t_sorted, dim=-1)
|
||||
return (cdf_s - cdf_t).abs().sum(dim=-1)
|
||||
|
||||
|
||||
def uld_distill_loss(
|
||||
student_logits,
|
||||
teacher_logits,
|
||||
*,
|
||||
config: ULDConfig,
|
||||
attention_mask=None,
|
||||
):
|
||||
"""Cross-tokenizer ULD distillation loss (v0.71.11 #236).
|
||||
|
||||
Computes a vocab-mismatch-tolerant distillation loss between
|
||||
``student_logits`` ``[B, T, Vs]`` and ``teacher_logits`` ``[B, T, Vt]``
|
||||
where ``Vs`` and ``Vt`` may differ. Two strategies:
|
||||
|
||||
- ``wasserstein``: softmax both, sort descending, pad to the common
|
||||
vocab length, take the 1D Wasserstein-1 (L1 of the CDF difference)
|
||||
between the sorted distributions. No alignment required — this is
|
||||
the Boizard et al. 2024 ULD surrogate.
|
||||
- ``topk_align``: take the top-``top_k`` probabilities of each model
|
||||
(rank-aligned, renormalised) and the same sorted-W1 between them.
|
||||
Distils only on the high-probability subset.
|
||||
|
||||
Differentiable w.r.t. the student logits (sort / topk both use
|
||||
gather). Returns a scalar mean over the (masked) token positions.
|
||||
"""
|
||||
import torch
|
||||
|
||||
if not isinstance(config, ULDConfig):
|
||||
raise TypeError(
|
||||
f"config must be ULDConfig, got {type(config).__name__}"
|
||||
)
|
||||
p_s = torch.softmax(student_logits, dim=-1)
|
||||
p_t = torch.softmax(teacher_logits, dim=-1)
|
||||
|
||||
if config.strategy == "topk_align":
|
||||
k = int(config.top_k)
|
||||
k_s = min(k, p_s.shape[-1])
|
||||
k_t = min(k, p_t.shape[-1])
|
||||
s_top, _ = torch.topk(p_s, k_s, dim=-1)
|
||||
t_top, _ = torch.topk(p_t, k_t, dim=-1)
|
||||
# Renormalise the truncated top-k so each sums to ~1.
|
||||
s_top = s_top / s_top.sum(dim=-1, keepdim=True).clamp(min=1e-12)
|
||||
t_top = t_top / t_top.sum(dim=-1, keepdim=True).clamp(min=1e-12)
|
||||
per_pos = _sorted_w1(s_top, t_top)
|
||||
else:
|
||||
# wasserstein — full sorted distributions.
|
||||
s_sorted, _ = torch.sort(p_s, dim=-1, descending=True)
|
||||
t_sorted, _ = torch.sort(p_t, dim=-1, descending=True)
|
||||
per_pos = _sorted_w1(s_sorted, t_sorted)
|
||||
|
||||
if attention_mask is not None:
|
||||
mask = attention_mask.to(per_pos.dtype)
|
||||
denom = mask.sum().clamp(min=1.0)
|
||||
return (per_pos * mask).sum() / denom
|
||||
return per_pos.mean()
|
||||
|
||||
|
||||
class ULDProjection:
|
||||
"""Callable wrapper around :func:`uld_distill_loss` (v0.71.11 #236).
|
||||
|
||||
Holds the frozen :class:`ULDConfig` and exposes ``__call__`` so the
|
||||
distill trainer can swap in the cross-tokenizer ULD loss without
|
||||
re-validating per step. There is no learned projection matrix — the
|
||||
ULD surrogate operates directly on sorted logit distributions, which
|
||||
is the whole point (no alignment matrix to fit).
|
||||
"""
|
||||
|
||||
def __init__(self, config: ULDConfig) -> None:
|
||||
if not isinstance(config, ULDConfig):
|
||||
raise TypeError(
|
||||
f"config must be ULDConfig, got {type(config).__name__}"
|
||||
)
|
||||
self.config = config
|
||||
|
||||
def __call__(self, student_logits, teacher_logits, *, attention_mask=None):
|
||||
return uld_distill_loss(
|
||||
student_logits,
|
||||
teacher_logits,
|
||||
config=self.config,
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
|
||||
|
||||
def build_uld_projection(config) -> "ULDProjection":
|
||||
"""Build the live cross-tokenizer ULD projection (v0.71.11 #236).
|
||||
|
||||
Lifts the v0.70.0 ``NotImplementedError`` stub. Validates the config
|
||||
type at the public boundary (fail-fast policy), then returns a
|
||||
:class:`ULDProjection` callable that computes the ULD distillation
|
||||
loss for a (student_logits, teacher_logits) pair.
|
||||
"""
|
||||
if not isinstance(config, ULDConfig):
|
||||
raise TypeError(
|
||||
f"config must be ULDConfig, got {type(config).__name__}"
|
||||
)
|
||||
raise NotImplementedError(
|
||||
f"Live ULD projection for strategy={config.strategy!r} is deferred "
|
||||
"to v0.70.1. v0.70.0 ships the schema + validators only."
|
||||
)
|
||||
return ULDProjection(config)
|
||||
|
|
|
|||
|
|
@ -665,21 +665,18 @@ class TestTrueWeightedCombine:
|
|||
|
||||
class TestStabilityCallbackEMA:
|
||||
def test_on_step_end_runs_ema_update(self):
|
||||
# v0.71.11 #160 — on_step_end now mutates the ref parameters in
|
||||
# place (no full state_dict / load_state_dict round-trip).
|
||||
import torch.nn as nn
|
||||
|
||||
from soup_cli.monitoring.grpo_stability_callback import GRPOStabilityCallback
|
||||
|
||||
class _Model:
|
||||
def __init__(self, val):
|
||||
self._val = val
|
||||
|
||||
def state_dict(self):
|
||||
return {"w": torch.full((2,), self._val)}
|
||||
|
||||
def load_state_dict(self, sd, strict=True):
|
||||
self._loaded = sd
|
||||
|
||||
cb = GRPOStabilityCallback(ref_model_ema_alpha=0.5)
|
||||
ref = _Model(0.0)
|
||||
pol = _Model(1.0)
|
||||
ref = nn.Linear(2, 2)
|
||||
pol = nn.Linear(2, 2)
|
||||
with torch.no_grad():
|
||||
ref.weight.fill_(0.0)
|
||||
pol.weight.fill_(1.0)
|
||||
cb._policy_model = pol
|
||||
cb._ref_model = ref
|
||||
|
||||
|
|
@ -688,8 +685,8 @@ class TestStabilityCallbackEMA:
|
|||
|
||||
state = _State()
|
||||
cb.on_step_end(args=None, state=state, control=None, model=pol)
|
||||
# Loaded back into ref: midpoint 0.5.
|
||||
assert torch.allclose(ref._loaded["w"], torch.full((2,), 0.5))
|
||||
# ref mutated in place to the midpoint 0.5.
|
||||
assert torch.allclose(ref.weight, torch.full_like(ref.weight, 0.5))
|
||||
|
||||
def test_on_step_end_logs_counters(self):
|
||||
from soup_cli.monitoring.grpo_stability_callback import GRPOStabilityCallback
|
||||
|
|
|
|||
|
|
@ -406,25 +406,25 @@ class TestRewardHackReport:
|
|||
|
||||
|
||||
class TestBuildRewardHackCallbackStub:
|
||||
"""Live trainer-callback wiring deferred to v0.70.1.
|
||||
"""Live in v0.71.11 #235 — the factory now returns a real callback.
|
||||
|
||||
The factory validates inputs at construction time and raises
|
||||
NotImplementedError with explicit v0.70.1 marker (mirrors v0.50.0
|
||||
apply_variant_loss policy).
|
||||
The factory still validates inputs at construction time (fail-fast).
|
||||
"""
|
||||
|
||||
def test_invalid_detector_rejected_before_deferred(self):
|
||||
def test_invalid_detector_rejected_before_build(self):
|
||||
from soup_cli.utils.reward_hacking import build_reward_hack_callback
|
||||
|
||||
# Validation runs BEFORE NotImplementedError.
|
||||
with pytest.raises(ValueError, match="not supported"):
|
||||
build_reward_hack_callback(detector="evil")
|
||||
|
||||
def test_deferred_v0701(self):
|
||||
from soup_cli.utils.reward_hacking import build_reward_hack_callback
|
||||
def test_live_returns_callback(self):
|
||||
from soup_cli.utils.reward_hacking import (
|
||||
RewardHackCallback,
|
||||
build_reward_hack_callback,
|
||||
)
|
||||
|
||||
with pytest.raises(NotImplementedError, match="v0.70.1"):
|
||||
build_reward_hack_callback(detector="info_rm")
|
||||
cb = build_reward_hack_callback(detector="info_rm")
|
||||
assert isinstance(cb, RewardHackCallback)
|
||||
|
||||
def test_bool_halt_on_hack_rejected(self):
|
||||
from soup_cli.utils.reward_hacking import build_reward_hack_callback
|
||||
|
|
|
|||
|
|
@ -235,7 +235,7 @@ class TestULDConfig:
|
|||
|
||||
|
||||
class TestBuildULDProjection:
|
||||
"""Deferred to v0.70.1 — validates config type then raises."""
|
||||
"""Live in v0.71.11 #236 — returns a ULDProjection; validates type."""
|
||||
|
||||
def test_non_config_rejected(self):
|
||||
from soup_cli.utils.uld import build_uld_projection
|
||||
|
|
@ -243,16 +243,15 @@ class TestBuildULDProjection:
|
|||
with pytest.raises(TypeError, match="ULDConfig"):
|
||||
build_uld_projection({"strategy": "wasserstein"}) # type: ignore[arg-type]
|
||||
|
||||
def test_deferred_v0701(self):
|
||||
from soup_cli.utils.uld import ULDConfig, build_uld_projection
|
||||
def test_live_returns_projection(self):
|
||||
from soup_cli.utils.uld import ULDConfig, ULDProjection, build_uld_projection
|
||||
|
||||
cfg = ULDConfig(
|
||||
strategy="wasserstein",
|
||||
student_vocab_size=32000,
|
||||
teacher_vocab_size=128256,
|
||||
)
|
||||
with pytest.raises(NotImplementedError, match="v0.70.1"):
|
||||
build_uld_projection(cfg)
|
||||
assert isinstance(build_uld_projection(cfg), ULDProjection)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -211,7 +211,7 @@ class TestMiniLLMConfig:
|
|||
|
||||
|
||||
class TestBuildMiniLLMCallback:
|
||||
"""Live trainer callback deferred to v0.70.1."""
|
||||
"""Live in v0.71.11 #237 — returns a MiniLLMCallback; validates type."""
|
||||
|
||||
def test_non_config_rejected(self):
|
||||
from soup_cli.utils.minillm import build_minillm_callback
|
||||
|
|
@ -219,12 +219,14 @@ class TestBuildMiniLLMCallback:
|
|||
with pytest.raises(TypeError, match="MiniLLMConfig"):
|
||||
build_minillm_callback({}) # type: ignore[arg-type]
|
||||
|
||||
def test_deferred(self):
|
||||
from soup_cli.utils.minillm import MiniLLMConfig, build_minillm_callback
|
||||
def test_live_returns_callback(self):
|
||||
from soup_cli.utils.minillm import (
|
||||
MiniLLMCallback,
|
||||
MiniLLMConfig,
|
||||
build_minillm_callback,
|
||||
)
|
||||
|
||||
cfg = MiniLLMConfig()
|
||||
with pytest.raises(NotImplementedError, match="v0.70.1"):
|
||||
build_minillm_callback(cfg)
|
||||
assert isinstance(build_minillm_callback(MiniLLMConfig()), MiniLLMCallback)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -245,7 +245,7 @@ class TestRLCheckpointState:
|
|||
|
||||
|
||||
class TestBuildRLCheckpointCallback:
|
||||
"""Live callback deferred to v0.70.1."""
|
||||
"""Live in v0.71.11 #238 — returns a callback; requires output_dir."""
|
||||
|
||||
def test_non_config_rejected(self):
|
||||
from soup_cli.utils.rl_checkpoint import build_rl_checkpoint_callback
|
||||
|
|
@ -253,15 +253,25 @@ class TestBuildRLCheckpointCallback:
|
|||
with pytest.raises(TypeError, match="RLCheckpointConfig"):
|
||||
build_rl_checkpoint_callback({"save_every_steps": 100}) # type: ignore[arg-type]
|
||||
|
||||
def test_deferred(self):
|
||||
def test_live_returns_callback(self):
|
||||
from soup_cli.utils.rl_checkpoint import (
|
||||
RLCheckpointCallback,
|
||||
RLCheckpointConfig,
|
||||
build_rl_checkpoint_callback,
|
||||
)
|
||||
|
||||
cfg = RLCheckpointConfig(save_every_steps=100)
|
||||
with pytest.raises(NotImplementedError, match="v0.70.1"):
|
||||
build_rl_checkpoint_callback(cfg)
|
||||
cb = build_rl_checkpoint_callback(cfg, output_dir="run", task="grpo")
|
||||
assert isinstance(cb, RLCheckpointCallback)
|
||||
|
||||
def test_requires_output_dir(self):
|
||||
from soup_cli.utils.rl_checkpoint import (
|
||||
RLCheckpointConfig,
|
||||
build_rl_checkpoint_callback,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="output_dir"):
|
||||
build_rl_checkpoint_callback(RLCheckpointConfig(save_every_steps=100))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -350,34 +350,50 @@ class TestBuildIterativeDPOPlan:
|
|||
|
||||
|
||||
class TestRunIterativeDPODeferred:
|
||||
"""Live in v0.71.11 #239 — runs the sample → score → pair → train loop."""
|
||||
|
||||
def test_non_plan_rejected(self):
|
||||
from soup_cli.utils.iterative_dpo import run_iterative_dpo
|
||||
|
||||
with pytest.raises(TypeError, match="IterativeDPOPlan"):
|
||||
run_iterative_dpo({"rounds": 1}) # type: ignore[arg-type]
|
||||
|
||||
def test_deferred(self):
|
||||
def test_live_runs_with_fakes(self, tmp_path, monkeypatch):
|
||||
import json
|
||||
|
||||
from soup_cli.utils.iterative_dpo import (
|
||||
IterativeDPOPlan,
|
||||
IterativeDPORound,
|
||||
IterativeDPOResult,
|
||||
build_iterative_dpo_plan,
|
||||
run_iterative_dpo,
|
||||
)
|
||||
|
||||
plan = IterativeDPOPlan(
|
||||
monkeypatch.chdir(tmp_path)
|
||||
(tmp_path / "p.jsonl").write_text(json.dumps({"prompt": "q"}), encoding="utf-8")
|
||||
plan = build_iterative_dpo_plan(
|
||||
base_model="m",
|
||||
reward_model="./rm",
|
||||
rounds=(
|
||||
IterativeDPORound(
|
||||
round_index=0,
|
||||
prompts_path="./p.jsonl",
|
||||
pairs_path="./r0.jsonl",
|
||||
adapter_path="./out/r0",
|
||||
pairs_count=10,
|
||||
),
|
||||
),
|
||||
reward_model="rm",
|
||||
prompts_path="p.jsonl",
|
||||
output_dir="out",
|
||||
rounds=1,
|
||||
pairs_per_round=10,
|
||||
)
|
||||
with pytest.raises(NotImplementedError, match="v0.70.1"):
|
||||
run_iterative_dpo(plan)
|
||||
|
||||
def fake_sample(**kwargs):
|
||||
return [["aa", "bbb"]]
|
||||
|
||||
def fake_score(**kwargs):
|
||||
return [1.0, 2.0]
|
||||
|
||||
def fake_train(**kwargs):
|
||||
import os
|
||||
|
||||
os.makedirs(kwargs["adapter_path"], exist_ok=True)
|
||||
|
||||
result = run_iterative_dpo(
|
||||
plan, sample_fn=fake_sample, score_fn=fake_score, train_fn=fake_train
|
||||
)
|
||||
assert isinstance(result, IterativeDPOResult)
|
||||
assert result.rounds_completed == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -446,11 +462,13 @@ class TestIterativeDPOCli:
|
|||
)
|
||||
assert result.exit_code == 2
|
||||
|
||||
def test_live_deferred_exits_3(self, tmp_path, monkeypatch):
|
||||
"""Without --plan-only, the deferred live runner exits 3."""
|
||||
def test_live_runner_bad_model_exits_1(self, tmp_path, monkeypatch):
|
||||
"""Without --plan-only, the live runner runs; a bad model exits 1."""
|
||||
from soup_cli.commands.iterative_dpo import app
|
||||
|
||||
monkeypatch.chdir(tmp_path)
|
||||
# Empty prompt rows → no prompts → default sample_fn tries to load
|
||||
# the (non-existent) model "m" → run fails → CLI exits 1 (NOT 3).
|
||||
(tmp_path / "prompts.jsonl").write_text("{}\n", encoding="utf-8")
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(
|
||||
|
|
@ -459,18 +477,18 @@ class TestIterativeDPOCli:
|
|||
"--base-model",
|
||||
"m",
|
||||
"--reward-model",
|
||||
"./rm",
|
||||
"rm",
|
||||
"--prompts",
|
||||
"./prompts.jsonl",
|
||||
"prompts.jsonl",
|
||||
"--output-dir",
|
||||
"./out",
|
||||
"out",
|
||||
"--rounds",
|
||||
"2",
|
||||
"--pairs-per-round",
|
||||
"100",
|
||||
],
|
||||
)
|
||||
assert result.exit_code == 3, (result.output, repr(result.exception))
|
||||
assert result.exit_code == 1, (result.output, repr(result.exception))
|
||||
|
||||
|
||||
class TestSourceWiring:
|
||||
|
|
|
|||
|
|
@ -332,11 +332,13 @@ class TestBuildEchoTrapCallbackDeferred:
|
|||
with pytest.raises(ValueError, match="bool"):
|
||||
build_echo_trap_callback(threshold=True)
|
||||
|
||||
def test_deferred(self):
|
||||
from soup_cli.utils.echo_trap import build_echo_trap_callback
|
||||
def test_live_returns_callback(self):
|
||||
from soup_cli.utils.echo_trap import (
|
||||
EchoTrapCallback,
|
||||
build_echo_trap_callback,
|
||||
)
|
||||
|
||||
with pytest.raises(NotImplementedError, match="v0.70.1"):
|
||||
build_echo_trap_callback(threshold=0.5)
|
||||
assert isinstance(build_echo_trap_callback(threshold=0.5), EchoTrapCallback)
|
||||
|
||||
def test_halt_must_be_bool(self):
|
||||
from soup_cli.utils.echo_trap import build_echo_trap_callback
|
||||
|
|
|
|||
|
|
@ -0,0 +1,862 @@
|
|||
"""v0.71.11 — GRPO / RL callbacks live wiring.
|
||||
|
||||
Closes #235 (reward-hack), #236 (ULD), #237 (MiniLLM), #238 (RL checkpoint),
|
||||
#239 (iterative-DPO), #240 (echo-trap), #159 (variant fallback warning),
|
||||
#160 (in-place GRPO EMA).
|
||||
|
||||
These tests lift the v0.70.0 deferred-stub family to live behaviour and
|
||||
exercise the math + wiring on CPU / tiny fakes (no GPU). The Step-6 smoke
|
||||
runs a real SmolLM2-135M GRPO loop separately.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
import soup_cli
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Shared helpers
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _FakeState:
|
||||
def __init__(self, global_step: int = 1):
|
||||
self.global_step = global_step
|
||||
self.log_history: list[dict] = []
|
||||
|
||||
|
||||
class _FakeControl:
|
||||
def __init__(self):
|
||||
self.should_training_stop = False
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Shared RL signal buffer
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRLSignalBuffer:
|
||||
def test_record_and_snapshot(self):
|
||||
from soup_cli.utils.rl_signal_buffer import RLSignalBuffer
|
||||
|
||||
buf = RLSignalBuffer()
|
||||
buf.record(
|
||||
func_name="reward", completions=["a", "b", "c", "d"], rewards=[1, 2, 3, 4]
|
||||
)
|
||||
snap = buf.snapshot()
|
||||
assert snap["completions"] == ["a", "b", "c", "d"]
|
||||
assert snap["rewards"] == [1.0, 2.0, 3.0, 4.0]
|
||||
assert "reward" in snap["per_func"]
|
||||
|
||||
def test_aggregate_sums_across_funcs(self):
|
||||
from soup_cli.utils.rl_signal_buffer import RLSignalBuffer
|
||||
|
||||
buf = RLSignalBuffer()
|
||||
buf.record(func_name="a", completions=["x", "y"], rewards=[1, 2])
|
||||
buf.record(func_name="b", completions=["x", "y"], rewards=[10, 20])
|
||||
snap = buf.snapshot()
|
||||
assert snap["rewards"] == [11.0, 22.0]
|
||||
assert set(snap["per_func"]) == {"a", "b"}
|
||||
|
||||
def test_non_finite_reward_dropped(self):
|
||||
from soup_cli.utils.rl_signal_buffer import RLSignalBuffer
|
||||
|
||||
buf = RLSignalBuffer()
|
||||
buf.record(func_name="r", completions=["a"], rewards=[float("nan")])
|
||||
snap = buf.snapshot()
|
||||
# NaN coerced to None → aggregate position is 0.0 (no finite value).
|
||||
assert snap["per_func"]["r"] == [None]
|
||||
|
||||
def test_conversational_completion_extracted(self):
|
||||
from soup_cli.utils.rl_signal_buffer import RLSignalBuffer
|
||||
|
||||
buf = RLSignalBuffer()
|
||||
buf.record(
|
||||
func_name="r",
|
||||
completions=[[{"role": "assistant", "content": "hello world"}]],
|
||||
rewards=[1.0],
|
||||
)
|
||||
snap = buf.snapshot()
|
||||
assert snap["completions"] == ["hello world"]
|
||||
|
||||
def test_wrap_preserves_name_and_captures(self):
|
||||
from soup_cli.utils.rl_signal_buffer import RLSignalBuffer, wrap_reward_funcs
|
||||
|
||||
buf = RLSignalBuffer()
|
||||
|
||||
def my_reward(prompts=None, completions=None, **kwargs):
|
||||
return [float(len(c)) for c in completions]
|
||||
|
||||
wrapped = wrap_reward_funcs(my_reward, buf)
|
||||
assert wrapped.__name__ == "my_reward"
|
||||
result = wrapped(prompts=["p"], completions=["aa", "bbb"])
|
||||
assert result == [2.0, 3.0]
|
||||
snap = buf.snapshot()
|
||||
assert snap["completions"] == ["aa", "bbb"]
|
||||
assert snap["rewards"] == [2.0, 3.0]
|
||||
|
||||
def test_wrap_list_shape_preserved(self):
|
||||
from soup_cli.utils.rl_signal_buffer import RLSignalBuffer, wrap_reward_funcs
|
||||
|
||||
buf = RLSignalBuffer()
|
||||
fns = [lambda completions=None, **k: [1.0]]
|
||||
wrapped = wrap_reward_funcs(fns, buf)
|
||||
assert isinstance(wrapped, list)
|
||||
assert len(wrapped) == 1
|
||||
|
||||
def test_capture_never_breaks_reward(self):
|
||||
from soup_cli.utils.rl_signal_buffer import RLSignalBuffer, wrap_reward_funcs
|
||||
|
||||
buf = RLSignalBuffer()
|
||||
|
||||
def reward(prompts=None, completions=None, **kwargs):
|
||||
return [1.0, 2.0]
|
||||
|
||||
wrapped = wrap_reward_funcs(reward, buf)
|
||||
# Bizarre completions that the normaliser can't handle must not raise.
|
||||
out = wrapped(prompts=None, completions=object())
|
||||
assert out == [1.0, 2.0]
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# #235 — reward-hack callback
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRewardHackCallback:
|
||||
def test_build_returns_callback_not_notimplemented(self):
|
||||
from soup_cli.utils.reward_hacking import (
|
||||
RewardHackCallback,
|
||||
build_reward_hack_callback,
|
||||
)
|
||||
|
||||
cb = build_reward_hack_callback(detector="info_rm")
|
||||
assert isinstance(cb, RewardHackCallback)
|
||||
|
||||
def test_build_rejects_unknown_detector(self):
|
||||
from soup_cli.utils.reward_hacking import build_reward_hack_callback
|
||||
|
||||
with pytest.raises(ValueError, match="not supported"):
|
||||
build_reward_hack_callback(detector="evil")
|
||||
|
||||
def test_build_rejects_non_bool_halt(self):
|
||||
from soup_cli.utils.reward_hacking import build_reward_hack_callback
|
||||
|
||||
with pytest.raises(TypeError, match="halt_on_hack"):
|
||||
build_reward_hack_callback(detector="info_rm", halt_on_hack="yes")
|
||||
|
||||
def test_info_rm_compute_signal(self):
|
||||
from soup_cli.utils.reward_hacking import build_reward_hack_callback
|
||||
|
||||
cb = build_reward_hack_callback(detector="info_rm")
|
||||
snap = {"rewards": [0.0, 0.0, 5.0, 5.0], "per_func": {}}
|
||||
sig = cb.compute_signal(snap)
|
||||
assert sig is not None and sig > 0.0
|
||||
|
||||
def test_info_rm_insufficient_data_returns_none(self):
|
||||
from soup_cli.utils.reward_hacking import build_reward_hack_callback
|
||||
|
||||
cb = build_reward_hack_callback(detector="info_rm")
|
||||
assert cb.compute_signal({"rewards": [1.0, 2.0], "per_func": {}}) is None
|
||||
|
||||
def test_observe_baseline_then_drop_to_hack(self):
|
||||
from soup_cli.utils.reward_hacking import build_reward_hack_callback
|
||||
|
||||
cb = build_reward_hack_callback(detector="info_rm")
|
||||
r0 = cb.observe_signal(10.0, step=1) # baseline separation 10
|
||||
assert r0.verdict == "OK"
|
||||
r1 = cb.observe_signal(2.0, step=2) # dropped 80% → HACK
|
||||
assert r1.verdict == "HACK"
|
||||
assert r1.signal == 2.0
|
||||
|
||||
def test_on_step_end_halts_on_hack(self):
|
||||
from soup_cli.utils.reward_hacking import build_reward_hack_callback
|
||||
from soup_cli.utils.rl_signal_buffer import RLSignalBuffer
|
||||
|
||||
buf = RLSignalBuffer()
|
||||
cb = build_reward_hack_callback(
|
||||
detector="info_rm", halt_on_hack=True, buffer=buf
|
||||
)
|
||||
state, control = _FakeState(1), _FakeControl()
|
||||
# Step 1 — high separation = baseline.
|
||||
buf.record(func_name="r", completions=["a"] * 4, rewards=[0, 0, 9, 9])
|
||||
cb.on_step_end(None, state, control)
|
||||
# Step 2 — bunched rewards = HACK.
|
||||
buf.record(func_name="r", completions=["a"] * 4, rewards=[5, 5, 5, 5])
|
||||
state.global_step = 2
|
||||
cb.on_step_end(None, state, control)
|
||||
assert control.should_training_stop is True
|
||||
assert any("reward_hack_verdict" in e for e in state.log_history)
|
||||
|
||||
def test_rm_ensemble_needs_two_funcs(self):
|
||||
from soup_cli.utils.reward_hacking import build_reward_hack_callback
|
||||
|
||||
cb = build_reward_hack_callback(detector="rm_ensemble")
|
||||
# One func → None.
|
||||
assert cb.compute_signal({"rewards": [], "per_func": {"a": [1.0, 2.0]}}) is None
|
||||
# Two funcs → divergence.
|
||||
sig = cb.compute_signal(
|
||||
{"rewards": [], "per_func": {"a": [1.0, 2.0], "b": [3.0, 0.0]}}
|
||||
)
|
||||
assert sig is not None and sig >= 0.0
|
||||
|
||||
def test_on_log_fallback_without_buffer(self):
|
||||
from soup_cli.utils.reward_hacking import build_reward_hack_callback
|
||||
|
||||
cb = build_reward_hack_callback(detector="info_rm", buffer=None)
|
||||
state, control = _FakeState(1), _FakeControl()
|
||||
cb.on_log(None, state, control, logs={"reward": 5.0, "reward_std": 0.1})
|
||||
assert cb.last_report() is not None
|
||||
|
||||
def test_compute_separation_from_stats(self):
|
||||
from soup_cli.utils.reward_hacking import compute_separation_from_stats
|
||||
|
||||
high = compute_separation_from_stats(5.0, 0.1)
|
||||
low = compute_separation_from_stats(5.0, 10.0)
|
||||
assert high > low
|
||||
|
||||
def test_separation_stats_rejects_bool(self):
|
||||
from soup_cli.utils.reward_hacking import compute_separation_from_stats
|
||||
|
||||
with pytest.raises(ValueError, match="bool"):
|
||||
compute_separation_from_stats(True, 1.0)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# #240 — echo-trap callback
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _FakeTokenizer:
|
||||
def encode(self, text, add_special_tokens=False):
|
||||
# deterministic id per whitespace token
|
||||
return [abs(hash(t)) % 1000 for t in text.split()]
|
||||
|
||||
|
||||
class TestEchoTrapCallback:
|
||||
def test_build_returns_callback_not_notimplemented(self):
|
||||
from soup_cli.utils.echo_trap import EchoTrapCallback, build_echo_trap_callback
|
||||
|
||||
cb = build_echo_trap_callback(threshold=0.5)
|
||||
assert isinstance(cb, EchoTrapCallback)
|
||||
|
||||
def test_build_rejects_bad_threshold(self):
|
||||
from soup_cli.utils.echo_trap import build_echo_trap_callback
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
build_echo_trap_callback(threshold=2.0)
|
||||
|
||||
def test_build_rejects_bool_threshold(self):
|
||||
from soup_cli.utils.echo_trap import build_echo_trap_callback
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
build_echo_trap_callback(threshold=True)
|
||||
|
||||
def test_build_rejects_non_bool_halt(self):
|
||||
from soup_cli.utils.echo_trap import build_echo_trap_callback
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
build_echo_trap_callback(threshold=0.5, halt_on_trap="yes")
|
||||
|
||||
def test_compute_signal_repetitive_high(self):
|
||||
from soup_cli.utils.echo_trap import build_echo_trap_callback
|
||||
|
||||
cb = build_echo_trap_callback(threshold=0.5)
|
||||
# "a a a a a" — every 2-gram repeats.
|
||||
snap = {"completions": ["a a a a a", "b b b b b"]}
|
||||
sig = cb.compute_signal(snap)
|
||||
assert sig is not None and sig > 0.5
|
||||
|
||||
def test_compute_signal_no_completions(self):
|
||||
from soup_cli.utils.echo_trap import build_echo_trap_callback
|
||||
|
||||
cb = build_echo_trap_callback(threshold=0.5)
|
||||
assert cb.compute_signal({"completions": []}) is None
|
||||
|
||||
def test_tokenizer_aware_path(self):
|
||||
from soup_cli.utils.echo_trap import build_echo_trap_callback
|
||||
|
||||
cb = build_echo_trap_callback(
|
||||
threshold=0.5, tokenizer_aware=True, tokenizer=_FakeTokenizer()
|
||||
)
|
||||
sig = cb.compute_signal({"completions": ["x x x x", "y y y y"]})
|
||||
assert sig is not None and sig >= 0.0
|
||||
|
||||
def test_on_step_end_halts_on_trap(self):
|
||||
from soup_cli.utils.echo_trap import build_echo_trap_callback
|
||||
from soup_cli.utils.rl_signal_buffer import RLSignalBuffer
|
||||
|
||||
buf = RLSignalBuffer()
|
||||
cb = build_echo_trap_callback(threshold=0.3, halt_on_trap=True, buffer=buf)
|
||||
buf.record(
|
||||
func_name="r",
|
||||
completions=["a a a a a a", "b b b b b b"],
|
||||
rewards=[1.0, 1.0],
|
||||
)
|
||||
state, control = _FakeState(1), _FakeControl()
|
||||
cb.on_step_end(None, state, control)
|
||||
assert control.should_training_stop is True
|
||||
assert any("echo_trap_verdict" in e for e in state.log_history)
|
||||
|
||||
def test_observe_classifies_ok(self):
|
||||
from soup_cli.utils.echo_trap import build_echo_trap_callback
|
||||
|
||||
cb = build_echo_trap_callback(threshold=0.5)
|
||||
report = cb.observe_signal(0.0, step=1, n_trajectories=3)
|
||||
assert report.verdict == "OK"
|
||||
assert report.trajectories_seen == 3
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# #238 — RL checkpoint callback
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _FakeSavableModel:
|
||||
def __init__(self):
|
||||
self.saved_to = None
|
||||
|
||||
def save_pretrained(self, path):
|
||||
import os
|
||||
|
||||
os.makedirs(path, exist_ok=True)
|
||||
with open(os.path.join(path, "adapter_model.safetensors"), "wb") as fh:
|
||||
fh.write(b"\x00")
|
||||
self.saved_to = path
|
||||
|
||||
|
||||
class TestRLCheckpointCallback:
|
||||
def test_build_requires_output_dir(self):
|
||||
from soup_cli.utils.rl_checkpoint import (
|
||||
RLCheckpointConfig,
|
||||
build_rl_checkpoint_callback,
|
||||
)
|
||||
|
||||
cfg = RLCheckpointConfig(save_every_steps=1)
|
||||
with pytest.raises(ValueError, match="output_dir"):
|
||||
build_rl_checkpoint_callback(cfg)
|
||||
|
||||
def test_build_rejects_non_config(self):
|
||||
from soup_cli.utils.rl_checkpoint import build_rl_checkpoint_callback
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
build_rl_checkpoint_callback({"save_every_steps": 1}, output_dir="x")
|
||||
|
||||
def test_save_checkpoint_writes_manifest(self, tmp_path, monkeypatch):
|
||||
import torch
|
||||
|
||||
from soup_cli.utils.rl_checkpoint import (
|
||||
RLCheckpointConfig,
|
||||
build_rl_checkpoint_callback,
|
||||
)
|
||||
|
||||
monkeypatch.chdir(tmp_path)
|
||||
cfg = RLCheckpointConfig(save_every_steps=1)
|
||||
cb = build_rl_checkpoint_callback(cfg, output_dir="run", task="grpo")
|
||||
model = _FakeSavableModel()
|
||||
opt = torch.optim.SGD([torch.nn.Parameter(torch.zeros(2))], lr=0.1)
|
||||
ckpt = cb.save_checkpoint(step=2, model=model, optimizer=opt)
|
||||
manifest = Path(ckpt) / "manifest.json"
|
||||
assert manifest.is_file()
|
||||
data = json.loads(manifest.read_text())
|
||||
assert data["step"] == 2 and data["task"] == "grpo"
|
||||
assert data["has_optimizer"] is True
|
||||
assert (Path(ckpt) / "optimizer.pt").is_file()
|
||||
|
||||
def test_prune_keeps_keep_last(self, tmp_path, monkeypatch):
|
||||
from soup_cli.utils.rl_checkpoint import (
|
||||
RLCheckpointConfig,
|
||||
build_rl_checkpoint_callback,
|
||||
)
|
||||
|
||||
monkeypatch.chdir(tmp_path)
|
||||
cfg = RLCheckpointConfig(save_every_steps=1, keep_last=2)
|
||||
cb = build_rl_checkpoint_callback(cfg, output_dir="run", task="grpo")
|
||||
for step in (1, 2, 3):
|
||||
cb.save_checkpoint(step=step, model=_FakeSavableModel(), optimizer=None)
|
||||
root = tmp_path / "run" / "rl-checkpoints"
|
||||
dirs = sorted(p.name for p in root.iterdir())
|
||||
assert dirs == ["step-2", "step-3"] # step-1 pruned
|
||||
|
||||
def test_on_step_end_cadence(self, tmp_path, monkeypatch):
|
||||
from soup_cli.utils.rl_checkpoint import (
|
||||
RLCheckpointConfig,
|
||||
build_rl_checkpoint_callback,
|
||||
)
|
||||
|
||||
monkeypatch.chdir(tmp_path)
|
||||
cfg = RLCheckpointConfig(save_every_steps=2)
|
||||
cb = build_rl_checkpoint_callback(cfg, output_dir="run", task="grpo")
|
||||
# step 1 → no save; step 2 → save.
|
||||
cb.on_step_end(None, _FakeState(1), _FakeControl(), model=_FakeSavableModel())
|
||||
assert not (tmp_path / "run" / "rl-checkpoints").exists()
|
||||
cb.on_step_end(None, _FakeState(2), _FakeControl(), model=_FakeSavableModel())
|
||||
assert (tmp_path / "run" / "rl-checkpoints" / "step-2").is_dir()
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# #236 — ULD
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestULD:
|
||||
def test_build_returns_projection(self):
|
||||
from soup_cli.utils.uld import ULDConfig, ULDProjection, build_uld_projection
|
||||
|
||||
proj = build_uld_projection(
|
||||
ULDConfig(strategy="wasserstein", student_vocab_size=10, teacher_vocab_size=12)
|
||||
)
|
||||
assert isinstance(proj, ULDProjection)
|
||||
|
||||
def test_build_rejects_non_config(self):
|
||||
from soup_cli.utils.uld import build_uld_projection
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
build_uld_projection({"strategy": "wasserstein"})
|
||||
|
||||
def test_wasserstein_loss_different_vocab(self):
|
||||
import torch
|
||||
|
||||
from soup_cli.utils.uld import ULDConfig, uld_distill_loss
|
||||
|
||||
cfg = ULDConfig(strategy="wasserstein", student_vocab_size=8, teacher_vocab_size=12)
|
||||
s = torch.randn(2, 3, 8, requires_grad=True)
|
||||
t = torch.randn(2, 3, 12)
|
||||
loss = uld_distill_loss(s, t, config=cfg)
|
||||
assert torch.isfinite(loss)
|
||||
loss.backward()
|
||||
assert s.grad is not None
|
||||
|
||||
def test_topk_align_loss(self):
|
||||
import torch
|
||||
|
||||
from soup_cli.utils.uld import ULDConfig, uld_distill_loss
|
||||
|
||||
cfg = ULDConfig(
|
||||
strategy="topk_align", student_vocab_size=8, teacher_vocab_size=12, top_k=4
|
||||
)
|
||||
s = torch.randn(2, 3, 8, requires_grad=True)
|
||||
t = torch.randn(2, 3, 12)
|
||||
loss = uld_distill_loss(s, t, config=cfg)
|
||||
assert torch.isfinite(loss)
|
||||
loss.backward()
|
||||
|
||||
def test_identical_distributions_low_loss(self):
|
||||
import torch
|
||||
|
||||
from soup_cli.utils.uld import ULDConfig, uld_distill_loss
|
||||
|
||||
cfg = ULDConfig(strategy="wasserstein", student_vocab_size=8, teacher_vocab_size=8)
|
||||
logits = torch.randn(2, 3, 8)
|
||||
loss = uld_distill_loss(logits, logits.clone(), config=cfg)
|
||||
assert float(loss) < 1e-5
|
||||
|
||||
def test_attention_mask_applied(self):
|
||||
import torch
|
||||
|
||||
from soup_cli.utils.uld import ULDConfig, uld_distill_loss
|
||||
|
||||
cfg = ULDConfig(strategy="wasserstein", student_vocab_size=8, teacher_vocab_size=8)
|
||||
s = torch.randn(2, 3, 8)
|
||||
t = torch.randn(2, 3, 8)
|
||||
mask = torch.tensor([[1, 1, 0], [1, 0, 0]])
|
||||
loss = uld_distill_loss(s, t, config=cfg, attention_mask=mask)
|
||||
assert torch.isfinite(loss)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# #237 — MiniLLM
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestMiniLLM:
|
||||
def test_build_returns_callback(self):
|
||||
from soup_cli.utils.minillm import (
|
||||
MiniLLMCallback,
|
||||
MiniLLMConfig,
|
||||
build_minillm_callback,
|
||||
)
|
||||
|
||||
cb = build_minillm_callback(MiniLLMConfig(teacher_mix_ratio=0.5))
|
||||
assert isinstance(cb, MiniLLMCallback)
|
||||
|
||||
def test_build_rejects_non_config(self):
|
||||
from soup_cli.utils.minillm import build_minillm_callback
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
build_minillm_callback({})
|
||||
|
||||
def test_distill_term_finite_and_differentiable(self):
|
||||
import torch
|
||||
|
||||
from soup_cli.utils.minillm import MiniLLMConfig, minillm_distill_term
|
||||
|
||||
cfg = MiniLLMConfig(teacher_mix_ratio=0.5, length_normalize=True)
|
||||
s = torch.randn(2, 4, 16, requires_grad=True)
|
||||
t = torch.randn(2, 4, 16)
|
||||
labels = torch.randint(0, 16, (2, 4))
|
||||
labels[0, 0] = -100 # masked
|
||||
loss = minillm_distill_term(s, t, labels, config=cfg)
|
||||
assert torch.isfinite(loss)
|
||||
loss.backward()
|
||||
assert s.grad is not None
|
||||
|
||||
def test_teacher_mix_ratio_zero_gives_near_zero(self):
|
||||
import torch
|
||||
|
||||
from soup_cli.utils.minillm import MiniLLMConfig, minillm_distill_term
|
||||
|
||||
cfg = MiniLLMConfig(teacher_mix_ratio=0.0)
|
||||
s = torch.randn(2, 4, 16)
|
||||
t = torch.randn(2, 4, 16)
|
||||
labels = torch.randint(0, 16, (2, 4))
|
||||
loss = minillm_distill_term(s, t, labels, config=cfg)
|
||||
# ratio=0 → target = student_detached → reverse-KL ≈ 0.
|
||||
assert abs(float(loss)) < 1e-4
|
||||
|
||||
def test_anchor_term_with_file(self, tmp_path, monkeypatch):
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
from soup_cli.utils.minillm import MiniLLMConfig, build_minillm_callback
|
||||
|
||||
monkeypatch.chdir(tmp_path)
|
||||
anchor = tmp_path / "anchor.jsonl"
|
||||
anchor.write_text(
|
||||
"\n".join(json.dumps({"text": f"sentence number {i}"}) for i in range(4))
|
||||
)
|
||||
tok = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
|
||||
if tok.pad_token is None:
|
||||
tok.pad_token = tok.eos_token
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
"hf-internal-testing/tiny-random-gpt2"
|
||||
)
|
||||
cb = build_minillm_callback(
|
||||
MiniLLMConfig(pretrain_anchor_weight=0.1, pretrain_anchor_path="anchor.jsonl"),
|
||||
tokenizer=tok,
|
||||
)
|
||||
term = cb.anchor_term(model)
|
||||
assert term is not None
|
||||
assert torch.isfinite(term)
|
||||
|
||||
def test_anchor_term_disabled_returns_none(self):
|
||||
import torch.nn as nn
|
||||
|
||||
from soup_cli.utils.minillm import MiniLLMConfig, build_minillm_callback
|
||||
|
||||
cb = build_minillm_callback(MiniLLMConfig(teacher_mix_ratio=0.3))
|
||||
assert cb.anchor_term(nn.Linear(2, 2)) is None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# #239 — iterative DPO
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestIterativeDPO:
|
||||
def test_build_pairs_from_scored(self):
|
||||
from soup_cli.utils.iterative_dpo import build_pairs_from_scored
|
||||
|
||||
assert build_pairs_from_scored([("a", 1.0), ("b", 3.0)]) == ("b", "a")
|
||||
assert build_pairs_from_scored([("a", 1.0)]) is None
|
||||
assert build_pairs_from_scored([("a", 2.0), ("b", 2.0)]) is None
|
||||
|
||||
def test_run_iterative_dpo_with_fakes(self, tmp_path, monkeypatch):
|
||||
from soup_cli.utils.iterative_dpo import (
|
||||
IterativeDPOResult,
|
||||
build_iterative_dpo_plan,
|
||||
run_iterative_dpo,
|
||||
)
|
||||
|
||||
monkeypatch.chdir(tmp_path)
|
||||
prompts = tmp_path / "prompts.jsonl"
|
||||
prompts.write_text(
|
||||
"\n".join(json.dumps({"prompt": f"q{i}"}) for i in range(3))
|
||||
)
|
||||
|
||||
plan = build_iterative_dpo_plan(
|
||||
base_model="tiny",
|
||||
reward_model="rm",
|
||||
prompts_path="prompts.jsonl",
|
||||
output_dir="out",
|
||||
rounds=2,
|
||||
pairs_per_round=10,
|
||||
)
|
||||
|
||||
calls = {"sample_adapters": [], "score": 0, "train": []}
|
||||
|
||||
def fake_sample(*, base_model, adapter_path, prompts, num_samples,
|
||||
max_new_tokens, device):
|
||||
calls["sample_adapters"].append(adapter_path)
|
||||
return [[f"{p}-a", f"{p}-b"] for p in prompts]
|
||||
|
||||
def fake_score(*, reward_model, prompt, completions, device):
|
||||
calls["score"] += 1
|
||||
return [float(len(c)) for c in completions]
|
||||
|
||||
def fake_train(*, base_model, pairs_path, adapter_path):
|
||||
calls["train"].append((base_model, adapter_path))
|
||||
import os
|
||||
|
||||
os.makedirs(adapter_path, exist_ok=True)
|
||||
|
||||
result = run_iterative_dpo(
|
||||
plan, sample_fn=fake_sample, score_fn=fake_score, train_fn=fake_train
|
||||
)
|
||||
assert isinstance(result, IterativeDPOResult)
|
||||
assert result.rounds_completed == 2
|
||||
# Training ALWAYS starts from the plan's base (never an adapter dir).
|
||||
assert calls["train"][0][0] == "tiny"
|
||||
assert calls["train"][1][0] == "tiny"
|
||||
# Round 0 samples from base (None adapter); round 1 from round-0 adapter.
|
||||
assert calls["sample_adapters"][0] is None
|
||||
assert calls["sample_adapters"][1].endswith("round-00/adapter")
|
||||
# pairs written
|
||||
assert (tmp_path / "out" / "round-00" / "pairs.jsonl").is_file()
|
||||
|
||||
def test_run_rejects_non_plan(self):
|
||||
from soup_cli.utils.iterative_dpo import run_iterative_dpo
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
run_iterative_dpo({"rounds": 1})
|
||||
|
||||
def test_cli_plan_only_still_exits_zero(self, tmp_path, monkeypatch):
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from soup_cli.commands.iterative_dpo import app
|
||||
|
||||
monkeypatch.chdir(tmp_path)
|
||||
(tmp_path / "p.jsonl").write_text(json.dumps({"prompt": "q"}))
|
||||
runner = CliRunner()
|
||||
res = runner.invoke(
|
||||
app,
|
||||
[
|
||||
"--base-model", "b", "--reward-model", "rm",
|
||||
"--prompts", "p.jsonl", "--output-dir", "o",
|
||||
"--rounds", "1", "--pairs-per-round", "10", "--plan-only",
|
||||
],
|
||||
)
|
||||
assert res.exit_code == 0, res.output
|
||||
|
||||
def test_cli_runs_with_monkeypatched_runner(self, tmp_path, monkeypatch):
|
||||
from typer.testing import CliRunner
|
||||
|
||||
import soup_cli.utils.iterative_dpo as idpo
|
||||
from soup_cli.commands.iterative_dpo import app
|
||||
from soup_cli.utils.iterative_dpo import IterativeDPOResult
|
||||
|
||||
monkeypatch.chdir(tmp_path)
|
||||
(tmp_path / "p.jsonl").write_text(json.dumps({"prompt": "q"}))
|
||||
|
||||
def fake_run(plan, **kwargs):
|
||||
return IterativeDPOResult(
|
||||
rounds_completed=1, final_adapter="o/round-00/adapter",
|
||||
per_round_pairs=(1,),
|
||||
)
|
||||
|
||||
monkeypatch.setattr(idpo, "run_iterative_dpo", fake_run)
|
||||
runner = CliRunner()
|
||||
res = runner.invoke(
|
||||
app,
|
||||
[
|
||||
"--base-model", "b", "--reward-model", "rm",
|
||||
"--prompts", "p.jsonl", "--output-dir", "o",
|
||||
"--rounds", "1", "--pairs-per-round", "10",
|
||||
],
|
||||
)
|
||||
assert res.exit_code == 0, res.output
|
||||
assert "Done" in res.output
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# #159 — GRPO variant fallback warning
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _FakeGRPOBase:
|
||||
"""Minimal stand-in for trl.GRPOTrainer for the variant subclass."""
|
||||
|
||||
class _Args:
|
||||
beta = 0.1
|
||||
|
||||
def __init__(self):
|
||||
self.args = self._Args()
|
||||
self.super_called = 0
|
||||
|
||||
def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
|
||||
self.super_called += 1
|
||||
return 0.0
|
||||
|
||||
|
||||
class TestGrpoVariantFallbackWarning:
|
||||
def test_fallback_warns_once(self, caplog):
|
||||
from soup_cli.trainer.grpo import make_grpo_trainer_variant
|
||||
|
||||
cls = make_grpo_trainer_variant(_FakeGRPOBase, "gspo")
|
||||
inst = cls()
|
||||
# inputs missing per-token logps → fallback path.
|
||||
with caplog.at_level(logging.WARNING, logger="soup_cli.trainer.grpo"):
|
||||
inst.compute_loss(None, {})
|
||||
inst.compute_loss(None, {})
|
||||
warnings = [r for r in caplog.records if "fell back" in r.message]
|
||||
assert len(warnings) == 1 # one-shot
|
||||
assert inst.super_called == 2 # but fallback happened both times
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# #160 — in-place GRPO EMA
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestGrpoEmaInPlace:
|
||||
def test_in_place_blends_toward_policy(self):
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from soup_cli.monitoring.grpo_stability_callback import update_ema_in_place
|
||||
|
||||
ref = nn.Linear(3, 3)
|
||||
pol = nn.Linear(3, 3)
|
||||
with torch.no_grad():
|
||||
ref.weight.fill_(0.0)
|
||||
pol.weight.fill_(1.0)
|
||||
update_ema_in_place(ref, pol, 0.25)
|
||||
# ref = 0.75*0 + 0.25*1 = 0.25
|
||||
assert torch.allclose(ref.weight, torch.full_like(ref.weight, 0.25))
|
||||
|
||||
def test_returns_updated_count(self):
|
||||
import torch.nn as nn
|
||||
|
||||
from soup_cli.monitoring.grpo_stability_callback import update_ema_in_place
|
||||
|
||||
# nn.Linear(2, 2) has weight + bias → 2 shared params updated.
|
||||
n = update_ema_in_place(nn.Linear(2, 2), nn.Linear(2, 2), 0.5)
|
||||
assert n == 2
|
||||
|
||||
def test_zero_overlap_returns_zero(self):
|
||||
import torch.nn as nn
|
||||
|
||||
from soup_cli.monitoring.grpo_stability_callback import update_ema_in_place
|
||||
|
||||
# Disjoint parameter names → no overlap → count 0.
|
||||
class _A(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.alpha = nn.Linear(2, 2)
|
||||
|
||||
class _B(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.beta = nn.Linear(2, 2)
|
||||
|
||||
assert update_ema_in_place(_A(), _B(), 0.5) == 0
|
||||
|
||||
def test_rejects_bool_alpha(self):
|
||||
import torch.nn as nn
|
||||
|
||||
from soup_cli.monitoring.grpo_stability_callback import update_ema_in_place
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
update_ema_in_place(nn.Linear(2, 2), nn.Linear(2, 2), True)
|
||||
|
||||
def test_rejects_out_of_range_alpha(self):
|
||||
import torch.nn as nn
|
||||
|
||||
from soup_cli.monitoring.grpo_stability_callback import update_ema_in_place
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
update_ema_in_place(nn.Linear(2, 2), nn.Linear(2, 2), 1.5)
|
||||
|
||||
def test_shape_mismatch_skipped(self):
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from soup_cli.monitoring.grpo_stability_callback import update_ema_in_place
|
||||
|
||||
ref = nn.Linear(3, 3)
|
||||
pol = nn.Linear(2, 2) # different shapes for the same param name
|
||||
with torch.no_grad():
|
||||
ref.weight.fill_(7.0)
|
||||
update_ema_in_place(ref, pol, 0.5)
|
||||
# shape mismatch → ref untouched.
|
||||
assert torch.allclose(ref.weight, torch.full_like(ref.weight, 7.0))
|
||||
|
||||
def test_callback_warns_once_on_zero_overlap(self, caplog):
|
||||
import logging
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
from soup_cli.monitoring.grpo_stability_callback import (
|
||||
GRPOStabilityCallback,
|
||||
)
|
||||
|
||||
class _A(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.alpha = nn.Linear(2, 2)
|
||||
|
||||
class _B(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.beta = nn.Linear(2, 2)
|
||||
|
||||
cb = GRPOStabilityCallback(ref_model_ema_alpha=0.5)
|
||||
cb._policy_model = _A()
|
||||
cb._ref_model = _B()
|
||||
with caplog.at_level(logging.WARNING):
|
||||
cb.on_step_end(args=None, state=None, control=None, model=cb._policy_model)
|
||||
cb.on_step_end(args=None, state=None, control=None, model=cb._policy_model)
|
||||
warnings = [r for r in caplog.records if "0 shared parameters" in r.message]
|
||||
assert len(warnings) == 1 # one-shot
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Source wiring + patch invariants
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSourceWiring:
|
||||
def _read(self, rel: str) -> str:
|
||||
root = Path(__file__).resolve().parent.parent
|
||||
return (root / "src" / "soup_cli" / rel).read_text(encoding="utf-8")
|
||||
|
||||
def test_grpo_wires_rl_callbacks(self):
|
||||
src = self._read("trainer/grpo.py")
|
||||
assert "attach_rl_callbacks" in src
|
||||
assert "wrap_reward_funcs" in src
|
||||
|
||||
def test_distill_wires_uld_and_minillm(self):
|
||||
src = self._read("trainer/distill.py")
|
||||
assert "build_uld_projection" in src
|
||||
assert "build_minillm_callback" in src
|
||||
|
||||
def test_stability_callback_uses_in_place_ema(self):
|
||||
src = self._read("monitoring/grpo_stability_callback.py")
|
||||
assert "update_ema_in_place" in src
|
||||
# the old full-state_dict round-trip should be gone from on_step_end.
|
||||
assert "self._ref_model.load_state_dict(ref_sd" not in src
|
||||
|
||||
def test_no_top_level_torch_in_new_utils(self):
|
||||
for rel in (
|
||||
"utils/rl_signal_buffer.py",
|
||||
"utils/reward_hacking.py",
|
||||
"utils/echo_trap.py",
|
||||
"utils/uld.py",
|
||||
"utils/minillm.py",
|
||||
"utils/iterative_dpo.py",
|
||||
"utils/rl_checkpoint.py",
|
||||
):
|
||||
src = self._read(rel)
|
||||
assert "\nimport torch" not in src, rel
|
||||
assert "\nfrom torch" not in src, rel
|
||||
|
||||
|
||||
class TestPatchInvariants:
|
||||
def test_version_bumped(self):
|
||||
parts = soup_cli.__version__.split(".")
|
||||
assert (int(parts[0]), int(parts[1]), int(parts[2])) >= (0, 71, 11)
|
||||
Loading…
Reference in New Issue