freesolo-flash-dev 1.0.8__tar.gz → 1.0.9__tar.gz

{freesolo_flash_dev-1.0.8 → freesolo_flash_dev-1.0.9}/.env.example

@@ -31,4 +31,4 @@ FREESOLO_BASE_URL=http://backend:8000
 # Modal app that serves every adapter on one GPU per base model, scaling to zero when
 # idle). `flash deploy`/`undeploy`/`chat` register/deregister/chat against it. Defaults
 # to the hosted Modal URL when unset.
-FREESOLO_SERVING_URL=https://clado-ai--freesolo-lora-serving.modal.run
+FREESOLO_SERVING_URL=https://serve.freesolo.co

freesolo_flash_dev-1.0.9/AB_PLAN.md

@@ -0,0 +1,86 @@
+# OPD rollout-memory utilization A/B plan
+
+## Question
+
+Does sizing the colocated vLLM executor from measured post-load free VRAM, while reserving the modeled OPD training peak and allocator margin, eliminate the Qwen3.5 4B negative-KV startup failure and improve rollout throughput without changing training behavior or causing an OOM?
+
+No paid GPU run is part of this change.
+
+## Frozen comparison
+
+- control: `origin/dev` at `8fdfc858`
+- treatment: branch `perf/opd-rollout-util`
+- hardware: one H100 80 GB for each arm, from the same RunPod GPU type and image
+- launch both arms concurrently on separate matched GPUs
+- use the same base checkpoint, dataset revision, environment revision, seed, teacher, token budget, batch order, and worker image inputs
+- use immutable commits for both arms
+
+An H100 80 GB is the smallest card for this A/B. The reported failure is a large-card utilization defect, with roughly 60 GB left idle and a negative-KV startup failure for Qwen3.5 4B. An RTX 4090 24 GB does not have enough residual memory to expose the intended utilization increase and therefore cannot demonstrate the optimization.
+
+## OPD configuration
+
+Use the existing reproducing environment and dataset with this training shape held constant:
+
+```toml
+model = "Qwen/Qwen3.5-4B"
+algorithm = "opd"
+seed = 42
+
+[environment]
+id = "<frozen-reproduction-environment>"
+
+[train]
+epochs = 1
+batch_size = 8
+group_size = 1
+max_context_tokens = 1536
+max_completion_tokens = 512
+lora_rank = 32
+teacher_model = "glm-5.2"
+```
+
+Use the same bounded number of optimizer steps in both arms. Do not tune either arm after launch.
+
+## Measurements
+
+Record at vLLM initialization and for every rollout step:
+
+1. measured free and total VRAM immediately before vLLM construction
+2. modeled OPD training peak reserve, allocator margin, rollout budget, and final `gpu_memory_utilization`
+3. vLLM executor budget in GiB and available KV-cache budget or cache blocks
+4. rollout batch size and generated tokens per second
+5. peak allocated and peak reserved VRAM across the OPD loss forward/backward
+6. initialization outcome, negative-KV/cache-block errors, CUDA OOMs, and worker health
+7. per-step OPD loss, reward or task metric already emitted by the environment, and completed optimizer steps
+
+## Gates
+
+### Feasibility and memory safety
+
+- treatment must initialize vLLM and complete every planned optimizer step with zero negative-KV/cache-block failures and zero CUDA OOMs
+- treatment executor budget must never exceed measured free VRAM minus the logged training reserve and allocator margin
+- treatment loss forward/backward peak must retain at least the allocator margin after accounting for the executor budget
+- if the control reproduces the startup crash, treatment must complete the same configuration; otherwise the A/B is inconclusive on crash recovery but can still evaluate utilization and throughput
+
+### Utilization and throughput
+
+- direction: treatment `gpu_memory_utilization` and KV budget must be higher than control on the H100
+- threshold: treatment executor budget must increase by at least 16 GiB or utilization by at least 0.20 absolute
+- direction: treatment rollout generated tokens per second must be no worse than control
+- threshold: when both arms complete, treatment median steady-state rollout throughput must improve by at least 10 percent; if control crashes, report treatment throughput without claiming a relative speedup
+
+### Training parity
+
+- both arms must consume the same prompts and completion-token budget and complete the same optimizer-step count
+- per-checkpoint OPD loss curves must remain within normal deterministic/runtime noise, with no sustained treatment regression above 2 percent relative to control
+- final task metric must not regress by more than 1 percentage point absolute
+- any change in gradients, optimizer state, sampled inputs, or training math invalidates the A/B
+
+## Decision
+
+Accept the optimization only if the treatment passes all memory-safety and training-parity gates and either:
+
+1. recovers the exact control startup failure, or
+2. meets both the executor-budget and throughput thresholds when both arms complete.
+
+Reject or revise if treatment OOMs, violates the protected-memory budget, changes the training curve beyond the parity threshold, or fails to increase the H100 rollout budget materially.

{freesolo_flash_dev-1.0.8 → freesolo_flash_dev-1.0.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 1.0.8
+Version: 1.0.9
 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
@@ -186,12 +186,12 @@ for serving. If the run is not deployed yet, `/v1/runs/<run_id>/chat` returns `4
 with a hint to deploy first.
 
 Operators can also call the serving backend directly after the adapter is registered.
-The default serving backend is `https://clado-ai--freesolo-lora-serving.modal.run`, and
+The default serving backend is `https://serve.freesolo.co`, and
 operators can point Flash at another backend by setting `FREESOLO_SERVING_URL`.
 Use that same base URL when calling the backend directly; pass the run id as `model`:
 
 ```bash
-export FREESOLO_SERVING_URL=https://clado-ai--freesolo-lora-serving.modal.run
+export FREESOLO_SERVING_URL=https://serve.freesolo.co
 
 curl -X POST "$FREESOLO_SERVING_URL/v1/chat/completions" \
   -H "Content-Type: application/json" \

{freesolo_flash_dev-1.0.8 → freesolo_flash_dev-1.0.9}/README.md

@@ -129,12 +129,12 @@ for serving. If the run is not deployed yet, `/v1/runs/<run_id>/chat` returns `4
 with a hint to deploy first.
 
 Operators can also call the serving backend directly after the adapter is registered.
-The default serving backend is `https://clado-ai--freesolo-lora-serving.modal.run`, and
+The default serving backend is `https://serve.freesolo.co`, and
 operators can point Flash at another backend by setting `FREESOLO_SERVING_URL`.
 Use that same base URL when calling the backend directly; pass the run id as `model`:
 
 ```bash
-export FREESOLO_SERVING_URL=https://clado-ai--freesolo-lora-serving.modal.run
+export FREESOLO_SERVING_URL=https://serve.freesolo.co
 
 curl -X POST "$FREESOLO_SERVING_URL/v1/chat/completions" \
   -H "Content-Type: application/json" \

{freesolo_flash_dev-1.0.8 → freesolo_flash_dev-1.0.9}/flash/cli/__init__.py

@@ -162,7 +162,7 @@ class _FlashParser(_ThemedParser):
             f"train after publishing: `{CLI_NAME} env push --name my-env .`, "
             f"then `{CLI_NAME} train configs/sft.toml`",
             f"any command in depth: `{CLI_NAME} <command> --help`",
-            "docs: https://freesolo.co/docs",
+            "docs: https://docs.freesolo.co",
         ]
         return render.help_page(
             "managed LoRA post-training", usage, _HELP_GROUPS, _HELP_OPTIONS, footers

{freesolo_flash_dev-1.0.8 → freesolo_flash_dev-1.0.9}/flash/cli/training_doc.py

@@ -8,7 +8,7 @@ TRAINING_MD = r"""# TRAINING.md — how to actually improve a model with Flash
 > `flash env setup` dropped this file next to your `environment.py` and `configs/`.
 > It is the playbook Freesolo's own training agents follow to turn a *finished*
 > run into a model that *measurably improved*. The mechanics live in the hosted
-> docs (https://freesolo.co/docs); this file is the judgment that sits on top of them.
+> docs (https://docs.freesolo.co); this file is the judgment that sits on top of them.
 
 A run that reaches `done` is **not** the same as a run that worked. Submitting a run
 is not a result. The whole job is to design the learning signal, read what the run
@@ -713,5 +713,5 @@ flash export --adapter-id <run-id> --repository <you>/<repo>  # export final ada
 flash export --adapter-id <run-id>/step-N --repository <you>/<repo>  # export a checkpoint
 ```
 
-See the full reference at https://freesolo.co/docs.
+See the full reference at https://docs.freesolo.co.
 """

{freesolo_flash_dev-1.0.8 → freesolo_flash_dev-1.0.9}/flash/engine/vram.py

@@ -27,6 +27,9 @@ _BYTES_PER_PARAM = {
 _BASE_OVERHEAD_GB = 4.0
 _ACT_COEF = 0.12
 _SFT_PER_DEVICE_BS_DEFAULT = 4
+_VOCAB_DEFAULT = 248_320
+_OPD_TRAINING_RESERVE_MULTIPLIER = 1.15
+_OPD_ALLOCATOR_MARGIN_MAX_GB = 6.0
 OPD_CE_CHUNK_SIZE = 64
 _OPD_CE_PEAK_BYTES_PER_LOGIT = 16
 
@@ -126,6 +129,75 @@ def opd_loss_microbatch(params_b: float, prompts_per_step: int = 1, group_size:
     return max(1, min(total, default))
 
 
+def opd_training_peak_gb(
+    params_b: float,
+    seq_len: int,
+    *,
+    max_tokens: int | None = None,
+    prompts_per_step: int = 1,
+    group_size: int = 1,
+    thinking: bool = False,
+    vocab: int = _VOCAB_DEFAULT,
+    active_params_b: float | None = None,
+) -> float:
+    """Additional OPD loss-forward/backward memory above the resident student model."""
+    eff_b = float(active_params_b) if active_params_b else float(params_b)
+    width = math.sqrt(max(eff_b, 0.1))
+    loss_mb = opd_loss_microbatch(params_b, prompts_per_step, group_size)
+    activations = loss_mb * _ACT_COEF * (seq_len / 1024.0) * width
+    completion = opd_completion_len(max_tokens, thinking)
+    # the backward peak holds both full-sequence bf16 logits and fp32 completion rows plus gradients.
+    dense_logits = 2 * loss_mb * (seq_len * 2 + completion * 4) * vocab / 1e9
+    return activations + dense_logits
+
+
+def opd_training_reserve_gb(
+    params_b: float,
+    seq_len: int,
+    *,
+    max_tokens: int | None = None,
+    prompts_per_step: int = 1,
+    group_size: int = 1,
+    thinking: bool = False,
+    vocab: int = _VOCAB_DEFAULT,
+    active_params_b: float | None = None,
+) -> float:
+    """OPD loss-forward/backward peak with the runtime safety multiplier."""
+    return _OPD_TRAINING_RESERVE_MULTIPLIER * opd_training_peak_gb(
+        params_b,
+        seq_len,
+        max_tokens=max_tokens,
+        prompts_per_step=prompts_per_step,
+        group_size=group_size,
+        thinking=thinking,
+        vocab=vocab,
+        active_params_b=active_params_b,
+    )
+
+
+def opd_allocator_margin_gb(total_vram_gb: float) -> float:
+    """Allocator fragmentation slack protected by the OPD startup budget."""
+    return max(2.0, min(_OPD_ALLOCATOR_MARGIN_MAX_GB, float(total_vram_gb) * 0.05))
+
+
+def opd_post_init_reserve_gb(params_b: float, lora_rank: int) -> float:
+    """Pending LoRA gradient, AdamW state, and persistent post-init growth above model residency.
+
+    OPD trains all linear layers with bf16 LoRA parameters. PyTorch AdamW lazily allocates two
+    parameter-dtype moment tensors on the first step, so each estimated trainable parameter needs
+    2 bytes for its pending gradient plus 4 bytes for optimizer state. Use total model parameters for
+    the rank-linear geometry so MoE experts are not undercounted, then reserve the measured 35B-class
+    first-generation and post-init growth with a size-scaled 12 GB ceiling.
+    """
+    model_params_b = max(0.1, float(params_b))
+    rank = max(1, int(lora_rank))
+    trainable_params_b = (rank / 16.0) * (0.05 + model_params_b / 150.0)
+    gradient_gb = trainable_params_b * 2.0
+    adamw_state_gb = trainable_params_b * 4.0
+    persistent_growth_gb = max(1.0, min(12.0, model_params_b * 0.35))
+    return gradient_gb + adamw_state_gb + persistent_growth_gb
+
+
 def _resident_kv_gb(
     params_b: float | None, vllm_max_len: int, num_generations: int = 8, fp8_kv: bool = False
 ) -> float:
@@ -224,7 +296,6 @@ _VLLM_COLOCATE_FLOOR_GB = 28.0
 # 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).
 _SFT_LOGITS_BYTES_PER_ELEM = 16.0

{freesolo_flash_dev-1.0.8 → freesolo_flash_dev-1.0.9}/flash/engine/worker/hf.py

@@ -197,6 +197,7 @@ class _OptionalUpload:
     label: str
     staged_dir: str
     run: Callable[[], None]
+    on_coalesce: Callable[[_OptionalUpload], None] | None = None
 
 
 class _SingleSlotUploader:
@@ -209,12 +210,18 @@ class _SingleSlotUploader:
         self._next_sequence = 0
         self._thread: threading.Thread | None = None
 
-    def enqueue(self, label: str, staged_dir: str, run: Callable[[], None]) -> None:
+    def enqueue(
+        self,
+        label: str,
+        staged_dir: str,
+        run: Callable[[], None],
+        on_coalesce: Callable[[_OptionalUpload], None] | None = None,
+    ) -> None:
         replaced: _OptionalUpload | None = None
         thread: threading.Thread | None = None
         with self._condition:
             self._next_sequence += 1
-            task = _OptionalUpload(self._next_sequence, label, staged_dir, run)
+            task = _OptionalUpload(self._next_sequence, label, staged_dir, run, on_coalesce)
             replaced, self._pending = self._pending, task
             if self._thread is None:
                 thread = threading.Thread(
@@ -229,7 +236,11 @@ class _SingleSlotUploader:
                 f"[upload] coalesced optional {replaced.label} into newer {label} "
                 f"(sequence {task.sequence})"
             )
-            shutil.rmtree(replaced.staged_dir, ignore_errors=True)
+            if replaced.on_coalesce is not None:
+                # the coalesce hook takes over cleanup of the replaced staged tree.
+                replaced.on_coalesce(replaced)
+            else:
+                shutil.rmtree(replaced.staged_dir, ignore_errors=True)
         if thread is not None:
             thread.start()
 
@@ -264,8 +275,68 @@ class _SingleSlotUploader:
             return True
 
 
+class _FifoUploader:
+    """run uploads one at a time in FIFO order, never dropping a pending task.
+
+    unlike _SingleSlotUploader this keeps every enqueued upload, so per-step artifacts that must
+    each land (deployable adapters) are not coalesced away by a newer save.
+    """
+
+    def __init__(self) -> None:
+        self._condition = threading.Condition()
+        self._in_flight: _OptionalUpload | None = None
+        self._pending: list[_OptionalUpload] = []
+        self._next_sequence = 0
+        self._thread: threading.Thread | None = None
+
+    def enqueue(self, label: str, staged_dir: str, run: Callable[[], None]) -> None:
+        thread: threading.Thread | None = None
+        with self._condition:
+            self._next_sequence += 1
+            self._pending.append(_OptionalUpload(self._next_sequence, label, staged_dir, run))
+            if self._thread is None:
+                thread = threading.Thread(
+                    target=self._drain, name="flash-deployable-uploader", daemon=True
+                )
+                self._thread = thread
+            self._condition.notify_all()
+        if thread is not None:
+            thread.start()
+
+    def _drain(self) -> None:
+        while True:
+            with self._condition:
+                if not self._pending:
+                    self._thread = None
+                    self._condition.notify_all()
+                    return
+                task = self._pending.pop(0)
+                self._in_flight = task
+            try:
+                task.run()
+            except Exception as e:
+                print(f"deployable upload warn ({task.label}): {sanitize_diagnostic(e, limit=500)}")
+            finally:
+                shutil.rmtree(task.staged_dir, ignore_errors=True)
+                with self._condition:
+                    self._in_flight = None
+                    self._condition.notify_all()
+
+    def flush(self, timeout_s: float) -> bool:
+        """wait for the in-flight and all pending FIFO uploads to finish."""
+        deadline = time.monotonic() + max(0.0, timeout_s)
+        with self._condition:
+            while self._in_flight is not None or self._pending:
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    return False
+                self._condition.wait(remaining)
+            return True
+
+
 _OPTIONAL_CHECKPOINT_UPLOADER = _SingleSlotUploader()
 _OPTIONAL_AUX_UPLOADER = _SingleSlotUploader()
+_OPTIONAL_DEPLOYABLE_UPLOADER = _FifoUploader()
 _DEBUG_UPLOAD_LOCK = threading.Lock()
 
 
@@ -331,8 +402,10 @@ def flush_optional_uploads(timeout_s: float = _OPTIONAL_UPLOAD_FLUSH_TIMEOUT_S)
     started = time.monotonic()
     checkpoints_flushed = _OPTIONAL_CHECKPOINT_UPLOADER.flush(timeout_s)
     remaining = max(0.0, timeout_s - (time.monotonic() - started))
+    deployables_flushed = _OPTIONAL_DEPLOYABLE_UPLOADER.flush(remaining)
+    remaining = max(0.0, timeout_s - (time.monotonic() - started))
     aux_flushed = _OPTIONAL_AUX_UPLOADER.flush(remaining)
-    return checkpoints_flushed and aux_flushed
+    return checkpoints_flushed and deployables_flushed and aux_flushed
 
 
 def upload_debug_jsonl(name: str, rows: list[dict], *, keep_last: int = 200) -> None:
@@ -985,14 +1058,37 @@ def make_checkpoint_upload_callback(save_at_steps=()):
                 ckpt_dir, f"checkpoint-{step}"
             )
         except Exception as e:
-            print(f"[ckpt] step {step} snapshot warn: {sanitize_diagnostic(e, limit=500)}")
+            # surface the miss explicitly rather than logging a soft warning and continuing as if
+            # the periodic save reached hf.
+            print(
+                f"[ckpt] step {step} snapshot failed; step not published: "
+                f"{sanitize_diagnostic(e, limit=500)}"
+            )
             return
+
+        def _publish_coalesced_deployable(
+            replaced: _OptionalUpload,
+            step: int = step,
+            staged_checkpoint: str = staged_checkpoint,
+        ) -> None:
+            # a newer optional save coalesced this resume checkpoint away; still publish this step's
+            # small durable deployable through the non-coalescing fifo path (which owns the staged
+            # tree cleanup) so per-step deployables are never dropped.
+            _OPTIONAL_DEPLOYABLE_UPLOADER.enqueue(
+                f"coalesced deployable step {step}",
+                replaced.staged_dir,
+                lambda: publish_deployable_checkpoint(
+                    staged_checkpoint, step, _provenance_ready=True, _emit_heartbeat=False
+                ),
+            )
+
         _OPTIONAL_CHECKPOINT_UPLOADER.enqueue(
             f"checkpoint step {step}",
             staged_dir,
             lambda: _upload(
                 step, staged_checkpoint, provenance_ready=True, emit_heartbeat=False
             ),
+            on_coalesce=_publish_coalesced_deployable,
         )
 
     class _CheckpointUpload(TrainerCallback):

{freesolo_flash_dev-1.0.8 → freesolo_flash_dev-1.0.9}/flash/engine/worker/opd.py

@@ -819,10 +819,16 @@ def run_opd():
 
     # resolve the model/tokenizer halt set once for vllm sampling and termination classification.
     generation_eos_ids = _generation_eos_ids(model, tok)
-    vllm_kwargs = _opd_vllm_kwargs(model_id, knobs, seq_cap)
     lora_rank = _opd_lora_rank(
         model, getattr(_w.JOB_SPEC.train, "lora_rank", 32) if _w.JOB_SPEC else 32
     )
+    vllm_kwargs = _opd_vllm_kwargs(
+        model_id,
+        knobs,
+        seq_cap,
+        prompts_per_step=prompts_per_step,
+        lora_rank=lora_rank,
+    )
     print(
         f"[opd] rollout backend: colocated vLLM model={rollout_model_source} "
         f"ctx={seq_cap} lora_rank={lora_rank} util={vllm_kwargs['gpu_memory_utilization']:.2f} "

{freesolo_flash_dev-1.0.8 → freesolo_flash_dev-1.0.9}/flash/engine/worker/opd_vllm.py

@@ -92,7 +92,7 @@ def _startup_oom_error(
     free_gb: float,
     total_gb: float,
     requested_util: float,
-    margin_gb: float,
+    reserve_gb: float,
     rollout_batch_size: int,
 ) -> BaseException:
     requested_gb = float(requested_util) * float(total_gb)
@@ -100,8 +100,8 @@ def _startup_oom_error(
         "Free memory on device cuda:0 "
         f"({free_gb:.2f}/{total_gb:.2f} GiB) on startup is less than desired GPU "
         f"memory utilization ({requested_util:.6f} -> {requested_gb:.2f} GiB) "
-        f"with {margin_gb:.2f} GiB startup margin after reducing OPD rollout_batch_size "
-        f"to {rollout_batch_size}. Retry on a larger GPU."
+        f"with {reserve_gb:.2f} GiB reserved for the OPD training peak and allocator "
+        f"after reducing OPD rollout_batch_size to {rollout_batch_size}. Retry on a larger GPU."
     )
     try:
         import torch
@@ -118,7 +118,34 @@ def _decode_only_compilation_config() -> dict[str, Any]:
     }
 
 
-def opd_vllm_kwargs(model_id: str, knobs: Any, seq_cap: int) -> dict[str, Any]:
+def _sizing_lora_rank(knobs: Any, default: int = 32) -> int:
+    rank = getattr(knobs, "lora_rank", None)
+    if rank is None:
+        try:
+            from flash.engine.worker._pkg import W as _w
+
+            train = getattr(getattr(_w, "JOB_SPEC", None), "train", None)
+            rank = getattr(train, "lora_rank", None)
+        except Exception:
+            rank = None
+    try:
+        return max(1, int(rank if rank is not None else default))
+    except (TypeError, ValueError):
+        return default
+
+
+def _rollout_uplift_validated(params_b: float | None, lora_rank: int) -> bool:
+    return params_b is not None and float(params_b) <= 4.7 and int(lora_rank) <= 32
+
+
+def opd_vllm_kwargs(
+    model_id: str,
+    knobs: Any,
+    seq_cap: int,
+    *,
+    prompts_per_step: int | None = None,
+    lora_rank: int | None = None,
+) -> dict[str, Any]:
     """Direct vLLM LLM(...) kwargs mirroring the GRPO colocate rollout tuning."""
     kwargs: dict[str, Any] = {
         "gpu_memory_utilization": 0.10,
@@ -155,18 +182,32 @@ def opd_vllm_kwargs(model_id: str, knobs: Any, seq_cap: int) -> dict[str, Any]:
 
     if card_gb > 0:
         try:
-            from flash.catalog import MODELS
+            from flash.catalog import MODELS, vocab_size_for
             from flash.engine.vram import (
                 colocate_kv_util,
+                opd_allocator_margin_gb,
+                opd_post_init_reserve_gb,
                 opd_rollout_concurrency,
+                opd_training_reserve_gb,
                 resolve_params_b,
             )
 
             info = MODELS.get(model_id)
             params_b = resolve_params_b(model_id)
+            sizing_params_b = float(params_b or 1.0)
+            resolved_lora_rank = (
+                _sizing_lora_rank(knobs)
+                if lora_rank is None
+                else max(1, int(lora_rank))
+            )
+            resolved_prompts_per_step = (
+                getattr(knobs, "prompts_per_step", 1)
+                if prompts_per_step is None
+                else prompts_per_step
+            )
             active_b = float(getattr(info, "active_params_b", 0.0) or 0.0) if info else 0.0
             target_concurrency = opd_rollout_concurrency(
-                getattr(knobs, "prompts_per_step", 1),
+                resolved_prompts_per_step,
                 getattr(knobs, "group_size", 1),
             )
             rollout_concurrency = target_concurrency
@@ -184,12 +225,26 @@ def opd_vllm_kwargs(model_id: str, knobs: Any, seq_cap: int) -> dict[str, Any]:
 
             util = _util_for(rollout_concurrency)
             if free_gb > 0:
-                # vLLM rejects startup unless its whole executor budget is free right now. OPD builds
-                # the rollout engine after the student + warm-start adapter are resident, so choose the
-                # largest rollout chunk that fits the observed residual memory instead of reserving for
-                # every prompt in the optimizer step.
-                startup_margin_gb = max(1.0, min(4.0, card_gb * 0.03))
-                max_startup_util = max(0.0, (free_gb - startup_margin_gb) / max(1.0, card_gb))
+                # free memory is measured after the student and adapter are resident. protect the
+                # modeled loss forward/backward peak plus allocator slack, then give the remaining
+                # memory to the rollout executor instead of leaving it idle behind the generic cap.
+                training_reserve_gb = opd_training_reserve_gb(
+                    sizing_params_b,
+                    int(seq_cap),
+                    max_tokens=getattr(knobs, "max_completion", None),
+                    prompts_per_step=resolved_prompts_per_step,
+                    group_size=getattr(knobs, "group_size", 1),
+                    vocab=vocab_size_for(model_id),
+                    active_params_b=active_b,
+                )
+                post_init_reserve_gb = opd_post_init_reserve_gb(
+                    sizing_params_b, resolved_lora_rank
+                )
+                allocator_margin_gb = opd_allocator_margin_gb(card_gb)
+                protected_gb = training_reserve_gb + post_init_reserve_gb + allocator_margin_gb
+                rollout_budget_gb = max(0.0, free_gb - protected_gb)
+                max_startup_util = rollout_budget_gb / max(1.0, card_gb)
+                lean_trimmed = False
                 while rollout_concurrency > 1 and util > max_startup_util:
                     rollout_concurrency -= 1
                     util = _util_for(rollout_concurrency)
@@ -198,16 +253,61 @@ def opd_vllm_kwargs(model_id: str, knobs: Any, seq_cap: int) -> dict[str, Any]:
                         "[opd] reduced vLLM rollout batch "
                         f"{target_concurrency}->{rollout_concurrency} to fit startup memory "
                         f"(free={free_gb:.1f} GiB, request={util * card_gb:.1f} GiB, "
-                        f"margin={startup_margin_gb:.1f} GiB)"
+                        f"protected={protected_gb:.1f} GiB)"
                     )
                 if util > max_startup_util:
-                    startup_oom = _startup_oom_error(
-                        free_gb=free_gb,
-                        total_gb=card_gb,
-                        requested_util=util,
-                        margin_gb=startup_margin_gb,
-                        rollout_batch_size=rollout_concurrency,
-                    )
+                    # concurrency bottomed out but the resident-KV floor (vLLM weight copy + _KV_CAP)
+                    # keeps the minimal executor above the protected budget. before hard-OOMing, trim
+                    # the allocator slack to the old ~1 GiB startup margin so configs that fit before
+                    # the protected-budget change are not newly rejected; the modeled training peak and
+                    # post-init reserve stay intact.
+                    lean_protected_gb = training_reserve_gb + post_init_reserve_gb + 1.0
+                    lean_startup_util = max(0.0, free_gb - lean_protected_gb) / max(1.0, card_gb)
+                    if util <= lean_startup_util:
+                        print(
+                            "[opd] trimmed allocator margin "
+                            f"{allocator_margin_gb:.1f}->1.0 GiB to fit rollout executor at "
+                            f"concurrency {rollout_concurrency} "
+                            f"(free={free_gb:.1f} GiB, request={util * card_gb:.1f} GiB)"
+                        )
+                        protected_gb = lean_protected_gb
+                        max_startup_util = lean_startup_util
+                        lean_trimmed = True
+                        # the concurrency loop above bottomed out against the stricter full
+                        # budget; re-seek how many sequences now fit under the relaxed lean
+                        # ceiling so the lean path is not frozen at the reduced floor.
+                        while (
+                            rollout_concurrency < target_concurrency
+                            and _util_for(rollout_concurrency + 1) <= max_startup_util
+                        ):
+                            rollout_concurrency += 1
+                            util = _util_for(rollout_concurrency)
+                    else:
+                        startup_oom = _startup_oom_error(
+                            free_gb=free_gb,
+                            total_gb=card_gb,
+                            requested_util=util,
+                            reserve_gb=protected_gb,
+                            rollout_batch_size=rollout_concurrency,
+                        )
+                if (
+                    startup_oom is None
+                    and not lean_trimmed
+                    and _rollout_uplift_validated(params_b, resolved_lora_rank)
+                ):
+                    # 0.80 is a final ceiling for the validated qwen3.5-4b/rank-32 size class; larger
+                    # or higher-rank models keep the existing colocate ceiling until gpu validation.
+                    # skip uplift when we had to lean-trim: the lean budget only leaves ~1 GiB of
+                    # headroom for the training peak, so spending it on the rollout executor risks OOM.
+                    util = max(util, min(0.80, max_startup_util))
+                print(
+                    "[opd] vLLM memory budget "
+                    f"free={free_gb:.1f}/{card_gb:.1f} GiB "
+                    f"training_reserve={training_reserve_gb:.1f} GiB "
+                    f"post_init_reserve={post_init_reserve_gb:.1f} GiB "
+                    f"allocator_margin={allocator_margin_gb:.1f} GiB "
+                    f"rollout={util * card_gb:.1f} GiB util={util:.3f}"
+                )
             kwargs["gpu_memory_utilization"] = util
             kwargs["max_num_seqs"] = rollout_concurrency
             kwargs["rollout_batch_size"] = rollout_concurrency

{freesolo_flash_dev-1.0.8 → freesolo_flash_dev-1.0.9}/flash/lora_rank.py

@@ -13,6 +13,11 @@ from typing import TYPE_CHECKING, Any
 
 from flash.catalog import serving_context_cap, serving_lora_rank_cap
 
+
+class ServingPreflightError(ValueError):
+    """Serving-path preflight rejection. Homed here (dependency-light) so unstructured
+    preparation can raise/catch it without importing the heavy flash.serve.preflight module."""
+
 if TYPE_CHECKING:
     from flash.spec import JobSpec