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

{freesolo_flash_dev-1.0.7 → 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.7 → freesolo_flash_dev-1.0.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 1.0.7
+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.7 → 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.7 → 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.7 → 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
@@ -277,13 +277,16 @@ spending another GPU run:
 - **Judge the trend, not a single number.** The proof of training is the curve:
   loss falling (SFT) or `reward` rising over steps (GRPO). Record the base/early
   value and the final value. A flat or noisy trend with no improvement is not success.
-- **Read the model's outputs, not just the metrics.** A rising reward can come from
-  reward-hacking or a degenerate output the reward still credits — metrics alone never
-  establish that the model got better. Flash does not expose training-time rollouts
-  through the CLI (`flash log` gives you the metric trend and the worker's console/error
-  logs, not the sampled generations), so to read real outputs **deploy the adapter and
-  probe it**: `flash deploy <run-id>` then `flash chat <run-id> -m "..."` on at least a
-  few real inputs, including ones it should get wrong.
+- **Read the model's outputs, not just the metrics.** A rising reward (or falling loss)
+  can come from reward-hacking or a degenerate output the metric still credits — metrics
+  alone never establish that the model got better. For GRPO and OPD, `flash log` surfaces a
+  handful of full (untruncated) sample completions at the heartbeat cadence — GRPO shows
+  each completion's reward, OPD its distillation loss — with the first sample-bearing update
+  forced through, so you can catch skipped reasoning or a parroted prompt placeholder by
+  step 1-2. These are bounded diagnostics, not every rollout or a held-out
+  evaluation, so still **deploy the adapter and probe it**: `flash deploy <run-id>` then
+  `flash chat <run-id> -m "..."` on at least a few real inputs, including ones it should
+  get wrong.
 
   ```bash
   flash status <run-id>            # state + accrued cost
@@ -710,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.7 → freesolo_flash_dev-1.0.9}/flash/cost/analytical.py

@@ -166,7 +166,9 @@ def seconds_per_step(config: RunConfig, gpu: str) -> float:
         # sequence (see _opd_step_shape), not completion-only, or long-prompt opd is underquoted.
         completions, seq_tokens = _opd_step_shape(n)
         gen_s = (GRPO_GEN_FLOPS_PER_TOKEN_PER_PARAM * params * seq_tokens) / (peak * MFU_DECODE)
-        update_s = (OPD_UPDATE_FLOPS_PER_TOKEN_PER_PARAM * params * seq_tokens) / (peak * update_mfu)
+        update_s = (OPD_UPDATE_FLOPS_PER_TOKEN_PER_PARAM * params * seq_tokens) / (
+            peak * update_mfu
+        )
         teacher_lat = teacher_seconds_per_completion()
         # run_opd's primary path scores a step's completions CONCURRENTLY over Fireworks with a fan-out
         # cap of the step's OWN completion count (prompts_per_step * group_size, opd.py Phase 2), so
@@ -191,6 +193,7 @@ def seconds_per_step(config: RunConfig, gpu: str) -> float:
     )  # ceil: a partial wave still costs one latency
     return overhead + gen_s + reward_s + update_s
 
+
 def sft_seconds_for_tokens(config: RunConfig, gpu: str, train_tokens: float) -> float:
     """SFT steady-state wall time for an actual token count on ``gpu``."""
     n = config.normalized()
@@ -276,7 +279,11 @@ def _notes(
 def estimate_cost(config: RunConfig, *, wall_cap_s: float = DEFAULT_WALL_CAP_S) -> CostEstimate:
     """Deterministic pre-flight cost calculation."""
     # Billing cap: mirror the runner's max(60, max_wall_seconds) floor so a sub-60s cap isn't underpriced.
-    cap_s = max(60.0, float(config.max_wall_seconds)) if config.max_wall_seconds is not None else wall_cap_s
+    cap_s = (
+        max(60.0, float(config.max_wall_seconds))
+        if config.max_wall_seconds is not None
+        else wall_cap_s
+    )
     # Vast market duration filter: price against offers that outlast the run, using the SAME semantics
     # ``usable_offers`` applies at LAUNCH (not the 60s-floored billing cap_s) — a non-positive wall means
     # NO filter, a positive one is floored at 60s by usable_offers itself:
@@ -331,7 +338,9 @@ def estimate_cost(config: RunConfig, *, wall_cap_s: float = DEFAULT_WALL_CAP_S)
     compile_s = compile_seconds(config, gpu)
     raw_train = compile_s + config.steps * sps + required_save_s
     if not config.is_grpo and config.train_tokens is not None:
-        raw_train = compile_s + sft_seconds_for_tokens(config, gpu, config.train_tokens) + required_save_s
+        raw_train = (
+            compile_s + sft_seconds_for_tokens(config, gpu, config.train_tokens) + required_save_s
+        )
     sps = raw_train / config.steps
 
     # The cap is on total elapsed wall; setup is reported but not billed, so only training
@@ -367,9 +376,11 @@ def estimate_cost(config: RunConfig, *, wall_cap_s: float = DEFAULT_WALL_CAP_S)
         train_seconds=train,
         wall_clock_seconds=wall,
         wall_capped=wall_capped,
+        gpu_count=config.gpu_count,
         # total_usd is the customer gpu charge. the platform-owned teacher spend is itemized
-        # only as a diagnostic and is not passed through to the customer.
-        total_usd=train / 3600.0 * hourly,
+        # only as a diagnostic and is not passed through to the customer. an n-card job occupies
+        # n cards for the billed training wall, so the charge scales linearly with gpu_count.
+        total_usd=train / 3600.0 * hourly * config.gpu_count,
         teacher_api_usd=teacher_api_usd,
         notes=_notes(config, raw_train, wall_capped, cap_s),
     )

{freesolo_flash_dev-1.0.7 → freesolo_flash_dev-1.0.9}/flash/cost/spec.py

@@ -88,19 +88,17 @@ def _on_policy_prompts_per_step(spec, examples: int) -> int:
 def _sft_realized_batch(spec) -> int:
     from flash.catalog import resolve_vocab_size
     from flash.engine.recipe import RECIPE
-    from flash.engine.vram import resolve_params_b, sft_logits_fused, sft_realized_batch
+    from flash.engine.vram import sft_chunked_nll_enabled, sft_realized_batch
 
     t = spec.train
     requested_batch = int(t.batch_size) if t.batch_size is not None else RECIPE.sft.effective_batch
     sft_seq = _sft_seq_len(spec)
-    # Resolve params_b via the shared helper (catalog stat else HF safetensors for an open model) —
-    # the SAME resolution the worker's run_sft uses. The fused-CE decision (and thus the big-vocab
-    # micro-batch cap) hinges on the >=3B threshold, so an uncataloged >=3B model must not be priced
-    # as <3B (which would flip fused off, change the realized batch via the cap, and misprice the
-    # step count). Best-effort: no network -> None -> the prior <3B (cap-on) behavior.
-    sft_fused = sft_logits_fused(
-        resolve_params_b(spec.model, revision=spec.model_revision), sft_seq
-    )
+    # mirror the worker's validated chunked-nll model set so pricing uses the same realized batch.
+    # every allowlisted model is a gdn hybrid; at packable context lengths the worker may force
+    # per-device=1, so price the requested batch rather than the larger chunked-nll rounded batch.
+    sft_fused = sft_chunked_nll_enabled(spec.model)
+    if sft_fused and sft_seq <= 16_384:
+        return requested_batch
     return sft_realized_batch(
         requested_batch,
         seq_len=sft_seq,
@@ -167,6 +165,7 @@ def runconfig_from_spec(spec) -> RunConfig:
         exact_type=g.exact_type,
         model_revision=spec.model_revision,
         disk_gb=float(getattr(g, "disk_gb", 0.0) or 0.0),
+        gpu_count=g.count,
         max_wall_seconds=g.max_wall_seconds,
         environment=spec.environment.id or None,
         save_at_steps=t.save_at_steps,

{freesolo_flash_dev-1.0.7 → freesolo_flash_dev-1.0.9}/flash/cost/types.py

@@ -46,6 +46,9 @@ class RunConfig:
     # Spec gpu.disk_gb, carried so an exact-auto quote allocates at the run's real disk floor (parity
     # with the launch allocate call), keeping the persisted quote aligned with the pinned hardware.
     disk_gb: float = 0.0
+    # Spec gpu.count: cards the job occupies. total cost scales linearly with it (n cards for the
+    # billed training wall); 1 = the historical single-gpu quote.
+    gpu_count: int = 1
 
     def __post_init__(self) -> None:
         object.__setattr__(self, "method", normalize_algorithm(self.method))
@@ -71,6 +74,14 @@ class RunConfig:
         object.__setattr__(self, "model_revision", self.model_revision.strip())
         if self.steps < 1:
             raise ValueError(f"steps must be >= 1, got {self.steps}")
+        if isinstance(self.gpu_count, bool) or not isinstance(self.gpu_count, int):
+            raise TypeError("gpu_count must be an integer")
+        if self.gpu_count < 1:
+            raise ValueError(f"gpu_count must be >= 1, got {self.gpu_count}")
+        # upper bound mirrors GpuSpec's 1..8 so a direct RunConfig(gpu_count=...) cannot price a
+        # count the spec layer would reject.
+        if self.gpu_count > 8:
+            raise ValueError(f"gpu_count must be <= 8, got {self.gpu_count}")
         # Reject 0/negative positive-only knobs (bogus quote). max_wall_seconds is NOT here: the
         # runner floors it to max(60, ...) and estimate_cost mirrors that, so a non-positive cap
         # is accepted (floored to 60s), not rejected.
@@ -173,6 +184,8 @@ class CostEstimate:
     wall_clock_seconds: float
     wall_capped: bool
     total_usd: float
+    # cards the job occupies; total_usd already reflects n-card billing. 1 = single-gpu quote.
+    gpu_count: int = 1
     # opd only: external fireworks teacher token spend (0.0 for sft/grpo). billed by fireworks
     # to the platform-managed teacher key (users don't supply one), tracked separately from the
     # platform-billed gpu charge, so it is not part of total_usd and is shown as its own itemized
@@ -192,9 +205,9 @@ class CostEstimate:
         """Multi-line itemized breakdown for CLI output."""
         lines = [
             f"Run        : {self.model_id}  [{self.method.upper()}, {self.steps} steps]",
-            f"GPU        : {self.gpu} "
+            f"GPU        : {f'{self.gpu_count}x ' if self.gpu_count > 1 else ''}{self.gpu} "
             f"({self.gpu_vram_gb} GB; run needs >= {self.required_vram_gb} GB) "
-            f"@ ${self.gpu_hourly_usd:.2f}/hr",
+            f"@ ${self.gpu_hourly_usd:.2f}/hr{' per card' if self.gpu_count > 1 else ''}",
             f"Setup      : {self.setup_seconds / 60:.1f} min (cold start: boot + deps + model load"
             + (" + vLLM init" if self.method == "grpo" else "")
             + "; not billed)",

{freesolo_flash_dev-1.0.7 → freesolo_flash_dev-1.0.9}/flash/diagnostics.py

@@ -39,3 +39,14 @@ def sanitize_diagnostic(value: Any, *, limit: int = 2000) -> str:
     text = _BEARER_RE.sub("Bearer <redacted>", text)
     text = _SECRET_KEY_RE.sub(lambda match: f"{match.group(1)}{match.group(2)}<redacted>", text)
     return text[: max(0, int(limit))]
+
+
+def neutralize_control_chars(value: Any) -> str:
+    """Escape terminal control characters while preserving newlines as separators."""
+    text = str(value)
+    return "".join(
+        char
+        if char == "\n" or 0x20 <= ord(char) < 0x7F or ord(char) >= 0xA0
+        else f"\\x{ord(char):02x}"
+        for char in text
+    )

{freesolo_flash_dev-1.0.7 → 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,21 +296,30 @@ _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
-# Single source of truth for the SFT fused-CE gate. Keep the historical constant names because
-# engine.worker.perf and tests import them.
+_SFT_CHUNKED_NLL_TOKENS = 256
+# shared thresholds for the independent liger and worker memory gates.
 _LIGER_MIN_PARAMS_B = 3.0
 _LIGER_LONG_CTX_TOKENS = 2048
 
+# trl 1.6 chunked_nll is validated against these model families. other models keep plain nll
+# until their output-head and backbone traversal are covered by the same parity tests.
+_SFT_CHUNKED_NLL_MODELS = frozenset(
+    {
+        "Qwen/Qwen3.5-0.8B",
+        "Qwen/Qwen3.5-2B",
+        "Qwen/Qwen3.5-4B",
+        "Qwen/Qwen3.5-9B",
+        "Qwen/Qwen3.6-35B-A3B",
+    }
+)
+
 
-def sft_logits_fused(params_b: float | None, seq_len: int) -> bool:
-    """True when the worker fuses SFT cross-entropy (>=3B model OR >=2048-token context)."""
-    if seq_len >= _LIGER_LONG_CTX_TOKENS:
-        return True
-    return (params_b or 0.0) >= _LIGER_MIN_PARAMS_B
+def sft_chunked_nll_enabled(model_id: str) -> bool:
+    """whether the sft worker uses trl's dense-logit-free chunked nll path."""
+    return model_id in _SFT_CHUNKED_NLL_MODELS
 
 
 def sft_logits_per_device_cap(seq_len: int, vocab: int) -> int:
@@ -384,14 +465,15 @@ def estimate_vram_gb(
         dense_image_logits = (seq_len * 4 + completion * 8) * vocab / 1e9
         logits = max(chunked_logits, dense_image_logits)
         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.
+    # direct callers default to the conservative plain-nll estimate. model_required_vram_gb passes
+    # the worker's validated chunked-nll decision explicitly when it knows the model identity.
     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.
-    logits = 0.0 if fused else pd * seq_len * vocab * _SFT_LOGITS_BYTES_PER_ELEM / 1e9
+    # plain nll retains every sequence position. chunked nll retains at most one 256-token vocab
+    # projection at a time, independent of micro-batch and context once the chunk is full.
+    projected_tokens = min(pd * seq_len, _SFT_CHUNKED_NLL_TOKENS) if fused else pd * seq_len
+    logits = projected_tokens * vocab * _SFT_LOGITS_BYTES_PER_ELEM / 1e9
     return base + activations + logits
 
 
@@ -465,10 +547,10 @@ def sft_gc_off_peak_gb(
     lora_rank: int = 32,
     quant: str = "bf16",
 ) -> float:
-    """Estimated peak VRAM (GB) for a FUSED-CE LoRA SFT step with gradient checkpointing OFF: the
-    resident weights + optimizer/base + the no-recompute activations held across ALL ``num_layers``.
-    Fused CE (chalk FLCE) is assumed, so there is no ``[B, T, vocab]`` logits term (the thing that
-    made GC-off impossible at a 248k vocab). Unknown architecture dims -> ``inf`` (caller keeps GC on).
+    """Estimated peak VRAM (GB) for a dense-logit-free LoRA SFT step with gradient checkpointing OFF:
+    the resident weights + optimizer/base + the no-recompute activations held across ALL ``num_layers``.
+    Chunked or fused CE is assumed, so there is no ``[B, T, vocab]`` logits term (the thing that made
+    GC-off impossible at a 248k vocab). Unknown architecture dims -> ``inf`` (caller keeps GC on).
 
     MoE: the activation backbone scales with the model's real ``hidden`` x ``num_layers`` (geometry),
     NOT params_b -- the ~3B-active expert FFN is already folded into ``_GC_OFF_ACT_K``. ``weights``
@@ -497,7 +579,7 @@ def sft_grad_checkpoint_can_disable(
     quant: str = "bf16",
     margin_gb: float = 18.0,
 ) -> bool:
-    """True when a FUSED-CE LoRA SFT step fits a ``card_vram_gb`` card WITHOUT gradient checkpointing,
+    """True when a dense-logit-free LoRA SFT step fits a ``card_vram_gb`` card WITHOUT gradient checkpointing,
     so GC -- a ~+33% recompute tax on every step -- can be turned off for the speed win.
 
     Conservative by construction: an unknown card / unknown architecture dims, or a peak that doesn't
@@ -716,12 +798,10 @@ def model_required_vram_gb(
     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
+    # trl chunked_nll removes ignored positions before the lm head and projects valid tokens in
+    # bounded chunks. size validated qwen sft jobs without a dense [batch, seq, vocab] term; models
+    # outside that validated set keep the conservative plain-nll estimate and micro-batch cap.
+    sft_fused_ce = None if is_grpo else sft_chunked_nll_enabled(model_id)
     if info is not None:
         params_b = info.params_b
         if model_revision:

{freesolo_flash_dev-1.0.7 → freesolo_flash_dev-1.0.9}/flash/engine/worker/gpu_setup.py

@@ -109,11 +109,56 @@ def finalize_alloc_conf_for_sleep() -> None:
         print("[alloc] auto-conf skipped:", e)
 
 
+def _is_cudagraph_capture_failure(exc: BaseException) -> bool:
+    markers = (
+        "capture_model",
+        "cuda graph",
+        "cudagraph",
+        "graph capture",
+    )
+    current: BaseException | None = exc
+    seen: set[int] = set()
+    while current is not None and id(current) not in seen:
+        seen.add(id(current))
+        if any(marker in f"{type(current).__name__}: {current}".lower() for marker in markers):
+            return True
+        traceback = current.__traceback__
+        while traceback is not None:
+            frame = traceback.tb_frame
+            location = f"{frame.f_code.co_filename}:{frame.f_code.co_name}".lower()
+            if any(marker in location for marker in markers):
+                return True
+            traceback = traceback.tb_next
+        current = current.__cause__ or current.__context__
+    return False
+
+
+def _release_failed_cudagraph_capture(exc: BaseException) -> None:
+    current: BaseException | None = exc
+    seen: set[int] = set()
+    while current is not None and id(current) not in seen:
+        seen.add(id(current))
+        next_exc = current.__cause__ or current.__context__
+        current.__traceback__ = None
+        current = next_exc
+    try:
+        import gc
+
+        gc.collect()
+        import torch
+
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+    except Exception as cleanup_exc:
+        print(f"[rl][warn] cuda graph fallback cleanup failed: {cleanup_exc}")
+
+
 def patch_trl_colocate_llm_kwargs(
     *,
     kv_cache_dtype: str | None = None,
     max_num_batched_tokens: int | None = None,
     enforce_eager: bool | None = None,
+    compilation_config: dict | None = None,
     attention_backend: str | None = None,
     mm_encoder_attn_backend: str | None = None,
     reasoning_parser: str | None = None,
@@ -144,6 +189,8 @@ def patch_trl_colocate_llm_kwargs(
     * ``enforce_eager=True``: run the rollout engine in pure eager mode (no torch.compile, no CUDA-graph
       capture) — the SUPPORTED replacement for the removed ``VLLM_TORCH_COMPILE_LEVEL=0`` env. Dodges
       the vLLM 0.19.1 aot_compile (Ampere sm86) + Triton slot-mapping (graph-capture) crashes off B200.
+    * ``compilation_config``: select a narrower vLLM compilation profile, such as decode-only CUDA
+      graphs with torch.compile/AOT disabled.
     * ``attention_backend="FLASHINFER"|"TRITON_ATTN"``: pin a PTX-independent decoder attention backend
       (consumer-Blackwell sm120) — the SUPPORTED replacement for the removed ``VLLM_ATTENTION_BACKEND``
       env. ``EngineArgs`` coerces the bare member name through ``AttentionConfig.validate_backend_before``.
@@ -171,6 +218,8 @@ def patch_trl_colocate_llm_kwargs(
         new_overrides["max_num_batched_tokens"] = int(max_num_batched_tokens)
     if enforce_eager is not None:
         new_overrides["enforce_eager"] = bool(enforce_eager)
+    if compilation_config is not None:
+        new_overrides["compilation_config"] = dict(compilation_config)
     if attention_backend is not None:
         new_overrides["attention_backend"] = attention_backend
     if mm_encoder_attn_backend is not None:
@@ -194,7 +243,9 @@ def patch_trl_colocate_llm_kwargs(
         _vg._flash_llm_overrides = overrides
     overrides.update(new_overrides)
     if getattr(_vg, "_flash_llm_kwargs_patched", False):
-        print(f"[rl] colocate vLLM LLM kwargs extended (wrapper already installed): {new_overrides}")
+        print(
+            f"[rl] colocate vLLM LLM kwargs extended (wrapper already installed): {new_overrides}"
+        )
         return True
     # TRL binds the module-global ``LLM`` only under ``if is_vllm_available():`` (it's the symbol the
     # colocate ``self.llm = LLM(...)`` references). If vLLM isn't importable here there is nothing to
@@ -205,11 +256,31 @@ def patch_trl_colocate_llm_kwargs(
         return False
 
     def _patched_LLM(*args, **kwargs):
-        # Override TRL's hardcoded/absent values with ours (the kwargs TRL/env can't set). Read the
+        # override trl's hardcoded/absent values with ours (the kwargs trl/env can't set). read the
         # accumulated dict live so kwargs registered by later calls are applied too.
         kwargs.update(_vg._flash_llm_overrides)
         print(f"[rl] colocate vLLM LLM(...) kwargs override applied: {_vg._flash_llm_overrides}")
-        return _orig_LLM(*args, **kwargs)
+        compilation = kwargs.get("compilation_config") or {}
+        decode_graphs = (
+            not kwargs.get("enforce_eager", False)
+            and compilation.get("cudagraph_mode") == "FULL_DECODE_ONLY"
+        )
+        try:
+            return _orig_LLM(*args, **kwargs)
+        except Exception as capture_exc:
+            if not decode_graphs or not _is_cudagraph_capture_failure(capture_exc):
+                raise
+            _release_failed_cudagraph_capture(capture_exc)
+            eager_kwargs = dict(kwargs)
+            eager_kwargs["enforce_eager"] = True
+            eager_kwargs.pop("compilation_config", None)
+            _vg._flash_llm_overrides["enforce_eager"] = True
+            _vg._flash_llm_overrides.pop("compilation_config", None)
+            print(
+                f"[rl][warn] decode-only cuda graph capture failed ({capture_exc}); "
+                "retrying colocate vLLM initialization once in eager mode"
+            )
+            return _orig_LLM(*args, **eager_kwargs)
 
     _vg.LLM = _patched_LLM
     _vg._flash_llm_kwargs_patched = True