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

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

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

{freesolo_flash_dev-0.2.47 → 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 ""
     )

{freesolo_flash_dev-0.2.47 → 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
@@ -397,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
@@ -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.47 → freesolo_flash_dev-0.2.48}/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.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,

{freesolo_flash_dev-0.2.47 → freesolo_flash_dev-0.2.48}/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.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
@@ -59,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.47 → 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.47 → 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