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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 0.2.44
+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
@@ -51,7 +51,8 @@ Description-Content-Type: text/markdown
 
 # Flash
 
-Managed LoRA post-training service: SFT and GRPO on managed Flash GPUs.
+Managed LoRA post-training service: SFT, GRPO, and on-policy distillation (from a
+Fireworks-hosted GLM teacher) on managed Flash GPUs.
 The allocator picks the cheapest validated managed GPU class that fits the run.
 
 ## Scope
@@ -79,8 +80,9 @@ The allocator picks the cheapest validated managed GPU class that fits the run.
 - `flash/providers/` — managed GPU substrate code (pricing, GPU classes, durable
   submit/poll, preflight) behind the `base.Provider` protocol, with an
   `allocator.py` that picks the cheapest fitting managed GPU class
-- `flash/engine/` — the on-GPU worker (TRL + colocated vLLM rollouts) and the
-  shared recipe; SFT targets and RL rewards route through the active environment
+- `flash/engine/` — the on-GPU worker (TRL + colocated vLLM rollouts; distillation
+  scores on-policy student samples against a remote GLM teacher) and the shared
+  recipe; SFT targets and RL rewards route through the active environment
   (task-specific grading lives with its example, not in the engine)
 - `flash/envs/` — environment machinery: registry and the adapter that loads
   Freesolo SDK environments onto the worker's interface

{freesolo_flash_dev-0.2.44 → freesolo_flash_dev-0.2.46}/README.md

@@ -1,6 +1,7 @@
 # Flash
 
-Managed LoRA post-training service: SFT and GRPO on managed Flash GPUs.
+Managed LoRA post-training service: SFT, GRPO, and on-policy distillation (from a
+Fireworks-hosted GLM teacher) on managed Flash GPUs.
 The allocator picks the cheapest validated managed GPU class that fits the run.
 
 ## Scope
@@ -28,8 +29,9 @@ The allocator picks the cheapest validated managed GPU class that fits the run.
 - `flash/providers/` — managed GPU substrate code (pricing, GPU classes, durable
   submit/poll, preflight) behind the `base.Provider` protocol, with an
   `allocator.py` that picks the cheapest fitting managed GPU class
-- `flash/engine/` — the on-GPU worker (TRL + colocated vLLM rollouts) and the
-  shared recipe; SFT targets and RL rewards route through the active environment
+- `flash/engine/` — the on-GPU worker (TRL + colocated vLLM rollouts; distillation
+  scores on-policy student samples against a remote GLM teacher) and the shared
+  recipe; SFT targets and RL rewards route through the active environment
   (task-specific grading lives with its example, not in the engine)
 - `flash/envs/` — environment machinery: registry and the adapter that loads
   Freesolo SDK environments onto the worker's interface

{freesolo_flash_dev-0.2.44 → freesolo_flash_dev-0.2.46}/flash/catalog.py

@@ -6,7 +6,24 @@ import math
 from dataclasses import asdict, dataclass
 from typing import Any
 
-ALGORITHMS = ("sft", "grpo")
+ALGORITHMS = ("sft", "grpo", "opd")
+
+# Algorithms whose training step SAMPLES on-policy student completions (a rollout): GRPO and OPD,
+# versus SFT which trains on a fixed dataset. Single source of truth for the "has rollouts" concept
+# the cost model, VRAM sizing and worker allocator branch on — import ``samples_on_policy`` instead
+# of hand-rolling the ("grpo", "opd") tuple at each site.
+_ON_POLICY_ALGORITHMS = frozenset({"grpo", "opd"})
+
+
+def samples_on_policy(algorithm: str) -> bool:
+    """True when a training step samples on-policy student completions (GRPO or OPD).
+
+    Tolerant of the ``rl`` phase-name alias for grpo (``JobSpec.phase``) so callers that pass a
+    phase rather than an algorithm resolve identically."""
+    algo = (algorithm or "").strip().lower()
+    if algo == "rl":  # phase-name alias for grpo
+        algo = "grpo"
+    return algo in _ON_POLICY_ALGORITHMS
 
 
 def normalize_algorithm(value: str) -> str:
@@ -145,7 +162,7 @@ MODELS: dict[str, ModelInfo] = {
         params="1.2B dense (Llama arch)",
         params_b=1.2,
         vocab_size=130_560,
-        algos=("sft", "grpo"),
+        algos=("sft", "grpo", "opd"),
         min_vram_gb=12,
         recommended_gpu="RTX 4090",
         serving=ServingCapacity(
@@ -164,7 +181,7 @@ MODELS: dict[str, ModelInfo] = {
         params="0.9B (text-only fine-tune)",
         params_b=0.9,
         vocab_size=248_320,
-        algos=("sft", "grpo"),
+        algos=("sft", "grpo", "opd"),
         min_vram_gb=12,
         recommended_gpu="RTX 4090",
         serving=ServingCapacity(
@@ -183,7 +200,7 @@ MODELS: dict[str, ModelInfo] = {
         params="2.3B (text-only fine-tune)",
         params_b=2.3,
         vocab_size=248_320,
-        algos=("sft", "grpo"),
+        algos=("sft", "grpo", "opd"),
         min_vram_gb=16,
         recommended_gpu="RTX 4090",
         serving=ServingCapacity(
@@ -201,7 +218,7 @@ MODELS: dict[str, ModelInfo] = {
         params="4.7B (text-only fine-tune)",
         params_b=4.7,
         vocab_size=248_320,
-        algos=("sft", "grpo"),
+        algos=("sft", "grpo", "opd"),
         min_vram_gb=32,
         recommended_gpu="RTX 5090",
         serving=ServingCapacity(
@@ -223,7 +240,7 @@ MODELS: dict[str, ModelInfo] = {
         params="9.7B (text-only fine-tune)",
         params_b=9.7,
         vocab_size=248_320,
-        algos=("sft", "grpo"),
+        algos=("sft", "grpo", "opd"),
         min_vram_gb=48,
         # NOT QLoRA: peft bnb merge during GRPO rollout diverges trainer precision -> TRL ratio collapses to 0.
         grpo_min_vram_gb=80,
@@ -255,7 +272,7 @@ MODELS: dict[str, ModelInfo] = {
         num_layers=40,
         hidden_size=2048,
         vocab_size=248_320,
-        algos=("sft", "grpo"),
+        algos=("sft", "grpo", "opd"),
         min_vram_gb=141,
         # Floor to 100 GB so SFT lands on H200, not the thin-margin consumer Blackwell or 80 GB H100.
         sft_min_vram_gb=100,
@@ -403,8 +420,9 @@ def _resolve_open_model(model_id: str, algo: str, gpu: str | None) -> ModelInfo:
 def validate_model_for_algorithm(model_id: str, algorithm: str) -> ModelInfo:
     info = get_model(model_id)
     algo = normalize_algorithm(algorithm)
-    required = "grpo" if algo == "grpo" else "sft"
-    if required not in info.algos:
+    # Each algorithm gates on its own capability entry (sft/grpo/opd), so a model must
+    # explicitly opt into opd via its algos tuple — same contract as sft/grpo.
+    if algo not in info.algos:
         allowed = ", ".join(info.algos)
         raise ValueError(f"{model_id} supports {allowed}, not {algo}")
     return info

{freesolo_flash_dev-0.2.44 → freesolo_flash_dev-0.2.46}/flash/cli/commands.py

@@ -210,11 +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.
 
-Both SFT and GRPO train off this file:
+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) 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
@@ -487,6 +490,33 @@ def cmd_env_setup(args) -> int:
             "# 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"
         )
+    opd = Path("configs/opd.toml")
+    if not opd.exists():
+        # 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 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\n\n'
+            "# Environment: upload this project folder with\n"
+            "# `flash env push --name my-env .`, then paste the returned id below.\n"
+            "# The GLM teacher key is platform-managed — nothing to set up or export.\n"
+            "[environment]\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"
+        )
     training = Path("TRAINING.md")
     if not training.exists():
         # Explicit UTF-8: TRAINING_MD has non-ASCII chars that raise UnicodeEncodeError under a non-UTF-8 locale.
@@ -496,6 +526,7 @@ def cmd_env_setup(args) -> int:
         "dataset/train.jsonl",
         "configs/sft.toml",
         "configs/rl.toml",
+        "configs/opd.toml",
         "TRAINING.md",
     ]
     if render.styled():
@@ -655,6 +686,13 @@ def _log_follow_progress(status: dict | None, fallback_state: str) -> tuple[str,
         step = heartbeat.get("step")
         if step is not None:
             parts.append(f"step={step}")
+        # live heartbeat age so a long quiet phase reads as "alive, throttled" not "frozen".
+        # minute granularity: the non-TTY follow path prints a line whenever this string changes,
+        # so a seconds-precision age would emit one line per poll.
+        ts = heartbeat.get("ts")
+        if isinstance(ts, (int, float)) and ts > 0:
+            mins = int(max(0.0, time.time() - ts) // 60)
+            parts.append(f"hb={mins}m" if mins else "hb=<1m")
     realized = status.get("realized_cost_usd")
     if realized is not None:
         if isinstance(realized, (int, float)):
@@ -800,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
@@ -817,7 +867,7 @@ def cmd_checkpoints(args) -> int:
     if not checkpoints:
         message = (
             f"no deployable checkpoints for {args.run_id} yet "
-            "(RL streams one per save interval; SFT-only runs have none)."
+            "(RL/opd stream one per save interval; SFT-only runs have none)."
         )
         if render.styled():
             print(render.empty("checkpoints", "0 deployable", message))

{freesolo_flash_dev-0.2.44 → freesolo_flash_dev-0.2.46}/flash/cli/render.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import os
 import sys
+import time
 
 from flash._channel import CLI_NAME
 
@@ -520,6 +521,47 @@ def _humanize_ts(value) -> str | None:
     return datetime.datetime.fromtimestamp(value, datetime.UTC).strftime("%Y-%m-%d %H:%M UTC")
 
 
+def _humanize_age(value) -> str | None:
+    """Format an epoch seconds value as a compact age ("3m ago"), None for non-numbers."""
+    if not isinstance(value, (int, float)) or value <= 0:
+        return None
+    secs = max(0.0, time.time() - value)
+    if secs < 90:
+        return f"{int(secs)}s ago"
+    if secs < 5400:
+        return f"{int(secs // 60)}m ago"
+    return f"{secs / 3600:.1f}h ago"
+
+
+# heartbeat age past which the panel reminds that quiet is normal: worker uploads are throttled
+# (240s quiet phases, up to 900s mid-training), so a frozen ts is usually NOT a dead worker.
+_HB_QUIET_HINT_AFTER_S = 300.0
+
+
+def _heartbeat_pairs(obj: dict) -> list[tuple[str, str]]:
+    """Worker heartbeat rows for the status panel: stage, step, age, and a quiet-is-normal hint."""
+    hb = obj.get("last_heartbeat")
+    if not isinstance(hb, dict) or not hb.get("stage"):
+        return []
+    worker = str(hb["stage"])
+    step = hb.get("step")
+    if step is not None:
+        worker += f" · step {step}"
+    if hb.get("liveness"):
+        worker += " · alive ping"
+    pairs = [("worker", worker)]
+    ts = hb.get("ts")
+    age = _humanize_age(ts)
+    if age:
+        running = str(obj.get("state") or "") == "running"
+        if running and isinstance(ts, (int, float)) and (time.time() - ts) > _HB_QUIET_HINT_AFTER_S:
+            age += _dim(
+                "  (heartbeat uploads are throttled; quiet is not dead - check flash log -f <run-id>)"
+            )
+        pairs.append(("heartbeat", age))
+    return pairs
+
+
 def run_status(obj: dict) -> str:
     """A curated status panel for `flash status`, with the full JSON below for completeness."""
     spec = obj.get("spec") or {}
@@ -534,6 +576,7 @@ def run_status(obj: dict) -> str:
     realized = obj.get("realized_cost_usd")
     if realized is not None:
         pairs.append(("realized", money(realized)))
+    pairs += _heartbeat_pairs(obj)
     pairs += [
         ("created", _humanize_ts(obj.get("created_at"))),
         ("updated", _humanize_ts(obj.get("updated_at"))),
@@ -661,6 +704,16 @@ def cost_panel(est) -> str:
         ("wall clock", f"{est.wall_clock_hours:.2f} h"),
         ("billable", f"{est.billable_hours:.2f} h  {_dim('(training only)')}"),
     ]
+    # opd teacher spend is itemized but NOT part of total_usd (billed by Fireworks on the user's key).
+    # Mirror CostEstimate.breakdown() so the styled panel doesn't silently drop it (cursor[bot]).
+    if getattr(est, "teacher_api_usd", 0) > 0:
+        pairs.append(
+            (
+                "teacher api",
+                f"{money(est.teacher_api_usd, 2)}  "
+                f"{_dim('(Fireworks GLM token spend on your key — billed by Fireworks, NOT in TOTAL)')}",
+            )
+        )
     panel = _kv(pairs)
     total = f"  {_paint('TOTAL'.ljust(10), _GRAY, '1')} {_paint(_glyph('·', '-'), _FAINT)} {_paint(f'${est.total_usd:.2f}', _TEAL, '1')}"
     out = f"{header('train', 'pre-flight cost estimate')}\n{panel}\n{_rule()}\n{total}"

{freesolo_flash_dev-0.2.44 → freesolo_flash_dev-0.2.46}/flash/cli/training_doc.py

@@ -19,8 +19,8 @@ actually produced, and decide — honestly — whether the model got better.
 ## Using Flash
 
 Flash is a **managed** training service with a thin CLI/client. You author an
-environment (the task + its reward), publish it, and submit SFT or GRPO runs from a
-TOML config. Flash allocates the cheapest fitting managed GPU class, runs the job,
+environment (the task + its reward), publish it, and submit SFT, GRPO, or on-policy
+**distillation** runs from a TOML config. Flash allocates the cheapest fitting managed GPU class, runs the job,
 streams logs back, and serves the result. You never handle infrastructure credentials —
 you authenticate once with a freesolo API key, and everything below is a `flash` CLI
 command.
@@ -42,6 +42,7 @@ environment.py          # the task: how to prompt the model and how to score it
 dataset/train.jsonl     # training rows, one JSON object per line: {"input": ..., "output": ...}
 configs/sft.toml        # an SFT run config
 configs/rl.toml         # a GRPO (RL) run config
+configs/opd.toml        # an on-policy distillation run config
 TRAINING.md             # this file
 ```
 
@@ -92,7 +93,7 @@ edit to `environment.py` or `dataset/` so the managed run uses your change.
 
 ```toml
 model = "Qwen/Qwen3.5-4B"   # see `flash models`
-algorithm = "sft"           # "sft" (supervised) or "grpo" (RL)
+algorithm = "sft"           # "sft" (supervised), "grpo" (RL), or "opd" (on-policy distillation)
 # thinking = true           # opt-in reasoning mode, for models that support it
 
 [environment]
@@ -373,6 +374,51 @@ 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
+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.
+
+- **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.
+
+```toml
+# configs/opd.toml
+model = "Qwen/Qwen3.5-4B"
+algorithm = "opd"
+
+[environment]
+id = "your-org/my-env"
+
+[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
+```
+
+---
+
 ## GRPO knobs that matter
 
 Set these in `[train]`. Each is `None` by default — the worker's tuned recipe fills

{freesolo_flash_dev-0.2.44 → freesolo_flash_dev-0.2.46}/flash/client/runtime_secrets.py

@@ -5,6 +5,13 @@ from __future__ import annotations
 import os
 from pathlib import Path
 
+# The global runtime-secret allow-list: keys collected/gated/injected for EVERY run regardless of
+# algorithm. Only WANDB_API_KEY belongs here (optional experiment logging, useful to all arms). The
+# opd teacher key (FIREWORKS_API_KEY) is NOT global: opd auto-declares it as a required environment
+# secret in spec_from_dict, and each gate unions the spec's declared secrets on top of this set
+# (client collection below; server create gate server/_deps.py; worker-env injection
+# providers/runpod/train/deps.py). Keeping it out of the default means SFT/GRPO workers never
+# receive the teacher key they don't need.
 DEFAULT_RUNTIME_SECRET_KEYS = frozenset({"WANDB_API_KEY"})
 
 

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

@@ -20,6 +20,8 @@ from .facts import (
     model_quant,
     pick_gpu,
     reward_seconds_per_completion,
+    teacher_seconds_per_completion,
+    teacher_token_cost_usd,
     total_params_b,
 )
 from .types import CostEstimate, RunConfig
@@ -28,6 +30,9 @@ from .types import CostEstimate, RunConfig
 SFT_FLOPS_PER_TOKEN_PER_PARAM = 6.0  # forward (2) + backward (4)
 GRPO_GEN_FLOPS_PER_TOKEN_PER_PARAM = 2.0  # autoregressive rollout forward
 GRPO_UPDATE_FLOPS_PER_TOKEN_PER_PARAM = 8.0  # policy fwd+bwd (6) + frozen-ref fwd (2)
+# OPD's update is policy fwd+bwd ONLY (6): the teacher's per-token logprobs come from the
+# Fireworks API, so there is NO local frozen-reference forward (GRPO's extra 2).
+OPD_UPDATE_FLOPS_PER_TOKEN_PER_PARAM = 6.0
 
 # Model-FLOPs utilization (fraction of peak sustained), calibrated against real RunPod
 # wall clock. LoRA + small batches sit well below dense-pretraining MFU.
@@ -39,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
@@ -77,6 +87,14 @@ def setup_seconds(config: RunConfig) -> float:
     return s
 
 
+def _opd_step_shape(n: RunConfig) -> tuple[int, int]:
+    """(completions per step, prompt+completion tokens per step) for one OPD step, from a NORMALIZED
+    config. completions = batch x group; each is billed over the FULL n.seq_len (prompt+completion,
+    not completion-only) since the loss forward runs model(prompt_ids + student_ids)."""
+    completions = n.batch_size * n.group_size
+    return completions, completions * n.seq_len
+
+
 def seconds_per_step(config: RunConfig, gpu: str) -> float:
     """Steady-state wall time for one optimizer step on ``gpu``."""
     n = config.normalized()
@@ -85,6 +103,22 @@ def seconds_per_step(config: RunConfig, gpu: str) -> float:
     params = active_params_b(n.model_id) * 1e9
     peak = gpu_tflops(gpu) * 1e12  # FLOP/s
 
+    if n.is_opd:
+        # 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.
+        completions, seq_tokens = _opd_step_shape(n)
+        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'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:
         flops = SFT_FLOPS_PER_TOKEN_PER_PARAM * params * (n.batch_size * n.seq_len)
         return flops / (peak * MFU_SFT_TRAIN)
@@ -95,7 +129,9 @@ def seconds_per_step(config: RunConfig, gpu: str) -> float:
     gen_s = (GRPO_GEN_FLOPS_PER_TOKEN_PER_PARAM * params * gen_tokens) / (peak * MFU_DECODE)
     update_s = (GRPO_UPDATE_FLOPS_PER_TOKEN_PER_PARAM * params * gen_tokens) / (peak * MFU_TRAIN)
     latency = reward_seconds_per_completion(n.reward_seconds_per_completion)
-    reward_s = math.ceil(completions / REWARD_CONCURRENCY) * latency  # ceil: a partial wave still costs one latency
+    reward_s = (
+        math.ceil(completions / REWARD_CONCURRENCY) * latency
+    )  # ceil: a partial wave still costs one latency
     return gen_s + reward_s + update_s
 
 def sft_seconds_for_tokens(config: RunConfig, gpu: str, train_tokens: float) -> float:
@@ -114,7 +150,9 @@ def select_gpu(config: RunConfig, *, max_wall_seconds: float = 0.0) -> tuple[str
     it considers every fitting class, validated or not — so the estimate reflects the cheapest
     card that *could* run the job. The live allocator restricts to the validated pool, so the
     actually-provisioned class can be pricier than this. Catalog sizing is offline/deterministic."""
-    total_params_b(config.model_id)  # catalog-only: reject a non-catalog model before any (HF) sizing
+    total_params_b(
+        config.model_id
+    )  # catalog-only: reject a non-catalog model before any (HF) sizing
     need = required_vram_gb(
         config.model_id,
         config.method,
@@ -125,12 +163,22 @@ def select_gpu(config: RunConfig, *, max_wall_seconds: float = 0.0) -> tuple[str
     return gpu, need
 
 
-def _notes(config: RunConfig, raw_train_s: float, wall_capped: bool, cap_s: float) -> tuple[str, ...]:
+def _notes(
+    config: RunConfig, raw_train_s: float, wall_capped: bool, cap_s: float
+) -> tuple[str, ...]:
     n = config.normalized()
     notes: list[str] = []
     if (quant := model_quant(n.model_id)) != "bf16":
         notes.append(f"{quant}: smaller VRAM footprint -> cheaper GPU class fits")
-    if n.is_grpo:
+    if n.is_opd:
+        comps, _ = _opd_step_shape(n)
+        tsec = teacher_seconds_per_completion()
+        notes.append(
+            f"opd step = student rollout of {n.batch_size}x{n.group_size}={comps} completions "
+            f"@ {n.completion_len} tok + GLM teacher scoring ({tsec:.2f}s/completion) + policy "
+            "update (no local reference forward)"
+        )
+    elif n.is_grpo:
         comps = n.batch_size * n.group_size
         rsec = reward_seconds_per_completion(n.reward_seconds_per_completion)
         notes.append(
@@ -187,6 +235,18 @@ def estimate_cost(config: RunConfig, *, wall_cap_s: float = DEFAULT_WALL_CAP_S)
     train = max(0.0, cap_s - setup) if wall_capped else raw_train
     wall = setup + train
 
+    # OPD: add the external Fireworks teacher token spend. The teacher echo-scores every
+    # sampled completion (input ~ prompt+completion per completion), so bill INPUT tokens over the
+    # EFFECTIVE (wall-capped) step count — not the uncapped `steps` — so a wall-capped run's teacher
+    # bill tracks the GPU time it is actually billed for.
+    teacher_api_usd = 0.0
+    if config.is_opd:
+        n = config.normalized()
+        effective_steps = (train / sps) if sps > 0 else config.steps
+        _, tokens_per_step = _opd_step_shape(n)
+        teacher_input_tokens = effective_steps * tokens_per_step
+        teacher_api_usd = teacher_token_cost_usd(teacher_input_tokens, 0.0, config.teacher_model)
+
     return CostEstimate(
         model_id=config.model_id,
         method=config.method,
@@ -201,6 +261,11 @@ 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 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.44 → freesolo_flash_dev-0.2.46}/flash/cost/facts.py

@@ -159,3 +159,46 @@ def reward_seconds_per_completion(override: float | None = None) -> float:
     if override is not None:
         return max(0.0, override)
     return AVG_REWARD_SECONDS_PER_COMPLETION
+
+
+# Fireworks serverless pricing for the on-policy-distillation GLM teacher, $/1M tokens as
+# (input, output). Static like gpu_hourly_usd so cost estimation is deterministic/offline and needs
+# 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.
+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
+
+
+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)."""
+    from flash.engine.recipe import RECIPE
+
+    default = TEACHER_USD_PER_1M[RECIPE.opd.teacher_model]
+    return TEACHER_USD_PER_1M.get(teacher_model or RECIPE.opd.teacher_model, default)
+
+
+def teacher_token_cost_usd(
+    input_tokens: float, output_tokens: float = 0.0, teacher_model: str = ""
+) -> float:
+    """External teacher-API dollar cost for a token count. Deterministic; no network."""
+    inp, outp = teacher_price_per_1m(teacher_model)
+    return (max(0.0, input_tokens) * inp + max(0.0, output_tokens) * outp) / 1_000_000.0
+
+
+def teacher_seconds_per_completion(override: float | None = None) -> float:
+    """Per-completion teacher-scoring latency (s): the explicit override, else the average."""
+    if override is not None:
+        return max(0.0, override)
+    return AVG_TEACHER_SECONDS_PER_COMPLETION