world-model-optimizer 0.2.0__py3-none-any.whl

wmo/distill/tripwire.py

@@ -0,0 +1,411 @@
+"""Degeneration tripwires: entropy and generation length, relative to baseline.
+
+Reverse-KL distillation can degenerate in ways a KL curve cannot show, because
+KL falls while the policy collapses. Two failures are on record in sibling
+lanes:
+
+- Length/mode collapse: generation length fell 50x (2,866 to about 50 tokens),
+  task accuracy fell 0.813 to 0.596, and entropy breached the lane's
+  pre-registered floor at steps 17 to 23; the run was killed at step 30/200.
+- The mirror pathology, reported by a second lane: "pure KL gives EOS no
+  gradient, so the student never learns to stop", i.e. episodes that run to
+  the turn or token cap instead of submitting.
+
+This module measures the two statistics that make both visible and decides
+warn/kill against a baseline THIS run measured, from data the training step
+already holds (the student's own recorded spans): no extra sampling, no extra
+spend.
+
+Why relative and never absolute. The sibling lane's rule was "entropy < 0.2
+nats means collapse". Our healthy, untrained baseline is 0.181 nats/token
+(`PROBE_BASELINE_ENTROPY_NATS`, pooled over 356,122 sampled tokens of a live
+48-episode TerminalBench-2 probe on Super-120B base weights at temperature
+0.7), so that rule would fire at step 0, before a single gradient step.
+Terminal-command tokens are much more predictable than math reasoning, and a
+sampled-token entropy estimate taken at temperature 0.7 is biased downward
+besides. An always-firing tripwire gets muted, so an absolute threshold here
+would be worse than no threshold at all: every bound in `TripwireConfig` is a
+fraction of the measured baseline, and it must stay that way.
+
+Why batch-pooled and never per-episode. On that same probe the healthiest
+single episode already sits at entropy 0.082 nats/token and episode lengths
+span 349 to 30,869 sampled tokens. Any per-episode rule (or a mean of
+per-episode ratios) would fire on healthy runs constantly, so both statistics
+pool over the whole batch: total sampled logprob over total sampled tokens, and
+total sampled tokens over episodes.
+
+What the defaults cost in false alarms. Resampling the probe's own 47 healthy
+episodes into batches (200,000 draws per shape) gives the pooled ratio a
+completely stationary policy would report:
+
+    episodes/batch   length ratio p0.1 / min   entropy ratio p0.1 / min
+    32 (8 tasks x 4)      0.62 / 0.47               0.86 / 0.81
+    16                    0.50 / 0.36               0.80 / 0.70
+     8                    0.37 / 0.21               0.71 / 0.57
+
+At the headline batch shape the length WARN (0.5) has probability 1e-5 per step
+and the length KILL (0.25) never occurred in 200,000 draws; no draw at any shape
+put the entropy ratio under 0.5, let alone the 0.3 kill bound. Requiring
+`kill_consecutive_steps` in a row squares an already negligible number. The
+tripwire tightens as the batch shrinks, though, so a step whose batch was
+decimated by infrastructure failures is the one case worth reading beside its
+`trials` and `empty_span_trials` counts before believing a length warning.
+
+Both directions, since 2026-07-26. The bounds above are floors, and for a while
+they were the only bounds, on the argument that a turn cap already bounds the
+runaway direction. A measured run refuted that: under a 20-turn training cap the
+pooled length still reached 5.33x baseline and entropy 1.68x, because the cap
+bounds turns while each turn grew (172 turns pinned the 12,288-token output cap
+in one step). Nothing fired, and the damage landed somewhere neither floor
+watches — episodes over the context budget are dropped whole, so 22 of 64
+episodes and a third of the trainable datums disappeared while every metric the
+tripwire reads looked "fine, just larger". Each metric now answers to a ceiling
+(`*_mult`) as well, and the ceilings are set from that run's own trajectory:
+see `TripwireConfig` for which observation fixed each number.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Sequence
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from wmo.distill.config import TripwireConfig
+from wmo.distill.tokens import TrialRecord
+
+logger = logging.getLogger(__name__)
+
+TripwireMetric = Literal["entropy_per_token", "mean_generation_tokens"]
+"""The two statistics under tripwire, named exactly as the metrics-row keys."""
+
+TripwireLevel = Literal["warn", "kill"]
+
+TripwireDirection = Literal["floor", "ceiling"]
+"""Which side of the baseline a breach is on.
+
+`floor` is collapse (the ratio fell to a fraction); `ceiling` is runaway (it
+rose to a multiple). Both are the same pathology seen from opposite ends, and a
+run has been lost to each: a sibling lane's length fell 50x, and this lane's
+length rose 5.33x until context overflow was discarding a third of every batch.
+"""
+
+
+class PolicyHealth(BaseModel):
+    """One rollout batch's batch-pooled entropy proxy and generation length."""
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    episodes: int = Field(ge=0)
+    """Episodes that recorded at least one sampled token.
+
+    Episodes with no spans are excluded from BOTH denominators on purpose: a
+    trial that died at sandbox creation or lost its transport generated nothing,
+    so counting it as a zero-length episode would drag the length statistic down
+    for an infrastructure outage and could kill a healthy policy. The
+    all-empty-batch abort (`MAX_CONSECUTIVE_EMPTY_STEPS`) owns that failure
+    mode, and `RolloutStats.empty_span_trials` counts it in the same row."""
+
+    sampled_tokens: int = Field(ge=0)
+    """Sampled (loss-position) tokens pooled over those episodes."""
+
+    entropy_per_token: float | None
+    """`-mean(sampled_logprobs)` over every sampled token in the batch.
+
+    The tokens were drawn from the policy that scored them, so `-log p(x)` with
+    `x ~ p` is an unbiased single-sample estimator of that policy's entropy
+    `H(p) = E_{x~p}[-log p(x)]`; pooling over the batch averages many such
+    samples. Two caveats to read it with:
+
+    - Temperature. Rollouts sample at `sampling.temperature` (0.7 for the
+      headline run), which concentrates the distribution the tokens come from,
+      so this sits BELOW the T=1 entropy: treat it as a lower bound and never
+      compare it against a number measured at another temperature.
+    - It is a policy-shape statistic, not a per-token uncertainty: a batch
+      dominated by long, highly predictable tool output reads low without
+      anything being wrong, which is exactly why the tripwire compares it only
+      against this run's own baseline.
+
+    None when the batch recorded no sampled token at all."""
+
+    mean_generation_tokens: float | None
+    """Sampled tokens per episode, pooled (`sampled_tokens / episodes`).
+
+    None when no episode recorded a sampled token."""
+
+
+def policy_health(records: Sequence[TrialRecord]) -> PolicyHealth:
+    """Pool one rollout batch's recorded spans into the two tripwire statistics.
+
+    Measured from the raw `TrialRecord` spans rather than from the built
+    `TrainDatum`s, though the datum builder's loss positions are exactly these
+    sampled tokens (its loss mask is 1.0 on sampled tokens and 0.0 on context).
+    The reason is drops: `build_datums` discards whole episodes that overflow
+    the context budget or exceed `train.max_datum_tokens`, which are precisely
+    the LONGEST episodes, so measuring generation length after the drops would
+    deflate it in exactly the tail that matters and could kill a healthy run.
+
+    Args:
+        records: The batch's assembled trial records.
+
+    Returns:
+        The batch-pooled health; both statistics are None when no episode
+        recorded a sampled token (an all-empty batch, whose own abort path
+        handles it).
+    """
+    episodes = 0
+    sampled_tokens = 0
+    logprob_total = 0.0
+    for record in records:
+        episode_tokens = sum(len(span.sampled_token_ids) for span in record.spans)
+        if episode_tokens == 0:
+            continue
+        episodes += 1
+        sampled_tokens += episode_tokens
+        logprob_total += sum(sum(span.sampled_logprobs) for span in record.spans)
+    if episodes == 0 or sampled_tokens == 0:
+        return PolicyHealth(
+            episodes=episodes,
+            sampled_tokens=sampled_tokens,
+            entropy_per_token=None,
+            mean_generation_tokens=None,
+        )
+    return PolicyHealth(
+        episodes=episodes,
+        sampled_tokens=sampled_tokens,
+        entropy_per_token=-logprob_total / sampled_tokens,
+        mean_generation_tokens=sampled_tokens / episodes,
+    )
+
+
+class TripwireBaseline(BaseModel):
+    """The reference the tripwires compare every later step against.
+
+    Captured from the run's FIRST training step and persisted in the run's
+    checkpoint manifest, because a resumed run that re-baselines against an
+    already-degenerated policy has a blind tripwire: the collapse becomes the
+    new normal and nothing fires again.
+    """
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    step: int = Field(ge=0)
+    """The 0-based training step measured, so a resumed run can tell whether
+    the current step IS the baseline step (which never fires against itself)."""
+
+    entropy_per_token: float = Field(gt=0.0)
+    mean_generation_tokens: float = Field(gt=0.0)
+    episodes: int = Field(ge=1)
+    """Episodes behind the measurement, for reading the row's confidence."""
+
+    sampled_tokens: int = Field(ge=1)
+
+
+def capture_baseline(step: int, health: PolicyHealth) -> TripwireBaseline | None:
+    """Turn a first step's measured health into the run's baseline.
+
+    Args:
+        step: The 0-based training step being measured.
+        health: That step's batch-pooled health.
+
+    Returns:
+        The baseline, or None when this step cannot serve as one: no episode
+        sampled anything, or a statistic came out non-positive (an entropy of
+        exactly 0.0 would make every later ratio undefined). The caller leaves
+        the tripwire unarmed and tries again at the next step.
+    """
+    entropy = health.entropy_per_token
+    length = health.mean_generation_tokens
+    if entropy is None or length is None or entropy <= 0.0 or length <= 0.0:
+        return None
+    return TripwireBaseline(
+        step=step,
+        entropy_per_token=entropy,
+        mean_generation_tokens=length,
+        episodes=health.episodes,
+        sampled_tokens=health.sampled_tokens,
+    )
+
+
+class TripwireBreach(BaseModel):
+    """One metric under one threshold at one step."""
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    metric: TripwireMetric
+    level: TripwireLevel
+    direction: TripwireDirection
+    baseline: float
+    value: float
+    ratio: float
+    """`value / baseline`; the tripwire is a bound on this, never on `value`."""
+
+    threshold_frac: float
+    """The configured bound the ratio reached: a fraction for a `floor` breach,
+    a multiple for a `ceiling` one."""
+
+    config_field: str
+    """The `[tripwire]` key that set `threshold_frac`, so a breach message says
+    exactly which knob to move."""
+
+    def describe(self) -> str:
+        """One line naming the metric, its baseline, the value, and the ratio."""
+        crossed = "at or under" if self.direction == "floor" else "at or above"
+        return (
+            f"{self.metric} {self.value:.4g} is {self.ratio:.2f}x this run's baseline "
+            f"{self.baseline:.4g} (measured at its first training step), {crossed} the "
+            f"{self.level} threshold {self.threshold_frac:.2f} ({self.config_field})"
+        )
+
+
+def metric_ratio(value: float | None, baseline: float | None) -> float | None:
+    """`value / baseline`, or None when either side is missing or unusable."""
+    if value is None or baseline is None or baseline <= 0.0:
+        return None
+    return value / baseline
+
+
+class _MetricReading(BaseModel):
+    """One metric's current value beside the baseline and bounds it answers to."""
+
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    metric: TripwireMetric
+    value: float | None
+    baseline: float
+    warn_frac: float
+    kill_frac: float
+    warn_field: str
+    kill_field: str
+    warn_mult: float
+    kill_mult: float
+    warn_mult_field: str
+    kill_mult_field: str
+
+
+def _readings(
+    cfg: TripwireConfig, baseline: TripwireBaseline, health: PolicyHealth
+) -> list[_MetricReading]:
+    """Both metrics wired to their configured bounds, in report order."""
+    return [
+        _MetricReading(
+            metric="entropy_per_token",
+            value=health.entropy_per_token,
+            baseline=baseline.entropy_per_token,
+            warn_frac=cfg.entropy_warn_frac,
+            kill_frac=cfg.entropy_kill_frac,
+            warn_field="tripwire.entropy_warn_frac",
+            kill_field="tripwire.entropy_kill_frac",
+            warn_mult=cfg.entropy_warn_mult,
+            kill_mult=cfg.entropy_kill_mult,
+            warn_mult_field="tripwire.entropy_warn_mult",
+            kill_mult_field="tripwire.entropy_kill_mult",
+        ),
+        _MetricReading(
+            metric="mean_generation_tokens",
+            value=health.mean_generation_tokens,
+            baseline=baseline.mean_generation_tokens,
+            warn_frac=cfg.length_warn_frac,
+            kill_frac=cfg.length_kill_frac,
+            warn_field="tripwire.length_warn_frac",
+            kill_field="tripwire.length_kill_frac",
+            warn_mult=cfg.length_warn_mult,
+            kill_mult=cfg.length_kill_mult,
+            warn_mult_field="tripwire.length_warn_mult",
+            kill_mult_field="tripwire.length_kill_mult",
+        ),
+    ]
+
+
+def evaluate_breaches(
+    cfg: TripwireConfig, baseline: TripwireBaseline, health: PolicyHealth
+) -> list[TripwireBreach]:
+    """Compare one step's health against the baseline bounds, both sides.
+
+    A metric with no measurement (None) is skipped rather than treated as zero:
+    a batch that sampled nothing is an infrastructure failure, and the
+    all-empty-batch abort owns it. The boundary belongs to the breach side (a
+    ratio exactly at the fraction or the multiple fires), so a threshold set to a
+    value that was actually observed is never silently inert.
+
+    Kill is checked before warn on BOTH sides before either warn is considered,
+    so a metric that has blown through a kill bound is never reported as a mere
+    warning. A ratio cannot breach a floor and a ceiling at once (the validator
+    keeps every fraction <= 1 < every multiple), so the two directions cannot
+    collide.
+
+    Args:
+        cfg: The run's `[tripwire]` section (its `enabled` flag is the caller's
+            business; this function always reports what it sees).
+        baseline: The run's captured baseline.
+        health: The step's batch-pooled health.
+
+    Returns:
+        At most one breach per metric, the more severe level when both bounds
+        are crossed; an empty list when nothing breached.
+    """
+    breaches: list[TripwireBreach] = []
+    for reading in _readings(cfg, baseline, health):
+        ratio = metric_ratio(reading.value, reading.baseline)
+        if reading.value is None or ratio is None:
+            continue
+        level: TripwireLevel
+        direction: TripwireDirection
+        if ratio <= reading.kill_frac:
+            level, direction = "kill", "floor"
+            threshold, field = reading.kill_frac, reading.kill_field
+        elif ratio >= reading.kill_mult:
+            level, direction = "kill", "ceiling"
+            threshold, field = reading.kill_mult, reading.kill_mult_field
+        elif ratio <= reading.warn_frac:
+            level, direction = "warn", "floor"
+            threshold, field = reading.warn_frac, reading.warn_field
+        elif ratio >= reading.warn_mult:
+            level, direction = "warn", "ceiling"
+            threshold, field = reading.warn_mult, reading.warn_mult_field
+        else:
+            continue
+        breaches.append(
+            TripwireBreach(
+                metric=reading.metric,
+                level=level,
+                direction=direction,
+                baseline=reading.baseline,
+                value=reading.value,
+                ratio=ratio,
+                threshold_frac=threshold,
+                config_field=field,
+            )
+        )
+    return breaches
+
+
+def health_summary(health: PolicyHealth, baseline: TripwireBaseline | None) -> str:
+    """The progress-line fragment for one step's health.
+
+    Args:
+        health: The step's batch-pooled health.
+        baseline: The run's baseline, when one is armed.
+
+    Returns:
+        A compact "entropy X (Rx baseline), gen Y tok/episode (Rx)" fragment,
+        with `n/a` in place of a statistic the batch could not measure and the
+        ratios omitted until a baseline exists.
+    """
+    parts: list[str] = []
+    entropy_ratio = metric_ratio(
+        health.entropy_per_token, baseline.entropy_per_token if baseline else None
+    )
+    length_ratio = metric_ratio(
+        health.mean_generation_tokens, baseline.mean_generation_tokens if baseline else None
+    )
+    entropy = health.entropy_per_token
+    length = health.mean_generation_tokens
+    parts.append("entropy/token n/a" if entropy is None else f"entropy/token {entropy:.4f}")
+    if entropy_ratio is not None:
+        parts[-1] += f" ({entropy_ratio:.2f}x baseline)"
+    parts.append("gen tokens/episode n/a" if length is None else f"gen tokens/episode {length:.0f}")
+    if length_ratio is not None:
+        parts[-1] += f" ({length_ratio:.2f}x)"
+    return ", ".join(parts)

wmo/distill/xtoken/byte_offsets.py

@@ -0,0 +1,152 @@
+"""Exact per-token byte offsets for a token sequence, with no re-encoding.
+
+Cross-tokenizer chunk alignment needs to know which BYTES of the decoded text
+each token covers, on both the student side (the exact sampled ids, which
+TITO forbids re-encoding) and the teacher side. Re-encoding is not an option
+for the student: sampling routinely emits a NON-canonical BPE segmentation of
+its own text (measured on the headline run's real sinks, `'S' + 'olver'` where
+canonical BPE gives `'Solver'`), so `encode(decode(ids)) != ids` for 14% of
+spans covering 42% of sampled tokens. Gating on that round trip would discard
+those tokens as "unstable" when nothing about them is unstable.
+
+Byte-level BPE vocabularies (the GPT-2 lineage, which both Qwen3.6 and
+GLM-5.2 use) make the exact answer available directly: every ordinary token's
+surface form is a reversible byte-level encoding of the bytes it covers, so
+per-token byte lengths come from the vocabulary alone. `span_byte_ends`
+reconstructs the span's bytes that way and VERIFIES the reconstruction against
+the tokenizer's own decode before returning, so a vocabulary that does not
+work this way (a SentencePiece tokenizer that normalizes text, say) is
+reported as unusable instead of silently producing wrong offsets.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Protocol
+
+logger = logging.getLogger(__name__)
+
+
+class ByteOffsetTokenizer(Protocol):
+    """The tokenizer slice byte-offset reconstruction needs.
+
+    HuggingFace fast tokenizers satisfy this structurally, as do the small
+    deterministic fakes in the tests (extra defaulted parameters do not
+    matter).
+    """
+
+    def convert_ids_to_tokens(self, ids: list[int]) -> list[str | None]:
+        """The surface form of each token id; None for ids outside the vocab."""
+        ...
+
+    def decode(self, token_ids: list[int]) -> str:
+        """Decode token ids back to text."""
+        ...
+
+
+def _byte_decoder() -> dict[str, int]:
+    """The GPT-2 byte-level BPE surface-character to raw-byte map.
+
+    Byte-level BPE renders each of the 256 possible bytes as exactly one
+    printable character so a BPE vocabulary can be built over text. The map is
+    a fixed property of the scheme, not of any model, so it is derived here
+    rather than read from a tokenizer.
+    """
+    printable = list(range(33, 127)) + list(range(161, 173)) + list(range(174, 256))
+    surface = list(printable)
+    spare = 0
+    for value in range(256):
+        if value not in printable:
+            printable.append(value)
+            surface.append(256 + spare)
+            spare += 1
+    return {chr(code): value for value, code in zip(printable, surface, strict=True)}
+
+
+BYTE_DECODER = _byte_decoder()
+"""Surface character to raw byte, for byte-level BPE token surface forms."""
+
+
+def token_bytes(surface: str) -> bytes | None:
+    """The raw bytes one byte-level BPE token surface form stands for.
+
+    Args:
+        surface: A token's surface form from `convert_ids_to_tokens`.
+
+    Returns:
+        The bytes the token covers, or None when the surface form contains a
+        character outside the byte-level alphabet. That is the normal case for
+        special and added tokens (`<|im_end|>`), whose surface form is their
+        literal text; callers fall back to encoding that text directly.
+    """
+    try:
+        return bytes(BYTE_DECODER[char] for char in surface)
+    except KeyError:
+        return None
+
+
+def span_byte_ends(
+    tokenizer: ByteOffsetTokenizer, token_ids: list[int]
+) -> tuple[list[int], bytes] | None:
+    """Cumulative byte-end offset per token, plus the span's decoded bytes.
+
+    Entry i is the byte offset just past token i, so token i covers
+    `bytes[ends[i - 1] : ends[i]]` (with `ends[-1] == 0` implied for i = 0).
+    The result is exact for the tokenizer's ACTUAL ids: nothing is re-encoded,
+    so a non-canonical sampled segmentation is handled like any other.
+
+    Args:
+        tokenizer: The tokenizer that produced `token_ids`.
+        token_ids: The span's token ids, in order.
+
+    Returns:
+        `(byte_ends, span_bytes)` when reconstruction agrees with the
+        tokenizer's own decode, else None. None means this tokenizer does not
+        expose reversible per-token bytes (it normalizes text, or a token id
+        is outside the vocabulary), so the caller must treat the span as
+        unscoreable rather than guess at offsets.
+    """
+    if not token_ids:
+        return [], b""
+    surfaces = tokenizer.convert_ids_to_tokens(list(token_ids))
+    if len(surfaces) != len(token_ids):
+        logger.warning(
+            "tokenizer returned %d surface form(s) for %d token id(s); cannot "
+            "reconstruct byte offsets for this span",
+            len(surfaces),
+            len(token_ids),
+        )
+        return None
+    pieces: list[bytes] = []
+    for index, surface in enumerate(surfaces):
+        if surface is None:
+            logger.warning(
+                "token id %d at position %d has no surface form (outside the "
+                "vocabulary); cannot reconstruct byte offsets for this span",
+                token_ids[index],
+                index,
+            )
+            return None
+        raw = token_bytes(surface)
+        pieces.append(surface.encode("utf-8") if raw is None else raw)
+    span_bytes = b"".join(pieces)
+    # The reconstruction is only trustworthy if it reproduces what the
+    # tokenizer itself decodes. A normalizing tokenizer (SentencePiece
+    # whitespace rewriting, NFKC) fails here, which is exactly the signal the
+    # caller needs: alignment by byte offset is invalid for that vocabulary.
+    decoded = tokenizer.decode(list(token_ids))
+    if span_bytes != decoded.encode("utf-8"):
+        logger.warning(
+            "per-token byte reconstruction (%d bytes) does not match the "
+            "tokenizer's decode (%d bytes); this vocabulary does not expose "
+            "reversible per-token bytes, so the span cannot be byte-aligned",
+            len(span_bytes),
+            len(decoded.encode("utf-8")),
+        )
+        return None
+    ends: list[int] = []
+    total = 0
+    for piece in pieces:
+        total += len(piece)
+        ends.append(total)
+    return ends, span_bytes