freesolo-flash-dev 0.2.42__tar.gz → 0.2.44__tar.gz

{freesolo_flash_dev-0.2.42 → freesolo_flash_dev-0.2.44}/Dockerfile.worker

@@ -69,7 +69,12 @@ RUN pip install --no-cache-dir \
 # miscomputed with Triton>=3.4 (fla #640) — the fix is fla's TILELANG backend (correct there), so
 # we bake the pinned tilelang==0.1.11 + apache-tvm-ffi==0.1.11 (0.1.12 double-registers TVM-FFI ->
 # `import tilelang` aborts) and KEEP fla on every arch; flash.engine.worker._ensure_fla_fastpath_on_hopper
-# ensures the tilelang backend is live on sm90 before any model import. Validated win vs the old
+# ensures the tilelang backend is live on sm90 before any model import. The INVERSE holds on
+# B200/sm100: THERE the tilelang chunk_bwd_dqkwg miscomputes gradients while the Triton path is
+# correct — and the pinned fla default-enables tilelang wherever it imports (i.e. everywhere,
+# since it is baked here for Hopper), so the worker opts sm100 out via FLA_TILELANG=0
+# (engine.worker.perf._force_fla_triton_gdn_on_sm100; upstream default-gates tilelang to
+# Hopper-only since fla #975). Validated win vs the old
 # pure-PyTorch fallback (H100 SXM, Qwen3.5 hidden-2560 LoRA): seq4096 4.2x faster / 1.6x less mem,
 # seq16384 12.6x / 1.9x, loss matches to 1.8e-4. fla is from git: the PyPI wheel is a broken stub
 # (no fla.modules). The git URL is PINNED to a commit (not the moving default branch) for

{freesolo_flash_dev-0.2.42 → freesolo_flash_dev-0.2.44}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 0.2.42
+Version: 0.2.44
 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

{freesolo_flash_dev-0.2.42 → freesolo_flash_dev-0.2.44}/flash/cli/commands.py

@@ -783,12 +783,32 @@ def cmd_runs(args) -> int:
 
 
 def cmd_cancel(args) -> int:
-    status = client_from_config().cancel_run(args.run_id)
+    client = client_from_config()
+    status = client.cancel_run(args.run_id)
     payload = {"run_id": args.run_id, "state": status["state"]}
+    # A cancelled run is not necessarily worthless: every completed save interval already streamed
+    # a deployable checkpoint, even though the run shows adapter_ref=null / cost=0. Surface the
+    # surviving steps here so the run isn't discarded unseen. Best-effort: cancel never fails on it.
+    checkpoints: list[dict] = []
+    if payload["state"] == "cancelled":
+        try:
+            checkpoints = client.checkpoints(args.run_id)
+        except Exception:
+            checkpoints = []
     if render.styled():
         print(render.cancelled(payload))
     else:
         print(json.dumps(payload, indent=2))
+    if checkpoints:
+        top = max(c["step"] for c in checkpoints)
+        # stderr in the plain path so the machine-readable stdout JSON stays untouched.
+        out = sys.stdout if render.styled() else sys.stderr
+        msg = (
+            f"{len(checkpoints)} deployable checkpoint(s) survive this cancel — list with "
+            f"`flash checkpoints {args.run_id}`, deploy one with "
+            f"`flash deploy {args.run_id}/step-{top}`."
+        )
+        print(render.note(msg) if render.styled() else msg, file=out)
     return 0
 
 

{freesolo_flash_dev-0.2.42 → freesolo_flash_dev-0.2.44}/flash/engine/worker/__init__.py

@@ -100,6 +100,7 @@ from flash.engine.worker.perf import (
     _estimate_params,
     _flash_attn_3_available,
     _flash_attn_available,
+    _force_fla_triton_gdn_on_sm100,
     _GpuPeakSampler,
     _liger_default_for_model,
     _memory_mode,
@@ -108,6 +109,7 @@ from flash.engine.worker.perf import (
     _peak_gpu_gb,
     _remove_fla_from_disk,
     _reset_peak_gpu,
+    _restrict_fla_gdn_autotune_on_blackwell,
     _sdpa_cudnn_ctx,
     free_gpu,
     fused_optim_name,
@@ -265,9 +267,15 @@ def main():
                 raise RetriableInfraError(
                     f"DONE present but metrics.json unreadable after retries (transient HF): {last_err}"
                 )
+        # BEFORE any model import / fla dispatch: on sm100 the baked tilelang GDN backend
+        # computes wrong gradients — opt out so fla uses its (correct-there) Triton path.
+        _force_fla_triton_gdn_on_sm100()
         _ensure_fla_fastpath_on_hopper()
         # Must run AFTER fla fast path (may reinstall tilelang) and BEFORE model/vLLM import.
         _neutralize_tilelang_cudart_stub()
+        # AFTER the fla fast path (which may (re)install fla), BEFORE any model import / GDN
+        # launch: restrict fla's Blackwell GDN bwd autotune to grad-correct configs (fla #913).
+        _restrict_fla_gdn_autotune_on_blackwell()
         heartbeat("boot", gpu=gpu_diagnostics(include_torch=False))
         finalize_alloc_conf_for_sleep()
         load_mega_cache()
@@ -344,6 +352,7 @@ __all__ = [
     "_finalize",
     "_flash_attn_3_available",
     "_flash_attn_available",
+    "_force_fla_triton_gdn_on_sm100",
     "_grpo_is_no_op_failure",
     "_grpo_resume_already_complete",
     "_hf_upload",
@@ -359,6 +368,7 @@ __all__ = [
     "_remove_fla_from_disk",
     "_reset_peak_gpu",
     "_resolve_adapter_ref",
+    "_restrict_fla_gdn_autotune_on_blackwell",
     "_sdpa_cudnn_ctx",
     "assert_adapter_delta_nonzero",
     "assert_adapter_load_clean",

{freesolo_flash_dev-0.2.42 → freesolo_flash_dev-0.2.44}/flash/engine/worker/hf.py

@@ -310,11 +310,13 @@ def publish_deployable_checkpoint(
     return None
 
 
-# Daemon upload threads are killed on worker exit; on_train_end blocks up to this long to flush.
-_CKPT_FLUSH_TIMEOUT_S = 300.0
-# Retry/backoff for the lock-timeout fallback publish (rides out a transient concurrent-commit 409).
+# Retry/backoff for the deployable-adapter publish (rides out a transient concurrent-commit 409).
 _CKPT_FLUSH_RETRIES = 3
 _CKPT_FLUSH_BACKOFF_S = 1.0
+# Retry/backoff for each synchronous checkpoint upload. `on_save` BLOCKS the training loop on the
+# upload, so a transient HF error is retried until the step lands rather than costing the step.
+_CKPT_UPLOAD_RETRIES = 3
+_CKPT_UPLOAD_BACKOFF_S = 5.0
 
 
 def _latest_checkpoint_dir(output_dir: str) -> tuple[int, str] | None:
@@ -368,72 +370,14 @@ def _prune_stale_resume_checkpoints(keep_step: int) -> None:
             api.delete_folder(path_in_repo=folder, repo_id=_w.HF_REPO, repo_type="dataset")
 
 
-class _UploadPump:
-    """Coalescing single-slot upload queue: at most one upload in flight, at most one queued.
-
-    ONE lock guards two fields — ``uploading`` and ``queued`` (newest wins). Every producer path
-    is enqueue-then-pump and the upload thread's finally is done()-then-pump, so a save that
-    lands while the in-flight upload is draining can never be lost (no check-then-set window
-    between separate locks).
-    """
-
-    def __init__(self, start_upload) -> None:
-        self._cond = threading.Condition()
-        self._start_upload = start_upload  # (step, ckpt_dir) -> spawns thread; must call done()
-        self.uploading = False
-        self.queued: tuple[int, str] | None = None
-
-    def enqueue(self, step: int, ckpt_dir: str) -> None:
-        """Queue `step` (replacing any older queued step) and start it if idle."""
-        with self._cond:
-            busy = self.uploading
-            self.queued = (step, ckpt_dir)
-        if busy:
-            print(
-                f"[ckpt] upload busy; queued step {step} (uploads when the in-flight one finishes)"
-            )
-        self.pump()
-
-    def pump(self) -> None:
-        """Atomically start the queued upload iff nothing is in flight."""
-        with self._cond:
-            if self.uploading or self.queued is None:
-                return
-            step, ckpt_dir = self.queued
-            self.queued = None
-            gone = not os.path.isdir(ckpt_dir)
-            if not gone:
-                self.uploading = True
-        if gone:  # pruned by the trainer's save_total_limit meanwhile
-            print(f"[ckpt] queued step {step} checkpoint dir is gone; dropping")
-            return
-        self._start_upload(step, ckpt_dir)
-
-    def done(self) -> None:
-        """Mark the in-flight upload finished, then pump any queued step."""
-        with self._cond:
-            self.uploading = False
-            self._cond.notify_all()
-        self.pump()
-
-    def take_queued(self) -> tuple[int, str] | None:
-        """Remove and return the queued step, if any (for a synchronous train-end flush)."""
-        with self._cond:
-            queued, self.queued = self.queued, None
-            return queued
-
-    def wait_idle(self, timeout: float) -> bool:
-        """Wait up to `timeout` seconds for the in-flight upload (if any) to finish."""
-        with self._cond:
-            return self._cond.wait_for(lambda: not self.uploading, timeout)
-
-
 def make_checkpoint_upload_callback():
     """Return a TrainerCallback that streams each save to HF and publishes deployable per-step adapters.
 
-    Uploads run in a background daemon thread; a save that fires while one is in flight is queued
-    (newest wins) and uploaded next. on_train_end flushes the queue and the final deployable
-    synchronously before exit.
+    Uploads are SYNCHRONOUS: on_save blocks the training loop until the checkpoint is durably on
+    HF. A save therefore never returns while its upload is still running, so there is never a
+    second upload in flight when the next save fires — the "upload busy" contention (and the
+    coalescing/dropping it used to cause) cannot happen by construction. Every `save_every` step is
+    guaranteed to upload, retrying transient HF errors instead of skipping.
     """
     from transformers import TrainerCallback
 
@@ -491,72 +435,57 @@ def make_checkpoint_upload_callback():
         _prune_stale_resume_checkpoints(step)
         _w.heartbeat("checkpoint_uploaded", step=step)
 
-    def _start_upload(step: int, ckpt_dir: str) -> None:
-        """Upload `step` in a background thread; the pump is marked done (and re-pumped) after."""
+    uploaded_steps: set[int] = set()
 
-        def _upload():
+    def _upload_with_retries(step: int, ckpt_dir: str) -> bool:
+        """Run the full upload for `step`, retrying transient failures; True once it lands."""
+        for attempt in range(_CKPT_UPLOAD_RETRIES):
             try:
                 _upload_once(step, ckpt_dir)
+                uploaded_steps.add(step)
+                return True
             except Exception as e:
-                print("ckpt upload warn:", e)
-            finally:
-                pump.done()
-
-        threading.Thread(target=_upload, daemon=True).start()
-
-    # `save_every` must not be advisory: a save that fires while a prior upload is still in
-    # flight is QUEUED (newest wins) and uploaded as soon as the in-flight one finishes, instead
-    # of being dropped with "upload busy; skipping step N". Coalescing to the newest pending step
-    # keeps at most one waiter, so a slow uplink can't pile up threads, while the registered step
-    # list stays as dense as the uplink allows.
-    pump = _UploadPump(_start_upload)
+                if attempt + 1 < _CKPT_UPLOAD_RETRIES:
+                    print(
+                        f"[ckpt] step {step} upload retry {attempt + 1}/{_CKPT_UPLOAD_RETRIES}: {e}"
+                    )
+                    time.sleep(_CKPT_UPLOAD_BACKOFF_S * (attempt + 1))
+                else:
+                    print(
+                        f"[ckpt] step {step} upload FAILED after {_CKPT_UPLOAD_RETRIES} attempts: {e}"
+                    )
+                    with contextlib.suppress(Exception):
+                        _w.heartbeat("checkpoint_upload_failed", step=step)
+        return False
 
     class _CheckpointUpload(TrainerCallback):
         def on_save(self, args, state, control, **kwargs):
+            # SYNCHRONOUS on the training thread: the trainer blocks here until this checkpoint is
+            # uploaded. `save_total_limit=1` rotation (which runs earlier, inside this same save's
+            # `_save_checkpoint`) deletes the PREVIOUS checkpoint, never the current one — and the
+            # next save can't begin until we return — so the dir is stable for the whole upload and
+            # no snapshot/queue is needed. Because the upload finishes before the save returns,
+            # there is never a concurrent upload for a later save to be "busy" against.
             if not _w.HF_REPO:
                 return
             step = int(state.global_step)
             ckpt_dir = os.path.join(args.output_dir, f"checkpoint-{step}")
             if not os.path.isdir(ckpt_dir):
                 return
-            pump.enqueue(step, ckpt_dir)
+            _upload_with_retries(step, ckpt_dir)
 
         def on_train_end(self, args, state, control, **kwargs):
-            # Daemon threads are killed on exit; flush synchronously here.
+            # Safety net for a final checkpoint the trainer wrote without an on_save (e.g. a
+            # load_best_model_at_end / end-of-training save). Synchronous on_save already uploaded
+            # every step it saw, so this only publishes an as-yet-unuploaded latest step.
             if not _w.HF_REPO:
                 return
             latest = _latest_checkpoint_dir(args.output_dir)
             if latest is None:
                 return
             step, ckpt_dir = latest
-            # A queued save the pump never started must not be silently dropped: take it out
-            # of the pump (so a racing done()->pump() can't double-start it) and run it here.
-            queued = pump.take_queued()
-            idle = pump.wait_idle(_CKPT_FLUSH_TIMEOUT_S)
-            if queued is not None and os.path.isdir(queued[1]):
-                q_step, q_dir = queued
-                if idle:
-                    try:
-                        _upload_once(q_step, q_dir)
-                    except Exception as e:
-                        print("ckpt upload warn:", e)
-                else:
-                    # Another upload is still pruning latest-only resume checkpoints. Publish only
-                    # the durable per-step adapter here so concurrent prune passes cannot delete a
-                    # newer resume checkpoint while it is being flushed.
-                    try:
-                        _publish_deployable_recombined(q_dir, q_step, with_retry=True)
-                    except Exception as e:
-                        print("ckpt deployable flush warn:", e)
-                if q_step == step:
-                    return  # the queued flush already published the final deployable
-            if idle:
-                _publish_deployable_recombined(ckpt_dir, step)
-            else:
-                print(
-                    f"[ckpt] upload still in flight after {_CKPT_FLUSH_TIMEOUT_S:.0f}s; publishing "
-                    f"final deployable (step {step}) alongside it"
-                )
-                _publish_deployable_recombined(ckpt_dir, step, with_retry=True)
+            if step in uploaded_steps:
+                return
+            _upload_with_retries(step, ckpt_dir)
 
     return _CheckpointUpload()

{freesolo_flash_dev-0.2.42 → freesolo_flash_dev-0.2.44}/flash/engine/worker/kernel_warmup.py

@@ -77,7 +77,9 @@ def _require_gpu():
         if not torch.cuda.is_available():
             _log("no CUDA device visible — kernel warmup must run on a GPU builder; nothing baked")
             return None
-        _log(f"GPU: {torch.cuda.get_device_name(0)} ({_torch_sm(torch)}), torch {torch.__version__}")
+        _log(
+            f"GPU: {torch.cuda.get_device_name(0)} ({_torch_sm(torch)}), torch {torch.__version__}"
+        )
         return torch
     except Exception as e:
         _log(f"torch unavailable ({e}); cannot warm kernels")
@@ -128,6 +130,13 @@ def warm_flash_attn(torch) -> bool:
 def warm_fla_gdn(torch) -> bool:
     """Compile flash-linear-attention's Gated-DeltaNet chunk kernels (Qwen3.5/3.6 hybrid path)."""
     try:
+        # sm100: tilelang's GDN bwd miscomputes grads — warm the Triton path the worker will run.
+        try:
+            from flash.engine.worker.perf import _force_fla_triton_gdn_on_sm100
+
+            _force_fla_triton_gdn_on_sm100()
+        except Exception as e:
+            _log(f"sm100 fla tilelang opt-out skipped: {e}")
         # fla #640: GDN chunk_bwd needs tilelang on Hopper; Triton path miscomputes/raises there.
         try:
             from flash.engine.worker.perf import _ensure_fla_fastpath_on_hopper
@@ -135,6 +144,14 @@ def warm_fla_gdn(torch) -> bool:
             _ensure_fla_fastpath_on_hopper()
         except Exception as e:
             _log(f"hopper fla fast-path setup skipped: {e}")
+        # fla #913: on Blackwell, warm the SAME restricted autotune space the worker boot will
+        # run (grad-correct configs only) so the baked cache holds the config that actually runs.
+        try:
+            from flash.engine.worker.perf import _restrict_fla_gdn_autotune_on_blackwell
+
+            _restrict_fla_gdn_autotune_on_blackwell()
+        except Exception as e:
+            _log(f"blackwell fla autotune restriction skipped: {e}")
         from fla.ops.gated_delta_rule import chunk_gated_delta_rule
 
         b, h, t, d = 1, 4, 256, 64

{freesolo_flash_dev-0.2.42 → freesolo_flash_dev-0.2.44}/flash/engine/worker/perf/__init__.py

@@ -131,7 +131,7 @@ def _find_real_libcudart() -> str | None:
             with open("/proc/self/maps") as f:
                 for line in f:
                     if base in line and "/" in line:
-                        p = line[line.index("/"):].rstrip()
+                        p = line[line.index("/") :].rstrip()
                         if os.path.basename(p).startswith(base) and os.path.exists(p):
                             return os.path.realpath(p)
         except OSError:
@@ -208,6 +208,138 @@ def _neutralize_tilelang_cudart_stub() -> None:
         print(f"[worker] libcudart stub neutralize failed: {e}", flush=True)
 
 
+def _force_fla_triton_gdn_on_sm100() -> None:
+    """Opt fla OUT of its tilelang GDN backend on B200 (sm100) via ``FLA_TILELANG=0``.
+
+    The worker image bakes tilelang for the HOPPER fast path (fla #640), but the pinned fla
+    default-enables its tilelang backend WHEREVER tilelang imports. On sm100, tilelang's
+    ``chunk_bwd_dqkwg`` computes WRONG GRADIENTS — measured on a real B200 with this exact image
+    at the production Qwen3.5/3.6 GDN call shapes (H==HV after transformers' head repeat):
+    dq/dk rel-err ~0.72, dg ~1.28 vs the fp32 reference, deterministic, in bf16 AND fp32, while
+    the same shapes on H200 are correct to ~4e-3. This is the root cause of the B200 35B-A3B SFT
+    incident (grad_norm ~1e8 from the first logged step; loss flat at LR 2e-5/5e-5, collapse at
+    1e-4 — garbage gradient DIRECTIONS no LR can fix, while H200 trained healthily at 1e-4).
+
+    ``FLA_TILELANG=0`` flips fla's dispatch to its Triton chunk_bwd_dqkwg — verified correct on
+    sm100 at the same shapes (~4e-3). Upstream has since stopped default-enabling tilelang
+    off-Hopper (fla #975 gates the default to sm90 + Triton>=3.4) and an explicit env opt-out
+    wins on every fla version, so this env set reproduces upstream's current gate on our pin and
+    stays a correct no-op across a future pin bump.
+
+    * ``setdefault``: an explicitly pre-set FLA_TILELANG (e.g. testing a fixed tilelang) wins.
+    * sm100 only: sm90 NEEDS tilelang (fla #640, the fast-path installer owns it); sm89/sm120
+      train healthily today under the pin's default and upstream's next-pin gate flips them to
+      Triton anyway — no evidence to justify changing them from here.
+    * Env is read live per-dispatch by fla, but set it at boot BEFORE any model import so no
+      kernel ever launches under the old default.
+    """
+    try:
+        import torch
+
+        if not (torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0)):
+            return
+    except Exception:
+        return
+    prior = os.environ.get("FLA_TILELANG")
+    if prior is not None:
+        print(
+            f"[blackwell] FLA_TILELANG={prior!r} pre-set; respecting it "
+            "(default here is 0 on sm100: tilelang chunk_bwd_dqkwg miscomputes grads)",
+            flush=True,
+        )
+        return
+    os.environ["FLA_TILELANG"] = "0"
+    print(
+        "[blackwell] sm100 (B200): FLA_TILELANG=0 -> fla GDN uses the Triton chunk_bwd_dqkwg "
+        "(tilelang backend computes wrong gradients on sm100; measured dq/dk ~0.72, dg ~1.28 "
+        "rel-err; upstream default-gates tilelang to Hopper since fla #975)",
+        flush=True,
+    )
+
+
+def _restrict_fla_gdn_autotune_on_blackwell() -> None:
+    """Restrict fla's GDN ``prepare_wy_repr_bwd_kernel`` autotune space on Blackwell (sm100/sm120).
+
+    Defense-in-depth for the Triton GDN backward path — the path sm100 defaults to once
+    ``_force_fla_triton_gdn_on_sm100`` opts out of tilelang. Upstream reports that Blackwell can
+    select unstable Triton configs for this kernel during autotuning (fla #913) and restricted
+    the space to the B200-validated config (num_warps=2, num_stages=4) in fla #1000 (ac6c648) —
+    NEWER than the worker's pinned fla SHA — so apply the same restriction in-process at boot,
+    before any GDN kernel launch. (Not the 35B-A3B incident root cause — that was tilelang, see
+    above — but there is no reason to keep autotuning over configs upstream deemed unstable on
+    the exact card we train on.)
+
+    * No-op on non-Blackwell arches, when fla is absent (pure-PyTorch delta path), or when the
+      pinned fla already carries the upstream fix (filtering an already-restricted space).
+    * Fail-CLOSED like the Hopper fast path: if fla is present on Blackwell but the restriction
+      cannot be applied (autotuner layout changed — e.g. an unreviewed pin bump), physically
+      remove fla so transformers falls back to the pure-PyTorch delta rule. Slow-but-correct
+      beats a run that burns GPU-hours computing garbage gradients.
+    * sm120 is included to match upstream's IS_NVIDIA_BLACKWELL gate: 5090 smokes trained fine on
+      the unrestricted space, but that is autotune luck (the winning config varies per shape/run),
+      not evidence of safety.
+    """
+    try:
+        import torch
+
+        if not (torch.cuda.is_available() and torch.cuda.get_device_capability()[0] in (10, 12)):
+            return
+    except Exception:
+        return
+    try:
+        import importlib.util
+
+        if importlib.util.find_spec("fla") is None:
+            return  # no fla -> transformers already uses the pure-PyTorch delta path
+    except Exception:
+        return
+
+    def _fail_closed(reason: str) -> None:
+        _removed, _still = _remove_fla_from_disk()
+        print(
+            f"[blackwell] fla GDN autotune restriction FAILED ({reason}) -> DISABLING fla "
+            f"(removed {len(_removed)} copy(ies); still_importable={_still}); "
+            "pure-PyTorch delta fallback (fla #913: unrestricted autotune miscomputes grads)",
+            flush=True,
+        )
+
+    try:
+        from fla.ops.gated_delta_rule import wy_fast
+
+        tuner = getattr(wy_fast, "prepare_wy_repr_bwd_kernel", None)
+        for _ in range(8):  # unwrap decorator layers (heuristics/cache) down to the Autotuner
+            if tuner is None or hasattr(tuner, "configs"):
+                break
+            tuner = getattr(tuner, "fn", None)
+        configs = getattr(tuner, "configs", None)
+        if not configs:
+            _fail_closed("prepare_wy_repr_bwd_kernel autotuner not found")
+            return
+        keep = [
+            c
+            for c in configs
+            if getattr(c, "num_warps", None) == 2 and getattr(c, "num_stages", None) == 4
+        ]
+        if not keep:
+            _fail_closed("no B200-validated config (num_warps=2/num_stages=4) in autotune space")
+            return
+        if len(keep) == len(configs):
+            print(
+                "[blackwell] fla GDN autotune already restricted (pinned fla carries fla #1000)",
+                flush=True,
+            )
+            return
+        tuner.configs = keep
+        print(
+            f"[blackwell] fla GDN prepare_wy_repr_bwd autotune restricted {len(configs)} -> "
+            f"{len(keep)} config(s) (num_warps=2/num_stages=4; fla #913/#1000 grad miscompute)",
+            flush=True,
+        )
+    except Exception as e:
+        with contextlib.suppress(Exception):
+            _fail_closed(f"{type(e).__name__}: {e}")
+
+
 def _ensure_fla_fastpath_on_hopper() -> None:
     """Enable fla's tilelang GDN backend on Hopper (sm90); no-op elsewhere.
 
@@ -334,6 +466,7 @@ __all__ = [
     "_flash_attn_3_available",
     "_flash_attn_available",
     "_float_or_none",
+    "_force_fla_triton_gdn_on_sm100",
     "_int_or_none",
     "_liger_default_for_model",
     "_memory_mode",
@@ -344,6 +477,7 @@ __all__ = [
     "_query_nvidia_processes",
     "_remove_fla_from_disk",
     "_reset_peak_gpu",
+    "_restrict_fla_gdn_autotune_on_blackwell",
     "_round_gb_from_mib",
     "_sdpa_cudnn_ctx",
     "free_gpu",

{freesolo_flash_dev-0.2.42 → freesolo_flash_dev-0.2.44}/flash/engine/worker/sft.py

@@ -99,6 +99,21 @@ def _model_arch_dims(model_id: str) -> tuple[int, int]:
         return c_hidden, c_layers
 
 
+def select_sft_examples(train, max_examples, seed):
+    """Pick the SFT sample: the first ``max_examples`` rows of the dataset (file order), shuffled.
+
+    The slice happens BEFORE the shuffle so ``max_examples`` is a deterministic prefix fence,
+    not a random subsample. A train.jsonl that carries fully-labeled SFT rows first and
+    prompt-only (empty-output) GRPO rows after can cap SFT to the labeled head — an empty
+    completion can never be shuffled into the SFT sample and teach the model to emit nothing.
+    """
+    if max_examples > 0:
+        train = train[:max_examples]
+    rng = random.Random(seed)
+    rng.shuffle(train)
+    return train
+
+
 def run_sft():
     from datasets import Dataset
     from transformers import AutoTokenizer
@@ -122,12 +137,7 @@ def run_sft():
         val = getattr(_t, name, None) if _t else None
         return val if val is not None else default
 
-    train = env.dataset()
-    rng = random.Random(_w.SEED)
-    rng.shuffle(train)
-    max_examples = int(_train_opt("max_examples", 0) or 0)
-    if max_examples > 0:
-        train = train[:max_examples]
+    train = select_sft_examples(env.dataset(), int(_train_opt("max_examples", 0) or 0), _w.SEED)
     texts = []
     multiturn_targets = 0
     for ex in train:

{freesolo_flash_dev-0.2.42 → freesolo_flash_dev-0.2.44}/flash/providers/_worker.py

@@ -41,7 +41,9 @@ WORKER_DEPS = [
     "wandb>=0.17",
     # fla from git: PyPI wheel is a broken stub missing fla.modules. SHA-pinned for reproducibility;
     # keep in lockstep with Dockerfile.worker. fla kept on ALL arches — worker ensures tilelang
-    # backend on sm90 before model import (fla #640: chunk_bwd miscompute with Triton>=3.4 on Hopper).
+    # backend on sm90 before model import (fla #640: chunk_bwd miscompute with Triton>=3.4 on Hopper)
+    # and OPTS OUT of tilelang on sm100 (B200) where tilelang's chunk_bwd_dqkwg miscomputes grads
+    # (worker _force_fla_triton_gdn_on_sm100; upstream default-gates tilelang to Hopper since fla #975).
     "flash-linear-attention @ git+https://github.com/fla-org/flash-linear-attention.git@f0e213dbd8b5fb90c3c7eca869ac1706d5377139",
     # tilelang version-pinned in lockstep with Dockerfile.worker + perf.py runtime reinstall.
     "tilelang==0.1.11",
@@ -110,13 +112,14 @@ def _effective_worker_env(spec=None) -> dict[str, str]:
 # Chalk is published to PUBLIC PyPI on every version bump (chalk .github/workflows/publish.yml).
 # Pin the exact PUBLIC version so worker + kernel-cache-bake pods install with NO GitHub auth — the
 # chalk repo is INTERNAL, so a git+https default made installs fail wherever GITHUB_TOKEN is absent
-# (bake pod, tokenless workers). freesolo-chalk 0.5.2 = chalk dev->main sync #44 (standalone, Liger
-# fully replaced; adds arch-TUNED gdn_gated_rmsnorm + lora on sm89/sm90/sm100, MoE grouped-GEMM,
-# serving surface, MiniCPM/Llama kernels). Bump DEFAULT_CHALK_VERSION when the merged kernel surface moves.
-DEFAULT_CHALK_VERSION = "0.5.2"
+# (bake pod, tokenless workers). freesolo-chalk 0.5.3 = chalk dev->main sync #67 (adds the decode
+# serving kernels rope + GDN gated-rmsnorm; GPU-A/B'd launch tuning: flce chunk_mult 16, rope
+# num_warps=2, gdn_conv warps formula, lora fwd autotune candidate; OP_MEAS_OVERRIDE; flce
+# chunk_mult tradeoff docs). Bump DEFAULT_CHALK_VERSION when the merged kernel surface moves.
+DEFAULT_CHALK_VERSION = "0.5.3"
 DEFAULT_CHALK_SPEC = f"freesolo-chalk=={DEFAULT_CHALK_VERSION}"
 # Provenance only (kernel-cache fingerprint / traceability): the chalk commit this version was cut from.
-LATEST_CHALK_MAIN_SHA = "37b2fb9e1a123ceb902ba6ffc4ed182182c6e7d7"
+LATEST_CHALK_MAIN_SHA = "c63f0367e01795006bc0e91e2540ec85e740f19b"
 
 
 def chalk_extra_pip(spec=None) -> list[str]:

{freesolo_flash_dev-0.2.42 → freesolo_flash_dev-0.2.44}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "freesolo-flash-dev"
-version = "0.2.42"
+version = "0.2.44"
 description = "Flash — managed LoRA post-training (SFT/GRPO) for Freesolo environments, driven by the `flash` CLI"
 readme = "README.md"
 # RunPod Flash supports Python 3.11-3.12 (not 3.13 yet).
@@ -106,7 +106,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 = "0.2.42"
+version = "0.2.44"
 
 [tool.pytest.ini_options]
 pythonpath = ["."]