freesolo-flash-dev 0.2.56__tar.gz → 0.2.57__tar.gz

{freesolo_flash_dev-0.2.56 → freesolo_flash_dev-0.2.57}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 0.2.56
+Version: 0.2.57
 Summary: Flash — managed LoRA post-training (SFT/GRPO/OPD) for Freesolo environments, driven by the `flash` CLI
 Project-URL: Homepage, https://github.com/freesolo-co/flash
 Project-URL: Repository, https://github.com/freesolo-co/flash

{freesolo_flash_dev-0.2.56 → freesolo_flash_dev-0.2.57}/flash/engine/recipe.py

@@ -169,6 +169,17 @@ class OPDConfig:
     # that pushes up log P(eos) at the position a rollout naturally terminated, restoring the stop
     # signal the alignment cannot express. 0 disables it; overridable via [train].opd_eos_loss_coef.
     eos_loss_coef: float = 0.5
+    # student-side entropy floor: a hinge on the per-sample MEAN completion entropy (nats), added to
+    # the reverse-KL loss. reverse-KL is mode-seeking and progressively over-sharpens small students
+    # (<=2b) until greedy (temp=0) serving falls into a repetition loop that never emits EOS — the
+    # teacher-gated fix (EOPD, arxiv 2603.07079) is unimplementable here (Fireworks echo-scoring
+    # returns no teacher distribution, and teacher/student are cross-tokenizer), so the STUDENT's own
+    # distribution — local, same-vocab, already computed — is regularized instead. the hinge is
+    # self-limiting: zero loss and zero gradient while mean entropy sits at/above the floor, so runs
+    # that don't collapse pay nothing. 0 disables the term (mean entropy is still logged);
+    # overridable via [train].opd_entropy_floor_coef / [train].opd_entropy_floor.
+    entropy_floor_coef: float = 0.0
+    entropy_floor: float = 0.5
 
 
 @dataclass(frozen=True)

{freesolo_flash_dev-0.2.56 → freesolo_flash_dev-0.2.57}/flash/engine/worker/opd.py

@@ -158,6 +158,12 @@ class OpdKnobs:
     # over shared decoded-text spans cannot supervise the zero-width stop token, so this restores the
     # stop signal; 0 disables it. [train].opd_eos_loss_coef, else RECIPE.opd.eos_loss_coef.
     eos_loss_coef: float = RECIPE.opd.eos_loss_coef
+    # student-side entropy floor (see _entropy_floor_term): hinge weight and target (nats) on the
+    # per-sample mean completion entropy, countering reverse-KL over-sharpening on small students.
+    # coef 0 disables the loss term; the mean entropy is still computed (detached) for logging.
+    # [train].opd_entropy_floor_coef / opd_entropy_floor, else RECIPE.opd defaults.
+    entropy_floor_coef: float = RECIPE.opd.entropy_floor_coef
+    entropy_floor: float = RECIPE.opd.entropy_floor
     save_every: int = 0
     max_length: int = 0
     # Student on-policy sampling stops at these delimiters (parity with GRPO), so the teacher never
@@ -224,6 +230,19 @@ def _resolve_opd_knobs() -> OpdKnobs:
                 _eos if (_eos := opt("opd_eos_loss_coef", None)) is not None else d.eos_loss_coef
             ),
         ),
+        # like eos_loss_coef: explicit 0 is valid (disable), only None falls back to the recipe.
+        entropy_floor_coef=max(
+            0.0,
+            float(
+                _ec
+                if (_ec := opt("opd_entropy_floor_coef", None)) is not None
+                else d.entropy_floor_coef
+            ),
+        ),
+        entropy_floor=max(
+            0.0,
+            float(_ef if (_ef := opt("opd_entropy_floor", None)) is not None else d.entropy_floor),
+        ),
         save_every=int(opt("save_every", 0) or 20),
         max_length=int(opt("max_context_tokens", 0) or 0),
         stop_sequences=tuple(getattr(t, "stop_sequences", ()) or ()),
@@ -315,7 +334,9 @@ def _student_model(model_id, model_init_kwargs, device):
 
     # init_model is the base id (fresh run). VL runs load the full multimodal model so all-linear LoRA
     # sees every target; non-VL runs keep the lighter causal-LM loader.
-    base = model_cls.from_pretrained(init_model, trust_remote_code=True, **model_init_kwargs).to(device)
+    base = model_cls.from_pretrained(init_model, trust_remote_code=True, **model_init_kwargs).to(
+        device
+    )
     return get_peft_model(base, init_peft), str(init_model)
 
 
@@ -369,7 +390,8 @@ def run_opd():
         f"[opd] gkd (groupwise reverse-KL) teacher={knobs.teacher_model} "
         f"epochs={knobs.epochs} warm_start={warm_start or 'none'} "
         f"mode={'multi-turn' if multi_turn else 'single-turn'} "
-        f"eos_loss_coef={knobs.eos_loss_coef}"
+        f"eos_loss_coef={knobs.eos_loss_coef} "
+        f"entropy_floor_coef={knobs.entropy_floor_coef} entropy_floor={knobs.entropy_floor}"
     )
 
     # The GLM teacher key is a platform-owned credential the control plane injects into the worker
@@ -621,12 +643,24 @@ def run_opd():
     # EMA of the per-rollout truncation rate — the runaway signal that scales terminal-EOS reinforcement
     # (see _runaway_eos_scale). Starts at 0 (assume the student stops fine); rises if rollouts run away.
     trunc_ema = 0.0
+    # EMA of the per-rollout EMPTY-completion rate — the OVER-correction signal (see _overcorrection_damp).
+    # Terminal-EOS reinforcement raises the SHARED eos logit, so pushing it too hard makes the student
+    # emit eos FIRST (empty completion -> empty serve). trunc_ema can't see this: an empty rollout is
+    # truncated=False, so it reads as "terminated fine" and keeps the coef pinned high; empty_ema damps
+    # the coef back down as empties appear, closing the feedback loop.
+    empty_ema = 0.0
     granularity_sum = 0.0
     granularity_n = 0
     # Running mean of the reinforced terminal log P(eos): should RISE over a run as the student learns
     # to stop, the counterpart to truncated_rollouts FALLING. Surfaced in train_meta.
     eos_logprob_sum = 0.0
     eos_logprob_n = 0
+    # Running mean of the per-sample student completion entropy (nats) and how many distilled samples
+    # tripped the entropy-floor hinge — the in-band view of reverse-KL entropy collapse (see
+    # _entropy_floor_term). Surfaced in train_meta; the per-step mean goes to wandb/heartbeats.
+    entropy_sum = 0.0
+    entropy_n = 0
+    entropy_floor_active = 0
     opt_steps = (
         0  # optimizer steps actually applied (< steps if any iteration had no teacher signal)
     )
@@ -688,9 +722,16 @@ def run_opd():
         sample (scaled 1/accum_target for gradient accumulation), advance the step aggregates, and
         refresh the stall clock. Called after vLLM generation and teacher scoring; ``samples_seen`` is
         advanced by the caller once per generated rollout so its timing is unchanged."""
-        nonlocal teacher_ok, teacher_transient, teacher_error, truncated_rollouts, step_loss, step_cov
+        nonlocal \
+            teacher_ok, \
+            teacher_transient, \
+            teacher_error, \
+            truncated_rollouts, \
+            step_loss, \
+            step_cov
         nonlocal granularity_sum, granularity_n, generated_tokens, teacher_input_tokens, nseq
-        nonlocal eos_logprob_sum, eos_logprob_n, trunc_ema
+        nonlocal eos_logprob_sum, eos_logprob_n, trunc_ema, empty_ema
+        nonlocal entropy_sum, entropy_n, entropy_floor_active, step_ent_sum, step_ent_n
         if r.teacher_status == "ok":
             teacher_ok += 1
         elif r.teacher_status == "transient":
@@ -702,6 +743,10 @@ def run_opd():
         # Update the runaway EMA once per accounted rollout (truncated or not) so terminal-EOS
         # reinforcement tracks the student's live termination health (see _runaway_eos_scale).
         trunc_ema += (1.0 - _EOS_TRUNC_EMA_DECAY) * ((1.0 if r.truncated else 0.0) - trunc_ema)
+        # An empty completion terminates (truncated=False), so trunc_ema reads it as healthy; track it
+        # separately so the controller can back OFF when its own EOS push starts collapsing rollouts.
+        _is_empty = r.loss is None and _sample_skip_reason(r) == "empty_completion"
+        empty_ema += (1.0 - _EOS_TRUNC_EMA_DECAY) * ((1.0 if _is_empty else 0.0) - empty_ema)
         if r.loss is None:
             reason = _sample_skip_reason(r)
             skip_counts[reason] += 1
@@ -721,6 +766,13 @@ def run_opd():
         if r.eos_logprob is not None:
             eos_logprob_sum += r.eos_logprob
             eos_logprob_n += 1
+        if r.entropy is not None:
+            entropy_sum += r.entropy
+            entropy_n += 1
+            step_ent_sum += r.entropy
+            step_ent_n += 1
+            if knobs.entropy_floor_coef > 0 and r.entropy < knobs.entropy_floor:
+                entropy_floor_active += 1
         nseq += 1
         # Non-liveness progress ping WITHIN the step (rationale in _opd_progress).
         _opd_progress(opt_steps, nseq)
@@ -747,6 +799,7 @@ def run_opd():
     max_teacher_workers = _opd_teacher_workers(
         knobs.prompts_per_step * knobs.group_size, teacher_batch_size
     )
+
     # fields= carries opt_steps on the liveness thread's opd_step pings: opd_step is upload-throttled,
     # so a stepless liveness ping could win the slot and overwrite the main thread's stepped heartbeat,
     # leaving actual_steps_run to floor a cancelled run to 1 step (codex[bot]).
@@ -770,10 +823,15 @@ def run_opd():
             for no_signal_attempt in range(1, max_no_signal_attempts + 1):
                 it = step  # data-slice + display index for THIS rollout attempt
                 step += 1  # advance up front so the nseq==0 retry path can't spin forever
-                batch = [examples[(it * prompts_per_step + i) % len(examples)] for i in range(prompts_per_step)]
+                batch = [
+                    examples[(it * prompts_per_step + i) % len(examples)]
+                    for i in range(prompts_per_step)
+                ]
                 accum_target = max(1, prompts_per_step * group)
                 step_loss = 0.0
                 step_cov = 0.0
+                step_ent_sum = 0.0
+                step_ent_n = 0
                 nseq = 0
                 step_skip_counts: Counter[str] = Counter()
                 # Step pipeline. vLLM produces on-policy rollouts, scorable completions are immediately
@@ -848,6 +906,7 @@ def run_opd():
                         loss_microbatch_size,
                         backward_scale=1.0 / _accum_target,
                         runaway_rate=trunc_ema,
+                        overcorr_rate=empty_ema,
                     )
                     opd_phase_seconds["loss_backward"] += time.perf_counter() - loss_started
                     opd_phase_counts["loss_batches"] += 1
@@ -991,6 +1050,8 @@ def run_opd():
                     opd_phase_counts["vllm_syncs"] += 1
             avg_loss = step_loss / nseq
             avg_cov = step_cov / nseq
+            # None when no distilled sample this step carried an entropy (bare test doubles).
+            avg_ent = (step_ent_sum / step_ent_n) if step_ent_n else None
             loss_curve.append(avg_loss)
             coverage_curve.append(avg_cov)
             _w.heartbeat(
@@ -1004,6 +1065,7 @@ def run_opd():
                 step=opt_steps,
                 loss=avg_loss,
                 coverage=avg_cov,
+                entropy=avg_ent,
                 gpu=gpu_diagnostics(include_torch=False),
                 force=True,
             )
@@ -1012,11 +1074,17 @@ def run_opd():
                 with contextlib.suppress(Exception):
                     import wandb
 
-                    wandb.log({"opd/loss": avg_loss, "opd/coverage": avg_cov}, step=opt_steps)
+                    wandb_metrics = {"opd/loss": avg_loss, "opd/coverage": avg_cov}
+                    if avg_ent is not None:
+                        # falls over a collapsing run; should flatten at/above opd_entropy_floor
+                        # once the floor is enabled (see _entropy_floor_term).
+                        wandb_metrics["opd/entropy"] = avg_ent
+                    wandb.log(wandb_metrics, step=opt_steps)
             if it % 10 == 0:
                 print(
                     f"[opd] step {it + 1}/{steps} loss={avg_loss:.4f} "
                     f"coverage={avg_cov:.0%} seqs={nseq}"
+                    + (f" entropy={avg_ent:.3f}" if avg_ent is not None else "")
                 )
             # Checkpoint on OPTIMIZER-step count, not the loop index: a `step-N` artifact must
             # contain N real updates (skipped no-signal iterations don't advance opt_steps), else a
@@ -1122,6 +1190,17 @@ def run_opd():
             "eos_loss_coef": knobs.eos_loss_coef,
             "eos_reinforced_samples": eos_logprob_n,
             "mean_eos_logprob": (eos_logprob_sum / eos_logprob_n) if eos_logprob_n else None,
+            # Final EMA of the empty-completion rate (_overcorrection_damp input). A non-trivial value
+            # means the controller detected its own EOS push collapsing rollouts and damped the coef.
+            "final_empty_rate_ema": round(empty_ema, 5),
+            # Student entropy floor (see _entropy_floor_term): the knobs, the run-mean per-sample
+            # completion entropy (nats), and how many distilled samples tripped the hinge. A
+            # collapsing run shows mean_student_entropy FALLING step over step (opd/entropy in
+            # wandb); with the floor enabled it should flatten at/above opd_entropy_floor.
+            "entropy_floor_coef": knobs.entropy_floor_coef,
+            "entropy_floor": knobs.entropy_floor,
+            "mean_student_entropy": (entropy_sum / entropy_n) if entropy_n else None,
+            "entropy_floor_active_samples": entropy_floor_active,
             "teacher_transient_failures": teacher_transient,
             "teacher_errors": teacher_error,
             "no_signal_resamples": no_signal_resamples,
@@ -1169,7 +1248,9 @@ def run_opd():
                 else None
             ),
             "opd_rollout_pipeline_max_chunks": (
-                _opd_rollout_pipeline_max_chunks(prompts_per_step * group) if not multi_turn else None
+                _opd_rollout_pipeline_max_chunks(prompts_per_step * group)
+                if not multi_turn
+                else None
             ),
             "opd_teacher_workers": max_teacher_workers,
             "opd_teacher_batch_size": teacher_batch_size,
@@ -1218,6 +1299,10 @@ class SampleResult:
     # Detached log P(eos) at the reinforced terminal position, when the EOS behaviour-cloning term
     # was applied (else None). Surfaced as mean_eos_logprob so a run's rising termination is visible.
     eos_logprob: float | None = None
+    # Detached mean completion entropy (nats) of this sample's student predictive distributions
+    # (distilled samples only, else None). Surfaced as opd/entropy + mean_student_entropy so the
+    # reverse-KL entropy collapse behind the small-model greedy repetition loop is visible in-band.
+    entropy: float | None = None
 
 
 @dataclass(frozen=True)
@@ -1422,7 +1507,9 @@ def _score_many(
                 for _ in pendings
             ]
         return [_ScoreResult(teacher_toks=toks, status="ok") for toks in scored]
-    return [_ScoreResult(status="error", error="teacher batch attempts exhausted") for _ in pendings]
+    return [
+        _ScoreResult(status="error", error="teacher batch attempts exhausted") for _ in pendings
+    ]
 
 
 @dataclass(frozen=True)
@@ -1508,12 +1595,10 @@ def _gkd_loss_from_logps(sp_t, groups, kl_coef=1.0):
     )
     student_group_logsum = sp_det.new_zeros(len(prepared.group_lengths))
     student_group_logsum.index_add_(0, group_ids_t, sp_det.index_select(0, flat_idx_t))
-    teacher_logsum_t = torch.tensor(
-        prepared.teacher_logsums, dtype=sp_t.dtype, device=sp_t.device
+    teacher_logsum_t = torch.tensor(prepared.teacher_logsums, dtype=sp_t.dtype, device=sp_t.device)
+    coeffs = (
+        kl_coef * (student_group_logsum - teacher_logsum_t) / group_lengths_t.to(dtype=sp_t.dtype)
     )
-    coeffs = kl_coef * (
-        student_group_logsum - teacher_logsum_t
-    ) / group_lengths_t.to(dtype=sp_t.dtype)
     coeff_vec = coeffs.index_select(0, group_ids_t)
     sp_sel = sp_t.index_select(0, flat_idx_t)
     return (coeff_vec * sp_sel).mean()
@@ -1609,6 +1694,73 @@ def _eos_reinforce_term(
     return (-float(eos_coef)) * logp, float(logp.detach())
 
 
+# Row-chunk size for _entropy_floor_term's full-vocab passes: 128 rows x ~150k vocab x 3 fp32
+# temporaries ~= 250 MB transient, comfortably inside the existing dense-logits headroom.
+_ENTROPY_CHUNK_ROWS = 128
+
+
+def _entropy_floor_term(rows, coef: float, floor: float):
+    """Student-side entropy floor: hinge the per-sample MEAN completion entropy (nats) at ``floor``.
+
+    Reverse-KL distillation is mode-seeking: on small students (<=2b) it progressively over-sharpens
+    the next-token distribution over training steps until greedy (temperature=0) serving — the
+    deploy default — falls into a repetition attractor and never emits EOS (observed on the 2b:
+    step-20 checkpoint clean, step-63 loops on 80% of a 200-row eval; severity scales with
+    rank x steps / model size). The principled fix gates the KL direction on the TEACHER's per-token
+    entropy (EOPD, arxiv 2603.07079), but flash's teacher path cannot express it: Fireworks
+    echo-scoring returns only realized-token scalars (no top-k, no entropy) and teacher/student are
+    cross-tokenizer. So regularize the STUDENT's own distribution — local, same-vocab, already
+    computed for the reverse-KL loss — to keep its entropy from collapsing.
+
+    The hinge is on the per-sample MEAN, not per token: individual tokens are often legitimately
+    near-deterministic (syntax, format, ``</think>``), and a per-token floor would push every
+    confident token up, fighting correct behaviour everywhere. The observed pathology is
+    distribution-wide drift (mean entropy falls monotonically over OPD steps), which the mean hinge
+    targets exactly. Like the terminal-EOS controller above, it is SELF-LIMITING: zero loss and zero
+    gradient while the mean sits at/above the floor, so students that don't collapse (4b/9b, or the
+    2b's early steps) pay nothing.
+
+    ``rows`` is the same ``[completion_len, vocab]`` logits slice the reverse-KL loss consumes.
+    Returns ``(term_or_none, mean_entropy)``: the differentiable addend
+    ``coef * (floor - mean_entropy)`` when the hinge is active (else ``None``), plus the detached
+    mean entropy in nats for logging (always computed, even with coef 0, so a run's entropy
+    trajectory is visible in-band — serving exposes no logprobs to measure it after the fact).
+
+    MEMORY: entropy needs the full-vocab distribution, and an unchunked fp32 log_softmax + exp over
+    ``[completion_len, vocab]`` would materialize multiple ~GB temporaries per sample — an
+    unbudgeted addition to the dense-logits peak _opd_loss_microbatch_size was sized for, EVEN WITH
+    THE FEATURE DISABLED (codex[bot]). Both passes therefore run in _ENTROPY_CHUNK_ROWS-row chunks,
+    bounding the detached (logging) pass to a ~chunk x vocab transient. The AUTOGRAD pass still
+    retains ~2 fp32 ``[completion_len, vocab]`` buffers for backward — but only for samples the
+    hinge is actively penalizing (coef > 0 AND mean entropy below the floor), i.e. the opt-in
+    experimental case, never the default path."""
+    import torch
+    import torch.nn.functional as F
+
+    if rows is None or rows.shape[0] == 0:
+        return None, None
+    n = rows.shape[0]
+
+    def _chunked_entropy_sum(grad: bool):
+        total = None
+        for i in range(0, n, _ENTROPY_CHUNK_ROWS):
+            lp = F.log_softmax(rows[i : i + _ENTROPY_CHUNK_ROWS].float(), dim=-1)
+            part = -(lp.exp() * lp).sum()
+            if not grad:
+                part = float(part)
+            total = part if total is None else total + part
+        return total
+
+    # detached pass first: cheap logging value, and the gate that decides whether the (autograd)
+    # penalty pass is needed at all — at/above the floor the hinge is identically zero.
+    with torch.no_grad():
+        mean_entropy = _chunked_entropy_sum(grad=False) / n
+    if coef <= 0 or mean_entropy >= floor:
+        return None, mean_entropy
+    ent = _chunked_entropy_sum(grad=True) / n
+    return float(coef) * (float(floor) - ent), mean_entropy
+
+
 # Terminal-EOS reinforcement is a PROPORTIONAL CONTROLLER, not a constant push. #482 added the EOS
 # behaviour-cloning term on every distilled (cleanly-terminated) rollout unconditionally; because it
 # raises the SHARED eos output logit — which governs P(eos) at EVERY position, not just the terminal
@@ -1619,7 +1771,9 @@ def _eos_reinforce_term(
 # rollouts, then an all-empty collapse by step ~55/63. Fix: scale the coef by how much the student is
 # CURRENTLY failing to terminate (an EMA of the per-rollout truncation rate), so a run that stops
 # reliably gets ~zero EOS push and cannot ratchet, while genuine runaway still drives full reinforcement.
-_EOS_RUNAWAY_LO = 0.02  # truncation rate <= this: student stops fine -> NO reinforcement (no ratchet)
+_EOS_RUNAWAY_LO = (
+    0.02  # truncation rate <= this: student stops fine -> NO reinforcement (no ratchet)
+)
 _EOS_RUNAWAY_HI = 0.15  # truncation rate >= this: full opd_eos_loss_coef (clear runaway to correct)
 _EOS_TRUNC_EMA_DECAY = 0.97  # per-rollout EMA (~30-rollout window): responsive yet smoothed
 
@@ -1636,6 +1790,30 @@ def _runaway_eos_scale(runaway_rate: float) -> float:
     return (runaway_rate - _EOS_RUNAWAY_LO) / (_EOS_RUNAWAY_HI - _EOS_RUNAWAY_LO)
 
 
+# _runaway_eos_scale ramps the coef UP on runaway, but on its own it CANNOT stop an empty-collapse when
+# runaway is extreme (base-model / cold start): truncation stays high the whole run so the coef is pinned
+# at full, and an empty completion is truncated=False so it reads as "terminated" and never lowers the
+# ramp. Result (measured on a base-model 35B-thinking run, coef 0.5): 0 runaway but 58% of SAMPLED serves
+# come back empty. So damp the coef by an EMA of the empty-completion rate — the direct symptom of the
+# over-correction — giving two-sided negative feedback: reinforce EOS when running away, but SHUT THE
+# PUSH OFF the moment the student starts emitting eos first. Empties are always a failure (never
+# task-length like a truncation), so damp aggressively.
+_EOS_OVERCORR_LO = 0.01  # empty-rate EMA <= this: no damping (full reinforcement per _runaway_eos_scale)
+_EOS_OVERCORR_HI = 0.05  # empty-rate EMA >= this: reinforcement fully suppressed (student over-stops)
+
+
+def _overcorrection_damp(empty_rate: float) -> float:
+    """Fraction of the runaway-scaled EOS coef to KEEP given the CURRENT empty-completion rate (EMA).
+    1.0 while the student isn't collapsing to empty (rate <= _EOS_OVERCORR_LO); ramps linearly to 0.0 by
+    _EOS_OVERCORR_HI so the terminal-EOS push shuts off before it ratchets P(eos | position 0) past the
+    sampling threshold. Composes multiplicatively with _runaway_eos_scale."""
+    if empty_rate <= _EOS_OVERCORR_LO:
+        return 1.0
+    if empty_rate >= _EOS_OVERCORR_HI:
+        return 0.0
+    return 1.0 - (empty_rate - _EOS_OVERCORR_LO) / (_EOS_OVERCORR_HI - _EOS_OVERCORR_LO)
+
+
 def _resolve_samples_batched(
     model,
     tok,
@@ -1646,6 +1824,7 @@ def _resolve_samples_batched(
     *,
     backward_scale: float | None = None,
     runaway_rate: float = 1.0,
+    overcorr_rate: float = 0.0,
 ) -> list[SampleResult]:
     import torch
 
@@ -1657,11 +1836,16 @@ def _resolve_samples_batched(
     # don't change across the chunk. getattr defaults keep bare test stand-ins (SimpleNamespace knobs)
     # and eos-less fake tokenizers working with reinforcement OFF.
     eos_coef = float(getattr(knobs, "eos_loss_coef", 0.0) or 0.0)
-    # Controller: reinforce the terminal EOS only in proportion to the student's CURRENT failure to
-    # terminate (see _runaway_eos_scale). A cleanly-terminating run -> ~0 push -> the shared eos logit
-    # can't ratchet into an empty-collapse; genuine runaway still drives full reinforcement. Default
-    # runaway_rate=1.0 keeps bare callers (tests) at full strength = pre-controller behaviour.
-    eos_coef *= _runaway_eos_scale(runaway_rate)
+    # Two-sided controller: reinforce the terminal EOS in proportion to the student's CURRENT failure to
+    # terminate (_runaway_eos_scale), then DAMP by its current empty-completion rate (_overcorrection_damp)
+    # so the push shuts off if it starts collapsing rollouts to empty. A cleanly-terminating run -> ~0
+    # push; genuine runaway -> full push UNLESS that push is already producing empties. Defaults
+    # runaway_rate=1.0 / overcorr_rate=0.0 keep bare callers (tests) at full strength = pre-controller.
+    eos_coef *= _runaway_eos_scale(runaway_rate) * _overcorrection_damp(overcorr_rate)
+    # student entropy floor (see _entropy_floor_term). getattr defaults keep bare test stand-ins
+    # working with the penalty OFF; the detached mean entropy is computed either way for logging.
+    ent_coef = float(getattr(knobs, "entropy_floor_coef", 0.0) or 0.0)
+    ent_floor = float(getattr(knobs, "entropy_floor", 0.0) or 0.0)
     eos_stop_sequences = tuple(getattr(knobs, "stop_sequences", ()) or ())
     eos_ids = _generation_eos_ids(model, tok) if eos_coef > 0 else frozenset()
     eos_primary = _primary_eos_id(tok, eos_ids) if eos_coef > 0 else None
@@ -1726,13 +1910,9 @@ def _resolve_samples_batched(
             seqs = [list(p.prompt_ids) + list(p.student_ids) for p in chunk]
             max_len = max(len(seq) for seq in seqs)
             input_ids = torch.full((len(seqs), max_len), pad_id, dtype=torch.long, device=device)
-            attention_mask = torch.zeros(
-                (len(seqs), max_len), dtype=torch.long, device=device
-            )
+            attention_mask = torch.zeros((len(seqs), max_len), dtype=torch.long, device=device)
             for row, seq in enumerate(seqs):
-                input_ids[row, : len(seq)] = torch.tensor(
-                    seq, dtype=torch.long, device=device
-                )
+                input_ids[row, : len(seq)] = torch.tensor(seq, dtype=torch.long, device=device)
                 attention_mask[row, : len(seq)] = 1
             logits = _forward_logits(model, input_ids, attention_mask)
             losses = []
@@ -1771,6 +1951,12 @@ def _resolve_samples_batched(
                         if eos_out is not None:
                             eos_term, eos_logprob = eos_out
                             loss = loss + eos_term
+                    # entropy floor on the SAME rows the reverse-KL consumed; added to distilled
+                    # samples only (loss is not None), so skip accounting is unchanged. the detached
+                    # mean entropy is logged even when the penalty is off/inactive.
+                    ent_term, entropy = _entropy_floor_term(rows, ent_coef, ent_floor)
+                    if ent_term is not None:
+                        loss = loss + ent_term
                     loss_for_result = loss
                     if backward_scale is not None:
                         losses.append(loss)
@@ -1783,6 +1969,7 @@ def _resolve_samples_batched(
                         teacher_tokens=p.teacher_tokens,
                         group_granularity=p.group_granularity,
                         eos_logprob=eos_logprob,
+                        entropy=entropy,
                     )
             if losses:
                 (sum(losses) * float(backward_scale)).backward()

{freesolo_flash_dev-0.2.56 → freesolo_flash_dev-0.2.57}/flash/engine/worker/sft.py

@@ -196,8 +196,12 @@ def run_sft():
     from flash.engine.vram import sft_grad_accum
 
     sft_lr = _train_opt("learning_rate", RECIPE.sft.learning_rate)
+    # Total SFT sequence budget (prompt + completion), honoring [train] max_context_tokens like the
+    # cost/preflight path (flash.cost.spec._sft_seq_len); unset -> the tuned recipe cap. Read the
+    # RENAMED knob: reading a stale "max_length" here made getattr silently return the 1024 default.
     sft_max_len = _train_opt(
-        "max_length", RECIPE.sft.max_seq_len_thinking if _w.THINKING else RECIPE.sft.max_seq_len
+        "max_context_tokens",
+        RECIPE.sft.max_seq_len_thinking if _w.THINKING else RECIPE.sft.max_seq_len,
     )
     with liveness_heartbeat("sft_pretokenizing"):
         texts, _pretok, _dropped = _pretokenize_completion_only(texts, tok, sft_max_len)

{freesolo_flash_dev-0.2.56 → freesolo_flash_dev-0.2.57}/flash/schema/__init__.py

@@ -208,6 +208,8 @@ _TRAIN_KEYS = frozenset(
         "advantage_clip",
         "thinking_length_penalty_coef",
         "opd_eos_loss_coef",
+        "opd_entropy_floor_coef",
+        "opd_entropy_floor",
         "teacher_model",
         "stop_sequences",
         "structured_outputs",
@@ -216,6 +218,7 @@ _TRAIN_KEYS = frozenset(
     }
 )
 
+
 def spec_from_dict(raw: dict[str, Any], run_id: str | None = None) -> JobSpec:
     # Only reject table-valued unknowns — callers pass harmless scalar flags like dry_run alongside spec.
     unknown = sorted(k for k in set(raw) - _TOP_LEVEL_KEYS if isinstance(raw[k], dict))
@@ -356,6 +359,10 @@ def spec_from_dict(raw: dict[str, Any], run_id: str | None = None) -> JobSpec:
             ),
             # OPD-only terminal-EOS reinforcement weight; 0 disables. None -> recipe default (0.5).
             opd_eos_loss_coef=_train_float(train_raw, "opd_eos_loss_coef", minimum=0.0),
+            # OPD-only student entropy floor (hinge weight + target nats); coef 0 disables the term.
+            # None -> recipe defaults (0.0 / 0.5, i.e. off until validated).
+            opd_entropy_floor_coef=_train_float(train_raw, "opd_entropy_floor_coef", minimum=0.0),
+            opd_entropy_floor=_train_float(train_raw, "opd_entropy_floor", minimum=0.0),
             # OPD-only managed teacher alias, validated against the allow-list; "" -> default GLM 5.2.
             teacher_model=_train_teacher(train_raw),
             stop_sequences=_train_stops(train_raw),

{freesolo_flash_dev-0.2.56 → freesolo_flash_dev-0.2.57}/flash/serve/deploy.py

@@ -11,6 +11,7 @@ from dataclasses import asdict, dataclass
 import httpx
 
 from flash._logging import get_logger
+from flash.engine.structured_outputs import parse_structured_outputs
 from flash.lora_rank import rank_from_adapter_config
 
 logger = get_logger(__name__)
@@ -263,6 +264,34 @@ def adapter_artifact_lora_rank(hf_repo: str, subfolder: str) -> int:
     return rank_from_adapter_config(config, source=f"{hf_repo}:{filename}")
 
 
+def _structured_outputs_body(
+    run_id: str, structured_outputs: str, *, thinking: bool
+) -> dict | None:
+    """The per-adapter structured-outputs serving default for the registration body, or None.
+
+    Returns the parsed canonical StructuredOutputsParams-kwargs dict when the run configured a
+    constraint AND is non-thinking; None otherwise (no default is registered). Non-thinking only:
+    serving sets no reasoning-parser, so a per-adapter grammar would bind from the very first
+    generated token and force a thinking adapter's ``<think>`` phase to open the schema (e.g. ``{``)
+    — strictly worse than serving it unconstrained. Training stays compatible with thinking via a
+    DEFERRED grammar (``reasoning_parser="deepseek_r1"`` holds the constraint until ``</think>``),
+    but serving has no such deferral, so we skip the serve-time default for thinking runs rather
+    than break reasoning. Raises ValueError on a corrupt spec — a wiring bug, since
+    ``[train].structured_outputs`` is already canonicalized at run creation (schema/fields.py), so
+    this cannot fire for a real run; it fails loudly rather than ever silently serving unconstrained."""
+    if not structured_outputs:
+        return None
+    if thinking:
+        logger.info(
+            "adapter %s trained with structured_outputs but thinking=True; not registering a "
+            "serving guided-decoding default — serving has no reasoning-parser deferral, so it "
+            "would bind the grammar from token 0 and suppress the <think> phase",
+            run_id,
+        )
+        return None
+    return parse_structured_outputs(structured_outputs)
+
+
 def deploy_adapter(
     run_id: str,
     model: str,
@@ -272,9 +301,18 @@ def deploy_adapter(
     dry_run: bool = False,
     lora_rank: int = 32,
     thinking: bool = False,
+    structured_outputs: str = "",
     org_id: str | None = None,
 ) -> Deployment:
-    """Register the trained adapter with the freesolo serving app."""
+    """Register the trained adapter with the freesolo serving app.
+
+    ``structured_outputs`` is the run's ``[train].structured_outputs`` spec (canonical
+    StructuredOutputsParams-kwargs JSON, or "" for none). When set, it is registered as the
+    adapter's per-adapter guided-decoding DEFAULT so every serve request is constrained by the
+    SAME grammar the adapter trained under — otherwise a structured-outputs run drifts at
+    unconstrained serving (train/serve exposure bias: it over-generates arrays that truncate to
+    invalid JSON and emits keys the training grammar had masked). See ``_structured_outputs_body``.
+    """
     validate_serving_lora_rank(model, lora_rank, rank_source="configured train.lora_rank")
     subfolder = f"{adapter_prefix}/adapter"
     if not dry_run:
@@ -298,6 +336,9 @@ def deploy_adapter(
         # Preserves thinking parity: without this, Qwen3.5 defaults to thinking ON regardless of training.
         "thinking": bool(thinking),
     }
+    so_default = _structured_outputs_body(run_id, structured_outputs, thinking=thinking)
+    if so_default is not None:
+        body["structured_outputs"] = so_default
     normalized_org_id = (org_id or "").strip()
     if normalized_org_id:
         body["org_id"] = normalized_org_id

{freesolo_flash_dev-0.2.56 → freesolo_flash_dev-0.2.57}/flash/server/routes/serving.py

@@ -386,6 +386,10 @@ def deploy(run_id: str, key: Annotated[dict, Depends(require_key)], payload: dic
             "lora_rank": spec.train.lora_rank,
             # a run trained with thinking serves with thinking (per-run parity)
             "thinking": spec.thinking,
+            # a run trained with structured_outputs serves under the SAME grammar (guided-decoding
+            # parity), so it doesn't drift at unconstrained serving; deploy_adapter registers it as
+            # the per-adapter default (non-thinking runs only — see _structured_outputs_body).
+            "structured_outputs": spec.train.structured_outputs,
             "org_id": deploy_org_id,
         }
         if dry_run:

{freesolo_flash_dev-0.2.56 → freesolo_flash_dev-0.2.57}/flash/spec.py

@@ -125,6 +125,11 @@ class TrainSpec:
     # OPD only: weight of the terminal-EOS behaviour-cloning term (recipe default when None). 0
     # disables it. Raise it for a student that keeps running past the length cap without emitting EOS.
     opd_eos_loss_coef: float | None = None
+    # OPD only: student-side entropy floor — weight and target (nats) of the hinge on the per-sample
+    # mean completion entropy (recipe defaults when None; coef 0 disables the term). counteracts
+    # reverse-KL over-sharpening on small students (the greedy repetition loop at temp=0).
+    opd_entropy_floor_coef: float | None = None
+    opd_entropy_floor: float | None = None
     # OPD only: the managed teacher to distil from, stored as its resolved Fireworks model id (parse
     # canonicalizes any alias / spaced / raw-id form via recipe.resolve_teacher, e.g. "glm-5.2" =>
     # "accounts/fireworks/models/glm-5p2"). "" => the recipe default (GLM 5.2).
@@ -218,6 +223,8 @@ class JobSpec:
                 advantage_clip=_opt_float(train.get("advantage_clip")),
                 thinking_length_penalty_coef=_opt_float(train.get("thinking_length_penalty_coef")),
                 opd_eos_loss_coef=_opt_float(train.get("opd_eos_loss_coef")),
+                opd_entropy_floor_coef=_opt_float(train.get("opd_entropy_floor_coef")),
+                opd_entropy_floor=_opt_float(train.get("opd_entropy_floor")),
                 teacher_model=str(train.get("teacher_model") or ""),
                 stop_sequences=_str_tuple(train.get("stop_sequences")),
                 structured_outputs=str(train.get("structured_outputs") or ""),

{freesolo_flash_dev-0.2.56 → freesolo_flash_dev-0.2.57}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "freesolo-flash-dev"
-version = "0.2.56"
+version = "0.2.57"
 description = "Flash — managed LoRA post-training (SFT/GRPO/OPD) for Freesolo environments, driven by the `flash` CLI"
 readme = "README.md"
 # RunPod Flash supports Python 3.11-3.12 (not 3.13 yet).
@@ -106,7 +106,7 @@ artifacts = ["flash/py.typed"]
 # auto-published by .github/workflows/publish-dev.yml: a push to `dev` whose version isn't yet on
 # PyPI publishes it; pushes that don't change it find it already published and no-op. So cutting a
 # dev-channel release = bump this (and [project].version) and merge to `dev`.
-version = "0.2.56"
+version = "0.2.57"
 
 [tool.pytest.ini_options]
 pythonpath = ["."]