freesolo-flash-dev 1.0.63__tar.gz → 1.0.65__tar.gz

{freesolo_flash_dev-1.0.63 → freesolo_flash_dev-1.0.65}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 1.0.63
+Version: 1.0.65
 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.63 → freesolo_flash_dev-1.0.65}/flash/engine/worker/rl_train.py

@@ -105,6 +105,17 @@ _MULTI_TURN_SCORE_BATCH_SIZE = 64
 _MULTI_TURN_SCORE_FLUSH_WAIT_S = 0.1
 _MULTI_TURN_SCORE_SHUTDOWN_WAIT_S = 5.0
 
+# the bridge's listen() backlog. verl opens one connection per episode and starts a whole step's
+# episodes at once, so the accept queue sees the entire rollout batch -- prompts_per_step *
+# group_size, 512 on the default 64x8 recipe -- connecting in a burst. socketserver's default of 5
+# overflows there and the kernel RESETS the excess, which surfaces client-side as
+# ConnectionResetError at getresponse() (the connect and send landed in the queue; the reset arrives
+# when the queue is dropped). bridge_post deliberately does not retry, so that reset kills the run.
+# the opd teacher bridge already carries this fix as _TEXT_TEACHER_REQUEST_BACKLOG; a fixed constant
+# would only move the cliff, so size the queue from the burst the caller actually generates.
+# linux caps the effective backlog at somaxconn (4096 here) and silently clamps beyond it.
+_REWARD_BRIDGE_MIN_REQUEST_BACKLOG = 128
+
 
 # --------------------------------------------------------------------------------------------
 # pure helpers (no verl, no gpu, no network) -- unit tested directly.
@@ -1229,14 +1240,20 @@ def render_reentrant_checkpointing_shim(reentrant: bool, *, multimodal: bool = F
     engines whose config may legitimately have checkpointing off, and enabling it there would turn
     on a feature verl chose to leave off.
 
-    on a MULTIMODAL run the same hook also restores vision input gradients. reentrant recompute
-    drops the backward for a checkpointed block when none of that block's inputs require grad, and
-    the vision tower's patch embeddings are exactly that case: the pixels are frozen inputs, so
-    without a forward hook marking the patch-embed output as requiring grad the visual modules
-    silently receive nothing while the language side trains normally and the run reports success
-    (``tests/test_multimodal_input_grads.py``). the retired trl path installed the same hook via a
-    trainer callback; verl has no callback surface, so it rides this shim instead -- and only here,
-    because non-reentrant recompute does not have the behaviour that makes it necessary
+    reentrant recompute drops the backward for a checkpointed block when none of that block's
+    inputs require grad. that hits the LANGUAGE side on every lora run: lora freezes the
+    embeddings, so the hidden states entering the first checkpointed decoder layer have
+    ``requires_grad=False`` and the whole segment -- containing every lora parameter -- receives no
+    gradient, while the run reports success and bills (GRAD-001). ``enable_input_require_grads()``
+    is what prevents it, and it is unconditional here because every flash rl run is lora.
+
+    on a MULTIMODAL run the same hook additionally restores VISION input gradients. the vision
+    tower's patch embeddings are the same failure in a place the language-side call does not
+    reach: the pixels are frozen inputs, so without a forward hook marking the patch-embed output
+    as requiring grad the visual modules silently receive nothing while the language side trains
+    normally (``tests/test_multimodal_input_grads.py``). the retired trl path installed that hook
+    via a trainer callback; verl has no callback surface, so it rides this shim instead -- and only
+    here, because non-reentrant recompute does not have the behaviour that makes it necessary
     (codex[bot]).
     """
     if not reentrant:
@@ -1290,6 +1307,13 @@ def _flash_reentrant_build_module(self):
     # only when verl actually enabled checkpointing: calling gradient_checkpointing_enable on a
     # model verl deliberately left uncheckpointed would turn it ON and change the memory profile.
     if getattr(self.model_config, "enable_gradient_checkpointing", False):
+        # the LANGUAGE-side counterpart of the vision hook below, and required for the same
+        # reason (GRAD-001): lora freezes the embeddings, so nothing entering the first
+        # checkpointed decoder layer requires grad and reentrant recompute drops the backward
+        # for the whole segment -- where every lora parameter lives. the vision hook only ever
+        # covered the patch embeddings on multimodal runs; text-only runs had no hook at all and
+        # trained nothing while reporting success.
+        module.enable_input_require_grads()
         module.gradient_checkpointing_enable(gradient_checkpointing_kwargs={{"use_reentrant": True}}){vision_hook}
         print("[rl-verl] reentrant gradient checkpointing is active", flush=True)
     return module
@@ -2071,12 +2095,15 @@ def multi_turn_child_env(inp: dict, *, reward_url: str, thinking: bool) -> dict[
     }
 
 
-def start_reward_server(score_by_index, *, example_count: int, multi_turn_bridge=None):
+def start_reward_server(
+    score_by_index, *, example_count: int, multi_turn_bridge=None, rollout_batch: int = 0
+):
     """start a localhost http reward server. score_by_index(index, solution_str) -> float.
 
     returns (server, base_url). the server runs in a daemon thread; call server.shutdown() when
     done. single-turn scoring lives at ``<base_url>/score``; when ``multi_turn_bridge`` is given,
-    its four episode routes are served alongside it from the same thread pool.
+    its four episode routes are served alongside it from the same thread pool. ``rollout_batch`` is
+    how many episodes verl starts at once (prompts_per_step * group_size) and sizes the accept queue.
     """
     # serialize scoring so the flash env sees sequential calls: a flash env is a plain python object
     # with no concurrency contract, and the retired trl path only ever called it from one thread.
@@ -2123,7 +2150,10 @@ def start_reward_server(score_by_index, *, example_count: int, multi_turn_bridge
             self.end_headers()
             self.wfile.write(body)
 
-    server = ThreadingHTTPServer(("127.0.0.1", 0), _Handler)
+    class _RewardBridgeHTTPServer(ThreadingHTTPServer):
+        request_queue_size = max(_REWARD_BRIDGE_MIN_REQUEST_BACKLOG, int(rollout_batch))
+
+    server = _RewardBridgeHTTPServer(("127.0.0.1", 0), _Handler)
     threading.Thread(target=server.serve_forever, daemon=True).start()
     port = server.server_address[1]
     return server, f"http://127.0.0.1:{port}"
@@ -3159,6 +3189,7 @@ def run_rl_train():
         _score,
         example_count=len(rollout_examples),
         multi_turn_bridge=multi_turn_bridge,
+        rollout_batch=int(inp["prompts_per_step"]) * int(inp["group_size"]),
     )
     # bound before the try so the finally can always ask whether it was started.
     resume_uploader: _VerlResumeUploader | None = None

{freesolo_flash_dev-1.0.63 → freesolo_flash_dev-1.0.65}/flash/engine/worker/sft_train.py

@@ -57,6 +57,10 @@ from flash.engine.worker.sft import (
 _SFT_LORAPLUS_RATIO = 16.0
 _LORAPLUS_READY_MARKER = "FLASH_LORAPLUS_READY"
 _VERL_STEP_RE = re.compile(r"(?:^|\s)step:(\d+)(?:\s|$)")
+# consecutive zero-grad-norm steps tolerated before the run is failed as untrainable (GRAD-001).
+# 2 is enough to separate a one-off fully-masked batch from a severed backward graph, and keeps
+# the wasted spend to a couple of steps rather than the whole run.
+_MAX_ZERO_GRAD_STEPS = 2
 
 _REQUIRED_OVERRIDE_KEYS = (
     "train_files",
@@ -535,6 +539,15 @@ _flash_original_build_module = _FlashFSDPEngine._build_module
 
 def _flash_build_reentrant_module(self):
     module = _flash_original_build_module(self)
+    # REQUIRED before enabling reentrant checkpointing on a lora model, and the reason every
+    # verl sft run silently trained nothing (GRAD-001). lora freezes the embeddings, so the
+    # hidden states entering the first checkpointed decoder layer have requires_grad=False.
+    # reentrant checkpointing needs at least one input tensor requiring grad to attach the
+    # recomputed forward to the graph; with none it warns "None of the inputs have
+    # requires_grad=True. Gradients will be None" and returns no gradient for the whole
+    # segment -- which is where every lora parameter lives. result was grad_norm 0.0 on every
+    # step while the run reported done and billed. opd already does this; sft did not.
+    module.enable_input_require_grads()
     module.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": True})
     return module
 
@@ -1340,6 +1353,9 @@ def run_sft_train(spec=None) -> None:
     ]
 
     progress = {"step": resume_step, "loss": None, "grad_norm": None, "lr": None}
+    # consecutive steps seen with grad_norm == 0.0 at a nonzero lr. one step can legitimately be
+    # zero (a fully-masked micro-batch), so require a short run of them before failing.
+    zero_grad_steps: list[int] = []
     loss_curve: list[float] = []
     train_tokens = sft_completed_train_tokens(
         total_tokens_per_epoch,
@@ -1379,6 +1395,21 @@ def run_sft_train(spec=None) -> None:
             progress["loss"] = loss
         if grad_norm is not None:
             progress["grad_norm"] = grad_norm
+            # a grad norm of exactly 0.0 while the lr is still nonzero means the backward pass
+            # produced nothing for every trainable parameter. that is never legitimate: it is a
+            # broken graph, not a small update. GRAD-001 shipped four runs that reported done and
+            # billed while training nothing, because this number was recorded and never read.
+            # fail the run instead of paying for a zero adapter that then deploys and serves.
+            if grad_norm == 0.0 and (learning_rate_value is None or learning_rate_value > 0.0):
+                zero_grad_steps.append(int(progress["step"] or 0))
+                if len(zero_grad_steps) >= _MAX_ZERO_GRAD_STEPS:
+                    raise RuntimeError(
+                        "verl reported train/grad_norm=0.0 with a nonzero learning rate on "
+                        f"{len(zero_grad_steps)} consecutive steps: no gradient is reaching the "
+                        "lora parameters, so this run would train nothing. see GRAD-001"
+                    )
+            else:
+                zero_grad_steps.clear()
         if learning_rate_value is not None:
             progress["lr"] = learning_rate_value
 
@@ -1402,6 +1433,14 @@ def run_sft_train(spec=None) -> None:
     return_code = 0
     if resume_step < update_horizon:
         watcher.start()
+        # completeness is only a meaningful question when training ran to the end. an on_line
+        # callback that raises (the zero-grad guard above, the lora+ guard) unwinds BEFORE
+        # return_code is assigned, so deriving the flag from return_code alone would leave it at
+        # its initial 0 and demand every required save from a run that stopped at step 2. the
+        # watcher would then raise "required saves were not durably published" from the finally
+        # and REPLACE the diagnosis with a downstream symptom. opd_train tracks the same flag for
+        # the same reason.
+        training_completed = False
         try:
             with liveness_heartbeat(
                 "sft_step",
@@ -1415,8 +1454,9 @@ def run_sft_train(spec=None) -> None:
                     on_line=on_line,
                     heartbeat=child_heartbeat,
                 )
+                training_completed = return_code == 0
         finally:
-            watcher.stop(require_complete=return_code == 0)
+            watcher.stop(require_complete=training_completed)
     train_wall = time.time() - train_started_at
     device_peak_gpu_gb = gpu_sampler.stop_gb()
     if return_code != 0:

{freesolo_flash_dev-1.0.63 → freesolo_flash_dev-1.0.65}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "freesolo-flash-dev"
-version = "1.0.63"
+version = "1.0.65"
 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).
@@ -123,7 +123,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.63"
+version = "1.0.65"
 
 [tool.pytest.ini_options]
 pythonpath = ["."]

{freesolo_flash_dev-1.0.63 → freesolo_flash_dev-1.0.65}/tests/test_rl_train.py

@@ -11,6 +11,7 @@ import math
 import os
 import re
 import shutil
+import socket
 import sys
 import textwrap
 import threading
@@ -1202,6 +1203,52 @@ def test_concurrent_scorers_are_serialized_for_the_env():
         server.shutdown()
 
 
+def test_reward_server_accept_queue_holds_a_whole_rollout_batch(monkeypatch):
+    # verl opens one connection per episode and starts a whole step at once, so the accept queue
+    # sees prompts_per_step * group_size connects in a burst. socketserver's default backlog of 5
+    # overflows there and the kernel RESETS the excess, which reaches the client as
+    # ConnectionResetError at getresponse(); bridge_post does not retry, so that kills the run.
+    rollout_batch = 512  # the default 64x8 recipe
+
+    # record the argument the server actually hands to listen(). asserting on request_queue_size
+    # alone cannot tell a working fix from a no-op: server_activate() reads that attribute once,
+    # so a value set after server_bind() would never reach the socket at all.
+    backlogs = []
+    real_listen = socket.socket.listen
+
+    def spy_listen(self, *args):
+        backlogs.append(args[0] if args else None)
+        return real_listen(self, *args)
+
+    monkeypatch.setattr(socket.socket, "listen", spy_listen)
+
+    server, _url = rl_train.start_reward_server(
+        lambda idx, solution: 1.0, example_count=8, rollout_batch=rollout_batch
+    )
+    try:
+        assert server.request_queue_size >= rollout_batch
+        assert backlogs, "server never called listen()"
+        assert backlogs[-1] is not None, "listen() was called with no backlog argument"
+        assert backlogs[-1] >= rollout_batch, f"listen() backlog is {backlogs[-1]}"
+    finally:
+        server.shutdown()
+        server.server_close()
+
+
+@pytest.mark.parametrize("rollout_batch", [0, 32, 64 * 8, 2048])
+def test_reward_bridge_backlog_never_falls_below_the_burst(rollout_batch):
+    # a fixed constant would only move the cliff, so the queue is sized from the caller's burst.
+    # an unspecified batch still keeps a floor well clear of socketserver's default of 5.
+    server, _url = rl_train.start_reward_server(
+        lambda idx, solution: 1.0, example_count=1, rollout_batch=rollout_batch
+    )
+    try:
+        assert server.request_queue_size >= max(128, rollout_batch)
+    finally:
+        server.shutdown()
+        server.server_close()
+
+
 # ------------------------------- reward parity -------------------------------
 class _BreakdownEnv:
     def scores_breakdown(self, graded, ex, state):
@@ -1699,6 +1746,48 @@ def test_the_reentrant_shim_installs_vision_input_grads_only_for_multimodal_runs
     assert rl_train.render_reentrant_checkpointing_shim(False, multimodal=True) == ""
 
 
+def test_the_reentrant_shim_enables_language_side_input_grads_before_checkpointing():
+    """GRAD-001: every rl run is lora, so lora freezes the embeddings and nothing entering the
+    first checkpointed decoder layer requires grad. reentrant recompute then drops the backward
+    for the whole segment -- where every lora parameter lives -- and the run reports success while
+    training nothing. the vision hook only ever covered the patch embeddings on multimodal runs;
+    text-only runs had no hook at all."""
+    calls = []
+
+    class FakeModule:
+        def enable_input_require_grads(self):
+            calls.append("require_grads")
+
+        def gradient_checkpointing_enable(self, **kwargs):
+            calls.append(("gc_enable", kwargs))
+
+    class FakeEngine:
+        def __init__(self, checkpointing):
+            self.model_config = SimpleNamespace(enable_gradient_checkpointing=checkpointing)
+
+        def _build_module(self):
+            return FakeModule()
+
+    source = rl_train.render_reentrant_checkpointing_shim(True)
+    start = source.index("def _flash_reentrant_build_module")
+    end = source.index("_FlashReentrantEngine._build_module = _flash_reentrant_build_module")
+    namespace = {"_flash_reentrant_original_build_module": FakeEngine._build_module}
+    exec(compile(source[start:end], "sitecustomize.py", "exec"), namespace)
+    build = namespace["_flash_reentrant_build_module"]
+
+    build(FakeEngine(checkpointing=True))
+    # order matters: enabling checkpointing first captures the graph before any input requires
+    # grad, so asserting mere presence would pass on a broken shim.
+    assert calls[0] == "require_grads"
+    assert calls[1] == ("gc_enable", {"gradient_checkpointing_kwargs": {"use_reentrant": True}})
+
+    # the guard is load-bearing: when verl left checkpointing OFF, touching the module at all
+    # would turn on a feature verl deliberately declined and change the memory profile.
+    calls.clear()
+    build(FakeEngine(checkpointing=False))
+    assert calls == []
+
+
 @contextlib.contextmanager
 def _vision_hook_installer(torch_module):
     """yield the rendered vision helper with a stand-in torch visible to it.
@@ -1831,6 +1920,10 @@ def test_the_reentrant_shim_flips_the_flag_and_leaves_uncheckpointed_models_alon
     class _Module:
         def __init__(self):
             self.kwargs = None
+            self.input_grads = False
+
+        def enable_input_require_grads(self):
+            self.input_grads = True
 
         def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
             self.kwargs = gradient_checkpointing_kwargs
@@ -1857,12 +1950,17 @@ def test_the_reentrant_shim_flips_the_flag_and_leaves_uncheckpointed_models_alon
             sys.modules.setdefault(name, types.ModuleType(name))
         sys.modules[module_stub.__name__] = module_stub
         exec(rl_train.render_reentrant_checkpointing_shim(True), {})
-        # checkpointing on -> the flag is put back to reentrant
+        # checkpointing on -> the flag is put back to reentrant, and the lora-frozen embeddings
+        # get input grads so the checkpointed segment actually produces a backward (GRAD-001)
         engine = _Engine(True)
-        assert engine._build_module().kwargs == {"use_reentrant": True}
+        built = engine._build_module()
+        assert built.kwargs == {"use_reentrant": True}
+        assert built.input_grads is True
         # checkpointing off -> untouched, so the shim cannot silently raise the memory profile
         off = _Engine(False)
-        assert off._build_module().kwargs is None
+        off_built = off._build_module()
+        assert off_built.kwargs is None
+        assert off_built.input_grads is False
     finally:
         for name, value in saved.items():
             if value is None:

{freesolo_flash_dev-1.0.63 → freesolo_flash_dev-1.0.65}/tests/test_sft_train.py

@@ -5,15 +5,18 @@ from __future__ import annotations
 import contextlib
 import importlib.util
 import os
+import re
 import sys
 import types
 from types import SimpleNamespace
 
 import pytest
 
+from flash.engine.worker.backend_common import parse_verl_metric
 from flash.engine.worker.sft import _pretokenize_completion_only
 from flash.engine.worker.sft_train import (
     _LORAPLUS_READY_MARKER,
+    _MAX_ZERO_GRAD_STEPS,
     _VERL_OPTIMIZER_IMPL,
     _VERL_OPTIMIZER_NAME,
     _build_verl_child_env,
@@ -478,6 +481,115 @@ def test_generated_sitecustomize_installs_linear_scheduler_and_required_loraplus
     assert scheduler_calls == [("optimizer", {"num_warmup_steps": 2, "num_training_steps": 20})]
 
 
+def test_reentrant_checkpointing_enables_input_grads_before_enabling_checkpointing():
+    """GRAD-001: lora freezes the embeddings, so nothing entering the first checkpointed layer
+    requires grad and reentrant checkpointing returns no gradient at all. the shim must call
+    enable_input_require_grads() BEFORE gradient_checkpointing_enable(), or every sft run
+    trains nothing while reporting done and billing."""
+    calls = []
+
+    class FakeModule:
+        def enable_input_require_grads(self):
+            calls.append("require_grads")
+
+        def gradient_checkpointing_enable(self, **kwargs):
+            calls.append(("gc_enable", kwargs))
+
+    class FakeEngine:
+        def _build_module(self):
+            return FakeModule()
+
+    source = _render_sft_sitecustomize(
+        seed=1,
+        loraplus_ratio=16,
+        save_at_steps=(),
+        total_steps=4,
+        reentrant_gradient_checkpointing=True,
+    )
+    # execute only the reentrant block: the surrounding shim imports verl/torch at module scope.
+    block = source[source.index("def _flash_build_reentrant_module") :]
+    block = block[: block.index("_FlashFSDPEngine._build_module = _flash_build_reentrant_module")]
+    namespace = {"_flash_original_build_module": FakeEngine._build_module}
+    exec(compile(block, "shim.py", "exec"), namespace)
+
+    namespace["_flash_build_reentrant_module"](FakeEngine())
+
+    # order matters: enabling checkpointing first would capture the graph before any input
+    # requires grad, so asserting mere presence would pass on a broken shim.
+    assert calls[0] == "require_grads"
+    assert calls[1] == ("gc_enable", {"gradient_checkpointing_kwargs": {"use_reentrant": True}})
+
+    # the non-reentrant path never patches _build_module at all, so it must not appear.
+    non_reentrant = _render_sft_sitecustomize(
+        seed=1,
+        loraplus_ratio=16,
+        save_at_steps=(),
+        total_steps=4,
+        reentrant_gradient_checkpointing=False,
+    )
+    assert "enable_input_require_grads" not in non_reentrant
+
+
+def test_zero_grad_norm_at_nonzero_lr_fails_the_run():
+    """GRAD-001: four runs reported done and charged while grad_norm was 0.0 on every step.
+    the number was parsed and recorded, never read. replay the real g4 log lines."""
+
+    def replay(lines):
+        zero_grad_steps: list[int] = []
+        for step, line in enumerate(lines):
+            grad_norm = parse_verl_metric(line, "train/grad_norm")
+            learning_rate_value = parse_verl_metric(line, "train/lr")
+            if grad_norm is None:
+                continue
+            if grad_norm == 0.0 and (learning_rate_value is None or learning_rate_value > 0.0):
+                zero_grad_steps.append(step)
+                if len(zero_grad_steps) >= _MAX_ZERO_GRAD_STEPS:
+                    raise RuntimeError(
+                        "verl reported train/grad_norm=0.0 with a nonzero learning rate on "
+                        f"{len(zero_grad_steps)} consecutive steps: no gradient is reaching "
+                        "the lora parameters, so this run would train nothing. see GRAD-001"
+                    )
+            else:
+                zero_grad_steps.clear()
+
+    # verbatim shape of the g4 gsm8k lines (flash-1785592071-e56cf3c6): loss barely moves on a
+    # replayed identical batch because nothing is learning.
+    broken = [
+        "step:1 - train/loss:0.5470 - train/grad_norm:0.0 - train/lr:0.0001",
+        "step:2 - train/loss:0.5437 - train/grad_norm:0.0 - train/lr:0.0001",
+        "step:3 - train/loss:0.5474 - train/grad_norm:0.0 - train/lr:0.0001",
+        "step:4 - train/loss:0.5444 - train/grad_norm:0.0 - train/lr:0.0001",
+    ]
+    with pytest.raises(RuntimeError, match=re.escape("grad_norm=0.0")):
+        replay(broken)
+
+    # a healthy run must not trip the guard.
+    replay(
+        [
+            "step:1 - train/loss:0.9 - train/grad_norm:1.4 - train/lr:0.0001",
+            "step:2 - train/loss:0.7 - train/grad_norm:0.9 - train/lr:0.0001",
+        ]
+    )
+
+    # an isolated zero is a legitimately fully-masked micro-batch, not a severed graph.
+    replay(
+        [
+            "step:1 - train/loss:0.9 - train/grad_norm:0.0 - train/lr:0.0001",
+            "step:2 - train/loss:0.7 - train/grad_norm:1.1 - train/lr:0.0001",
+            "step:3 - train/loss:0.6 - train/grad_norm:0.0 - train/lr:0.0001",
+        ]
+    )
+
+    # a decayed schedule reaching lr 0.0 produces a legitimate zero grad norm.
+    replay(
+        [
+            "step:1 - train/loss:0.5 - train/grad_norm:0.0 - train/lr:0.0",
+            "step:2 - train/loss:0.5 - train/grad_norm:0.0 - train/lr:0.0",
+            "step:3 - train/loss:0.5 - train/grad_norm:0.0 - train/lr:0.0",
+        ]
+    )
+
+
 def test_loraplus_shim_has_no_plain_lora_fallback():
     source = render_loraplus_shim(16)
     assert _LORAPLUS_READY_MARKER in source
@@ -615,7 +727,11 @@ def test_resume_credits_only_required_saves_that_are_durable(monkeypatch):
     assert sft_train._durable_required_save_steps((3, 5, 9), 5) == {3}
 
 
-def test_run_sft_train_orchestrates_exact_dataset_and_resume_accounting(monkeypatch):
+def _stub_sft_run(monkeypatch, *, save_at_steps=(), watcher_cls=None):
+    """monkeypatch every out-of-process dependency of run_sft_train and return (spec, captured).
+
+    the caller supplies its own ``run_verl_training`` fake, which is the only remaining seam.
+    """
     import flash.engine.worker as worker
     from flash.engine.worker import sft_train
 
@@ -630,7 +746,7 @@ def test_run_sft_train_orchestrates_exact_dataset_and_resume_accounting(monkeypa
             max_context_tokens=1024,
             max_examples=2,
             max_steps=2,
-            save_at_steps=(),
+            save_at_steps=save_at_steps,
             save_every=50,
             init_from_adapter="",
         ),
@@ -683,7 +799,7 @@ def test_run_sft_train_orchestrates_exact_dataset_and_resume_accounting(monkeypa
         def stop_gb(self):
             return 12.5
 
-    class Watcher:
+    class _DefaultWatcher:
         def __init__(self, **kwargs):
             self.processed_steps = set()
 
@@ -696,7 +812,25 @@ def test_run_sft_train_orchestrates_exact_dataset_and_resume_accounting(monkeypa
         def raise_if_failed(self):
             return None
 
+    Watcher = watcher_cls or _DefaultWatcher
+
     captured = {"heartbeats": [], "published": [], "uploads": []}
+    # run_sft_train imports AutoProcessor at data-loading time and transformers is not installed in
+    # the cpu test env. this used to pass only because some EARLIER test module left a transformers
+    # stub in sys.modules, so running this file alone failed -- stub it here so the test stands on
+    # its own. monkeypatch.setitem restores whatever was there (real module or nothing) afterwards.
+    monkeypatch.setitem(
+        sys.modules,
+        "transformers",
+        sys.modules.get("transformers")
+        or _module(
+            "transformers",
+            AutoProcessor=SimpleNamespace(from_pretrained=lambda *a, **k: None),
+            # datasets' dill serializer issubclass()-checks against this while writing the
+            # parquet, so it has to be a real class rather than a namespace attribute.
+            PreTrainedTokenizerBase=type("PreTrainedTokenizerBase", (), {}),
+        ),
+    )
     monkeypatch.delenv("WANDB_API_KEY", raising=False)
     monkeypatch.setattr(worker, "SEED", 7)
     monkeypatch.setattr(worker, "RUN_ID", "test-sft-verl-orchestration")
@@ -765,6 +899,14 @@ def test_run_sft_train_orchestrates_exact_dataset_and_resume_accounting(monkeypa
 
     monkeypatch.setattr(sft_train, "_export_checkpoint_adapter", fake_export)
 
+    return spec, captured
+
+
+def test_run_sft_train_orchestrates_exact_dataset_and_resume_accounting(monkeypatch):
+    from flash.engine.worker import sft_train
+
+    spec, captured = _stub_sft_run(monkeypatch)
+
     def fake_training(command, *, env, on_step, on_line, heartbeat):
         captured["command"] = command
         captured["child_env"] = env
@@ -798,6 +940,61 @@ def test_run_sft_train_orchestrates_exact_dataset_and_resume_accounting(monkeypa
     assert captured["meta"]["notes"]["realized_max_length"] > 0
 
 
+def test_a_guard_failure_is_not_replaced_by_the_watcher_completeness_error(monkeypatch):
+    """the zero-grad diagnosis must survive the finally block, not be overwritten by it.
+
+    an on_line guard raises from INSIDE run_verl_training, so return_code is never assigned and
+    keeps its initial 0. deriving require_complete from it would then demand every save_at_steps
+    entry from a run that died at step 2, and the watcher's "required saves were not durably
+    published" would unwind out of the finally in place of the real cause -- turning the one
+    error GRAD-001 exists to surface into a checkpointing red herring.
+    """
+    from flash.engine.worker import sft_train
+
+    stopped: list[bool] = []
+
+    class Watcher:
+        def __init__(self, **kwargs):
+            self.processed_steps = set()
+            self.required_steps = frozenset(kwargs.get("required_steps", ()))
+
+        def start(self):
+            return None
+
+        def raise_if_failed(self):
+            return None
+
+        def stop(self, *, require_complete):
+            stopped.append(require_complete)
+            if require_complete:
+                missing = sorted(self.required_steps - self.processed_steps)
+                if missing:
+                    raise RuntimeError(f"required saves were not durably published: {missing}")
+
+    # a required save the run never durably publishes: the guard raises on the same step, before
+    # the watcher has processed it. (the step is inside the 2-update horizon because
+    # validate_save_steps rejects anything beyond it at config time.)
+    spec, _ = _stub_sft_run(monkeypatch, save_at_steps=(2,), watcher_cls=Watcher)
+
+    def fake_training(command, *, env, on_step, on_line, heartbeat):
+        on_line(f"{_LORAPLUS_READY_MARKER} ratio=16 optimizer=AdamW\n")
+        for step in range(1, _MAX_ZERO_GRAD_STEPS + 1):
+            on_line(
+                f"step:{step} - train/loss:1.0 - train/grad_norm:0.0 - train/lr:5e-05 "
+                "- train/global_tokens:8\n"
+            )
+        raise AssertionError("the zero-grad guard should have stopped the run before this")
+
+    monkeypatch.setattr(sft_train, "run_verl_training", fake_training)
+
+    with pytest.raises(RuntimeError, match=re.escape("grad_norm=0.0")):
+        sft_train.run_sft_train(spec)
+
+    # the watcher still gets stopped -- it just is not asked to prove completeness for a run that
+    # never finished. without this, stop() raises and the grad_norm error never reaches the caller.
+    assert stopped == [False]
+
+
 def test_overrides_enable_fused_linear_ce_for_long_context():
     # 32k contexts must not materialize [tokens, vocab] logits; the fused torch-backend
     # linear-CE computes loss from hidden states in chunks (numerically exact CE).

{freesolo_flash_dev-1.0.63 → freesolo_flash_dev-1.0.65}/uv.lock

@@ -1025,7 +1025,7 @@ wheels = [
 
 [[package]]
 name = "freesolo-flash"
-version = "1.0.63"
+version = "1.0.65"
 source = { editable = "." }
 
 [package.optional-dependencies]