freesolo-flash-dev 0.2.45__tar.gz → 0.2.47__tar.gz

{freesolo_flash_dev-0.2.45 → freesolo_flash_dev-0.2.47}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 0.2.45
-Summary: Flash — managed LoRA post-training (SFT/GRPO) for Freesolo environments, driven by the `flash` CLI
+Version: 0.2.47
+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
 Author: Freesolo

{freesolo_flash_dev-0.2.45 → freesolo_flash_dev-0.2.47}/flash/cli/commands.py

@@ -210,14 +210,14 @@ A managed run should use the returned [environment] id from
 This starter implements a tiny "guess the secret number" game so you can see the
 episode hooks wired end-to-end. Replace it with your real task before a real run.
 
-SFT and GRPO train off this file; OPD (configs/opd.toml) does NOT support multi-turn:
+All three algorithms train off this file:
 - GRPO (configs/rl.toml) rolls out full episodes and optimizes `score_episode`.
 - SFT (configs/sft.toml) learns the gold trajectory. Provide it per row as
   `output = {"messages": [...]}` (a full assistant/tool trajectory) or a scalar
   `output` for a single gold assistant turn.
-- OPD (configs/opd.toml) distils one sampled completion per prompt and cannot drive the
-  turn loop, so it is single-turn only — `flash train configs/opd.toml` fails fast here.
-  Use OPD with a single-turn env (`flash env setup --single-turn`).
+- OPD (configs/opd.toml) rolls out each episode and distils EVERY assistant turn against
+  the Fireworks GLM teacher, conditioned on the transcript so far — the multi-turn
+  on-policy-distillation objective. The teacher key is platform-managed (nothing to set).
 """
 
 from __future__ import annotations
@@ -492,32 +492,27 @@ def cmd_env_setup(args) -> int:
         )
     opd = Path("configs/opd.toml")
     if not opd.exists():
-        # opd (on-policy distillation) is SINGLE-TURN only: it samples one completion per prompt and
-        # cannot drive a multi-turn/tool episode loop, so `flash train configs/opd.toml` FAILS FAST on a
-        # multi-turn env (run_opd rejects it at setup). Scaffold it in both modes for reference, but say
-        # so loudly in the multi-turn one so the config isn't a footgun.
+        # opd (on-policy distillation) works in BOTH modes: single-turn distils one sampled completion
+        # per prompt; multi-turn rolls out each episode and distils every assistant turn against the
+        # transcript so far. The scaffold differs only by a one-line note pointing at the mode.
         opd_multiturn_note = (
-            "# NOTE: opd is SINGLE-TURN only — it distils one sampled completion per prompt and cannot\n"
-            "# drive this multi-turn environment's episode loop, so `flash train configs/opd.toml` will\n"
-            "# fail fast here. Use configs/rl.toml (grpo) or configs/sft.toml for multi-turn; keep opd for\n"
-            "# a single-turn env (`flash env setup --single-turn`).\n\n"
+            "# NOTE: opd rolls out each episode and distils EVERY assistant turn (conditioned on the\n"
+            "# transcript so far) against the Fireworks GLM teacher — the multi-turn distillation path.\n\n"
             if multi_turn
             else ""
         )
         opd.write_text(
             f"{opd_multiturn_note}"
             'model = "Qwen/Qwen3.5-4B"\n'
-            'algorithm = "opd"   # on-policy distillation from a Fireworks GLM teacher (single-turn only)\n\n'
+            'algorithm = "opd"   # on-policy distillation from the managed GLM 5.2 teacher\n\n'
             "# Environment: upload this project folder with\n"
             "# `flash env push --name my-env .`, then paste the returned id below.\n"
-            "# FIREWORKS_API_KEY (the GLM teacher key) is read from your shell/.env at submit time.\n"
+            "# The GLM 5.2 teacher and key are platform-managed — nothing to set up or export.\n"
             "[environment]\n"
-            'id = ""\n'
-            'secrets = ["FIREWORKS_API_KEY"]\n\n'
+            'id = ""\n\n'
             "[train]\n"
             "steps = 100                     # opd is step-driven (like GRPO)\n"
             "lora_rank = 32\n"
-            '# teacher_model = "accounts/fireworks/models/glm-5p2"   # Fireworks GLM teacher (default)\n'
             "# GPU and HF artifacts are managed automatically by the platform: the GPU is\n"
             "# the cheapest fitting managed class, and artifacts live in a private environment-scoped repo.\n"
         )
@@ -842,13 +837,25 @@ def cmd_cancel(args) -> int:
     else:
         print(json.dumps(payload, indent=2))
     if checkpoints:
-        top = max(c["step"] for c in checkpoints)
+        # Best-effort hint (the cancel already succeeded), so never crash on a malformed checkpoint
+        # shape: coerce steps defensively — a dict missing 'step' or carrying a non-int must not raise
+        # a traceback here. Only surface the `step-N` deploy example when we recovered a real step.
+        steps = []
+        for c in checkpoints:
+            try:
+                steps.append(int(c["step"]))
+            except (KeyError, TypeError, ValueError):
+                continue
         # stderr in the plain path so the machine-readable stdout JSON stays untouched.
         out = sys.stdout if render.styled() else sys.stderr
-        msg = (
+        base = (
             f"{len(checkpoints)} deployable checkpoint(s) survive this cancel — list with "
-            f"`flash checkpoints {args.run_id}`, deploy one with "
-            f"`flash deploy {args.run_id}/step-{top}`."
+            f"`flash checkpoints {args.run_id}`"
+        )
+        msg = (
+            f"{base}, deploy one with `flash deploy {args.run_id}/step-{max(steps)}`."
+            if steps
+            else f"{base}."
         )
         print(render.note(msg) if render.styled() else msg, file=out)
     return 0

{freesolo_flash_dev-0.2.45 → freesolo_flash_dev-0.2.47}/flash/cli/training_doc.py

@@ -377,20 +377,28 @@ SFT is **epoch-driven** (`epochs`); GRPO is **step-driven** (`steps`).
 ## On-policy distillation (`algorithm = "opd"`)
 
 Pick distillation when a much stronger **teacher** model can grade your student's work
-token-by-token. The student samples on-policy (like GRPO), a Fireworks-hosted GLM teacher scores
+token-by-token. The student samples on-policy (like GRPO), the managed GLM 5.2 teacher scores
 each of *its own* completions, and a dense per-token loss teaches the student to match the teacher —
 far more sample-efficient than reward-based RL and with no reward to design. It is **step-driven**
 (`steps`) and produces a LoRA served exactly like SFT.
 
-- **Provide the teacher key.** Distillation needs `FIREWORKS_API_KEY` at submit time. List it under
-  `[environment] secrets = ["FIREWORKS_API_KEY"]` and export it in your shell / local `.env`; the
-  value travels out-of-band and is never stored in the spec or needed at serving time.
+- **No teacher key or model override to set up.** The GLM 5.2 teacher and its Fireworks key are platform-managed: the
+  service supplies its own key to every opd run, so there is nothing to export or declare — an opd
+  run submits like any other. Bring-your-own teacher models and keys are not supported; a
+  `FIREWORKS_API_KEY` in your shell is simply ignored. The key is never stored in the spec or
+  needed at serving time.
 - **The student (Qwen / MiniCPM / Kimi) and the teacher (GLM) have different tokenizers.** Flash
   bridges the vocabulary mismatch with **groupwise reverse-KL** (the collinear-ai *spider* / Tinker
   method): it aligns the two tokenizations by shared decoded-text spans and applies per-span reverse
   KL using only realized-token logprobs — no vocabulary projection, so it covers every token exactly
   and works for any student tokenizer. When the tokenizers happen to agree it reduces to plain
   per-token reverse KL (Thinking Machines, *On-Policy Distillation*). Nothing to configure.
+- **Works for multi-turn envs too.** Against an `EnvironmentMultiTurn`, opd rolls out each episode
+  (driving `step_episode` / observations just like GRPO) and distils EVERY assistant turn against the
+  teacher, each conditioned on the transcript up to that turn — the episode's total reverse-KL over
+  the student's generated tokens is the sum of its per-turn reverse-KLs. Env/observation tokens are
+  never distilled (they're context, not the student's output). Set `[train] max_length` to bound the
+  transcript; the teacher must cover it (GLM-5.2's context far exceeds the default budget).
 - **Judge it like SFT.** Distillation logs a falling per-token loss; a low loss alone is not proof.
   Keep a held-out split, `flash deploy` the adapter, and score it — confirm the student actually
   moved toward the teacher's behavior, not just its surface tokens.
@@ -402,12 +410,10 @@ algorithm = "opd"
 
 [environment]
 id = "your-org/my-env"
-secrets = ["FIREWORKS_API_KEY"]
 
 [train]
 steps = 100
 lora_rank = 32
-# teacher_model = "accounts/fireworks/models/glm-5p2"   # Fireworks GLM teacher (default)
 # kl_penalty_coef = 1.0                                 # reverse-KL scale
 ```
 

{freesolo_flash_dev-0.2.45 → freesolo_flash_dev-0.2.47}/flash/cost/analytical.py

@@ -99,7 +99,7 @@ def seconds_per_step(config: RunConfig, gpu: str) -> float:
     peak = gpu_tflops(gpu) * 1e12  # FLOP/s
 
     if n.is_opd:
-        # OPD step = on-policy student rollout (like GRPO) + remote teacher scoring (serial
+        # OPD step = on-policy student rollout (like GRPO) + remote teacher scoring (CONCURRENT
         # Fireworks round-trips, replaces reward grading) + policy update (fwd+bwd only, NO local
         # reference forward — the teacher is the API). Bill local compute on the FULL prompt+completion
         # sequence (see _opd_step_shape), not completion-only, or long-prompt opd is underquoted.
@@ -107,10 +107,12 @@ def seconds_per_step(config: RunConfig, gpu: str) -> float:
         gen_s = (GRPO_GEN_FLOPS_PER_TOKEN_PER_PARAM * params * seq_tokens) / (peak * MFU_DECODE)
         update_s = (OPD_UPDATE_FLOPS_PER_TOKEN_PER_PARAM * params * seq_tokens) / (peak * MFU_TRAIN)
         teacher_lat = teacher_seconds_per_completion()
-        # run_opd scores each completion serially (nested prompt/group loops each await
-        # teacher.score before the next), so the wall cost is the full serial sum — do NOT divide
-        # by a concurrency factor or the quote materially understates teacher-bound steps.
-        teacher_s = completions * teacher_lat
+        # run_opd's primary path scores a step's completions CONCURRENTLY over Fireworks with a fan-out
+        # cap of the step's OWN completion count (prompts_per_step * group_size, opd.py Phase 2), so
+        # every completion in a step is scored in ONE parallel wave — the teacher wall is a single
+        # latency, NOT the full serial sum (that describes only the CPU-test fallback that can't
+        # batch-generate). The teacher endpoint's rate limit is the real ceiling on this fan-out.
+        teacher_s = teacher_lat if completions else 0.0
         return gen_s + teacher_s + update_s
 
     if not n.is_grpo:
@@ -255,10 +257,10 @@ def estimate_cost(config: RunConfig, *, wall_cap_s: float = DEFAULT_WALL_CAP_S)
         train_seconds=train,
         wall_clock_seconds=wall,
         wall_capped=wall_capped,
-        # GPU (platform-billed) time only. OPD's teacher tokens are billed by Fireworks DIRECTLY to
-        # the user's FIREWORKS_API_KEY (a user-supplied [environment].secret), so folding
-        # teacher_api_usd into the charged total would bill it a second time — keep it itemized as a
-        # diagnostic instead (codex[bot]).
+        # GPU (platform-billed) time only. OPD's teacher tokens are billed by Fireworks to the
+        # platform-managed teacher key — never through this GPU charge — so folding teacher_api_usd
+        # into the charged total would bill it a second time. Keep it itemized as a diagnostic
+        # instead (codex[bot]).
         total_usd=train / 3600.0 * hourly,
         teacher_api_usd=teacher_api_usd,
         notes=_notes(config, raw_train, wall_capped, cap_s),

{freesolo_flash_dev-0.2.45 → freesolo_flash_dev-0.2.47}/flash/cost/facts.py

@@ -166,10 +166,9 @@ def reward_seconds_per_completion(override: float | None = None) -> float:
 # NO FIREWORKS_API_KEY. Source: https://fireworks.ai/models/fireworks/glm-5p2 lists
 # $1.40 / $0.14 / $4.40 (input / cached input / output) per 1M. opd echo-scores completions
 # (max_tokens=0), so only INPUT tokens are billed (the teacher never generates) — but the table keeps
-# both so a mispriced entry is obvious. glm-5p1 shares the GLM-5 serverless rate.
+# both so a mispriced entry is obvious.
 TEACHER_USD_PER_1M: dict[str, tuple[float, float]] = {
     "accounts/fireworks/models/glm-5p2": (1.40, 4.40),
-    "accounts/fireworks/models/glm-5p1": (1.40, 4.40),
 }
 # Fireworks echo-scoring round-trip per completion (wall time, concurrency-bound like reward grading).
 AVG_TEACHER_SECONDS_PER_COMPLETION = 2.0
@@ -178,11 +177,9 @@ AVG_TEACHER_SECONDS_PER_COMPLETION = 2.0
 def teacher_price_per_1m(teacher_model: str) -> tuple[float, float]:
     """(input, output) $/1M tokens for a teacher model.
 
-    The empty-string sentinel (an opd run that OMITS [train] teacher_model) resolves to the recipe's
-    default teacher (``RECIPE.opd.teacher_model``) — the SINGLE source of "who is the default
-    teacher" — so the quote uses that model's real rate with no forked default-price constant. An
-    unknown non-default teacher falls back to the default teacher's rate (spec-parse already rejects
-    an unpriced teacher_model, so this is only a defensive backstop)."""
+    OPD uses the recipe's fixed GLM 5.2 teacher. ``teacher_model`` is accepted only for internal
+    diagnostics/tests; user configs cannot override it. Unknown values fall back defensively to
+    the fixed teacher rate."""
     from flash.engine.recipe import RECIPE
 
     default = TEACHER_USD_PER_1M[RECIPE.opd.teacher_model]

{freesolo_flash_dev-0.2.45 → freesolo_flash_dev-0.2.47}/flash/cost/spec.py

@@ -110,9 +110,6 @@ def runconfig_from_spec(spec) -> RunConfig:
         provider="auto",
         max_wall_seconds=g.max_wall_seconds,
         environment=spec.environment.id or None,
-        # "" is the unset sentinel (RunConfig.teacher_model); facts.teacher_price_per_1m resolves it
-        # to the recipe default. Non-opd carries no teacher.
-        teacher_model=t.teacher_model if spec.algorithm == "opd" else "",
     )
 
 

{freesolo_flash_dev-0.2.45 → freesolo_flash_dev-0.2.47}/flash/cost/types.py

@@ -27,9 +27,8 @@ class RunConfig:
     thinking: bool = False
     # GRPO only: seconds to score one completion. None -> the single average grader latency.
     reward_seconds_per_completion: float | None = None
-    # OPD only: the Fireworks teacher id, so the teacher-token quote uses the right per-model rate.
-    # "" (the unset sentinel, matching TrainSpec.teacher_model) resolves to the recipe's default
-    # teacher (RECIPE.opd.teacher_model) at the pricing boundary (facts.teacher_price_per_1m).
+    # OPD only: internal Fireworks teacher id for teacher-token diagnostics. User configs cannot
+    # override this; an empty value resolves to the recipe's fixed GLM 5.2 teacher.
     teacher_model: str = ""
 
     max_wall_seconds: int | None = None  # wall cap (spec gpu.max_wall_seconds); None = 24h
@@ -150,8 +149,9 @@ class CostEstimate:
     wall_capped: bool
     total_usd: float
     # opd only: external Fireworks GLM teacher token spend (0.0 for sft/grpo). Billed by Fireworks
-    # DIRECTLY to the user's FIREWORKS_API_KEY, so it is NOT part of total_usd — shown as its own
-    # itemized diagnostic line only.
+    # to the platform-managed teacher key (users don't supply one), tracked separately from the
+    # platform-billed GPU charge — so it is NOT part of total_usd; shown as its own itemized
+    # diagnostic line only.
     teacher_api_usd: float = 0.0
     notes: tuple[str, ...] = ()
 
@@ -181,8 +181,8 @@ class CostEstimate:
         ]
         if self.teacher_api_usd > 0:
             lines.append(
-                f"Teacher API: ${self.teacher_api_usd:.2f} (Fireworks GLM token spend on your "
-                "FIREWORKS_API_KEY — billed by Fireworks, NOT included in TOTAL)"
+                f"Teacher API: ${self.teacher_api_usd:.2f} (Fireworks GLM token spend on the "
+                "platform-managed teacher key — tracked separately, NOT included in TOTAL)"
             )
         lines.append(f"TOTAL      : ${self.total_usd:.2f}")
         if self.notes:

{freesolo_flash_dev-0.2.45 → freesolo_flash_dev-0.2.47}/flash/engine/multiturn_rollout.py

@@ -27,6 +27,26 @@ class RolloutResult(TypedDict):
     reward: float
 
 
+class TurnRecord(TypedDict):
+    """One assistant turn of a multi-turn episode, as :func:`rollout_one_records` emits it.
+
+    ``prefix_ids`` are the student token ids of the whole transcript BEFORE this turn (initial prompt
+    + every prior assistant turn + every inter-turn env "glue"), i.e. the exact on-policy context the
+    student sampled this turn's completion after. ``context_messages`` is the parallel message list at
+    the same point (prompt + prior turns + prior env replies), for callers that render a separate
+    teacher/scoring prompt from it. ``gen`` is whatever the injected ``generate`` callable returned for
+    this turn (opaque to the driver beyond the four attributes it reads: ``completion_ids``,
+    ``completion_text``, ``truncated``, ``skip``) — e.g. OPD hands back its ``_GenResult`` so the
+    per-turn record feeds straight into the existing single-turn scoring/loss path. Distilling each
+    turn against ``prefix_ids`` is the multi-turn on-policy-distillation objective: the episode's total
+    reverse-KL over student-generated tokens is the sum of per-turn reverse-KLs, each conditioned on the
+    transcript so far."""
+
+    prefix_ids: list[int]
+    gen: object
+    context_messages: list[dict]
+
+
 _ROLLOUT_FIELDS: tuple[str, ...] = (
     "prompt_ids",
     "completion_ids",
@@ -99,6 +119,20 @@ def index_collisions(rows: list[dict], prompt_of: Callable[[dict], object]) -> i
     return len(rows) - len({_prompt_key(prompt_of(r)) for r in rows})
 
 
+def _dedup_seam_terminator(prev_completion_ids: list[int], glue: list[int]) -> list[int]:
+    """Collapse a duplicate turn terminator at the turn/env seam.
+
+    ``env_glue`` leads with the assistant turn's terminator (e.g. ``<|im_end|>``), which a
+    naturally-stopped assistant turn also keeps as its final token — so the raw stream would carry two.
+    When the previous turn's last token equals ``glue[0]`` drop the glue copy, keeping the assistant's
+    own token (real logprob / model-generated) so the stream has exactly one terminator per turn,
+    matching the chat template + SFT transcripts. Shared by :func:`rollout_one`,
+    :func:`_advance_after_turn`, and :func:`rollout_one_records` so the three paths can't drift."""
+    if glue and prev_completion_ids and prev_completion_ids[-1] == glue[0]:
+        return glue[1:]
+    return glue
+
+
 def rollout_one(
     *,
     example: dict,
@@ -152,13 +186,7 @@ def rollout_one(
         if active_env.rollout_done(state, max_turns):
             break
 
-        glue = env_glue(env_msgs)
-        # Collapse a duplicate turn terminator at the seam: env_glue leads with the assistant turn's
-        # terminator (e.g. <|im_end|>), which vLLM also keeps as the last token of a naturally-stopped
-        # turn — keep the assistant's own (env_mask=1, real logprob) over the env copy. See
-        # _advance_after_turn (the rollout_async twin) for the same fix.
-        if glue and completion_ids and completion_ids[-1] == glue[0]:
-            glue = glue[1:]
+        glue = _dedup_seam_terminator(completion_ids, env_glue(env_msgs))
         if token_budget is not None and len(completion_ids) + len(glue) > token_budget:
             break
         completion_ids.extend(glue)
@@ -176,6 +204,102 @@ def rollout_one(
     }
 
 
+def rollout_one_records(
+    *,
+    example: dict,
+    active_env,
+    render: Callable[[list, bool], list[int]],
+    generate: Callable[[list, int], object],
+    env_glue: Callable[[list], list[int]],
+    max_turns: int,
+    per_turn_max_tokens: int,
+    engine_max_len: int | None = None,
+    on_turn_generated: Callable[[], None] | None = None,
+) -> list[TurnRecord]:
+    """Drive one multi-turn episode and return a per-turn :class:`TurnRecord` list (NOT a flat masked
+    sequence like :func:`rollout_one`). Used by OPD, which distils each assistant turn as an independent
+    single-turn sample conditioned on the transcript so far.
+
+    The env turn loop is identical to :func:`rollout_one` (``new_rollout_state`` → generate →
+    ``record_model_turn`` → ``rollout_done`` → ``env_reply`` → glue), and the tokenizer-sensitive seam
+    dedup / engine-budget accounting reuse the same helpers, so the two paths can't drift. The
+    differences are all OPD-shaped:
+
+    - ``generate(prefix_ids, max_new)`` returns an OBJECT (opaque here) exposing ``completion_ids``
+      (``list[int] | None``), ``completion_text`` (``str``), ``truncated`` (``bool``) and ``skip``
+      (``bool``) — OPD passes its ``_GenResult`` from ``_generate_one`` (termination/trim/U+FFFD gates
+      already applied), so the record drops straight into the existing teacher-scoring + gkd-loss path.
+    - Every turn (including a truncated/empty one) is recorded so the caller counts it; a
+      truncated/skip turn ENDS the episode (a student that didn't terminate its turn, or emitted an
+      empty/invalid completion, can't meaningfully continue — and its bad turn is skipped, not
+      distilled, by the caller's ``_resolve_sample``).
+    - ``on_turn_generated`` (optional) fires AFTER each turn's generation to refresh the worker stall
+      clock and advance the sample counter — a many-turn episode is a long serial stretch of GPU
+      generates + teacher-less env steps that would otherwise emit no progress ping.
+
+    ``prefix_ids`` and ``context_messages`` are snapshotted BEFORE generation, so each record carries
+    the exact student-token and message context the teacher must condition on for that turn.
+    """
+    state = active_env.new_rollout_state(example)
+    initial_messages = state.get("prompt") or state.get("messages")
+    if not isinstance(initial_messages, list):
+        raise KeyError("multi-turn rollout state must include prompt or messages")
+    messages = [dict(m) for m in initial_messages]
+    prompt_ids = render(messages, True)
+    cur_ids = list(prompt_ids)  # invariant: cur_ids == prompt_ids + completion tokens so far
+    token_budget = (engine_max_len - len(prompt_ids) - 8) if engine_max_len else None
+    records: list[TurnRecord] = []
+
+    turns = 0
+    while True:
+        completion_so_far = len(cur_ids) - len(prompt_ids)
+        max_new = per_turn_max_tokens
+        if token_budget is not None:
+            remaining = token_budget - completion_so_far
+            if remaining <= 0:
+                break
+            max_new = min(max_new, remaining)
+        # Snapshot the student-token prefix and message context BEFORE generation: this is what the
+        # teacher must condition on to score THIS turn (the transcript up to, but excluding, the turn).
+        prefix_ids = list(cur_ids)
+        context_messages = [dict(m) for m in messages]
+        gen = generate(prefix_ids, max(1, max_new))
+        if on_turn_generated is not None:
+            on_turn_generated()
+        records.append(
+            {"prefix_ids": prefix_ids, "gen": gen, "context_messages": context_messages}
+        )
+        text = getattr(gen, "completion_text", "") or ""
+        active_env.record_model_turn(state, text)
+        messages.append({"role": "assistant", "content": text})
+        turns += 1
+        # A turn that didn't terminate naturally (truncated) or produced no usable text (skip) ends the
+        # episode: it's recorded (counted) but not distilled, and continuing from a broken turn is
+        # pointless (the student can't end its turn / said nothing).
+        if getattr(gen, "truncated", False) or getattr(gen, "skip", False):
+            break
+        asst_ids = getattr(gen, "completion_ids", None) or []
+        cur_ids.extend(asst_ids)
+        completion_so_far = len(cur_ids) - len(prompt_ids)
+        if token_budget is not None and completion_so_far >= token_budget:
+            break
+        if turns >= max_turns or active_env.rollout_done(state, max_turns):
+            break
+        env_msgs = active_env.env_reply(messages, state)
+        if not env_msgs:
+            break
+        messages.extend(env_msgs)
+        # Don't append glue if the env step finished — no next model turn.
+        if active_env.rollout_done(state, max_turns):
+            break
+        glue = _dedup_seam_terminator(asst_ids, env_glue(env_msgs))
+        if token_budget is not None and completion_so_far + len(glue) > token_budget:
+            break
+        cur_ids.extend(glue)
+
+    return records
+
+
 class _RolloutState:
     """Mutable per-rollout accumulator for :func:`rollout_async` (mirrors :func:`rollout_one` locals)."""
 
@@ -248,13 +372,7 @@ def _advance_after_turn(
     if active_env.rollout_done(r.state, max_turns):
         r.done = True
         return
-    glue = env_glue(env_msgs)
-    # env_glue's text starts at the assistant turn's terminator (e.g. <|im_end|>); vLLM also keeps that
-    # terminator as the final token of a naturally-stopped turn (it's in asst_ids). Collapse the duplicate
-    # at the seam so the stream has exactly one terminator per turn (matches the chat template + SFT
-    # transcripts), keeping the assistant's own token (env_mask=1, real logprob) over the env copy.
-    if glue and r.completion_ids and r.completion_ids[-1] == glue[0]:
-        glue = glue[1:]
+    glue = _dedup_seam_terminator(r.completion_ids, env_glue(env_msgs))
     if r.budget is not None and len(r.completion_ids) + len(glue) > r.budget:
         r.done = True
         return
@@ -429,6 +547,47 @@ def render_message_ids(tok, messages, add_generation_prompt: bool, *, thinking:
     return [int(t) for t in tok(text, add_special_tokens=False).input_ids]
 
 
+def make_env_glue(tok, *, thinking: bool, cache_size: int = 8192) -> Callable[[list], list[int]]:
+    """Build the inter-turn "glue" tokenizer used by both the GRPO rollout (:func:`build_rollout_func`)
+    and OPD's per-turn rollout (:func:`rollout_one_records`).
+
+    Given the env's reply messages for a turn, returns the token ids that sit BETWEEN the assistant
+    turn and the next assistant generation prompt (the turn terminator + the env/observation messages +
+    the next assistant header). Computed with the probe trick — render ``[{assistant: PROBE}, *env]``
+    with ``add_generation_prompt=True`` and take everything after the probe — so the history is never
+    re-rendered (Qwen3's template doesn't round-trip a re-rendered transcript). Results are LRU-cached
+    by the env messages. Raises ``ValueError`` if the model's chat template doesn't insert assistant
+    content verbatim (token-aligned multi-turn is then unsupported for that model)."""
+    cache = _LRUCache(cache_size)
+    probe = "flash-env-glue-probe"
+
+    def env_glue(env_messages: list) -> list[int]:
+        cache_key = json.dumps(env_messages, sort_keys=True, default=str)
+        cached = cache.get(cache_key)
+        if cached is not None:
+            return cached
+        text = tok.apply_chat_template(
+            [{"role": "assistant", "content": probe}, *env_messages],
+            add_generation_prompt=True,
+            tokenize=False,
+            enable_thinking=thinking,
+        )
+        first = text.find(probe)
+        if first == -1 or text.find(probe, first + len(probe)) != -1:
+            raise ValueError(
+                "multi-turn env_glue could not uniquely locate its probe in the rendered chat "
+                "template; this model's template does not insert assistant content verbatim, so "
+                "token-aligned multi-turn rollout is unsupported for it (use a single-turn/tool "
+                "env or a different model)."
+            )
+        glue_text = text[first + len(probe) :]
+        glue = [int(t) for t in tok(glue_text, add_special_tokens=False).input_ids]
+        cache.put(cache_key, glue)
+        return glue
+
+    return env_glue
+
+
 def _engine_vocab_size(engine) -> int | None:
     """Best-effort vocab size from the colocate vLLM engine, or None. Never raises."""
     try:
@@ -482,34 +641,7 @@ def build_rollout_func(
         _render_cache.put(cache_key, ids)
         return ids
 
-    _glue_cache = _LRUCache(8192)
-
-    def env_glue(env_messages: list) -> list[int]:
-        cache_key = json.dumps(env_messages, sort_keys=True, default=str)
-        cached = _glue_cache.get(cache_key)
-        if cached is not None:
-            return cached
-        # Render a probe assistant turn + env messages, then take everything after the probe
-        # to get the inter-turn glue without re-rendering history (Qwen3 doesn't round-trip).
-        probe = "flash-env-glue-probe"
-        text = tok.apply_chat_template(
-            [{"role": "assistant", "content": probe}, *env_messages],
-            add_generation_prompt=True,
-            tokenize=False,
-            enable_thinking=thinking,
-        )
-        first = text.find(probe)
-        if first == -1 or text.find(probe, first + len(probe)) != -1:
-            raise ValueError(
-                "multi-turn env_glue could not uniquely locate its probe in the rendered chat "
-                "template; this model's template does not insert assistant content verbatim, so "
-                "token-aligned multi-turn rollout is unsupported for it (use a single-turn/tool "
-                "env or a different model)."
-            )
-        glue_text = text[first + len(probe) :]
-        glue = [int(t) for t in tok(glue_text, add_special_tokens=False).input_ids]
-        _glue_cache.put(cache_key, glue)
-        return glue
+    env_glue = make_env_glue(tok, thinking=thinking)
 
     def rollout_func(prompts, trainer):
         engine = trainer.vllm_generation.llm

{freesolo_flash_dev-0.2.45 → freesolo_flash_dev-0.2.47}/flash/engine/recipe.py

@@ -50,9 +50,8 @@ class OPDConfig:
     """On-policy distillation: student samples, a remote teacher scores its tokens, and a groupwise
     reverse-KL loss (the collinear-ai spider / Tinker cross-tokenizer method) trains the student."""
 
-    # Fireworks-hosted teacher reached over the OpenAI-compatible API; overridable via
-    # [train].teacher_model. GLM-5.2 is a strong reasoning model whose per-token logprobs
-    # supervise the (much smaller) student.
+    # Fixed Fireworks-hosted GLM 5.2 teacher reached over the OpenAI-compatible API. This is
+    # platform-managed and intentionally not configurable from [train].
     teacher_model: str = "accounts/fireworks/models/glm-5p2"
     teacher_base_url: str = "https://api.fireworks.ai/inference/v1"
     # OPD is step-driven like GRPO (on-policy sampling), not epoch-driven like SFT.

{freesolo_flash_dev-0.2.45 → freesolo_flash_dev-0.2.47}/flash/engine/worker/__init__.py

@@ -35,6 +35,7 @@ from flash.engine.worker.decoding import (
 from flash.engine.worker.finalize import write_train_meta
 from flash.engine.worker.gpu_setup import (
     finalize_alloc_conf_for_sleep,
+    force_vit_sdpa_on_blackwell,
     force_vllm_backend_for_sm120,
     patch_trl_colocate_llm_kwargs,
 )
@@ -54,6 +55,8 @@ from flash.engine.worker.heartbeat import (
     _HB_TERMINAL_ONLY_INTERVAL_S,
     _HB_TERMINAL_STAGES,
     _HB_THROTTLED_STAGES,
+    _HB_TIGHT_LIVENESS_STAGES,
+    _HB_UPLOAD_LIVENESS_STAGES,
     _HB_UPLOAD_LOCK,
     _SFT_HEARTBEAT_INTERVAL_S,
     _STEP_GPU_DIAG_INTERVAL_S,
@@ -371,6 +374,8 @@ __all__ = [
     "_HB_TERMINAL_ONLY_INTERVAL_S",
     "_HB_TERMINAL_STAGES",
     "_HB_THROTTLED_STAGES",
+    "_HB_TIGHT_LIVENESS_STAGES",
+    "_HB_UPLOAD_LIVENESS_STAGES",
     "_HB_UPLOAD_LOCK",
     "_SFT_HEARTBEAT_INTERVAL_S",
     "_STEP_GPU_DIAG_INTERVAL_S",
@@ -412,6 +417,7 @@ __all__ = [
     "disable_liger_grpo_torch_compile",
     "error_artifact_name",
     "finalize_alloc_conf_for_sleep",
+    "force_vit_sdpa_on_blackwell",
     "force_vllm_backend_for_sm120",
     "free_gpu",
     "fused_optim_name",