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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 0.2.45
+Version: 0.2.46
 Summary: Flash — managed LoRA post-training (SFT/GRPO) 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.45 → freesolo_flash_dev-0.2.46}/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,28 +492,24 @@ 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 a Fireworks GLM 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 teacher key is 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"
@@ -842,13 +838,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.46}/flash/cli/training_doc.py

@@ -382,15 +382,22 @@ each of *its own* completions, and a dense per-token loss teaches the student to
 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 to set up.** The Fireworks key for the GLM teacher is 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 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,7 +409,6 @@ algorithm = "opd"
 
 [environment]
 id = "your-org/my-env"
-secrets = ["FIREWORKS_API_KEY"]
 
 [train]
 steps = 100

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

@@ -44,6 +44,11 @@ MFU_DECODE = 0.12  # batched vLLM rollout (decode is memory-bandwidth-bound)
 # wall is ceil(completions / slots) waves x latency, not completions x latency.
 REWARD_CONCURRENCY = 16.0
 
+# Teacher scoring is CONCURRENT too: run_opd's primary path fans a step's completions across up to
+# TEACHER_CONCURRENCY Fireworks calls (mirror of opd.py _TEACHER_SCORE_MAX_WORKERS), so the teacher
+# wall is ceil(completions / slots) waves x latency, not completions x latency.
+TEACHER_CONCURRENCY = 8.0
+
 # Cold-start overhead (seconds): container boot + deps + model load (+ vLLM init for GRPO).
 #
 # Calibrated against a real fresh-worker run (0.8B SFT, RTX 3090 @ $0.239/hr) whose elapsed wall
@@ -99,7 +104,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 +112,11 @@ 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, bounded by
+        # TEACHER_CONCURRENCY workers (opd.py _TEACHER_SCORE_MAX_WORKERS), so the teacher wall is
+        # ceil(completions / slots) waves x latency — NOT the full serial sum (that describes only the
+        # CPU-test fallback that can't batch-generate). ceil keeps a partial last wave at one latency.
+        teacher_s = math.ceil(completions / TEACHER_CONCURRENCY) * teacher_lat
         return gen_s + teacher_s + update_s
 
     if not n.is_grpo:
@@ -255,10 +261,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.46}/flash/cost/types.py

@@ -150,8 +150,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 +182,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.46}/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.46}/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",

{freesolo_flash_dev-0.2.45 → freesolo_flash_dev-0.2.46}/flash/engine/worker/gpu_setup.py

@@ -49,6 +49,45 @@ def force_vllm_backend_for_sm120() -> str | None:
     return backend
 
 
+def force_vit_sdpa_on_blackwell() -> bool:
+    """Force the VISION-encoder (ViT) attention backend to TORCH_SDPA on Blackwell (sm100/sm120).
+
+    Qwen3.6-35B-A3B is a VL model, so vLLM builds its vision tower even for a text-only GRPO rollout
+    and runs a ViT attention during engine init/profiling. On Blackwell vLLM 0.19.1 routes the ViT to
+    its CUTE flash-attn (``vit_flash_attn_wrapper`` -> ``flash_attn_varlen_func`` ->
+    ``vllm.vllm_flash_attn.cute``), which is UNIMPORTABLE against every published ``nvidia-cutlass-dsl``:
+    the vendored cute references ``cutlass.cute.core.ThrMma`` (present only <=4.5.x) while
+    ``vit_attn_wrappers`` also imports ``cutlass._mlir_helpers`` (present only >=4.6.0) — the two symbols
+    never coexist, so the first ViT attention call aborts with
+    ``AttributeError: module 'cutlass.cute.core' has no attribute 'ThrMma'`` (4.6.0) or
+    ``ModuleNotFoundError: cutlass._mlir_helpers`` (<=4.5.2), crashing EVERY B200 GRPO rollout (a
+    version pin cannot fix it — measured 2026-07-07). ``get_vit_attn_backend`` honors
+    ``MultiModalConfig.mm_encoder_attn_backend`` UNCONDITIONALLY, and TORCH_SDPA is a supported ViT
+    backend on cc>=8.0, so pinning it sidesteps the CUTE import entirely (the LM/decoder attention is
+    unaffected — that path is FLASHINFER/flash-attn, chosen separately).
+
+    * Injected as a colocate ``LLM(...)`` kwarg (``EngineArgs.mm_encoder_attn_backend``, str-typed),
+      composing with the other overrides; must run BEFORE ``GRPOTrainer.__init__`` builds the engine.
+    * No-op off Blackwell and on non-multimodal models (a text-only model simply never builds a ViT).
+    Returns True iff the override was injected."""
+    try:
+        import torch
+
+        if not torch.cuda.is_available() or torch.cuda.get_device_capability(0)[0] not in (10, 12):
+            return False
+    except Exception as e:
+        print("[rl] ViT-SDPA Blackwell probe skipped:", e)
+        return False
+    ok = patch_trl_colocate_llm_kwargs(mm_encoder_attn_backend="TORCH_SDPA")
+    if ok:
+        print(
+            "[rl] Blackwell (sm100/sm120): mm_encoder_attn_backend=TORCH_SDPA on the colocate rollout "
+            "engine (vLLM 0.19.1 ViT CUTE flash-attn is unimportable vs every nvidia-cutlass-dsl: "
+            "cute.core.ThrMma <=4.5.x XOR cutlass._mlir_helpers >=4.6.0 -> crashes the B200 rollout)"
+        )
+    return ok
+
+
 def finalize_alloc_conf_for_sleep() -> None:
     """Sync PYTORCH_ALLOC_CONF with the resolved GRPO sleep mode (RL only).
 
@@ -76,6 +115,7 @@ def patch_trl_colocate_llm_kwargs(
     max_num_batched_tokens: int | None = None,
     enforce_eager: bool | None = None,
     attention_backend: str | None = None,
+    mm_encoder_attn_backend: str | None = None,
 ) -> bool:
     """Inject vLLM ``LLM(...)`` kwargs into TRL's colocated rollout engine that can't be expressed via
     ``GRPOConfig`` or the (pinned) vLLM env registry — must run BEFORE ``GRPOTrainer.__init__`` builds
@@ -105,6 +145,10 @@ def patch_trl_colocate_llm_kwargs(
     * ``attention_backend="FLASHINFER"|"TRITON_ATTN"``: pin a PTX-independent decoder attention backend
       (consumer-Blackwell sm120) — the SUPPORTED replacement for the removed ``VLLM_ATTENTION_BACKEND``
       env. ``EngineArgs`` coerces the bare member name through ``AttentionConfig.validate_backend_before``.
+    * ``mm_encoder_attn_backend="TORCH_SDPA"``: force the VISION-encoder (ViT) attention backend for a
+      multimodal model (Qwen3.6-35B-A3B is VL) off vLLM's CUTE flash-attn path — see
+      ``force_vit_sdpa_on_blackwell``. Maps to ``MultiModalConfig.mm_encoder_attn_backend``, which
+      ``get_vit_attn_backend`` honors unconditionally.
 
     Repeated calls COMPOSE: run_rl injects the attention backend, the KV/prefill knobs, and eager mode
     at three SEPARATE points, so each call merges its kwargs into one accumulated module-level override
@@ -120,6 +164,8 @@ def patch_trl_colocate_llm_kwargs(
         new_overrides["enforce_eager"] = bool(enforce_eager)
     if attention_backend is not None:
         new_overrides["attention_backend"] = attention_backend
+    if mm_encoder_attn_backend is not None:
+        new_overrides["mm_encoder_attn_backend"] = mm_encoder_attn_backend
     if not new_overrides:
         return False
     try:

{freesolo_flash_dev-0.2.45 → freesolo_flash_dev-0.2.46}/flash/engine/worker/heartbeat.py

@@ -38,12 +38,21 @@ _HB_SETUP_LIVENESS_STAGES = frozenset(
         "rl_finalizing",
     }
 )
+# Mid-training the per-step checkpoint upload runs SYNCHRONOUSLY on the trainer thread (dev #445),
+# which freezes the trainer's global_step — so the ``rl_step``/``sft_step`` liveness daemon can only
+# emit bare liveness pings for the whole upload, and those DON'T advance the provider's stall clock.
+# A ``checkpoint_uploading`` keepalive daemon (liveness_heartbeat(keepalive=True)) wraps the upload to
+# keep that clock fed; it rides the SAME tight, throttled cadence as a setup stage. Kept OUT of
+# _HB_SETUP_LIVENESS_STAGES so that set's "setup-phase" meaning (and its tests) stay honest.
+_HB_UPLOAD_LIVENESS_STAGES = frozenset({"checkpoint_uploading"})
+# Liveness stages that ride the tighter setup-liveness upload interval (setup + mid-train upload).
+_HB_TIGHT_LIVENESS_STAGES = _HB_SETUP_LIVENESS_STAGES | _HB_UPLOAD_LIVENESS_STAGES
 # Throttled to avoid blowing the 128/hr HF commit cap; terminal transitions are never throttled. Every
-# setup-liveness stage is throttled (⊂) PLUS the per-step training stages: opd_filtering_prompts alone
+# tight-liveness stage is throttled (⊂) PLUS the per-step training stages: opd_filtering_prompts alone
 # emits a REAL (non-liveness) heartbeat every scan tick — ~120/hr on a large split before model load —
 # so unthrottled the setup stages blow the cap; throttle them exactly like their sft_pretokenizing
-# analogue (codex[bot]).
-_HB_THROTTLED_STAGES = _HB_SETUP_LIVENESS_STAGES | frozenset({"rl_step", "sft_step", "opd_step"})
+# analogue (codex[bot]). checkpoint_uploading keepalive re-emits every 30s too, so it MUST be throttled.
+_HB_THROTTLED_STAGES = _HB_TIGHT_LIVENESS_STAGES | frozenset({"rl_step", "sft_step", "opd_step"})
 _HB_TERMINAL_STAGES = frozenset({"done", "already_done"})
 # 600s -> ~6 commits/hr; keeps stall detector alive without hitting the HF commit cap.
 _HB_TERMINAL_ONLY_INTERVAL_S = 600.0
@@ -137,7 +146,7 @@ def heartbeat(stage: str, *, liveness: bool = False, force: bool = False, **kw):
         else:
             throttled = stage in _HB_THROTTLED_STAGES
             interval_s = _w._HB_MIN_INTERVAL_S
-            if stage in _HB_SETUP_LIVENESS_STAGES:
+            if stage in _HB_TIGHT_LIVENESS_STAGES:
                 interval_s = min(interval_s, _w._HB_SETUP_LIVENESS_INTERVAL_S)
             upload_due = not throttled or (now - _w._HB_LAST_UPLOAD) >= interval_s
             # ``force`` bypasses the per-stage throttle (but not TERMINAL_ONLY mode, handled above) when
@@ -291,9 +300,18 @@ _STALL_DUMP_S = 1200.0
 
 
 @contextlib.contextmanager
-def liveness_heartbeat(stage, progress=None, fields=None, progress_step=False):
+def liveness_heartbeat(stage, progress=None, fields=None, progress_step=False, keepalive=False):
     """Emit liveness pings for ``stage`` while the wrapped block runs on the main thread.
 
+    ``keepalive``: the wrapped block is legitimate BLOCKING I/O with no per-step progress signal the
+    daemon can observe — a synchronous checkpoint/adapter upload (dev #445) or the finalize upload —
+    so EVERY tick emits a REAL (non-liveness) heartbeat instead of a bare ping. Bare liveness pings do
+    NOT advance the provider's stall clock (`_poll.surface_heartbeat` returns stage=None for them), so
+    without this a healthy multi-minute upload that outlasts STALL_AFTER_S (1500s) is wrongly killed
+    mid-save. Safe because a genuinely wedged upload surfaces as an EXCEPTION through its own retry
+    budget (ending this context), not an infinite silent hang — so it does not mask a real stall. Pair
+    with a throttled stage (see _HB_UPLOAD_LIVENESS_STAGES) so the 30s re-emit can't blow the HF cap.
+
     ``progress``: optional ``() -> float | None`` monotonic counter; advances emit a REAL heartbeat.
     ``fields``: optional ``() -> dict`` of EXTRA payload fields merged into every emission (liveness
     and progress alike). Use it to carry the billing/stall ``step`` on a stage the poller step-gates:
@@ -342,7 +360,9 @@ def liveness_heartbeat(stage, progress=None, fields=None, progress_step=False):
                     extra = fields() or {}
             if progress_step and last_val is not None:
                 extra["step"] = int(last_val)
-            _w.heartbeat(stage, liveness=not made_progress, gpu=gpu, **extra)
+            # keepalive: a legitimate blocking upload IS progress the daemon can't sample per-step, so
+            # force a REAL heartbeat every tick to keep the provider's stall clock fed (see docstring).
+            _w.heartbeat(stage, liveness=(not made_progress) and not keepalive, gpu=gpu, **extra)
             last_progress = float(getattr(_w, "_HB_LAST_PROGRESS_TS", 0.0) or 0.0)
             if not dumped and last_progress and (time.time() - last_progress) > _STALL_DUMP_S:
                 _dump_thread_stacks(f"{stage}: no progress for >{_STALL_DUMP_S:.0f}s")