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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 0.2.44
+Version: 0.2.45
 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.45}/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.45}/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.45}/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:
+SFT and GRPO train off this file; OPD (configs/opd.toml) does NOT support multi-turn:
 - 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`).
 """
 
 from __future__ import annotations
@@ -487,6 +490,37 @@ 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) 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_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"
+            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'
+            "# 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"
+            "[environment]\n"
+            'id = ""\n'
+            'secrets = ["FIREWORKS_API_KEY"]\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 +530,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 +690,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)):
@@ -817,7 +859,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.45}/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.45}/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,45 @@ 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.
+
+- **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.
+- **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.
+- **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"
+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
+```
+
+---
+
 ## 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.45}/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.45}/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.
@@ -77,6 +82,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 +98,21 @@ 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 (serial
+        # 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 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
+        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 +123,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 +144,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 +157,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 +229,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 +255,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 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]).
         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.45}/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

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

@@ -7,6 +7,7 @@ from __future__ import annotations
 
 import math
 
+from flash.catalog import samples_on_policy
 from flash.cost.analytical import estimate_cost
 from flash.cost.types import CostEstimate, RunConfig
 
@@ -77,6 +78,12 @@ def spec_steps(spec) -> int:
         if t.steps is not None:
             return max(1, int(t.steps))
         return RECIPE.rl.num_steps
+    if spec.algorithm == "opd":
+        # Step-driven like GRPO (on-policy sampling), NOT example-driven — so a opd quote never
+        # demands [train].max_examples (the SFT fallback below would raise).
+        if t.steps is not None:
+            return max(1, int(t.steps))
+        return RECIPE.opd.num_steps
     # max_examples is a CAP; 0 (like None) means "no cap" (worker trains the full dataset), so
     # don't let max_examples=0 price a single step.
     return _sft_steps_from_examples(spec, _sft_example_count(spec), apply_cap=True)
@@ -87,20 +94,25 @@ def runconfig_from_spec(spec) -> RunConfig:
     estimate covers a single job. The estimate doesn't pin a GPU -- it does its own cheapest-fit
     (provider="auto")."""
     t, g = spec.train, spec.gpu
-    is_grpo = spec.algorithm == "grpo"
+    # Both grpo and opd sample on-policy student completions, so both carry the rollout
+    # dimensions (completion length + group size) into the cost model.
+    has_rollout = samples_on_policy(spec.algorithm)
     return RunConfig(
         model_id=spec.model,
         method=spec.algorithm,
         steps=spec_steps(spec),
         seq_len=t.max_length,
-        completion_len=t.max_tokens if is_grpo else None,
+        completion_len=t.max_tokens if has_rollout else None,
         batch_size=t.batch_size,
-        group_size=t.group_size if is_grpo else None,
+        group_size=t.group_size if has_rollout else None,
         lora_rank=t.lora_rank,
         thinking=spec.thinking,
         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 "",
     )