freesolo-flash-dev 0.2.30__tar.gz → 0.2.31__tar.gz

{freesolo_flash_dev-0.2.30 → freesolo_flash_dev-0.2.31}/.github/workflows/bake-kernel-cache.yml

@@ -76,16 +76,23 @@ jobs:
               gpu_type_id: "NVIDIA RTX A6000",
               allowed_cuda: "",
             }
+          # sm89 (Ada): warm on the L40S -- a DATACENTER Ada card with reliable secure-cloud capacity --
+          # instead of the consumer RTX 4090 (scarce/cold on secure cloud, where the bake pod often never
+          # finished pulling the ~20GB image before the deadline). The cache is sm-keyed, so an L40S-warmed
+          # sm89 cache is valid on every sm89 card (4090/L40/L4). NB: verify the exact RunPod gpuTypeId.
           - {
               sm: sm89,
               arch: "8.9",
-              gpu_type_id: "NVIDIA GeForce RTX 4090",
+              gpu_type_id: "NVIDIA L40S",
               allowed_cuda: "",
             }
+          # sm90 (Hopper): warm on the H200 (141GB, sm90) -- better secure-cloud capacity than the
+          # heavily-contended H100 80GB HBM3 (which cold-pull-stalled the bake). Same sm90, so the cache
+          # is valid on every H100/H200. H200 is in flash's validated GPU catalog (providers/base.py).
           - {
               sm: sm90,
               arch: "9.0",
-              gpu_type_id: "NVIDIA H100 80GB HBM3",
+              gpu_type_id: "NVIDIA H200",
               allowed_cuda: "",
             }
           # Blackwell needs CUDA-13 hosts to JIT its PTX (matches min_cuda_for in the provider). Bake

{freesolo_flash_dev-0.2.30 → freesolo_flash_dev-0.2.31}/Dockerfile.worker

@@ -63,7 +63,7 @@ RUN pip install --no-cache-dir \
     "tilelang==0.1.11" \
     "apache-tvm-ffi==0.1.11" \
     "hf_transfer" \
-    "freesolo>=0.2.49" \
+    "freesolo>=0.2.52" \
     "runpod"
 # Gated-DeltaNet (Qwen3.5/3.6) runs on flash-linear-attention's Triton kernels. On consumer cards
 # (Ampere/Ada/Blackwell) the Triton path is correct. On Hopper (sm90) fla's GDN chunk_bwd is

{freesolo_flash_dev-0.2.30 → freesolo_flash_dev-0.2.31}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 0.2.30
+Version: 0.2.31
 Summary: Flash — managed LoRA post-training (SFT/GRPO) 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
@@ -18,19 +18,20 @@ Requires-Python: <3.13,>=3.11
 Provides-Extra: dev
 Requires-Dist: datasets>=2.19; extra == 'dev'
 Requires-Dist: fastapi; extra == 'dev'
-Requires-Dist: freesolo>=0.2.49; extra == 'dev'
+Requires-Dist: freesolo>=0.2.52; extra == 'dev'
 Requires-Dist: httpx>=0.27; extra == 'dev'
 Requires-Dist: huggingface-hub>=0.34; extra == 'dev'
 Requires-Dist: mypy>=1.13.0; extra == 'dev'
 Requires-Dist: pytest>=9.0.3; extra == 'dev'
 Requires-Dist: ruff>=0.6; extra == 'dev'
 Requires-Dist: runpod-flash; extra == 'dev'
+Requires-Dist: transformers<5.11,>=5.6; extra == 'dev'
 Requires-Dist: uvicorn; extra == 'dev'
 Provides-Extra: gpu
 Requires-Dist: accelerate>=1.4; extra == 'gpu'
 Requires-Dist: bitsandbytes>=0.49; extra == 'gpu'
 Requires-Dist: datasets>=2.19; extra == 'gpu'
-Requires-Dist: freesolo>=0.2.49; extra == 'gpu'
+Requires-Dist: freesolo>=0.2.52; extra == 'gpu'
 Requires-Dist: huggingface-hub>=0.34; extra == 'gpu'
 Requires-Dist: peft>=0.19; extra == 'gpu'
 Requires-Dist: torch==2.10.0; extra == 'gpu'
@@ -40,10 +41,11 @@ Requires-Dist: vllm==0.19.1; extra == 'gpu'
 Provides-Extra: server
 Requires-Dist: datasets>=2.19; extra == 'server'
 Requires-Dist: fastapi; extra == 'server'
-Requires-Dist: freesolo>=0.2.49; extra == 'server'
+Requires-Dist: freesolo>=0.2.52; extra == 'server'
 Requires-Dist: httpx>=0.27; extra == 'server'
 Requires-Dist: huggingface-hub>=0.34; extra == 'server'
 Requires-Dist: runpod-flash; extra == 'server'
+Requires-Dist: transformers<5.11,>=5.6; extra == 'server'
 Requires-Dist: uvicorn; extra == 'server'
 Description-Content-Type: text/markdown
 
@@ -101,10 +103,10 @@ uv run flash-server                      # control plane (operator-side, run onc
 
 The control plane owns provider credentials: `RUNPOD_API_KEY` is always required,
 plus the shared `HF_TOKEN`.
-The artifact repo is platform-managed and per-run (each run gets its own
-`Freesolo-Co/flashrun-<run_id>`, written by the operator `HF_TOKEN`); it is not a user
-knob and not an operator-wide env var. Clients authenticate with their freesolo API key
-(`flash login`).
+The artifact repo is platform-managed and environment-scoped (runs for the same environment share
+one private `Freesolo-Co/flashrun-<environment>-<hash>` repo, written by the operator `HF_TOKEN`);
+it is not a user knob and not an operator-wide env var. Clients authenticate with their freesolo
+API key (`flash login`).
 
 ## Release channels
 

{freesolo_flash_dev-0.2.30 → freesolo_flash_dev-0.2.31}/README.md

@@ -52,10 +52,10 @@ uv run flash-server                      # control plane (operator-side, run onc
 
 The control plane owns provider credentials: `RUNPOD_API_KEY` is always required,
 plus the shared `HF_TOKEN`.
-The artifact repo is platform-managed and per-run (each run gets its own
-`Freesolo-Co/flashrun-<run_id>`, written by the operator `HF_TOKEN`); it is not a user
-knob and not an operator-wide env var. Clients authenticate with their freesolo API key
-(`flash login`).
+The artifact repo is platform-managed and environment-scoped (runs for the same environment share
+one private `Freesolo-Co/flashrun-<environment>-<hash>` repo, written by the operator `HF_TOKEN`);
+it is not a user knob and not an operator-wide env var. Clients authenticate with their freesolo
+API key (`flash login`).
 
 ## Release channels
 

{freesolo_flash_dev-0.2.30 → freesolo_flash_dev-0.2.31}/docker/bake_kernel_cache.py

@@ -66,15 +66,19 @@ def main() -> int:
     ap = argparse.ArgumentParser(description="offload the kernel-cache warmup to a RunPod GPU")
     ap.add_argument("--arch", required=True, help="TORCH_CUDA_ARCH_LIST target, e.g. 9.0")
     ap.add_argument("--sm", required=True, help="sm tag, e.g. sm90 (must match the produced cache)")
-    ap.add_argument("--gpu-type-id", required=True, help="RunPod gpuTypeId, e.g. 'NVIDIA H100 80GB HBM3'")
+    ap.add_argument(
+        "--gpu-type-id", required=True, help="RunPod gpuTypeId, e.g. 'NVIDIA H100 80GB HBM3'"
+    )
     ap.add_argument("--image", default="ghcr.io/freesolo-co/flash-worker:cu128")
     ap.add_argument("--out", default="build/kernel_cache")
     # the warm pod only pulls the ~20GB image + writes the cache (no model download), so keep this
     # modest -- an over-large ask shrinks the eligible host pool and trips "machine does not have the
     # resources" on scarce classes (e.g. Blackwell sm120 on secure cloud).
     ap.add_argument("--container-disk-gb", type=int, default=60)
-    ap.add_argument("--deadline-min", type=int, default=45)
-    ap.add_argument("--run-id", default="", help="unique suffix for the temp repo (default: time+uuid)")
+    ap.add_argument("--deadline-min", type=int, default=90)
+    ap.add_argument(
+        "--run-id", default="", help="unique suffix for the temp repo (default: time+uuid)"
+    )
     ap.add_argument(
         "--allowed-cuda",
         default="",
@@ -189,6 +193,59 @@ def main() -> int:
                 shutil.move(s, d)
             shutil.rmtree(tmp, ignore_errors=True)
             rc = _verify(args.out, args.sm)
+        else:
+            # Non-success (timeout / pod_died): the pod-side out/ (STARTED marker, warmup.log) lives in
+            # the temp dataset the `finally` below deletes, so pull + print it FIRST. out/STARTED present
+            # = the warmup entrypoint ran (a timeout then means it hung / was slow); STARTED absent is
+            # AMBIGUOUS -- either the pod ran the wrong entrypoint, or the best-effort STARTED upload
+            # failed while the warmup is still running (the pod log disambiguates). warmup.log is only
+            # uploaded AFTER the warmup process returns, so it's present only if the warmup
+            # finished/was killed, not on a pure mid-run hang.
+            log(f"outcome={outcome}: pulling pod-side out/ for diagnostics before cleanup")
+            try:
+                dbg = os.path.join(args.out, ".dbg")
+                snapshot_download(
+                    repo_id=repo,
+                    repo_type="dataset",
+                    allow_patterns=["out/**"],
+                    local_dir=dbg,
+                    token=token,
+                )
+                src = os.path.join(dbg, "out")
+                if os.path.isdir(src) and os.listdir(src):
+                    started = os.path.isfile(os.path.join(src, "STARTED"))
+                    log(f"warmup entrypoint ran (out/STARTED present): {started}")
+                    for root, _, files in os.walk(src):
+                        for f in sorted(files):
+                            p = os.path.join(root, f)
+                            log(f"   out/{os.path.relpath(p, src)} ({os.path.getsize(p)} b)")
+                    wl = os.path.join(src, "warmup.log")
+                    if os.path.isfile(wl):
+                        log("--- warmup.log tail (last 60 lines) ---")
+                        with open(wl, errors="replace") as wlf:
+                            for line in wlf.read().splitlines()[-60:]:
+                                log(f"   | {line}")
+                    else:
+                        log(
+                            "no warmup.log (warmup never returned -> mid-run hang or still running at deadline)"
+                        )
+                else:
+                    # Ambiguous: out/STARTED is uploaded best-effort BEFORE warmup, and the cache tree
+                    # only lands AFTER warmup returns. So an empty out/ means EITHER the entrypoint
+                    # never ran (wrong CMD / docker_args) OR the STARTED upload failed while warmup is
+                    # still mid-run. The pod retries that upload 3x, so a dropped marker is unlikely and
+                    # wrong-entrypoint is the probable cause -- but don't assert it. The only proof is
+                    # the pod's OWN console (its "[bake] uploaded out/STARTED" / WARNING lines), which
+                    # this helper can't fetch (pod terminated below) -> read it in the RunPod dashboard.
+                    log(
+                        "pod produced NO out/ -> warmup entrypoint probably never ran "
+                        "(wrong CMD / docker_args); the only other cause is a failed out/STARTED "
+                        "upload mid-run (unlikely, it retries 3x) -- confirm via the pod console "
+                        "in the RunPod dashboard (this CI log does not capture pod stdout)"
+                    )
+                shutil.rmtree(dbg, ignore_errors=True)
+            except Exception as e:
+                log(f"diagnostic fetch failed (ignore): {str(e)[:160]}")
     finally:
         if pod_id:
             try:

{freesolo_flash_dev-0.2.30 → freesolo_flash_dev-0.2.31}/docker/bake_pod_entry.py

@@ -28,6 +28,33 @@ def main() -> int:
     arch = os.environ.get("BAKE_ARCH", "")
     api = HfApi(token=token)
 
+    # Upload a STARTED marker FIRST -- before the code download / chalk install / warmup -- so a CI-side
+    # timeout can tell whether this entrypoint even ran. Present on a timeout = the warmup started (so it
+    # hung or was slow); absent = the pod ran the wrong entrypoint, OR this upload itself failed while
+    # the warmup is still running. Retry a few times so a transient HF blip doesn't drop the marker and
+    # make the helper misread a slow run as a wrong entrypoint.
+    started_uploaded = False
+    for attempt in range(3):
+        try:
+            api.upload_file(
+                path_or_fileobj=io.BytesIO(f"started arch={arch}\n".encode()),
+                path_in_repo="out/STARTED",
+                repo_id=repo,
+                repo_type="dataset",
+            )
+            started_uploaded = True
+            print("[bake] uploaded out/STARTED", flush=True)
+            break
+        except Exception as e:
+            # keep the "out/STARTED" token in the line so it greps the same as the success/WARNING lines
+            print(f"[bake] out/STARTED upload attempt {attempt + 1}/3 failed: {e}", flush=True)
+            if attempt < 2:  # no point sleeping after the last attempt
+                time.sleep(3)
+    if not started_uploaded:
+        # The marker never landed; warn loudly so a later "no out/STARTED" is read as upload-failed,
+        # not as a wrong-entrypoint run.
+        print("[bake] WARNING: out/STARTED never uploaded; warmup still proceeds", flush=True)
+
     # the flash code the helper uploaded (upload_code -> code/flash); run the warmup from it.
     snapshot_download(
         repo_id=repo,

{freesolo_flash_dev-0.2.30 → freesolo_flash_dev-0.2.31}/flash/catalog.py

@@ -29,6 +29,18 @@ DEFAULT_GPU = "RTX 5090"
 _DEFAULT_VOCAB_SIZE = 248_320
 
 
+@dataclass(frozen=True)
+class ServingCapacity:
+    gpu: str
+    max_loras: int
+    max_lora_rank: int
+    max_model_len: int
+    serve_model_id: str = ""
+    max_num_seqs: int = 0
+    max_num_batched_tokens: int = 0
+    gpu_memory_utilization: float = 0.0
+
+
 @dataclass(frozen=True)
 class ModelInfo:
     id: str
@@ -61,6 +73,10 @@ class ModelInfo:
     notes: str = ""
     # 0 = platform default (64 GB) suffices. Runner raises gpu.disk_gb to at least this.
     min_disk_gb: int = 0
+    # Deployment capacity of the external freesolo multi-LoRA serving app. This is separate from
+    # Flash's training GPU recommendation above; serving uses Modal/vLLM and sizes hot LoRA buffers
+    # by max_loras x max_lora_rank at engine init.
+    serving: ServingCapacity | None = None
     # "none" / "hybrid" (Qwen3-style) / "always" (can't disable) / "unknown" (open-model policy)
     thinking: str = "none"
     vocab_size: int = _DEFAULT_VOCAB_SIZE
@@ -79,7 +95,20 @@ class ModelInfo:
     hidden_size: int = 0
 
     def to_dict(self) -> dict[str, Any]:
-        return asdict(self)
+        data = asdict(self)
+        serving = data.get("serving")
+        if serving is None:
+            data.pop("serving", None)
+        else:
+            for key in (
+                "serve_model_id",
+                "max_num_seqs",
+                "max_num_batched_tokens",
+                "gpu_memory_utilization",
+            ):
+                if serving.get(key) in ("", 0, 0.0, None):
+                    serving.pop(key, None)
+        return data
 
 
 DEFAULT_MODEL = "Qwen/Qwen3.5-4B"
@@ -94,6 +123,7 @@ MODELS: dict[str, ModelInfo] = {
         algos=("sft", "grpo"),
         min_vram_gb=12,
         recommended_gpu="RTX 4090",
+        serving=ServingCapacity(gpu="L4", max_loras=16, max_lora_rank=64, max_model_len=32768),
         thinking="hybrid",
         notes="On-device class SLM (131k ctx); standard Llama architecture.",
     ),
@@ -106,6 +136,7 @@ MODELS: dict[str, ModelInfo] = {
         algos=("sft", "grpo"),
         min_vram_gb=12,
         recommended_gpu="RTX 4090",
+        serving=ServingCapacity(gpu="L4", max_loras=16, max_lora_rank=64, max_model_len=32768),
         thinking="hybrid",
         notes="Smallest Qwen3.5; cheap smoke/dev runs with the modern arch.",
     ),
@@ -118,6 +149,7 @@ MODELS: dict[str, ModelInfo] = {
         algos=("sft", "grpo"),
         min_vram_gb=16,
         recommended_gpu="RTX 4090",
+        serving=ServingCapacity(gpu="L4", max_loras=16, max_lora_rank=64, max_model_len=32768),
         thinking="hybrid",
     ),
     "Qwen/Qwen3.5-4B": ModelInfo(
@@ -129,6 +161,15 @@ MODELS: dict[str, ModelInfo] = {
         algos=("sft", "grpo"),
         min_vram_gb=32,
         recommended_gpu="RTX 5090",
+        serving=ServingCapacity(
+            gpu="L4",
+            serve_model_id="lovedheart/Qwen3.5-4B-FP8",
+            max_loras=64,
+            max_lora_rank=32,
+            max_model_len=8192,
+            max_num_seqs=8,
+            gpu_memory_utilization=0.98,
+        ),
         thinking="hybrid",
         notes="Current-gen 4B. GRPO uses the sleep-mode memory recipe (hybrid arch needs "
         "extra engine state-cache); fused DeltaNet kernels ship in the default stack.",
@@ -145,6 +186,15 @@ MODELS: dict[str, ModelInfo] = {
         grpo_min_vram_gb=80,
         quant="bf16",
         recommended_gpu="A100 PCIe",
+        serving=ServingCapacity(
+            gpu="L4",
+            serve_model_id="lovedheart/Qwen3.5-9B-FP8",
+            max_loras=44,
+            max_lora_rank=32,
+            max_model_len=8192,
+            max_num_seqs=8,
+            gpu_memory_utilization=0.98,
+        ),
         thinking="hybrid",
         notes="bf16 LoRA. ~19 GB of weights; SFT fits a 48 GB card, while colocated GRPO "
         "(two bf16 copies + KV + the 248k-vocab fp32 logits) needs an 80 GB-class card "
@@ -174,6 +224,16 @@ MODELS: dict[str, ModelInfo] = {
         sleep_unsupported=True,
         quant="bf16",
         recommended_gpu="H200",
+        serving=ServingCapacity(
+            gpu="A100-80GB",
+            serve_model_id="Qwen/Qwen3.6-35B-A3B-FP8",
+            max_loras=12,
+            max_lora_rank=32,
+            max_model_len=8192,
+            max_num_seqs=8,
+            max_num_batched_tokens=4096,
+            gpu_memory_utilization=0.98,
+        ),
         thinking="hybrid",
         min_disk_gb=200,
         notes="MoE (35B total / ~3B active), bf16 LoRA. SFT runs on the 141 GB H200 (the ~70 GB "
@@ -199,6 +259,24 @@ def get_model(model_id: str) -> ModelInfo:
         ) from exc
 
 
+def serving_lora_rank_cap(model: str | ModelInfo | None) -> int | None:
+    """Return the model's serving LoRA rank cap, or None when Flash has no local cap.
+
+    Serving capacity is model-specific: small serving models currently allow rank 64, while larger
+    serving paths can cap at rank 32. Unknown/open-policy models intentionally return None instead
+    of inheriting a global fallback.
+    """
+    if isinstance(model, ModelInfo):
+        info = model
+    elif isinstance(model, str) and model.strip():
+        info = MODELS.get(model.strip())
+    else:
+        info = None
+    if info is None or info.serving is None:
+        return None
+    return int(info.serving.max_lora_rank)
+
+
 def vocab_size_for(model_id: str) -> int:
     """Curated vocab_size for a model, or the safe default for open-model-policy entries."""
     info = MODELS.get(model_id)

{freesolo_flash_dev-0.2.30 → freesolo_flash_dev-0.2.31}/flash/cli/commands.py

@@ -338,7 +338,11 @@ def cmd_env_list(args) -> int:
 
 
 def _cmd_train_cost(args) -> int:
-    """Print the pre-flight USD cost for the config and exit without submitting."""
+    """`flash train --cost`: print the pre-flight USD cost for the config and exit (no submit).
+
+    Catalog-only and deterministic; SFT uses the actual local training-token count when the env
+    and tokenizer are importable. An uncapped SFT run must be able to count the env's train split,
+    otherwise it errors instead of guessing a dataset size."""
     from flash.cost import estimate_cost
 
     spec = spec_from_file(

{freesolo_flash_dev-0.2.30 → freesolo_flash_dev-0.2.31}/flash/cli/render.py

@@ -592,7 +592,8 @@ def cost_panel(est) -> str:
         ),
         (
             "setup",
-            f"{est.setup_seconds / 60:.1f} min  {_dim(f'(cold start: boot + deps + model load{setup_extra})')}",
+            f"{est.setup_seconds / 60:.1f} min  "
+            f"{_dim(f'(cold start: boot + deps + model load{setup_extra}; not billed)')}",
         ),
         ("per step", f"{est.seconds_per_step:.2f} s"),
         (
@@ -601,6 +602,7 @@ def cost_panel(est) -> str:
             + (_paint("  [capped at wall-clock limit]", _RED) if est.wall_capped else ""),
         ),
         ("wall clock", f"{est.wall_clock_hours:.2f} h"),
+        ("billable", f"{est.billable_hours:.2f} h  {_dim('(training only)')}"),
     ]
     panel = _kv(pairs)
     total = f"  {_paint('TOTAL'.ljust(10), _GRAY, '1')} {_paint(_glyph('·', '-'), _FAINT)} {_paint(f'${est.total_usd:.2f}', _TEAL, '1')}"

{freesolo_flash_dev-0.2.30 → freesolo_flash_dev-0.2.31}/flash/cli/training_doc.py

@@ -302,9 +302,12 @@ Pick SFT when you already have good answers and want the model to imitate them.
   more.
 - **SFT is a great warm start for GRPO.** SFT first to teach the format and a competent
   baseline, then GRPO to optimize past it. Across that lineage keep the **same base
-  model and the same `lora_rank` / `lora_alpha`** — `init_from_adapter` loads a LoRA
-  adapter specific to one base model and one adapter shape, so mixing sizes is an
-  invalid shape mismatch.
+  model**. For text-only continued adapters, keep the same adapter shape. For VL
+  warm-starts, Flash trains a fresh GRPO LoRA and rank-stacks it with the SFT LoRA
+  for deployment, so `SFT rank + GRPO rank` must stay within the selected model's
+  effective serving `max_lora_rank`. That cap comes from the serving/model policy:
+  some small serving models allow rank 64, while larger serving paths can cap at rank
+  32. Flash preflights the rank-stacked deploy rank against that model-specific cap.
 
 ```toml
 # configs/rl.toml — warm-start GRPO from the SFT run's adapter
@@ -314,8 +317,8 @@ algorithm = "grpo"
 # paste the full adapter_ref `flash status <sft-run-id>` prints, verbatim
 # (shape: <owner>/<repo>:sft/<run-id> — the owner/repo prefix is required)
 init_from_adapter = "your-org/your-repo:sft/<sft-run-id>"
-lora_rank = 32     # must match the SFT run
-lora_alpha = 64    # must match the SFT run
+lora_rank = 16     # for VL warm-starts, SFT rank + GRPO rank must fit the effective serving cap
+lora_alpha = 32
 ```
 
 SFT is **epoch-driven** (`epochs`); GRPO is **step-driven** (`steps`).

{freesolo_flash_dev-0.2.30 → freesolo_flash_dev-0.2.31}/flash/cost/__init__.py

@@ -1,4 +1,5 @@
-"""Flash training-cost estimator: deterministic pre-flight estimate of wall-clock hours x $/hr."""
+"""Flash training-cost estimator: a deterministic, equation-based pre-flight estimate
+(``estimate_cost``) of cost = training-only GPU hours x market $/hr. No output multiplier."""
 
 from __future__ import annotations
 

{freesolo_flash_dev-0.2.30 → freesolo_flash_dev-0.2.31}/flash/cost/analytical.py

@@ -1,6 +1,9 @@
-"""The analytical cost model: total = wall-clock hours x GPU $/hr, where wall = cold-start
-setup + steps x per-step time (a FLOPs/MFU estimate). GRPO splits each step into a vLLM
-rollout + reward grading + policy/reference update."""
+"""The analytical cost model: total = training-only GPU hours x GPU $/hr.
+
+Elapsed wall clock still includes cold-start setup + steps x per-step time, but setup/cold-start
+is reported as non-billable. GRPO splits each step into a vLLM rollout + reward grading +
+policy/reference update.
+"""
 
 from __future__ import annotations
 
@@ -38,11 +41,11 @@ REWARD_CONCURRENCY = 16.0
 
 # Cold-start overhead (seconds): container boot + deps + model load (+ vLLM init for GRPO).
 #
-# Calibrated against a real fresh-worker run (0.8B SFT, RTX 3090 @ $0.239/hr) whose billed wall
+# Calibrated against a real fresh-worker run (0.8B SFT, RTX 3090 @ $0.239/hr) whose elapsed wall
 # was ~708s for only ~26 priced steps -- i.e. cold start, not training, dominated. A fresh worker
 # spent ~12.5 min in `sft_model_load` alone (download + checkpoint deserialize + GPU placement +
-# framework/CUDA init), so the MODEL-LOAD term -- not boot/deps -- is the dominant cost of a short
-# job. MODEL_LOAD_BASE_S is the fixed (size-independent) load/init overhead; the download term on
+# framework/CUDA init), so the MODEL-LOAD term -- not boot/deps -- dominates a short job's elapsed
+# time. MODEL_LOAD_BASE_S is the fixed (size-independent) load/init overhead; the download term on
 # top of it scales with checkpoint size, so bigger models pay a longer cold start.
 WORKER_BOOT_S = 120.0  # container pull + start
 DEPS_INSTALL_S = 90.0  # pip/uv resolve + install
@@ -64,9 +67,9 @@ def _fmt_duration(seconds: float) -> str:
 
 
 def setup_seconds(config: RunConfig) -> float:
-    """Cold-start wall time billed before the first optimizer step: container boot + deps + model
-    load (a fixed deserialize/placement/init base + a size-scaled download), plus vLLM init for
-    GRPO. The model-load term dominates a short job's bill (see the constants above)."""
+    """Cold-start wall time before the first optimizer step: container boot + deps + model load
+    (a fixed deserialize/placement/init base + a size-scaled download), plus vLLM init for GRPO.
+    This elapsed setup time is reported but not included in customer-facing cost."""
     model_load = MODEL_LOAD_BASE_S + download_weight_gb(config.model_id) / DOWNLOAD_RATE_GBPS
     s = WORKER_BOOT_S + DEPS_INSTALL_S + model_load
     if config.is_grpo:
@@ -96,6 +99,15 @@ def seconds_per_step(config: RunConfig, gpu: str) -> float:
     return 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()
+    params = active_params_b(n.model_id) * 1e9
+    peak = gpu_tflops(gpu) * 1e12
+    flops = SFT_FLOPS_PER_TOKEN_PER_PARAM * params * train_tokens
+    return flops / (peak * MFU_SFT_TRAIN)
+
+
 def select_gpu(config: RunConfig) -> tuple[str, int]:
     """(chosen GPU class, required VRAM GB): the cheapest fitting class for the cost.
 
@@ -128,6 +140,8 @@ def _notes(config: RunConfig, raw_train_s: float, wall_capped: bool, cap_s: floa
             + (f", env {n.environment}" if n.environment else "")
             + ") + policy+reference update"
         )
+    elif n.train_tokens is not None:
+        notes.append(f"SFT priced on {n.train_tokens:,} actual train tokens")
     notes.append(f"GPU sized with {vram_headroom() - 1:.0%} VRAM headroom; static GPU $/hr")
     if wall_capped:
         notes.append(
@@ -141,14 +155,18 @@ def estimate_cost(config: RunConfig, *, wall_cap_s: float = DEFAULT_WALL_CAP_S)
     """Deterministic pre-flight cost calculation."""
     gpu, need = select_gpu(config)
     hourly = gpu_hourly_usd(gpu, provider=config.provider)
-    # Mirror the runner's max(60, max_wall_seconds) floor so a sub-60s cap isn't underpriced.
+    # Mirror the runner's max(60, max_wall_seconds) floor so elapsed wall time is not undercounted.
     cap_s = max(60.0, float(config.max_wall_seconds)) if config.max_wall_seconds is not None else wall_cap_s
 
     setup = setup_seconds(config)
     sps = seconds_per_step(config, gpu)
     raw_train = config.steps * sps
+    if not config.is_grpo and config.train_tokens is not None:
+        raw_train = sft_seconds_for_tokens(config, gpu, config.train_tokens)
+        sps = raw_train / config.steps
 
-    # The cap is on total wall; setup is billed too, so clamp training to fit it.
+    # The cap is on total elapsed wall; setup is reported but not billed, so only training
+    # contributes to total_usd.
     wall_capped = (setup + raw_train) > cap_s
     setup = min(setup, cap_s)
     train = max(0.0, cap_s - setup) if wall_capped else raw_train
@@ -168,6 +186,6 @@ def estimate_cost(config: RunConfig, *, wall_cap_s: float = DEFAULT_WALL_CAP_S)
         train_seconds=train,
         wall_clock_seconds=wall,
         wall_capped=wall_capped,
-        total_usd=wall / 3600.0 * hourly,
+        total_usd=train / 3600.0 * hourly,
         notes=_notes(config, raw_train, wall_capped, cap_s),
     )

{freesolo_flash_dev-0.2.30 → freesolo_flash_dev-0.2.31}/flash/cost/facts.py

@@ -9,9 +9,11 @@ from flash.providers.base import GPU_INFO, GpuClass, providers_for
 GPU_COMPUTE_TFLOPS: dict[str, float] = {
     "RTX 4090": 165.0,
     "RTX 5090": 210.0,
-    "RTX A6000": 155.0,
     "A100 PCIe": 312.0,
     "A100 SXM": 312.0,
+    # A100 SXM 40GB: same SMs/tensor cores as the 80GB A100 SXM, less HBM only (Lambda-only class).
+    # Without this, removing A6000 leaves the 33-40 GB Lambda pick falling back to _DEFAULT_TFLOPS.
+    "A100 SXM 40GB": 312.0,
     "H100": 990.0,
     # H200: same SMs/tensor cores as H100, more HBM only.
     "H200": 990.0,
@@ -56,7 +58,9 @@ def pick_gpu(required_vram_gb: int, *, provider: str | None = None) -> str:
     candidates = [g for g in GPU_INFO.values() if g.vram_gb >= required_vram_gb and _selectable(g)]
     if not candidates:
         raise ValueError(f"no GPU class fits >= {required_vram_gb} GB")
-    best = min(candidates, key=lambda g: (gpu_hourly_usd(g.name, provider=provider), g.vram_gb, g.name))
+    best = min(
+        candidates, key=lambda g: (gpu_hourly_usd(g.name, provider=provider), g.vram_gb, g.name)
+    )
     return best.name