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

freesolo_flash_dev-1.0.8/AB_PLAN.md

@@ -0,0 +1,45 @@
+# GRPO decode-only CUDA graph A/B plan
+
+## Objective
+
+Measure whether decode-only CUDA graphs speed up the colocated GRPO rollout engine without changing training behavior. This plan does not authorize a paid run.
+
+## Frozen comparison
+
+- Control: `origin/dev` at `8fdfc8580b9058d8c40289d4fa1a2e6e3b7bced2`. On RTX 5090/sm120 with vLLM 0.19.1, this selects `enforce_eager=True`.
+- Treatment: the committed head of `perf/grpo-cudagraphs`. On the same GPU and vLLM version, this selects `enforce_eager=False` with `mode=0` and `cudagraph_mode=FULL_DECODE_ONLY`.
+- Hardware: two separate, matched RunPod RTX 5090 workers from the same GPU offer, image, region, CPU/RAM class, and storage configuration. Do not substitute unmatched cards between arms.
+- Model: `Qwen/Qwen3.5-0.8B`, which fits within an RTX 4090-class memory budget.
+- Environment: the real `entropy-R1` hub environment, pinned to one immutable `freesolo-co/environment-hub` commit before launch.
+- Seed: 42 for both arms.
+- Budget: 30 optimizer updates after a two-update exact-path smoke. Use the same prompts per step, group size, maximum context, maximum completion length, LoRA settings, optimizer settings, reward implementation, dataset order, and checkpoint schedule.
+- Concurrency: launch both paid arms together only after both two-update smokes pass. Do not reuse a worker across arms.
+
+## Preflight and token-parity gate
+
+1. Resolve and record the exact Flash commit, environment-hub commit, model revision, container digest, vLLM version, CUDA/driver versions, and RunPod machine identifiers for each arm.
+2. Run the same fixed prompt batch through both immutable rollout engines with greedy decoding and the same maximum token count.
+3. Require identical token IDs, finish reasons, and decoded text for every fixed prompt. Any difference is a correctness failure and stops the A/B.
+4. Run two GRPO updates per arm. Require successful graph capture or vLLM's eager fallback, finite loss/reward values, and no OOM before scaling to 30 updates.
+
+## Metrics
+
+Measure rollout work separately from training work where telemetry permits.
+
+- Primary performance: generated rollout tokens per rollout-engine second, excluding engine startup, graph capture, checkpointing, reward evaluation, and optimizer time.
+- Secondary performance: median rollout wall time per update and end-to-end median step wall time. Report graph-capture startup time separately rather than amortizing it away.
+- Correctness: fixed-prompt rollout token parity, reward curve, training loss curve, completion length, finish-reason distribution, and invalid/empty completion rate.
+- Safety: peak allocated and reserved GPU memory, OOM count, graph-capture failures, eager-fallback events, worker restarts, and failed updates.
+
+## Decision thresholds
+
+The expected direction is faster treatment decode with unchanged learning behavior.
+
+- Performance works: treatment improves steady-state rollout tokens/second by at least 10% over control. A 10% to 40% improvement is the expected range. Confirm the direction with at least a 9% reduction in median rollout wall time over updates 3 through 30.
+- End-to-end value: treatment reduces median total step wall time by at least 5%. Report a decode-only win below this threshold as technically real but operationally marginal.
+- Reward parity: treatment reward AUC over the 30 matched updates must remain within 5% relative or 0.02 absolute of control, whichever tolerance is larger. There must be no sustained three-update reward regression larger than 0.05 absolute.
+- Loss parity: treatment loss AUC must remain within 5% relative of control, with finite values at every update and no new divergence pattern.
+- Rollout parity: the fixed greedy prompt batch must have exact token-ID parity. During stochastic training, completion length, finish-reason, and empty/invalid-rate differences must each remain within 5% relative or one sample, whichever is larger.
+- Memory and reliability: no OOM, no worker restart, no failed update, and treatment peak reserved memory no more than 5% above control after graph capture.
+
+The treatment passes only if all correctness and reliability gates pass and the primary performance threshold is met. If vLLM falls back to eager, report the fallback and classify the performance result as inconclusive rather than a decode-graph success.

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 1.0.6
+Version: 1.0.8
 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-1.0.6 → freesolo_flash_dev-1.0.8}/flash/cli/training_doc.py

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

{freesolo_flash_dev-1.0.6 → freesolo_flash_dev-1.0.8}/flash/cost/analytical.py

@@ -14,8 +14,8 @@ from flash.providers.allocator import required_vram_gb, vram_headroom
 from .facts import (
     active_params_b,
     download_weight_gb,
+    effective_train_tflops,
     gpu_hourly_usd,
-    gpu_tflops,
     gpu_vram_gb,
     model_quant,
     pick_gpu,
@@ -150,7 +150,7 @@ def compile_seconds(config: RunConfig, gpu: str) -> float:
 def seconds_per_step(config: RunConfig, gpu: str) -> float:
     """Steady-state wall time for one optimizer step on ``gpu``."""
     n = config.normalized()
-    peak = gpu_tflops(gpu) * 1e12  # FLOP/s
+    peak = effective_train_tflops(gpu) * 1e12  # FLOP/s (realized training throughput; see facts)
     # An MoE's per-step wall scales with TOTAL params (routing + all-expert coordination + grouped
     # GEMM under-utilization), not the tiny active-param FLOPs; dense models keep active (== total).
     moe = _is_moe(n.model_id)
@@ -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()
@@ -198,7 +201,7 @@ def sft_seconds_for_tokens(config: RunConfig, gpu: str, train_tokens: float) ->
     moe = _is_moe(n.model_id)
     params = (total_params_b(n.model_id) if moe else active_params_b(n.model_id)) * 1e9
     mfu = MFU_SFT_TRAIN_MOE if moe else MFU_SFT_TRAIN
-    peak = gpu_tflops(gpu) * 1e12
+    peak = effective_train_tflops(gpu) * 1e12
     flops = SFT_FLOPS_PER_TOKEN_PER_PARAM * params * train_tokens
     return flops / (peak * mfu)
 
@@ -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.6 → freesolo_flash_dev-1.0.8}/flash/cost/facts.py

@@ -30,6 +30,31 @@ def gpu_tflops(name: str) -> float:
     return GPU_COMPUTE_TFLOPS.get(name, _DEFAULT_TFLOPS)
 
 
+# realized TRAINING throughput sits well below peak when a class's training kernels don't reach it.
+# b200 (sm100) has no arch-tuned training kernels yet and falls back to the same portable paths as
+# h200 (sm90), so its 2.25 pflops dense-bf16 peak does not materialize for training -- realized
+# throughput is h200-class (and frequently lower for rl/grpo). cap b200 at h200 so the analytical
+# cost model does not rank it as faster/cheaper than h200 on peak flops alone (it is not, and is
+# often slower). this is a conservative floor: it never prices b200 above h200-equivalent training
+# time, so it cannot over-charge, and it removes the "b200 looks cheapest" inversion at the source.
+# refine per-workload once real b200 training samples exist. the vram/serving paths keep the true
+# peak via gpu_tflops; only the training-time model uses this.
+_TRAIN_TFLOPS_CAP: dict[str, float] = {
+    "B200": 990.0,  # h200-class realized training throughput, not the 2250 dense-bf16 peak
+}
+
+
+def effective_train_tflops(name: str) -> float:
+    """Realized bf16 TRAINING throughput for the analytical cost model.
+
+    Equal to ``gpu_tflops`` (peak) except where a class's training kernels fall short of peak: b200
+    training falls back to portable kernels, so it is capped at h200-class throughput rather than
+    its raw 2.25 pflops peak."""
+    peak = gpu_tflops(name)
+    cap = _TRAIN_TFLOPS_CAP.get(name)
+    return min(peak, cap) if cap is not None else peak
+
+
 def gpu_hourly_usd(
     name: str,
     provider: str | None = None,

{freesolo_flash_dev-1.0.6 → freesolo_flash_dev-1.0.8}/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.6 → freesolo_flash_dev-1.0.8}/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.6 → freesolo_flash_dev-1.0.8}/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.6 → freesolo_flash_dev-1.0.8}/flash/engine/vram.py

@@ -27,6 +27,8 @@ _BYTES_PER_PARAM = {
 _BASE_OVERHEAD_GB = 4.0
 _ACT_COEF = 0.12
 _SFT_PER_DEVICE_BS_DEFAULT = 4
+OPD_CE_CHUNK_SIZE = 64
+_OPD_CE_PEAK_BYTES_PER_LOGIT = 16
 
 
 def _sft_per_device_bs() -> int:
@@ -112,7 +114,7 @@ def opd_rollout_concurrency(prompts_per_step: int = 1, group_size: int = 1) -> i
 
 
 def opd_loss_microbatch(params_b: float, prompts_per_step: int = 1, group_size: int = 1) -> int:
-    """Loss microbatch size used by OPD's dense-logit GKD backward.
+    """Loss microbatch size used by OPD's GKD backward.
 
     Keep this in lockstep with worker.opd._opd_loss_microbatch_size without importing the worker from
     the sizing path. Small/medium catalog models run up to four loss samples per forward; 35B-class
@@ -226,17 +228,27 @@ _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:
@@ -370,44 +382,27 @@ def estimate_vram_gb(
         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
-        # vocab logits for >=3B models, so a long-completion opd job (e.g. max_tokens=8192) would be
-        # sized for an under-capacity card and OOM.
-        #
-        # Mirror opd_rollout_seq_len / run_opd's completion resolution: explicit max_tokens, else the
-        # OPD recipe default (thinking uses the longer max_completion_len_thinking). A GRPO-style
-        # min(seq_len, 1024) fallback would UNDER-budget a thinking opd job (1536-token completions).
+        # text opd projects only completion-prediction hidden states and computes exact full-vocabulary
+        # ce in checkpointed chunks. image samples still use the dense full-sequence logits path and are
+        # processed one at a time. allocation does not know the dataset modality, so reserve the larger
+        # of one text ce chunk and one dense image loss forward/backward peak.
         completion = opd_completion_len(max_tokens, thinking)
-        # run_opd backprops a small loss microbatch. Budget that actual dense-logit microbatch, not
-        # the full rollout/teacher batch and not a single sequence.
         loss_mb = opd_loss_microbatch(params_b, batch_size, group_size)
         activations = loss_mb * _ACT_COEF * (seq_len / 1024.0) * width
-        # Dense-logit peak spans BOTH the forward and the loss BACKWARD (OPD has no fused CE):
-        #   - forward:  bf16 full-sequence logits [seq, vocab]        (seq * 2)
-        #               fp32 completion rows [completion, vocab] for logsumexp, saved for backward
-        #                                                             (completion * 4)
-        #   - backward: fp32 gradient of those completion rows [completion, vocab]
-        #                                                             (completion * 4)
-        #               bf16 gradient scattered back into the full logits [seq, vocab] to reach lm_head
-        #                                                             (seq * 2)
-        # All four coexist at the backward peak. The old formula counted only the two FORWARD buffers,
-        # so a long-completion / large-vocab (248k) opd job under-budgeted the loss backward by
-        # ~(completion*4 + seq*2)*vocab bytes and could route to a GPU that OOMs in OPD loss backward
-        # (codex[bot]). Mirror the SFT dense-logit sizing by budgeting the backward buffers too.
-        logits_fwd = loss_mb * (seq_len * 2 + completion * 4) * vocab
-        logits_bwd = loss_mb * (seq_len * 2 + completion * 4) * vocab
-        logits = (logits_fwd + logits_bwd) / 1e9
+        ce_rows = min(OPD_CE_CHUNK_SIZE, loss_mb * completion)
+        chunked_logits = ce_rows * vocab * _OPD_CE_PEAK_BYTES_PER_LOGIT / 1e9
+        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
 
 
@@ -481,10 +476,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``
@@ -513,7 +508,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
@@ -732,12 +727,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.6 → freesolo_flash_dev-1.0.8}/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