freesolo-flash-dev 1.0.7__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.7 → freesolo_flash_dev-1.0.8}/PKG-INFO

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

@@ -228,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:
@@ -384,14 +394,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 +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``
@@ -497,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
@@ -716,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.7 → 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

{freesolo_flash_dev-1.0.7 → freesolo_flash_dev-1.0.8}/flash/engine/worker/heartbeat.py

@@ -18,6 +18,7 @@ import time
 
 from flash.engine.worker._pkg import W as _w
 from flash.engine.worker.perf import gpu_diagnostics
+from flash.engine.worker.rollout_samples import sanitize_rollout_text
 
 # Setup-phase liveness stages: emitted from a 30s liveness thread WITH a progress callback during the
 # cold download / model-load / split-scan phase, kept on the tighter setup-liveness upload cadence
@@ -151,26 +152,20 @@ def heartbeat(stage: str, *, liveness: bool = False, force: bool = False, **kw):
             if stage in _HB_TIGHT_LIVENESS_STAGES:
                 interval_s = min(interval_s, _w._HB_SETUP_LIVENESS_INTERVAL_S)
             upload_due = not throttled or (now - _w._HB_LAST_UPLOAD) >= interval_s
-            # ``force`` bypasses the per-stage throttle (but not TERMINAL_ONLY mode, handled above) when
-            # THIS heartbeat's payload must be the one on record and a throttled drop would leave a STALE
-            # value committed -- opd's post-optimizer-step ping forces so a mid-step progress ping
-            # (carrying the previous opt_steps) that just claimed the slot can't suppress the stepped
-            # commit (else a cancel is billed from the stale step). Force is gated on STEP ADVANCE (this
-            # payload's ``step`` exceeds the last committed step) AND a per-force time floor measured from
-            # the last FORCED commit. The floor throttles only a SUB-FLOOR BURST of step advances (a
-            # tiny/fast OPD config landing many optimizer updates per minute, which unthrottled would blow
-            # the HF per-repo commit cap before the final adapter/DONE upload, codex[bot]); because the
-            # floor keys off the last forced commit -- not the last upload of any kind -- a force still
-            # punches through IMMEDIATELY after a liveness/mid-step ping stole the slot, and steps spaced
-            # farther apart than the floor (the normal teacher-round-trip-gated regime) each still commit
-            # exactly once, so a cancel there still bills the true latest step. The bounded cost is that a
-            # cancel DURING a sub-floor burst under-bills by up to one floor-window of steps -- acceptable
-            # (it favours the customer) and unavoidable once steps outrun the commit cap.
+            # ``force`` bypasses the per-stage throttle when this payload must be on record. it normally
+            # requires a step advance, but a sample-bearing payload may match the committed step because
+            # the liveness daemon can commit that step first without the samples. the per-force floor still
+            # coalesces fast bursts to protect the hf commit cap, while an unrelated liveness commit does
+            # not arm the floor and therefore cannot suppress the first sample-bearing payload.
             if force and not upload_due:
                 fstep = kw.get("step")
+                has_samples = bool(kw.get("sampled_completions"))
+                force_step_due = isinstance(fstep, (int, float)) and (
+                    fstep > _w._HB_LAST_COMMITTED_STEP
+                    or (has_samples and fstep == _w._HB_LAST_COMMITTED_STEP)
+                )
                 if (
-                    isinstance(fstep, (int, float))
-                    and fstep > _w._HB_LAST_COMMITTED_STEP
+                    force_step_due
                     and (now - _w._HB_LAST_FORCED_UPLOAD) >= _w._HB_FORCE_MIN_INTERVAL_S
                 ):
                     upload_due = True
@@ -193,6 +188,7 @@ def heartbeat(stage: str, *, liveness: bool = False, force: bool = False, **kw):
             _committed_step = kw.get("step")
             if isinstance(_committed_step, (int, float)) and _committed_step > _w._HB_LAST_COMMITTED_STEP:
                 _w._HB_LAST_COMMITTED_STEP = int(_committed_step)
+    payload_committed = False
     if upload_due:
         critical = _is_critical_stage(stage)
         lock_timeout = _HB_CRITICAL_UPLOAD_LOCK_TIMEOUT_S if critical else _HB_UPLOAD_LOCK_TIMEOUT_S
@@ -210,19 +206,27 @@ def heartbeat(stage: str, *, liveness: bool = False, force: bool = False, **kw):
                     # ``is False`` (not falsy) so a mock/None never trips the rollback.
                     _rollback_throttle_slot(my_claim, prev_last_upload, prev_last_step, prev_last_forced)
                     print(f"HEARTBEAT upload failed; rolled back throttle slot for {stage}")
-                elif not liveness:
-                    # this committed snapshot carried real progress; settle the progress-carry
-                    # latch up to the seq captured when the snapshot was built (max: a concurrent
-                    # newer real heartbeat that lost the upload race must stay pending).
-                    with _HB_LOCK:
-                        if my_progress_seq > _w._HB_PROGRESS_UPLOADED_SEQ:
-                            _w._HB_PROGRESS_UPLOADED_SEQ = my_progress_seq
+                else:
+                    payload_committed = True
+                    if not liveness:
+                        # this committed snapshot carried real progress; settle the progress-carry
+                        # latch up to the seq captured when the snapshot was built (max: a concurrent
+                        # newer real heartbeat that lost the upload race must stay pending).
+                        with _HB_LOCK:
+                            if my_progress_seq > _w._HB_PROGRESS_UPLOADED_SEQ:
+                                _w._HB_PROGRESS_UPLOADED_SEQ = my_progress_seq
             finally:
                 _HB_UPLOAD_LOCK.release()
         else:
             _rollback_throttle_slot(my_claim, prev_last_upload, prev_last_step, prev_last_forced)
             print(f"HEARTBEAT upload-lock busy >{lock_timeout}s; skipping commit for {stage}")
-    print("HEARTBEAT", snapshot)
+    console_payload = payload
+    if not payload_committed and payload.get("sampled_completions"):
+        console_payload = {
+            key: value for key, value in payload.items() if key != "sampled_completions"
+        }
+    print("HEARTBEAT", json.dumps(console_payload))
+    return payload_committed
 
 
 def _maybe_attach_gpu_diag(payload: dict, last_gpu_diag_at: float, now: float) -> float:
@@ -290,7 +294,58 @@ def _bounded_reward_metrics(metrics) -> dict[str, float]:
     return {n: surviving[n] for n in sorted(priority + rest)}
 
 
-def make_reward_heartbeat_callback(reward_metrics=None):
+# Exactly three samples per heartbeat, always. Mirrors rollout_samples._SAMPLE_LIMIT.
+_SAMPLE_LIMIT = 3
+
+
+def _sampled_completion_scalar(sample: dict) -> tuple[str, float] | None:
+    """Return the (key, finite value) of a sample's scalar: GRPO ``reward`` or OPD ``loss``."""
+    for key in ("reward", "loss"):
+        if key not in sample:
+            continue
+        try:
+            value = float(sample.get(key))
+        except (TypeError, ValueError):
+            return None
+        return (key, value) if math.isfinite(value) else None
+    return None
+
+
+def _bounded_sampled_completions(samples) -> list[dict]:
+    if not isinstance(samples, (list, tuple)):
+        return []
+    bounded: list[dict] = []
+    for sample in samples:
+        if not isinstance(sample, dict):
+            continue
+        prompt_tail = sample.get("prompt_tail")
+        completion = sample.get("completion")
+        if not isinstance(prompt_tail, str) or not isinstance(completion, str):
+            continue
+        scalar = _sampled_completion_scalar(sample)
+        if scalar is None:
+            continue
+        generated_at_step = sample.get("generated_at_step")
+        if generated_at_step is not None:
+            try:
+                generated_at_step = int(generated_at_step)
+            except (TypeError, ValueError):
+                continue
+        scalar_key, scalar_value = scalar
+        bounded.append(
+            {
+                "prompt_tail": sanitize_rollout_text(prompt_tail),
+                "completion": sanitize_rollout_text(completion),
+                scalar_key: scalar_value,
+                "generated_at_step": generated_at_step,
+            }
+        )
+        if len(bounded) >= _SAMPLE_LIMIT:
+            break
+    return bounded
+
+
+def make_reward_heartbeat_callback(reward_metrics=None, samples=None):
     """Return a TRL callback that streams per-step reward to the HF heartbeat channel."""
     from transformers import TrainerCallback
 
@@ -298,6 +353,7 @@ def make_reward_heartbeat_callback(reward_metrics=None):
         def __init__(self):
             self.reward_history = []
             self.last_gpu_diag_at = 0.0
+            self.sent_first_sample_heartbeat = False
             self.latest_reward_metrics: dict[str, float] = {}
 
         def latest_fields(self) -> dict:
@@ -326,9 +382,19 @@ def make_reward_heartbeat_callback(reward_metrics=None):
             self.latest_reward_metrics = _bounded_reward_metrics(latest_metrics)
             if self.latest_reward_metrics:
                 payload["reward_metrics"] = dict(self.latest_reward_metrics)
+            latest_samples = samples() if callable(samples) else samples
+            bounded_samples = _bounded_sampled_completions(latest_samples)
+            if bounded_samples:
+                payload["sampled_completions"] = bounded_samples
             now = time.monotonic()
             self.last_gpu_diag_at = _maybe_attach_gpu_diag(payload, self.last_gpu_diag_at, now)
-            _w.heartbeat("rl_step", **payload)
+            force_first_samples = bool(bounded_samples) and not self.sent_first_sample_heartbeat
+            if force_first_samples:
+                committed = _w.heartbeat("rl_step", force=True, **payload)
+                if committed:
+                    self.sent_first_sample_heartbeat = True
+            else:
+                _w.heartbeat("rl_step", **payload)
 
     return _RewardHeartbeat()