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

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

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

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

@@ -447,11 +447,11 @@ def cmd_env_setup(args) -> int:
         '# secrets = ["SERPAPI_API_KEY"]\n\n'
     )
     # `thinking = true` opts the run into reasoning mode. Reasoning shares the generation budget with
-    # the answer, so GRPO also gets a raised max_tokens. These strings are empty when reasoning is
-    # off, keeping the default scaffold byte-for-byte identical.
+    # the answer, so GRPO also gets a raised completion budget. These strings are empty when reasoning
+    # is off, keeping the default scaffold byte-for-byte identical.
     thinking_line = "thinking = true\n" if reasoning else ""
     rl_reasoning_train = (
-        "max_tokens = 2048  # reasoning shares this budget with the answer; raised so it isn't truncated\n"
+        "max_completion_tokens = 2048  # reasoning shares this budget with the answer; raised so it isn't truncated\n"
         if reasoning
         else ""
     )
@@ -504,16 +504,15 @@ def cmd_env_setup(args) -> int:
         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'
+            'algorithm = "opd"   # on-policy distillation from the managed GLM 5.2 teacher\n\n'
             "# Environment: upload this project folder with\n"
             "# `flash env push --name my-env .`, then paste the returned id below.\n"
-            "# The GLM teacher key is platform-managed — nothing to set up or export.\n"
+            "# The GLM 5.2 teacher and key are 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"
         )

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

@@ -201,15 +201,15 @@ spending another GPU run:
 | --- | --- | --- |
 | Environment id is blank or stale | `flash train --dry-run` fails, or the worker uses old reward/data | Run `flash env push --name my-env .` after every environment/data edit and paste the returned id into every config you submit. |
 | Local-only env path in config | Config validation says there is no local path mode | Publish first, then use the returned slug in `[environment] id`. `flash train` only runs published env ids, not local paths. |
-| Config knobs are in the wrong table | Validation rejects `[grpo]`, `[sft]`, or unknown `[train]` keys | Put `steps`, `epochs`, `group_size`, `max_tokens`, `temperature`, `max_length`, LoRA, and other training knobs under `[train]`. |
+| Config knobs are in the wrong table | Validation rejects `[grpo]`, `[sft]`, or unknown `[train]` keys | Put `steps`, `epochs`, `group_size`, `max_completion_tokens`, `temperature`, `max_context_tokens`, LoRA, and other training knobs under `[train]`. |
 | Trying to pin managed infrastructure | `gpu.type`, `train.hf_repo`, or `model_policy` changes do not do what you expected | Treat GPU choice, model policy, and the run artifact repo as managed. Tune the model, algorithm, environment, and `[train]` knobs instead. |
 | Secrets are not available on the worker | Reward code works locally but remote logs show missing API keys or auth failures | List secret names under `[environment] secrets = [...]`, export those env vars locally before submit, or put them in local `.env` / `.env.local`. Never put secret values in `[worker_env]` or hard-code them in the config. |
 | Wrong model / thinking setting | Config validation fails, or chat behavior does not match the run | Use `flash models`; set `thinking = true` only for supported models. Thinking is a training-time/run-level choice and serving preserves that parity, so `flash chat` does not expose an override flag. |
 | Thinking reward grades the wrong text | Rewards accidentally score hidden reasoning, or ignore reasoning you meant to inspect | By default, score the answer text. In thinking mode the response object is still string-compatible, but also exposes `.completion`, `.thinking`, and `.raw` when a reward intentionally needs those fields. |
 | All-zero or flat GRPO reward | `reward_mean` stays near 0 and outputs do not improve | Make the reward dense: give partial credit for parse/format/execution/correctness tiers, and log a separate clean `success` metric. Do not keep rerunning an all-zero reward. |
 | Reward rises but behavior is worse | Short, templated, malformed, or reward-hacked outputs score well | Deploy the adapter and probe real examples. Add hard validity gates before judge calls, penalize degenerate shortcuts, and judge the outcome rather than the surface string. |
-| Output is truncated | Correct-looking answers cut off mid-response or JSON is incomplete | Increase `max_tokens` for GRPO rollouts or `max_length` for SFT only after seeing truncation. Oversizing them by default just burns memory/cost. |
-| Infrastructure, CUDA, OOM, vLLM, or kernel failure | Run errors before useful metrics, often during setup/model load | Treat this as infrastructure pressure, not proof the model is too large. Read `flash log <run-id>`, reduce footprint (`max_length`, `max_tokens`, `group_size`) if needed, and let Flash retry/allocate another fitting GPU class. |
+| Output is truncated | Correct-looking answers cut off mid-response or JSON is incomplete | Increase `max_completion_tokens` for GRPO/OPD rollouts or `max_context_tokens` for total prompt+completion context only after seeing truncation. Oversizing them by default just burns memory/cost. |
+| Infrastructure, CUDA, OOM, vLLM, or kernel failure | Run errors before useful metrics, often during setup/model load | Treat this as infrastructure pressure, not proof the model is too large. Read `flash log <run-id>`, reduce footprint (`max_context_tokens`, `max_completion_tokens`, `group_size`) if needed, and let Flash retry/allocate another fitting GPU class. |
 | Run looks stuck after disconnecting | Terminal stopped streaming but the job may still be alive | Ctrl-C detaches. Use `flash log <run-id> --follow` to reattach, `flash log <run-id>` for the console/error output, or `flash cancel <run-id>` if you intentionally want to stop it. |
 | Final checkpoint regresses | Last step is worse than an earlier checkpoint | Run `flash checkpoints <run-id>`, deploy a specific step with `flash deploy <run-id>/step-N`, and compare with held-out probes before exporting or relying on the final adapter. |
 | Export fails before upload | CLI says no HuggingFace token | Pass `flash export --api-key hf_...`, or set `HF_TOKEN` in your shell, `.env`, or `.env.local`. Exports are private unless you pass `--public`. |
@@ -335,8 +335,8 @@ Pick SFT when you already have good answers and want the model to imitate them.
   split the run never trains on, then **deploy the adapter and score it on that split**
   (`flash deploy` + `flash chat`). If held-out quality stalls or drops while train loss
   keeps falling, reduce `epochs` or add more data — not more passes.
-- **Start `max_length` small and grow it on evidence.** Begin from the smallest
-  `max_length` that plausibly fits prompt + completion, and only raise it when you see
+- **Start `max_context_tokens` small and grow it on evidence.** Begin from the smallest
+  `max_context_tokens` that plausibly fits prompt + completion, and only raise it when you see
   truncation (outputs cut off mid-thought, degraded loss). A bigger context just costs
   more.
 - **For Qwen3.5 thinking multi-turn SFT, put reasoning only in the final assistant
@@ -377,15 +377,16 @@ SFT is **epoch-driven** (`epochs`); GRPO is **step-driven** (`steps`).
 ## On-policy distillation (`algorithm = "opd"`)
 
 Pick distillation when a much stronger **teacher** model can grade your student's work
-token-by-token. The student samples on-policy (like GRPO), a Fireworks-hosted GLM teacher scores
+token-by-token. The student samples on-policy (like GRPO), the managed GLM 5.2 teacher scores
 each of *its own* completions, and a dense per-token loss teaches the student to match the teacher —
 far more sample-efficient than reward-based RL and with no reward to design. It is **step-driven**
 (`steps`) and produces a LoRA served exactly like SFT.
 
-- **No teacher key to set up.** The Fireworks key for the GLM teacher is platform-managed: the
+- **No teacher key or model override to set up.** The GLM 5.2 teacher and its Fireworks key are platform-managed: the
   service supplies its own key to every opd run, so there is nothing to export or declare — an opd
-  run submits like any other. (Bring-your-own teacher 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.
+  run submits like any other. Bring-your-own teacher models and keys are not supported; a
+  `FIREWORKS_API_KEY` in your shell is simply ignored. The key is never stored in the spec or
+  needed at serving time.
 - **The student (Qwen / MiniCPM / Kimi) and the teacher (GLM) have different tokenizers.** Flash
   bridges the vocabulary mismatch with **groupwise reverse-KL** (the collinear-ai *spider* / Tinker
   method): it aligns the two tokenizations by shared decoded-text spans and applies per-span reverse
@@ -396,7 +397,7 @@ far more sample-efficient than reward-based RL and with no reward to design. It
   (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
+  never distilled (they're context, not the student's output). Set `[train] max_context_tokens` 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
@@ -413,7 +414,6 @@ 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
 ```
 
@@ -427,7 +427,7 @@ in a sensible value, so only override with a reason.
 | Knob | Convention |
 | --- | --- |
 | `group_size` | Completions sampled per prompt (default 8). More = more signal and more cost; drop to 4 to trim cost. The group needs *within-group variance* for an advantage to exist. |
-| `max_tokens` | Completion budget per rollout. Size it to the expected output length — too small silently truncates good answers and poisons the reward; too large just costs more. |
+| `max_completion_tokens` | Completion budget per rollout. Size it to the expected output length — too small silently truncates good answers and poisons the reward; too large just costs more. |
 | `temperature` | Rollout sampling temperature. Keep it near 1.0 for GRPO — too low collapses diversity (and the model can collapse within a few steps); raise it to widen exploration against uniform-reward groups. |
 | `kl_penalty_coef` | Keeps the trained model from drifting too far from the base. Raise it to anchor against entropy collapse; lower it for more freedom to move. |
 | `thinking_length_penalty_coef` | Per-reasoning-token reward deduction — curb overthinking, but watch it doesn't push the model into terse degeneracy. |
@@ -468,7 +468,7 @@ targeted fix rather than leaning on the reward gate to slowly select against it.
 | --- | --- | --- |
 | Repetition / looping collapse | the same phrase repeats until truncation | repetition or length penalty; lower `temperature` |
 | Overthinking / verbose reasoning | reasoning eats the whole token budget | `thinking_length_penalty_coef`; tighten the prompt |
-| Max-token truncation | answers cut off mid-thought | raise `max_tokens` / `max_length` |
+| Completion truncation | answers cut off mid-thought | raise `max_completion_tokens` / `max_context_tokens` |
 | Unparsed / over-escaped output | reward can't read the answer | robust parser; return `0.0` on parse fail; format gate |
 | Wrapper / markdown around structured output | prose around the JSON/answer | a format gate; `stop_sequences` |
 | Uniform-reward groups | every rollout in a group scores the same → no gradient | shape the reward for partial credit; raise `temperature` |
@@ -513,7 +513,7 @@ on a beyond-noise improvement.
 ## Treat crashes as infra, not model size
 
 > A CUDA / OOM / vLLM / kernel / infrastructure error is an **infrastructure** problem, not a
-> sign the model is too big. Lower `max_length`, `max_tokens`, or `group_size` to shrink
+> sign the model is too big. Lower `max_context_tokens`, `max_completion_tokens`, or `group_size` to shrink
 > the run's footprint and let the allocator retry onto the next fitting GPU class — do
 > **not** switch to a smaller model to make a crash disappear. That silently destroys
 > quality.

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

@@ -44,11 +44,6 @@ 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
@@ -78,11 +73,11 @@ def _fmt_duration(seconds: float) -> str:
 
 def setup_seconds(config: RunConfig) -> float:
     """Cold-start wall time before the first optimizer step: container boot + deps + model load
-    (a fixed deserialize/placement/init base + a size-scaled download), plus vLLM init for GRPO.
+    (a fixed deserialize/placement/init base + a size-scaled download), plus vLLM init for rollouts.
     This elapsed setup time is reported but not included in customer-facing cost."""
     model_load = MODEL_LOAD_BASE_S + download_weight_gb(config.model_id) / DOWNLOAD_RATE_GBPS
     s = WORKER_BOOT_S + DEPS_INSTALL_S + model_load
-    if config.is_grpo:
+    if config.has_rollout:
         s += VLLM_INIT_S
     return s
 
@@ -112,11 +107,12 @@ def seconds_per_step(config: RunConfig, gpu: str) -> float:
         gen_s = (GRPO_GEN_FLOPS_PER_TOKEN_PER_PARAM * params * seq_tokens) / (peak * MFU_DECODE)
         update_s = (OPD_UPDATE_FLOPS_PER_TOKEN_PER_PARAM * params * seq_tokens) / (peak * MFU_TRAIN)
         teacher_lat = teacher_seconds_per_completion()
-        # run_opd'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
+        # run_opd's primary path scores a step's completions CONCURRENTLY over Fireworks with a fan-out
+        # cap of the step's OWN completion count (prompts_per_step * group_size, opd.py Phase 2), so
+        # every completion in a step is scored in ONE parallel wave — the teacher wall is a single
+        # latency, NOT the full serial sum (that describes only the CPU-test fallback that can't
+        # batch-generate). The teacher endpoint's rate limit is the real ceiling on this fan-out.
+        teacher_s = teacher_lat if completions else 0.0
         return gen_s + teacher_s + update_s
 
     if not n.is_grpo:

{freesolo_flash_dev-0.2.46 → freesolo_flash_dev-0.2.48}/flash/cost/facts.py

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

{freesolo_flash_dev-0.2.46 → freesolo_flash_dev-0.2.48}/flash/cost/spec.py

@@ -24,8 +24,8 @@ def _sft_seq_len(spec) -> int:
 
     t = spec.train
     return (
-        int(t.max_length)
-        if t.max_length is not None
+        int(t.max_context_tokens)
+        if t.max_context_tokens is not None
         else (RECIPE.sft.max_seq_len_thinking if spec.thinking else RECIPE.sft.max_seq_len)
     )
 
@@ -101,8 +101,8 @@ def runconfig_from_spec(spec) -> RunConfig:
         model_id=spec.model,
         method=spec.algorithm,
         steps=spec_steps(spec),
-        seq_len=t.max_length,
-        completion_len=t.max_tokens if has_rollout else None,
+        seq_len=t.max_context_tokens,
+        completion_len=t.max_completion_tokens if has_rollout else None,
         batch_size=t.batch_size,
         group_size=t.group_size if has_rollout else None,
         lora_rank=t.lora_rank,
@@ -110,9 +110,6 @@ def runconfig_from_spec(spec) -> RunConfig:
         provider="auto",
         max_wall_seconds=g.max_wall_seconds,
         environment=spec.environment.id or None,
-        # "" is the unset sentinel (RunConfig.teacher_model); facts.teacher_price_per_1m resolves it
-        # to the recipe default. Non-opd carries no teacher.
-        teacher_model=t.teacher_model if spec.algorithm == "opd" else "",
     )
 
 

{freesolo_flash_dev-0.2.46 → freesolo_flash_dev-0.2.48}/flash/cost/types.py

@@ -17,19 +17,18 @@ class RunConfig:
     method: str  # "sft" | "grpo" | "opd"
     steps: int
 
-    # Engine context length (forwarded as [train].max_length, NOT prompt length). When unset the
+    # Engine context length (forwarded as [train].max_context_tokens, NOT prompt length). When unset the
     # GRPO default mirrors the worker's max(1024, max_prompt_len + completion); see normalized().
     seq_len: int | None = None
-    completion_len: int | None = None  # GRPO only (max_tokens)
+    completion_len: int | None = None  # GRPO/OPD only (max_completion_tokens)
     batch_size: int | None = None  # SFT effective batch / GRPO prompts_per_step
     group_size: int | None = None  # GRPO completions per prompt (G)
     lora_rank: int | None = None
     thinking: bool = False
     # GRPO only: seconds to score one completion. None -> the single average grader latency.
     reward_seconds_per_completion: float | None = None
-    # OPD only: the Fireworks teacher id, so the teacher-token quote uses the right per-model rate.
-    # "" (the unset sentinel, matching TrainSpec.teacher_model) resolves to the recipe's default
-    # teacher (RECIPE.opd.teacher_model) at the pricing boundary (facts.teacher_price_per_1m).
+    # OPD only: internal Fireworks teacher id for teacher-token diagnostics. User configs cannot
+    # override this; an empty value resolves to the recipe's fixed GLM 5.2 teacher.
     teacher_model: str = ""
 
     max_wall_seconds: int | None = None  # wall cap (spec gpu.max_wall_seconds); None = 24h
@@ -83,7 +82,7 @@ class RunConfig:
             comp = self.completion_len
             if comp is None:
                 comp = rc.max_completion_len_thinking if self.thinking else rc.max_completion_len
-            # Explicit pin wins; else mirror the allocator's rollout sizing of an unset max_length:
+            # Explicit pin wins; else mirror the allocator's rollout sizing of an unset context:
             # max(1024, max_prompt_len + completion), not bare max_prompt_len (which under-sizes).
             seq = (
                 self.seq_len
@@ -117,9 +116,9 @@ class RunConfig:
         if self.batch_size is not None:
             knobs["batch_size"] = self.batch_size
         if n.seq_len is not None:
-            knobs["max_length"] = n.seq_len
+            knobs["max_context_tokens"] = n.seq_len
         if n.completion_len is not None:
-            knobs["max_tokens"] = n.completion_len
+            knobs["max_completion_tokens"] = n.completion_len
         if n.group_size is not None:
             knobs["group_size"] = n.group_size
         return knobs

{freesolo_flash_dev-0.2.46 → freesolo_flash_dev-0.2.48}/flash/engine/recipe.py

@@ -36,7 +36,7 @@ class RLConfig:
     learning_rate: float = 1e-5
     max_prompt_len: int = 2048
     max_completion_len: int = 320
-    # 320 is too short for <think> blocks; overridable via [train].max_tokens.
+    # 320 is too short for <think> blocks; overridable via [train].max_completion_tokens.
     max_completion_len_thinking: int = 1536
     prompts_per_step: int = 64
     group_size: int = 8
@@ -50,9 +50,8 @@ class OPDConfig:
     """On-policy distillation: student samples, a remote teacher scores its tokens, and a groupwise
     reverse-KL loss (the collinear-ai spider / Tinker cross-tokenizer method) trains the student."""
 
-    # Fireworks-hosted teacher reached over the OpenAI-compatible API; overridable via
-    # [train].teacher_model. GLM-5.2 is a strong reasoning model whose per-token logprobs
-    # supervise the (much smaller) student.
+    # Fixed Fireworks-hosted GLM 5.2 teacher reached over the OpenAI-compatible API. This is
+    # platform-managed and intentionally not configurable from [train].
     teacher_model: str = "accounts/fireworks/models/glm-5p2"
     teacher_base_url: str = "https://api.fireworks.ai/inference/v1"
     # OPD is step-driven like GRPO (on-policy sampling), not epoch-driven like SFT.
@@ -60,7 +59,7 @@ class OPDConfig:
     learning_rate: float = 1e-5
     max_prompt_len: int = 1024
     max_completion_len: int = 512
-    # 512 is short for <think> traces; overridable via [train].max_tokens.
+    # 512 is short for <think> traces; overridable via [train].max_completion_tokens.
     max_completion_len_thinking: int = 1536
     prompts_per_step: int = 8
     # Student samples per prompt. 1 is enough for a direct KD loss (no group-relative baseline

{freesolo_flash_dev-0.2.46 → freesolo_flash_dev-0.2.48}/flash/engine/vram.py

@@ -96,6 +96,20 @@ def opd_rollout_seq_len(
     return int(max_length or max(1024, RECIPE.opd.max_prompt_len + completion))
 
 
+def opd_rollout_concurrency(prompts_per_step: int = 1, group_size: int = 1) -> int:
+    """Concurrent student generations in one OPD vLLM rollout batch."""
+
+    def _positive_int(value: object, default: int) -> int:
+        try:
+            if isinstance(value, bool):
+                return default
+            return max(1, int(value))
+        except (TypeError, ValueError):
+            return default
+
+    return _positive_int(prompts_per_step, 1) * _positive_int(group_size, 1)
+
+
 def _resident_kv_gb(
     params_b: float | None, vllm_max_len: int, num_generations: int = 8, fp8_kv: bool = False
 ) -> float:
@@ -188,6 +202,12 @@ def colocate_kv_util(
 _TRAIN_COEF = 0.27
 # Small-model colocated GRPO floor: 0.8B OOMs 20 GB; 2B OOMs 24 GB -> both need 32 GB tier.
 _VLLM_COLOCATE_FLOOR_GB = 28.0
+# OPD builds its resident vLLM engine after the HF/PEFT student is already loaded. A real
+# Qwen3.5-2B OPD run with vLLM failed vLLM startup on a 32 GB RTX 5090 despite the raw equation
+# estimating ~28 GB with headroom, because vLLM requires the requested executor budget to be free at
+# init time. Keep 2B+ OPD off <=40 GB classes so the allocator picks an 80 GB-class card instead of a
+# consumer GPU that passes the aggregate estimate but fails the vLLM free-memory preflight.
+_OPD_VLLM_COLOCATE_FLOOR_GB = 41.0
 _VOCAB_DEFAULT = 248_320
 _LOGITS_BUDGET_GB = 6.0
 # 16 B/elem: fp32 logits+grad + bf16 logits+grad + CE temp. 8 B/elem under-counts (live OOM confirmed).
@@ -328,6 +348,14 @@ def estimate_vram_gb(
         train = activations + logits
         return base + (max(rollout, train) if sleep_offload else rollout + train)
     if is_opd:
+        # OPD always samples the student through a resident colocated vLLM engine, while the HF/PEFT
+        # trainer model stays resident for the GKD loss forward/backward. Budget the second bf16
+        # weight copy plus a viable KV pool through the training step. ``use_vllm`` is intentionally
+        # ignored for OPD; it only exists for legacy GRPO sizing callers.
+        rollout_concurrency = opd_rollout_concurrency(batch_size, group_size)
+        rollout = weights + max(
+            _KV_CAP, _resident_kv_gb(eff_b, seq_len, rollout_concurrency, fp8_kv=fp8_kv)
+        )
         # opd's gkd loss forward materializes DENSE logits — there is NO fused cross-entropy: the
         # student forward yields full-sequence bf16 logits [seq, vocab], then the completion rows are
         # gathered in fp32 [completion, vocab] for the logsumexp. The SFT fused path budgets ZERO
@@ -357,8 +385,11 @@ def estimate_vram_gb(
         logits_fwd = (seq_len * 2 + completion * 4) * vocab
         logits_bwd = (seq_len * 2 + completion * 4) * vocab
         logits = (logits_fwd + logits_bwd) / 1e9
-        return base + activations + logits
-    fused = sft_logits_fused(params_b, seq_len) if sft_fused_ce is None else bool(sft_fused_ce)
+        return base + rollout + activations + logits
+    # Actual TRL SFT keeps fused CE disabled (see worker/sft.py), so dense logits materialize even
+    # for long-context / >=3B models. Callers can pass sft_fused_ce=True only for theoretical
+    # comparisons; the default must mirror the worker.
+    fused = False if sft_fused_ce is None else bool(sft_fused_ce)
     pd = sft_per_device(batch_size, seq_len=seq_len, vocab=vocab, fused=fused)
     activations = _ACT_COEF * pd * (seq_len / 1024.0) * width
     # Don't clamp to budget: pd=1 is irreducible and the logits can exceed the budget at near-2048 ctx.
@@ -516,7 +547,7 @@ def model_required_vram_gb(
         except (TypeError, ValueError):
             return default
 
-    max_tokens = _pos_int(_g(train, "max_tokens"), None)
+    max_tokens = _pos_int(_g(train, "max_completion_tokens"), None)
     _algo = (algorithm or "").lower()
     if _algo in ("grpo", "rl"):
         _default_len = grpo_rollout_seq_len(0, max_tokens, thinking)
@@ -525,10 +556,18 @@ def model_required_vram_gb(
         _default_len = opd_rollout_seq_len(0, max_tokens, thinking)
     else:
         _default_len = 1024
-    seq_len = _pos_int(_g(train, "max_length"), _default_len)
+    seq_len = _pos_int(_g(train, "max_context_tokens"), _default_len)
     lora_rank = _pos_int(_g(train, "lora_rank"), 32)
-    group_size = _pos_int(_g(train, "group_size"), 8)
-    batch_size = _pos_int(_g(train, "batch_size"), _sft_per_device_bs())
+    if _algo == "opd":
+        from flash.engine.recipe import RECIPE
+
+        batch_size_default = int(RECIPE.opd.prompts_per_step)
+        group_size_default = int(RECIPE.opd.group_size)
+    else:
+        batch_size_default = _sft_per_device_bs()
+        group_size_default = 8
+    group_size = _pos_int(_g(train, "group_size"), group_size_default)
+    batch_size = _pos_int(_g(train, "batch_size"), batch_size_default)
 
     def _need(
         params_b: float,
@@ -557,17 +596,21 @@ def model_required_vram_gb(
         return math.ceil(est * headroom)
 
     from flash.catalog import MODELS, vocab_size_for
-    from flash.engine.chalk_kernels import chalk_supports_fused_ce_model
 
     info = MODELS.get(model_id)
     model_vocab = vocab_size_for(model_id)
-    is_grpo = (algorithm or "").lower() in ("grpo", "rl")
-    sft_fused_ce = (
-        None
-        if is_grpo
-        else sft_logits_fused(resolve_params_b(model_id), seq_len)
-        and chalk_supports_fused_ce_model(model_id)
+    is_grpo = _algo in ("grpo", "rl")
+    is_opd = _algo == "opd"
+    is_vllm_rollout = is_grpo or is_opd
+    vllm_concurrency = (
+        opd_rollout_concurrency(batch_size, group_size) if is_opd else group_size
     )
+    # Worker parity: TRL SFT deliberately calls install_chalk_kernels(..., fused_ce=False), because
+    # SFTTrainer.compute_loss reads outputs.logits and crashes when fused CE returns logits=None. So
+    # every SFT run materializes dense logits, even for long context / >=3B models where the old
+    # allocator assumed a fused CE path. Size SFT with fused CE OFF up front or long-context Qwen SFT
+    # routes to consumer cards and OOMs on the first backward.
+    sft_fused_ce = None if is_grpo else False
     if info is not None:
         params_b = info.params_b  # curated, authoritative (required field) — no string parsing
         quant = getattr(info, "quant", "bf16") or "bf16"
@@ -620,29 +663,39 @@ def model_required_vram_gb(
             else:
                 floor += grpo_seq_escalation_gb(active_b or params_b, seq_len)
         need = max(need, floor)
-        if is_grpo and use_vllm:
+        if is_vllm_rollout and use_vllm:
             floor_gb = 24 if (params_b or 0.0) <= 1.0 else int(_VLLM_COLOCATE_FLOOR_GB)
+            if is_opd and (params_b or 0.0) >= 2.0:
+                floor_gb = max(floor_gb, int(_OPD_VLLM_COLOCATE_FLOOR_GB))
             need = max(need, floor_gb)
             # vLLM KV-cache init preflight: the card must leave a viable cache-block pool
             # under the colocate utilization cap, or the engine dies at init ("No available
             # memory for the cache blocks") on a card the training-peak estimate accepted.
             need = max(
                 need,
-                grpo_kv_floor_gb(params_b or 4.0, seq_len, group_size, active_params_b=active_b),
+                grpo_kv_floor_gb(
+                    params_b or 4.0,
+                    seq_len,
+                    vllm_concurrency,
+                    active_params_b=active_b,
+                ),
             )
         return need
     params_b = fetch_hf_params_b(model_id)
     if params_b is None:
         return 24
     # Size the uncataloged (model_policy="allow") fallback with the ACTUAL algorithm, not a hardcoded
-    # "grpo". An open-model OPD run was quoted/allocated as a colocated-vLLM GRPO job, so the new OPD
-    # dense-logit estimator (estimate_vram_gb's is_opd path) was never applied -- rejecting fitting runs
-    # or routing them to pricier GPUs. The cataloged path above already threads `algorithm` through
-    # _need; do the same here (codex[bot]). The grpo-only escalations below stay gated on is_grpo.
+    # "grpo". The cataloged path above already threads `algorithm` through _need; do the same here so
+    # open-model OPD uses its own dense-logit + colocated-vLLM estimator. The grpo-only sequence
+    # escalation stays gated on is_grpo.
     need = _need(params_b, algorithm, vocab=model_vocab)
     if is_grpo:
         need += grpo_seq_escalation_gb(params_b, seq_len)
-        need = max(need, grpo_kv_floor_gb(params_b, seq_len, group_size))
+    if is_vllm_rollout:
+        floor_gb = 24 if params_b <= 1.0 else int(_VLLM_COLOCATE_FLOOR_GB)
+        if is_opd and params_b >= 2.0:
+            floor_gb = max(floor_gb, int(_OPD_VLLM_COLOCATE_FLOOR_GB))
+        need = max(need, floor_gb, grpo_kv_floor_gb(params_b, seq_len, vllm_concurrency))
     return need
 
 

{freesolo_flash_dev-0.2.46 → freesolo_flash_dev-0.2.48}/flash/engine/worker/grpo.py

@@ -137,7 +137,7 @@ def grpo_overrides() -> dict:
     cfg = {
         "group_size": train.group_size,
         "temperature": train.temperature,
-        "max_tokens": train.max_tokens,
+        "max_tokens": train.max_completion_tokens,
         "kl_penalty_coef": train.kl_penalty_coef,
         "advantage_clip": train.advantage_clip,
         "thinking_length_penalty_coef": train.thinking_length_penalty_coef,
@@ -173,7 +173,7 @@ def resolve_grpo_sleep_mode() -> tuple[bool, int, float, bool]:
     spec = _w.JOB_SPEC
     train = spec.train if spec else None
     model_id = spec.model if spec else ""
-    ctx = int(train.max_length if train and train.max_length else 0)
+    ctx = int(train.max_context_tokens if train and train.max_context_tokens else 0)
     gcfg = _w.grpo_overrides()
     group_size = int(gcfg.get("group_size") or RECIPE.rl.group_size)
     lora_rank = int(train.lora_rank) if train and train.lora_rank else 32