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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 1.0.6
+Version: 1.0.7
 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.7}/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)
@@ -198,7 +198,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)
 

{freesolo_flash_dev-1.0.6 → freesolo_flash_dev-1.0.7}/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.7}/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
@@ -370,35 +372,17 @@ 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

{freesolo_flash_dev-1.0.6 → freesolo_flash_dev-1.0.7}/flash/engine/worker/opd.py

@@ -46,7 +46,7 @@ from flash.engine.structured_outputs import (
     parse_structured_outputs,
     reasoning_parser_for,
 )
-from flash.engine.vram import opd_completion_len
+from flash.engine.vram import OPD_CE_CHUNK_SIZE, opd_completion_len
 from flash.engine.worker._pkg import W as _w
 from flash.engine.worker.heartbeat import liveness_heartbeat
 from flash.engine.worker.hf import RequiredSaveError, _deployable_adapter_on_hf
@@ -1645,7 +1645,9 @@ def run_opd():
             "opd_teacher_workers": max_teacher_workers,
             "opd_teacher_batch_size": teacher_batch_size,
             "opd_loss_microbatch_size": loss_microbatch_size,
-            "opd_full_logits_batches": getattr(model, "_flash_opd_full_logits_batches", 0),
+            "opd_completion_only_batches": getattr(
+                model, "_flash_opd_completion_only_batches", 0
+            ),
             # Multi-turn: each assistant turn is distilled independently against its transcript-so-far
             # prefix, so a "sample" is a TURN, not a whole episode. Report the mode + turn ceiling and
             # the mean turns/episode so a run that collapsed to one-turn episodes (env never replied /
@@ -2193,6 +2195,88 @@ def _forward_logits(
         raise
 
 
+class _LmHeadBypass(Exception):
+    pass
+
+
+def _forward_completion_hidden_states(
+    model, input_ids, attention_mask, position_ids, *, logits_to_keep: int
+):
+    """Run the causal body through the normal model boundary but stop before the LM-head matmul."""
+    output_embeddings = model.get_output_embeddings()
+    input_embeddings = model.get_input_embeddings()
+    if output_embeddings is None or output_embeddings is input_embeddings:
+        raise RuntimeError("opd chunked ce requires a distinct output embedding module")
+
+    captured = {}
+    marker = object()
+
+    def _capture_hidden(_module, args):
+        if not args:
+            raise RuntimeError("opd lm head received no hidden states")
+        captured["hidden_states"] = args[0]
+        raise _LmHeadBypass(marker)
+
+    handle = output_embeddings.register_forward_pre_hook(_capture_hidden)
+    try:
+        _forward_logits(
+            model,
+            input_ids,
+            attention_mask,
+            position_ids=position_ids,
+            logits_to_keep=logits_to_keep,
+        )
+    except _LmHeadBypass as exc:
+        if not exc.args or exc.args[0] is not marker:
+            raise
+    else:
+        raise RuntimeError("opd could not intercept the lm head for chunked ce")
+    finally:
+        handle.remove()
+
+    hidden_states = captured.get("hidden_states")
+    if hidden_states is None or hidden_states.shape[1] != logits_to_keep:
+        raise RuntimeError("opd model did not honor completion-only logits_to_keep")
+    return hidden_states, output_embeddings
+
+
+def _chunked_token_logps(lm_head, hidden_states, token_ids, *, chunk_size=OPD_CE_CHUNK_SIZE):
+    """Compute exact realized-token logprobs with checkpointed vocabulary chunks."""
+    import torch
+    import torch.nn.functional as F
+    from torch.utils.checkpoint import checkpoint
+
+    if hidden_states.ndim != 2 or token_ids.ndim != 1:
+        raise ValueError("opd chunked ce expects [tokens, hidden] states and [tokens] ids")
+    if hidden_states.shape[0] != token_ids.shape[0]:
+        raise ValueError("opd chunked ce hidden-state and token counts differ")
+    if hidden_states.shape[0] == 0:
+        return hidden_states.new_empty((0,), dtype=torch.float32)
+
+    size = max(1, int(chunk_size))
+
+    def _logps(hidden_chunk, ids_chunk):
+        logits = lm_head(hidden_chunk).float()
+        return -F.cross_entropy(logits, ids_chunk, reduction="none")
+
+    chunks = []
+    for start in range(0, hidden_states.shape[0], size):
+        hidden_chunk = hidden_states[start : start + size]
+        ids_chunk = token_ids[start : start + size]
+        if torch.is_grad_enabled():
+            chunks.append(
+                checkpoint(
+                    _logps,
+                    hidden_chunk,
+                    ids_chunk,
+                    use_reentrant=False,
+                )
+            )
+        else:
+            chunks.append(_logps(hidden_chunk, ids_chunk))
+    return torch.cat(chunks)
+
+
 def _bump_model_counter(model, name: str, inc: int = 1) -> None:
     with contextlib.suppress(Exception):
         setattr(model, name, int(getattr(model, name, 0) or 0) + int(inc))
@@ -2286,23 +2370,46 @@ def _resolve_samples_batched(
         mb = max(1, int(microbatch))
         for start in range(0, len(text_prepared), mb):
             chunk = text_prepared[start : start + mb]
-            _bump_model_counter(model, "_flash_opd_full_logits_batches")
+            _bump_model_counter(model, "_flash_opd_completion_only_batches")
             seqs = [list(p.prompt_ids) + list(p.student_ids) for p in chunk]
             max_len = max(len(seq) for seq in seqs)
+            max_completion = max(len(p.student_ids) for p in chunk)
+            logits_to_keep = max_completion + 1
             input_ids = torch.full((len(seqs), max_len), pad_id, dtype=torch.long, device=device)
             attention_mask = torch.zeros((len(seqs), max_len), dtype=torch.long, device=device)
             for row, seq in enumerate(seqs):
-                input_ids[row, : len(seq)] = torch.tensor(seq, dtype=torch.long, device=device)
-                attention_mask[row, : len(seq)] = 1
-            logits = _forward_logits(model, input_ids, attention_mask)
-            losses = []
+                offset = max_len - len(seq)
+                input_ids[row, offset:] = torch.tensor(seq, dtype=torch.long, device=device)
+                attention_mask[row, offset:] = 1
+            position_ids = attention_mask.cumsum(dim=-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 0)
+            hidden_tail, lm_head = _forward_completion_hidden_states(
+                model,
+                input_ids,
+                attention_mask,
+                position_ids,
+                logits_to_keep=logits_to_keep,
+            )
+            selected_hidden = []
+            selected_ids = []
             for row, p in enumerate(chunk):
-                prompt_len = len(p.prompt_ids)
                 comp_len = len(p.student_ids)
-                rows = logits[row, prompt_len - 1 : prompt_len - 1 + comp_len]
-                loss = _gkd_loss_from_logits_rows(
-                    rows, p.student_ids, p.groups, kl_coef=knobs.kl_coef
+                selected_hidden.append(
+                    hidden_tail[row, logits_to_keep - comp_len - 1 : logits_to_keep - 1]
                 )
+                selected_ids.extend(int(token_id) for token_id in p.student_ids)
+            flat_logps = _chunked_token_logps(
+                lm_head,
+                torch.cat(selected_hidden, dim=0),
+                torch.tensor(selected_ids, dtype=torch.long, device=device),
+            )
+            losses = []
+            token_offset = 0
+            for p in chunk:
+                comp_len = len(p.student_ids)
+                sp_t = flat_logps[token_offset : token_offset + comp_len]
+                token_offset += comp_len
+                loss = _gkd_loss_from_logps(sp_t, p.groups, kl_coef=knobs.kl_coef)
                 if loss is None:
                     results[p.idx] = SampleResult(
                         teacher_status="ok",

{freesolo_flash_dev-1.0.6 → freesolo_flash_dev-1.0.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "freesolo-flash-dev"
-version = "1.0.6"
+version = "1.0.7"
 description = "Flash — managed LoRA post-training (SFT/GRPO/OPD) for Freesolo environments, driven by the `flash` CLI"
 readme = "README.md"
 # RunPod Flash supports Python 3.11-3.12 (not 3.13 yet).
@@ -116,7 +116,7 @@ artifacts = ["flash/py.typed"]
 # auto-published by .github/workflows/publish-dev.yml: a push to `dev` whose version isn't yet on
 # PyPI publishes it; pushes that don't change it find it already published and no-op. So cutting a
 # dev-channel release = bump this (and [project].version) and merge to `dev`.
-version = "1.0.6"
+version = "1.0.7"
 
 [tool.pytest.ini_options]
 pythonpath = ["."]

{freesolo_flash_dev-1.0.6 → freesolo_flash_dev-1.0.7}/tests/test_allocator.py

@@ -468,6 +468,38 @@ def test_estimator_logits_term_uses_max_tokens_and_caps_at_budget():
     assert e(2.0, "sft", seq_len=4096, max_tokens=256) == e(2.0, "sft", seq_len=4096, max_tokens=8192)
 
 
+def test_opd_vram_estimate_reserves_one_dense_image_loss_peak():
+    from flash.engine.vram import estimate_vram_gb as e
+
+    seq_len = 4096
+    completion = 512
+    vocab = 248_320
+    with_vocab = e(
+        2.0,
+        "opd",
+        seq_len=seq_len,
+        max_tokens=completion,
+        vocab=vocab,
+        batch_size=4,
+        group_size=4,
+    )
+    without_vocab = e(
+        2.0,
+        "opd",
+        seq_len=seq_len,
+        max_tokens=completion,
+        vocab=0,
+        batch_size=4,
+        group_size=4,
+    )
+    expected = (seq_len * 4 + completion * 8) * vocab / 1e9
+    assert with_vocab - without_vocab == pytest.approx(expected)
+
+    short = e(2.0, "opd", seq_len=seq_len, max_tokens=128, vocab=vocab)
+    long = e(2.0, "opd", seq_len=seq_len, max_tokens=2048, vocab=vocab)
+    assert long > short
+
+
 def test_open_model_opd_uses_opd_sizing_not_grpo(monkeypatch):
     """Regression (codex[bot], vram.py): for an uncataloged (model_policy='allow') model the open-model
     fallback hardcoded ``_need(params_b, 'grpo', ...)``, so an OPD run was sized as a colocated-vLLM GRPO
@@ -629,11 +661,11 @@ def test_observed_qwen2_opd_vllm_case_routes_off_32gb_cards(monkeypatch):
                 "lora_rank": 32,
                 "lora_alpha": 64,
             },
-            "RTX Pro 6000",
+            "A100 PCIe",
         ),
         (
-            # fp8 KV (cc >= 8.9) halves the resident rollout KV, so this fits the 141 GB H200 and no
-            # longer needs the 180 GB B200 — still far off the 32/40 GB consumer classes this guards.
+            # the dense image fallback grows with context and keeps this mixed-modality-safe route
+            # above the 96 gb class.
             "max_context_16384",
             {
                 "epochs": 1,
@@ -653,7 +685,7 @@ def test_observed_qwen2_opd_vllm_case_routes_off_32gb_cards(monkeypatch):
                 "lora_rank": 32,
                 "lora_alpha": 64,
             },
-            None,
+            "H200",
         ),
         (
             "group_size_8",
@@ -748,9 +780,8 @@ def test_observed_qwen4b_opd_vllm_startup_case_routes_off_40gb_cards(monkeypatch
     }
 
     need = required_vram_gb("Qwen/Qwen3.5-4B", "opd", train=train)
-    # Off the 40 GB class, but no longer forced onto the 141 GB H200: the colocated OPD rollout uses an
-    # fp8 KV cache on cc >= 8.9 hardware, and a run this size can only land on a modern (cc >= 8.9)
-    # card, so its KV pool is fp8-sized — it fits the cheaper 96 GB RTX Pro 6000 (Blackwell).
+    # the dense image fallback keeps this run above the 80 gb class while still fitting the 96 gb
+    # rtx pro 6000.
     assert 80 < need <= 96
 
     preview_gpu = provisional_gpu("Qwen/Qwen3.5-4B", "opd", train=train)

{freesolo_flash_dev-1.0.6 → freesolo_flash_dev-1.0.7}/tests/test_cost_estimate.py

@@ -480,3 +480,33 @@ def test_estimate_exact_pinned_provider_routes_through_disk_aware_allocate(monke
     assert seen["disk_gb"] == 200.0
     # and the quote reflects the live allocation rather than a static gpu_hourly_usd lookup.
     assert (est.provider, est.gpu, est.gpu_hourly_usd) == ("runpod", "H100", 3.29)
+
+
+def test_b200_not_cheaper_or_faster_than_h200_for_grpo(monkeypatch):
+    # regression: the estimator must not advertise b200 as faster/cheaper than h200 on peak flops.
+    # b200/sm100 training is h200-class (portable kernels), so at its higher $/hr b200 must never
+    # come out cheaper, and never faster, than h200 for the same run.
+    from types import SimpleNamespace
+
+    import flash.providers.allocator as allocator_mod
+    from flash.cost.facts import gpu_hourly_usd
+
+    def alloc_as(gpu):
+        rate = gpu_hourly_usd(gpu)
+
+        def fake_allocate(model_id, method, **kwargs):
+            return SimpleNamespace(gpu=gpu, provider="runpod", hourly_usd=rate, min_vram_gb=80)
+
+        return fake_allocate
+
+    monkeypatch.setattr(allocator_mod, "allocate", alloc_as("H200"))
+    h200 = estimate_cost(RunConfig("Qwen/Qwen3.5-4B", "grpo", 100, exact_type="H200"))
+    monkeypatch.setattr(allocator_mod, "allocate", alloc_as("B200"))
+    b200 = estimate_cost(RunConfig("Qwen/Qwen3.5-4B", "grpo", 100, exact_type="B200"))
+
+    assert gpu_hourly_usd("B200") > gpu_hourly_usd("H200")  # b200 is the pricier card
+    # same effective training throughput => b200 is no faster than h200 ...
+    assert b200.seconds_per_step == pytest.approx(h200.seconds_per_step)
+    assert b200.train_seconds == pytest.approx(h200.train_seconds)
+    # ... and at its higher $/hr, never cheaper.
+    assert b200.total_usd > h200.total_usd

{freesolo_flash_dev-1.0.6 → freesolo_flash_dev-1.0.7}/tests/test_cost_hardware.py

@@ -100,3 +100,22 @@ def test_pick_gpu_impossible_raises():
 def test_pick_gpu_auto_matches_default():
     assert pick_gpu(24, provider="auto") == pick_gpu(24)
     assert pick_gpu(12) == "RTX 4090"
+
+
+def test_effective_train_tflops_caps_b200_at_h200_class():
+    # b200/sm100 training falls back to portable kernels, so realized training throughput is
+    # h200-class, not the 2.25 pflops peak. the cost model must not treat b200 as faster than h200.
+    from flash.cost.facts import effective_train_tflops
+
+    assert gpu_tflops("B200") == 2250.0  # raw peak unchanged (vram/serving still use gpu_tflops)
+    assert effective_train_tflops("B200") == 990.0
+    assert effective_train_tflops("B200") == effective_train_tflops("H200")
+
+
+def test_effective_train_tflops_is_peak_for_uncapped_classes():
+    from flash.cost.facts import effective_train_tflops
+
+    for name in ("H100", "H200", "A100 SXM", "RTX 4090", "B200"):
+        # only b200 is capped; every other class keeps its peak.
+        expected = 990.0 if name == "B200" else gpu_tflops(name)
+        assert effective_train_tflops(name) == expected