diff --git a/README.md b/README.md index cebfcfd..07ff7b3 100644 --- a/README.md +++ b/README.md @@ -131,9 +131,12 @@ soup iterative-dpo \ soup train --config grpo.yaml \ --echo-trap-enabled \ --echo-trap-threshold 0.6 \ - --echo-trap-halt + --echo-trap-halt \ + --echo-trap-tokenizer-aware ``` +`--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. ## Why Soup? diff --git a/soup_cli/commands/train.py b/soup_cli/commands/train.py index 93ae409..b70d34c 100644 --- a/soup_cli/commands/train.py +++ b/soup_cli/commands/train.py @@ -155,6 +155,14 @@ def train( "Default deny (v0.36.0). Only enable if you trust the source." ), ), + echo_trap_tokenizer_aware: bool = typer.Option( + False, + "--echo-trap-tokenizer-aware", + help=( + "Use tokenizer-id n-grams for echo-trap scoring. Requires " + "training.echo_trap_enabled=true on grpo/ppo." + ), + ), profile_run: bool = typer.Option( False, "--profile", @@ -230,6 +238,17 @@ def train( console.print(f"[dim]Loading config from {config_path}...[/]") cfg = load_config(config_path) + # --- Echo-trap tokenizer-aware shortcut --- + if echo_trap_tokenizer_aware: + if not cfg.training.echo_trap_enabled: + console.print( + "[red]--echo-trap-tokenizer-aware requires " + "training.echo_trap_enabled=true[/]" + ) + raise typer.Exit(1) + cfg.training.echo_trap_tokenizer_aware = True + console.print("[green]Echo-trap tokenizer-aware scoring enabled[/]") + # --- --push-as / --hf-resume validation --- if push_as: from soup_cli.utils.hf import validate_repo_id @@ -440,6 +459,8 @@ def train( script_args.append("--wandb") if tensorboard: script_args.append("--tensorboard") + if echo_trap_tokenizer_aware: + script_args.append("--echo-trap-tokenizer-aware") if yes: script_args.append("--yes") argv = build_accelerate_argv( diff --git a/soup_cli/config/schema.py b/soup_cli/config/schema.py index 3c50888..a9cd46e 100644 --- a/soup_cli/config/schema.py +++ b/soup_cli/config/schema.py @@ -1366,6 +1366,15 @@ class TrainingConfig(BaseModel): "echo_trap_enabled=True. (v0.70.0)" ), ) + echo_trap_tokenizer_aware: bool = Field( + default=False, + description=( + "Use tokenizer-id n-grams for echo-trap scoring instead of " + "whitespace tokens. More sensitive to subword repetition but " + "bound to the active tokenizer vocabulary. Requires " + "echo_trap_enabled=True. (v0.70.x)" + ), + ) # ---- v0.70.0 Part D — Mid-epoch RL checkpoint ------------------------ rl_checkpoint_save_every_steps: Optional[int] = Field( @@ -1520,6 +1529,7 @@ class TrainingConfig(BaseModel): @field_validator( "echo_trap_enabled", "echo_trap_halt", + "echo_trap_tokenizer_aware", mode="before", ) @classmethod @@ -4035,15 +4045,23 @@ class SoupConfig(BaseModel): ``echo_trap_enabled`` is a silent no-op footgun — reject. """ tcfg = self.training - if not tcfg.echo_trap_enabled and not tcfg.echo_trap_halt: + if ( + not tcfg.echo_trap_enabled + and not tcfg.echo_trap_halt + and not tcfg.echo_trap_tokenizer_aware + ): return self - if not tcfg.echo_trap_enabled and tcfg.echo_trap_halt: + if not tcfg.echo_trap_enabled and ( + tcfg.echo_trap_halt or tcfg.echo_trap_tokenizer_aware + ): raise ValueError( - "echo_trap_halt=True requires echo_trap_enabled=True" + "echo_trap_halt / echo_trap_tokenizer_aware require " + "echo_trap_enabled=True" ) if self.task not in ("grpo", "ppo"): raise ValueError( - "echo_trap_enabled / echo_trap_halt are only valid on " + "echo_trap_enabled / echo_trap_halt / " + "echo_trap_tokenizer_aware are only valid on " f"task in {{'grpo', 'ppo'}}; got task={self.task!r}" ) if self.backend == "mlx": diff --git a/soup_cli/utils/echo_trap.py b/soup_cli/utils/echo_trap.py index 57bba15..6a771cf 100644 --- a/soup_cli/utils/echo_trap.py +++ b/soup_cli/utils/echo_trap.py @@ -14,9 +14,8 @@ trajectory collection. Security: - Pure-Python math (no torch import at module top). - Bool / NaN / Inf / range rejection on every numeric input. -- Tokens must be strings; non-str rejected loudly so a caller that - hands tensor ids (instead of decoded strings) gets an actionable - error rather than silently misbehaving. +- Whitespace-mode tokens must be strings; tokeniser-aware mode accepts + integer token ids through the dedicated ``*_tokenized`` helpers. - ``_MAX_BATCH_TRAJECTORIES = 100_000`` DoS cap (matches v0.55 / v0.65 / v0.66 cap policy). - ``_MAX_NGRAM_N = 32`` keeps the n-gram counter bounded. @@ -75,6 +74,41 @@ def _check_tokens(tokens: object) -> tuple[str, ...]: return tuple(iterator) +def _check_token_ids(token_ids: object) -> tuple[int, ...]: + if isinstance(token_ids, (str, bytes)): + raise TypeError("token_ids must be a sequence of ints, not str/bytes") + try: + iterator = list(token_ids) # type: ignore[arg-type] + except TypeError as exc: + raise TypeError( + f"token_ids must be iterable, got {type(token_ids).__name__}" + ) from exc + if len(iterator) > _MAX_TRAJECTORY_TOKENS: + raise ValueError( + f"trajectory has {len(iterator)} token ids, exceeds " + f"{_MAX_TRAJECTORY_TOKENS} cap" + ) + for idx, token_id in enumerate(iterator): + if isinstance(token_id, bool) or not isinstance(token_id, int): + raise TypeError( + f"token_ids[{idx}] must be int, got {type(token_id).__name__}" + ) + return tuple(iterator) + + +def _score_repetition(units: Sequence[object], *, ngram_n: int) -> float: + if len(units) < ngram_n: + return 0.0 + counts: dict[tuple[object, ...], int] = {} + for i in range(len(units) - ngram_n + 1): + gram = tuple(units[i : i + ngram_n]) + counts[gram] = counts.get(gram, 0) + 1 + if not counts: + return 0.0 + repeating = sum(1 for count in counts.values() if count > 1) + return repeating / len(counts) + + def score_trajectory_repetition(tokens: object, *, ngram_n: object = 2) -> float: """Per-trajectory repetition score. @@ -88,16 +122,24 @@ def score_trajectory_repetition(tokens: object, *, ngram_n: object = 2) -> float """ n = _check_ngram_n(ngram_n) tok = _check_tokens(tokens) - if len(tok) < n: - return 0.0 - counts: dict[tuple[str, ...], int] = {} - for i in range(len(tok) - n + 1): - gram = tok[i : i + n] - counts[gram] = counts.get(gram, 0) + 1 - if not counts: - return 0.0 - repeating = sum(1 for c in counts.values() if c > 1) - return repeating / len(counts) + return _score_repetition(tok, ngram_n=n) + + +def score_trajectory_repetition_tokenized( + token_ids: object, + *, + ngram_n: object = 2, +) -> float: + """Per-trajectory repetition score over tokenizer ids. + + This mirrors :func:`score_trajectory_repetition`, but operates on + integer token ids before decoding/whitespace splitting can hide + subword repetition. It is intentionally separate so the existing + string-token API keeps rejecting accidental tensor-id input. + """ + n = _check_ngram_n(ngram_n) + ids = _check_token_ids(token_ids) + return _score_repetition(ids, ngram_n=n) def score_echo_signal( @@ -135,6 +177,41 @@ def score_echo_signal( return sum(scores) / len(scores) +def score_echo_signal_tokenized( + trajectories: object, + *, + ngram_n: object = 2, +) -> float: + """Mean repetition score across token-id trajectories. + + Use this when the caller has access to the trainer tokenizer and can + pass ``tokenizer.encode(text)`` output rather than decoded strings. + """ + n = _check_ngram_n(ngram_n) + if isinstance(trajectories, (str, bytes)): + raise TypeError( + "trajectories must be a sequence of token-id sequences, not str/bytes" + ) + try: + batch = list(trajectories) # type: ignore[arg-type] + except TypeError as exc: + raise TypeError( + f"trajectories must be iterable, got " + f"{type(trajectories).__name__}" + ) from exc + if len(batch) > _MAX_BATCH_TRAJECTORIES: + raise ValueError( + f"batch has {len(batch)} trajectories, exceeds " + f"{_MAX_BATCH_TRAJECTORIES} cap" + ) + if not batch: + return 0.0 + scores: list[float] = [] + for traj in batch: + scores.append(score_trajectory_repetition_tokenized(traj, ngram_n=n)) + return sum(scores) / len(scores) + + def classify_echo_signal(signal: object) -> str: """Map a signal in ``[0, 1]`` to OK / WARN / TRAP. @@ -222,6 +299,7 @@ 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. @@ -244,6 +322,11 @@ def build_echo_trap_callback( 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 " @@ -260,10 +343,14 @@ __all__ = [ "build_echo_trap_callback", "classify_echo_signal", "score_echo_signal", + "score_echo_signal_tokenized", "score_trajectory_repetition", + "score_trajectory_repetition_tokenized", ] # Type aliases retained for the v0.70.1 wiring. TrajectoryTokens = Sequence[str] TrajectoryBatch = Iterable[TrajectoryTokens] +TokenIdTrajectory = Sequence[int] +TokenIdTrajectoryBatch = Iterable[TokenIdTrajectory] diff --git a/tests/test_v0700_part_f.py b/tests/test_v0700_part_f.py index 669f3cd..08accfc 100644 --- a/tests/test_v0700_part_f.py +++ b/tests/test_v0700_part_f.py @@ -19,7 +19,9 @@ class TestEchoTrapPublicSurface: from soup_cli.utils import echo_trap assert hasattr(echo_trap, "score_trajectory_repetition") + assert hasattr(echo_trap, "score_trajectory_repetition_tokenized") assert hasattr(echo_trap, "score_echo_signal") + assert hasattr(echo_trap, "score_echo_signal_tokenized") assert hasattr(echo_trap, "classify_echo_signal") assert hasattr(echo_trap, "EchoTrapReport") assert hasattr(echo_trap, "build_echo_trap_callback") @@ -151,6 +153,38 @@ class TestScoreEchoSignal: score_echo_signal(big, ngram_n=2) +class TestTokenizedEchoSignal: + def test_tokenized_repetition_catches_subword_echo_trap(self): + from soup_cli.utils.echo_trap import ( + classify_echo_signal, + score_echo_signal, + score_echo_signal_tokenized, + ) + + decoded_tokens = ["ha,", "ha.", "ha!", "ha?", "ha;", "ha:"] + repeated_token_ids = [101, 202, 101, 202, 101, 202, 101, 202] + + whitespace_score = score_echo_signal([decoded_tokens], ngram_n=2) + tokenized_score = score_echo_signal_tokenized([repeated_token_ids], ngram_n=2) + + assert classify_echo_signal(whitespace_score) == "OK" + assert classify_echo_signal(tokenized_score) == "TRAP" + + def test_tokenized_trajectory_rejects_non_int_ids(self): + from soup_cli.utils.echo_trap import score_trajectory_repetition_tokenized + + with pytest.raises(TypeError, match="token_ids"): + score_trajectory_repetition_tokenized([1, "2", 3], ngram_n=2) + with pytest.raises(TypeError, match="token_ids"): + score_trajectory_repetition_tokenized([1, True, 3], ngram_n=2) + + def test_tokenized_batch_rejects_str(self): + from soup_cli.utils.echo_trap import score_echo_signal_tokenized + + with pytest.raises(TypeError, match="token-id"): + score_echo_signal_tokenized("not ids", ngram_n=2) + + class TestClassifyEchoSignal: """OK / WARN / TRAP taxonomy (mirrors v0.26 / v0.56 / v0.70 Part A). @@ -310,6 +344,15 @@ class TestBuildEchoTrapCallbackDeferred: with pytest.raises(TypeError, match="halt"): build_echo_trap_callback(threshold=0.5, halt_on_trap="yes") # type: ignore[arg-type] + def test_tokenizer_aware_must_be_bool(self): + from soup_cli.utils.echo_trap import build_echo_trap_callback + + with pytest.raises(TypeError, match="tokenizer_aware"): + build_echo_trap_callback( + threshold=0.5, + tokenizer_aware="yes", # type: ignore[arg-type] + ) + # --------------------------------------------------------------------------- # Schema integration — TrainingConfig + SoupConfig @@ -324,6 +367,7 @@ class TestSchemaTrainingConfig: assert tcfg.echo_trap_enabled is False assert tcfg.echo_trap_threshold == 0.6 assert tcfg.echo_trap_halt is False + assert tcfg.echo_trap_tokenizer_aware is False def test_threshold_bounds(self): from pydantic import ValidationError @@ -352,6 +396,7 @@ data: training: echo_trap_enabled: true echo_trap_threshold: 0.55 + echo_trap_tokenizer_aware: true """ def test_grpo_accepted(self): @@ -359,6 +404,7 @@ training: cfg = load_config_from_string(self._yaml("grpo")) assert cfg.training.echo_trap_enabled is True + assert cfg.training.echo_trap_tokenizer_aware is True def test_ppo_accepted(self): from soup_cli.config.loader import load_config_from_string @@ -388,6 +434,22 @@ training: """ ) + def test_tokenizer_aware_without_enabled_rejected(self): + from soup_cli.config.loader import load_config_from_string + + with pytest.raises(ValueError, match="echo_trap_enabled"): + load_config_from_string( + """ +base: meta-llama/Llama-3.1-8B +task: grpo +data: + train: ./data/train.jsonl + format: chatml +training: + echo_trap_tokenizer_aware: true +""" + ) + # --------------------------------------------------------------------------- # Source wiring guards