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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 0.2.47
+Version: 0.2.49
 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

{freesolo_flash_dev-0.2.47 → freesolo_flash_dev-0.2.49}/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 ""
     )
@@ -469,7 +469,8 @@ def cmd_env_setup(args) -> int:
             "\n"
             f"{env_comment}"
             "[train]\n"
-            "steps = 150\n"
+            "epochs = 1\n"
+            "max_examples = 2  # rows to train on; the starter dataset has 2 (raise as your dataset grows)\n"
             f"{rl_reasoning_train}"
             "lora_rank = 32\n"
             "# GPU and HF artifacts are managed automatically by the platform: the GPU is\n"
@@ -511,7 +512,8 @@ def cmd_env_setup(args) -> int:
             "[environment]\n"
             'id = ""\n\n'
             "[train]\n"
-            "steps = 100                     # opd is step-driven (like GRPO)\n"
+            "epochs = 1\n"
+            "max_examples = 2  # rows to train on; the starter dataset has 2 (raise as your dataset grows)\n"
             "lora_rank = 32\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.47 → freesolo_flash_dev-0.2.49}/flash/cli/training_doc.py

@@ -103,7 +103,7 @@ id = "your-org/my-env"      # the id printed by `flash env push`
                                    # never stored in the spec
 
 [train]
-epochs = 1                  # SFT is epoch-driven; GRPO is step-driven (steps = N)
+epochs = 1                  # one pass over the retained train rows
 max_examples = 2            # rows to train on (the starter dataset has 2)
 lora_rank = 32
 lora_alpha = 64
@@ -114,7 +114,7 @@ GPU and HF artifacts are **fully managed** — do not pick `gpu.type` or set
 `train.hf_repo`; the allocator picks the cheapest validated managed GPU class that fits,
 and run artifacts are stored in a private environment-scoped repo with content-addressed
 Flash code snapshots. Compose or tweak configs without editing files: `--config
-extra.toml` (deep-merge) and `--set key=value` (e.g. `--set train.steps=300`).
+extra.toml` (deep-merge) and `--set key=value` (e.g. `--set train.epochs=3`).
 
 ### 4. Submit
 
@@ -201,20 +201,20 @@ 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 `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`. |
 | SFT loss improves but quality does not | Train loss falls while held-out behavior stalls or degrades | Keep a held-out split outside training. Deploy and score that split; if quality drops, reduce epochs or improve data instead of adding more passes. |
-| Cost surprises | A quick experiment uses more GPU time than intended | Start with `--dry-run` and `--cost`, cap steps/epochs for smoke tests, and scale only after reward/data wiring is proven. Setup time is reported for observability; customer cost is based on training-loop GPU time. |
+| Cost surprises | A quick experiment uses more GPU time than intended | Start with `--dry-run` and `--cost`, keep `epochs` and `max_examples` small for smoke tests, and scale only after reward/data wiring is proven. Setup time is reported for observability; customer cost is based on training-loop GPU time. |
 
 ---
 
@@ -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
@@ -370,7 +370,9 @@ lora_rank = 16     # for VL warm-starts, SFT rank + GRPO rank must fit the effec
 lora_alpha = 32
 ```
 
-SFT is **epoch-driven** (`epochs`); GRPO is **step-driven** (`steps`).
+SFT, GRPO, and OPD all accept **epoch-driven** configs (`epochs`). For GRPO/OPD,
+an epoch is one pass over the retained prompt pool after `max_examples` and prompt-budget filtering;
+optimizer-step counts are derived from those epochs.
 
 ---
 
@@ -379,8 +381,8 @@ SFT is **epoch-driven** (`epochs`); GRPO is **step-driven** (`steps`).
 Pick distillation when a much stronger **teacher** model can grade your student's work
 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.
+far more sample-efficient than reward-based RL and with no reward to design. It supports `epochs`
+like SFT/GRPO and produces a LoRA served exactly like SFT.
 
 - **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
@@ -397,7 +399,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
@@ -412,7 +414,8 @@ algorithm = "opd"
 id = "your-org/my-env"
 
 [train]
-steps = 100
+epochs = 1
+max_examples = 2
 lora_rank = 32
 # kl_penalty_coef = 1.0                                 # reverse-KL scale
 ```
@@ -427,7 +430,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. |
@@ -449,7 +452,7 @@ rollouts score somewhere between all-fail and all-pass — is where GRPO has the
 signal.
 
 - If nearly every prompt is solved (most groups score ~1.0): **increase difficulty** —
-  harder prompts, tighter format/reward, more steps.
+  harder prompts, tighter format/reward, more epochs, or more data.
 - If nearly nothing is solved (most groups score ~0.0): **decrease difficulty** —
   easier or few-shot prompts, a more lenient (denser) reward, or warm-start with SFT.
 - In between: good signal — keep iterating at this difficulty.
@@ -468,7 +471,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` |
@@ -500,8 +503,8 @@ on a beyond-noise improvement.
 
 ## Scale the evidence
 
-- **A smoke test is not proof.** A single-digit `steps`, a tiny dataset, or a handful
-  of rollouts only validates the wiring. Scale `steps` / `epochs`, the dataset size,
+- **A smoke test is not proof.** A single-digit derived step count, a tiny dataset, or a handful
+  of rollouts only validates the wiring. Scale `epochs`, the dataset size,
   and `group_size` to the model and the data you actually have before you trust a
   result. Don't cite budget alone as the reason for an underpowered run.
 - **Use the data you have.** Deliberately assign every usable row to training or to a
@@ -513,7 +516,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.47 → freesolo_flash_dev-0.2.49}/flash/cost/analytical.py

@@ -73,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
 

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

@@ -10,6 +10,7 @@ import math
 from flash.catalog import samples_on_policy
 from flash.cost.analytical import estimate_cost
 from flash.cost.types import CostEstimate, RunConfig
+from flash.engine.steps import on_policy_steps
 
 
 def _sft_epochs(spec) -> int:
@@ -19,13 +20,21 @@ def _sft_epochs(spec) -> int:
     return int(t.epochs) if t.epochs is not None else RECIPE.sft.num_epochs
 
 
+def _on_policy_epochs(spec) -> int:
+    from flash.engine.recipe import RECIPE
+
+    t = spec.train
+    default = RECIPE.rl.num_epochs if spec.algorithm == "grpo" else RECIPE.opd.num_epochs
+    return int(t.epochs) if t.epochs is not None else default
+
+
 def _sft_seq_len(spec) -> int:
     from flash.engine.recipe import RECIPE
 
     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)
     )
 
@@ -41,6 +50,43 @@ def _sft_example_count(spec) -> int:
     )
 
 
+def _on_policy_example_count(spec) -> int:
+    t = spec.train
+    pinned_examples = int(t.max_examples) if t.max_examples else 0
+    if pinned_examples > 0:
+        return pinned_examples
+    env_examples = _env_max_examples(spec)
+    if env_examples > 0:
+        return env_examples
+    return _on_policy_requested_prompts_per_step(spec)
+
+
+def _env_max_examples(spec) -> int:
+    params = getattr(getattr(spec, "environment", None), "params", {}) or {}
+    if not isinstance(params, dict):
+        return 0
+    try:
+        value = int(params.get("max_examples") or 0)
+    except (TypeError, ValueError):
+        return 0
+    return max(0, value)
+
+
+def _on_policy_requested_prompts_per_step(spec) -> int:
+    from flash.engine.recipe import RECIPE
+
+    t = spec.train
+    default = (
+        RECIPE.rl.prompts_per_step if spec.algorithm == "grpo" else RECIPE.opd.prompts_per_step
+    )
+    return max(1, int(t.batch_size) if t.batch_size is not None else default)
+
+
+def _on_policy_prompts_per_step(spec, examples: int) -> int:
+    requested = _on_policy_requested_prompts_per_step(spec)
+    return max(1, min(requested, max(1, int(examples))))
+
+
 def _sft_realized_batch(spec) -> int:
     from flash.catalog import vocab_size_for
     from flash.engine.recipe import RECIPE
@@ -68,22 +114,26 @@ def _sft_steps_from_examples(spec, examples: int, *, apply_cap: bool) -> int:
 
 
 def spec_steps(spec) -> int:
-    """Per-seed optimizer steps implied by a train spec (mirrors the worker). GRPO: ``train.steps``
-    (else recipe default). SFT: ``epochs x ceil(num_examples / realized_batch)`` capped by
-    ``max_steps``, where ``num_examples`` must be pinned by ``max_examples``."""
-    from flash.engine.recipe import RECIPE
+    """Per-seed optimizer steps implied by a train spec (mirrors the worker).
 
-    t = spec.train
+    SFT: ``epochs x ceil(num_examples / realized_batch)`` capped by ``max_steps``. GRPO/OPD:
+    ``epochs`` means passes over ``max_examples`` rows when pinned, otherwise the estimate uses one
+    prompt batch per epoch.
+    """
     if spec.algorithm == "grpo":
-        if t.steps is not None:
-            return max(1, int(t.steps))
-        return RECIPE.rl.num_steps
+        examples = _on_policy_example_count(spec)
+        return on_policy_steps(
+            epochs=_on_policy_epochs(spec),
+            prompt_count=examples,
+            prompts_per_step=_on_policy_prompts_per_step(spec, examples),
+        )
     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
+        examples = _on_policy_example_count(spec)
+        return on_policy_steps(
+            epochs=_on_policy_epochs(spec),
+            prompt_count=examples,
+            prompts_per_step=_on_policy_prompts_per_step(spec, examples),
+        )
     # 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)
@@ -101,8 +151,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,

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

@@ -17,10 +17,10 @@ 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
@@ -82,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
@@ -116,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.47 → freesolo_flash_dev-0.2.49}/flash/engine/recipe.py

@@ -36,11 +36,11 @@ 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
-    num_steps: int = 150
+    num_epochs: int = 1
     sampling_temperature: float = 1.0
     sampling_top_p: float = 1.0
 
@@ -54,17 +54,16 @@ class OPDConfig:
     # 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.
-    num_steps: int = 100
     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
     # as in GRPO); raise for more teacher-scored coverage per prompt at higher cost.
     group_size: int = 1
+    num_epochs: int = 1
     sampling_temperature: float = 1.0
     sampling_top_p: float = 1.0
     # Reverse-KL coefficient for the groupwise reverse-KL loss; scales the per-span

freesolo_flash_dev-0.2.49/flash/engine/steps.py

@@ -0,0 +1,18 @@
+"""Training-step derivation shared by worker and cost-estimate paths."""
+
+from __future__ import annotations
+
+import math
+
+
+def on_policy_steps(
+    *,
+    epochs: int,
+    prompt_count: int,
+    prompts_per_step: int,
+) -> int:
+    """Resolve GRPO/OPD optimizer steps from full passes over the retained prompt pool."""
+    prompt_count = int(prompt_count)
+    if prompt_count <= 0:
+        raise ValueError("cannot derive epoch-based steps without at least one retained prompt")
+    return max(1, math.ceil(prompt_count * int(epochs) / max(1, int(prompts_per_step))))

{freesolo_flash_dev-0.2.47 → freesolo_flash_dev-0.2.49}/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.47 → freesolo_flash_dev-0.2.49}/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