freesolo-flash-dev 1.0.71__tar.gz → 1.0.72__tar.gz

{freesolo_flash_dev-1.0.71 → freesolo_flash_dev-1.0.72}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 1.0.71
+Version: 1.0.72
 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.71 → freesolo_flash_dev-1.0.72}/flash/engine/worker/opd_train.py

@@ -38,6 +38,7 @@ from flash.engine.worker.backend_common import (
     parse_wandb_link,
     ray_num_cpus,
     render_wandb_link_shim,
+    resolve_blackwell_attention_backends,
     resolve_rollout_enforce_eager,
     resolve_verl_device_capability,
     resolve_verl_loggers,
@@ -1561,6 +1562,25 @@ def build_opd_overrides(config: dict) -> list[str]:
         # rollout.enforce_eager is a real verl field, so this is a plain override, not a '+' append.
         # the caller resolves it from the device capability; absent/false keeps verl's default.
         *(["actor_rollout_ref.rollout.enforce_eager=True"] if config.get("enforce_eager") else []),
+        # blackwell attention pins, resolved by the caller and absent off blackwell. both are real
+        # AsyncEngineArgs fields in the pinned vllm 0.19.1 and verl spreads engine_kwargs.vllm
+        # straight into them, so '+' appends under the existing struct exactly as on the grpo path.
+        *(
+            [
+                "+actor_rollout_ref.rollout.engine_kwargs.vllm.attention_backend="
+                f"{config['attention_backend']}"
+            ]
+            if config.get("attention_backend")
+            else []
+        ),
+        *(
+            [
+                "+actor_rollout_ref.rollout.engine_kwargs.vllm.mm_encoder_attn_backend="
+                f"{config['mm_encoder_attn_backend']}"
+            ]
+            if config.get("mm_encoder_attn_backend")
+            else []
+        ),
         # keep the rollout engine RESIDENT for models whose vLLM wake/reload HANGS (catalog
         # sleep_unsupported), exactly as the grpo path does. opd is NOT exempt: main_ppo_sync calls
         # checkpoint_manager.sleep_replicas() during init_workers and again around validation, which
@@ -2389,7 +2409,19 @@ def run_opd_train(spec=None) -> None:
     # oom. see resolve_rollout_enforce_eager for why this one knob is enough and cannot fight
     # verl's: vllm resolves enforce_eager LAST (config/vllm.py:1024), after the async server has set
     # cudagraph_mode, and forces both compilation mode and cudagraph_mode to NONE.
-    enforce_eager = resolve_rollout_enforce_eager(resolve_verl_device_capability(python_bin))
+    # one capability probe feeds both rollout decisions, as it does on the grpo path.
+    verl_cc = resolve_verl_device_capability(python_bin)
+    enforce_eager = resolve_rollout_enforce_eager(verl_cc)
+    # the same grpo/opd divergence as enforce_eager above, one knob over: blackwell needs both
+    # rollout attention backends pinned because vllm 0.19.1's defaults are wrong there, and opd
+    # never pinned them. the ViT default routes into a CUTE flash-attn that is unimportable against
+    # every published nvidia-cutlass-dsl, which aborts the engine with
+    # `RuntimeError: Worker failed with error 'module 'cutlass.cute.core' has no attribute
+    # 'ThrMma''` -- and a VL model builds its vision tower even for a text-only rollout, so this
+    # reaches text-only opd too. no-op off blackwell. see resolve_blackwell_attention_backends.
+    attention_backend, mm_encoder_attn_backend = resolve_blackwell_attention_backends(
+        python_bin, verl_cc
+    )
 
     plugin_path = os.path.join(shim_dir, "flash_opd_plugin.py")
     shutil.copy2(os.path.join(os.path.dirname(__file__), "opd_plugin.py"), plugin_path)
@@ -2489,6 +2521,8 @@ def run_opd_train(spec=None) -> None:
             "structured_outputs": structured_outputs,
             "fp8_kv": fp8_kv,
             "enforce_eager": enforce_eager,
+            "attention_backend": attention_backend,
+            "mm_encoder_attn_backend": mm_encoder_attn_backend,
             "sleep_unsupported": rollout_sleep_unsupported(model_id),
             "loggers": loggers,
         }

{freesolo_flash_dev-1.0.71 → freesolo_flash_dev-1.0.72}/flash/runner/__init__.py

@@ -29,7 +29,7 @@ from flash.opd_retry_contract import (
     require_opd_retry_contract_version,
 )
 from flash.providers._poll import _MAX_ATTEMPT_ID, _attempt_int
-from flash.spec import MANAGED_GPU_KEYS, JobSpec, effective_backend
+from flash.spec import MANAGED_GPU_KEYS, TRAINER_BACKEND, JobSpec
 
 _STATE_DIR = os.path.join(os.path.expanduser("~"), ".flash")
 RUNS_DIR = os.path.join(_STATE_DIR, "runs")
@@ -1207,7 +1207,7 @@ def _persist_effective_worker_spec(worker_spec: JobSpec) -> bool:
         "worker_spec": worker_spec.to_internal_dict(),
         "adapter_identity": adapter_identity,
         "preparation_digest": _preparation_digest(public_spec, worker_spec, adapter_identity),
-        "backend": effective_backend(public_spec),
+        "backend": TRAINER_BACKEND,
     }
     return _update(
         worker_spec.run_id,
@@ -1297,7 +1297,7 @@ def submit_job(
             "preparation_digest": _preparation_digest(
                 public_spec, worker_spec, prepared.adapter_identity
             ),
-            "backend": effective_backend(public_spec),
+            "backend": TRAINER_BACKEND,
         },
         # Snapshot the instance providers available at submit so a later handle-less recovery can fail
         # closed for any phantom-capable one whose creds were since dropped (see _confirm_run_clear).

{freesolo_flash_dev-1.0.71 → freesolo_flash_dev-1.0.72}/flash/spec.py

@@ -215,18 +215,10 @@ def validate_worker_env_reserved(worker_env: Any) -> None:
         )
 
 
-def effective_backend(spec: JobSpec) -> str:
-    """the training backend this spec will actually run, as the worker resolves it.
-
-    verl is the only backend: run_sft, run_opd and run_rl all delegate straight to it, so every
-    spec resolves "verl" whatever [worker_env] says.
-
-    kept as a function rather than folded into its call sites because the answer is load-bearing --
-    it picks the allocator conf (a spec that resolved "trl" would take expandable_segments and crash
-    its vLLM rollout outright, cumem.py CuMemAllocator assert) and gates multi-gpu, so the
-    resolution belongs in one place even now that it has one answer.
-    """
-    return "verl"
+# the trainer every run reports. run_sft, run_opd and run_rl all delegate straight to verl, so this
+# is a constant rather than a resolution: nothing in [worker_env] or the spec can select anything
+# else. recorded on effective_preparation so a stored run says which trainer produced it.
+TRAINER_BACKEND = "verl"
 
 
 def require_matching_seed(spec: JobSpec, seed: Any) -> int:

{freesolo_flash_dev-1.0.71 → freesolo_flash_dev-1.0.72}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "freesolo-flash-dev"
-version = "1.0.71"
+version = "1.0.72"
 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.71"
+version = "1.0.72"
 
 [tool.pytest.ini_options]
 pythonpath = ["."]

{freesolo_flash_dev-1.0.71 → freesolo_flash_dev-1.0.72}/tests/test_grpo_resident_sizing.py

@@ -288,7 +288,7 @@ def test_a_stale_backend_key_cannot_change_the_allocator(stale_backend):
     trainer that cannot run, and "trl" on grpo/opd would hand verl's rollout the expandable conf that
     trips the CuMemAllocator assert before step 1."""
     from flash.providers._worker import build_worker_env
-    from flash.spec import JobSpec, effective_backend
+    from flash.spec import JobSpec
 
     for algorithm, phase, expect_expandable in (
         ("grpo", "rl", False),
@@ -304,7 +304,6 @@ def test_a_stale_backend_key_cannot_change_the_allocator(stale_backend):
             }
         )
         env = build_worker_env(spec, 0)
-        assert effective_backend(spec) == "verl"
         for key in ("PYTORCH_ALLOC_CONF", "PYTORCH_CUDA_ALLOC_CONF"):
             assert ("expandable_segments" in env[key]) is expect_expandable, (
                 f"{algorithm} allocator moved under a stale {stale_backend!r} key"

{freesolo_flash_dev-1.0.71 → freesolo_flash_dev-1.0.72}/tests/test_multi_gpu_gate.py

@@ -42,33 +42,37 @@ def test_gate_helper_allows_single_gpu():
 
 
 @pytest.mark.parametrize("algorithm", ["sft", "grpo", "opd"])
-def test_gate_allows_multi_gpu_for_verl_backend(algorithm):
+def test_gate_allows_multi_gpu_on_a_provider_that_rents_n_cards(algorithm):
+    # no phase has a selector left: run_sft, run_rl and run_opd all delegate to verl, whose worker
+    # launches nproc-per-node == gpu.count ranks. so multi-gpu needs no opt-in beyond the provider.
     from flash import runner
 
-    # every verl worker launches nproc-per-node == gpu.count ranks, so all three shard.
-    assert runner._require_supported_gpu_count(_spec(4, algorithm, backend="verl")) is None
-
-
-@pytest.mark.parametrize("algorithm", ["sft", "grpo", "opd"])
-def test_gate_allows_multi_gpu_with_no_key_at_all(algorithm):
-    from flash import runner
-
-    # no phase has a selector: run_sft, run_rl and run_opd all delegate to verl, so multi-gpu needs
-    # no opt-in anywhere.
     assert runner._require_supported_gpu_count(_spec(4, algorithm)) is None
 
 
-@pytest.mark.parametrize("algorithm", ["sft", "grpo", "opd"])
-def test_a_stale_backend_key_cannot_downgrade_any_phase(algorithm):
+@pytest.mark.parametrize("stale", ["verl", "trl", "bogus"])
+@pytest.mark.parametrize("provider", ["runpod", "vast"])
+def test_a_stale_backend_key_changes_no_gate_outcome(stale, provider):
+    """A [worker_env] selector carried over from the trl era must not move this gate either way.
+
+    Asserted as "same outcome as the identical spec WITHOUT the key" rather than as a fixed verdict:
+    a bare `is None` on a runpod spec passes no matter what the gate does with the key, since the
+    provider alone already decides it. Pairing against the no-key baseline on BOTH an accepting and
+    a rejecting provider is what makes the claim able to fail -- it breaks the moment the gate starts
+    reading worker_env, in either direction.
+    """
     from flash import runner
-    from flash.spec import effective_backend
 
-    # a config carried over from when these phases had selectors must not resolve a backend that no
-    # longer exists: the worker would still run verl, and the gate would refuse it multi-gpu while
-    # the launcher handed it trl's expandable_segments alloc conf (which kills verl's CuMemAllocator).
-    spec = _spec(4, algorithm, backend="trl")
-    assert effective_backend(spec) == "verl"
-    assert runner._require_supported_gpu_count(spec) is None
+    def outcome(backend: str) -> str:
+        try:
+            runner._require_supported_gpu_count(
+                _spec(4, "grpo", backend=backend, provider=provider)
+            )
+        except ValueError as exc:
+            return f"rejected: {exc}"
+        return "allowed"
+
+    assert outcome(stale) == outcome("")
 
 
 def test_gate_rejects_multi_gpu_on_providers_that_ignore_count():
@@ -81,17 +85,6 @@ def test_gate_rejects_multi_gpu_on_providers_that_ignore_count():
             runner._require_supported_gpu_count(_spec(4, backend="verl", provider=provider))
 
 
-def test_effective_backend_is_verl_for_every_phase():
-    from flash.spec import effective_backend
-
-    # verl is the only backend. this stays pinned because the answer is load-bearing: it picks the
-    # allocator conf and gates multi-gpu, and a spec that resolved anything else would take
-    # expandable_segments and crash its vllm rollout on the CuMemAllocator assert.
-    for algorithm in ("sft", "grpo", "opd"):
-        assert effective_backend(_spec(1, algorithm)) == "verl"
-        assert effective_backend(_spec(1, algorithm, backend="trl")) == "verl"
-
-
 def test_submit_job_rejects_multi_gpu_at_boundary(monkeypatch):
     from flash import runner
 

{freesolo_flash_dev-1.0.71 → freesolo_flash_dev-1.0.72}/tests/test_opd_train.py

@@ -4641,10 +4641,57 @@ def test_the_resolved_eager_flag_reaches_the_opd_verl_config():
     # the string assertions above pass against a resolver whose answer is never carried into the
     # config, which is exactly the shape of the bug this fixes. pin the wiring.
     built = inspect.getsource(opd_train.run_opd_train)
-    assert "resolve_rollout_enforce_eager(resolve_verl_device_capability(python_bin))" in built
+    # one probe now feeds both eager and the attention pins, so the call is bound to a name rather
+    # than nested inline. assert the resolver still consumes THAT probe, not a second one.
+    assert "verl_cc = resolve_verl_device_capability(python_bin)" in built
+    assert "resolve_rollout_enforce_eager(verl_cc)" in built
     assert '"enforce_eager": enforce_eager,' in built
 
 
+def test_opd_pins_the_blackwell_attention_backends_like_grpo_does():
+    """VERL-156: opd's rollout died on B200 because it never pinned the ViT attention backend.
+
+    vllm 0.19.1 defaults the ViT to a CUTE flash-attn that is unimportable against every published
+    nvidia-cutlass-dsl, so the engine aborts with `cutlass.cute.core has no attribute 'ThrMma'`.
+    grpo has pinned both backends since the trl driver; opd never did, which is the same one-sided
+    fix that produced the enforce_eager bug directly above. two attempts on two fresh B200s
+    reproduced it byte-identically before this was found.
+    """
+    # off blackwell the resolver answers (None, None) and nothing may be emitted -- an unconditional
+    # pin would change the rollout on every other card.
+    off = build_opd_overrides(_config())
+    assert not any("attention_backend" in value for value in off)
+
+    on = build_opd_overrides(
+        _config(attention_backend="FLASHINFER", mm_encoder_attn_backend="TORCH_SDPA")
+    )
+    # '+' appends under the existing engine_kwargs.vllm struct: these are AsyncEngineArgs fields
+    # spread by verl, not declared RolloutConfig fields, so a bare assignment would not reach them.
+    assert "+actor_rollout_ref.rollout.engine_kwargs.vllm.attention_backend=FLASHINFER" in on, on
+    assert (
+        "+actor_rollout_ref.rollout.engine_kwargs.vllm.mm_encoder_attn_backend=TORCH_SDPA" in on
+    ), on
+
+
+def test_both_ray_rollouts_pin_blackwell_attention_from_the_same_resolver():
+    # the enforce_eager cross-trainer check below exists because a one-sided fix is how that bug
+    # shipped. the attention pins are the same divergence one knob over, so assert them the same
+    # way rather than trusting opd alone.
+    for module, source in (
+        ("rl_train", inspect.getsource(rl_train.run_rl_train)),
+        ("opd_train", inspect.getsource(opd_train.run_opd_train)),
+    ):
+        # the resolver must be CALLED and its second return value must be carried onward. the two
+        # trainers spell the handoff differently (rl_train passes it as a kwarg into
+        # _rl_runconfig, opd_train builds its config dict inline), so assert the binding both
+        # share rather than either one's syntax.
+        assert "resolve_blackwell_attention_backends(" in source, module
+        assert "attention_backend, mm_encoder_attn_backend = " in source, module
+        assert "mm_encoder_attn_backend=mm_encoder_attn_backend" in source or (
+            '"mm_encoder_attn_backend": mm_encoder_attn_backend,' in source
+        ), module
+
+
 def test_both_ray_rollouts_resolve_eager_from_the_same_hardware_probe():
     # asserted across BOTH trainers rather than in one file: grpo has resolved this since the trl
     # driver and opd did not, so opd captured 102 cuda graphs on an rtx 4090 and died on HOST ram
@@ -5384,19 +5431,17 @@ def test_opd_delegates_to_verl_with_no_selector_left(monkeypatch):
 
 
 def test_opd_spec_never_resolves_the_allocator_conf_that_kills_vllm(monkeypatch):
-    """An OPD spec must resolve backend "verl" and a NON-expandable allocator conf.
+    """An OPD spec must build a NON-expandable allocator conf.
 
-    This is the regression the TRL deletion exposed. `effective_backend` defaulted to "trl" for any
-    phase whose worker_env carried no selector, and _worker.py maps a non-verl OPD run to
-    PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True. verl leaves rollout.enable_sleep_mode
-    defaulted True, so the engine always builds a CuMemAllocator, which asserts outright on
-    expandable_segments (cumem.py, pytorch#147851) -- the run would die before step 1.
+    verl leaves rollout.enable_sleep_mode defaulted True, so the engine always builds a
+    CuMemAllocator, which asserts outright on expandable_segments (cumem.py, pytorch#147851) -- the
+    run would die before step 1.
 
     A [worker_env] selector must not be able to reintroduce it either: run_opd delegates to verl
-    unconditionally, so honoring the key would pick an allocator for a backend that cannot run.
+    unconditionally, so honoring a stale key would pick an allocator for a backend that cannot run.
     """
     from flash.providers._worker import build_worker_env
-    from flash.spec import JobSpec, effective_backend
+    from flash.spec import JobSpec
 
     def _spec(worker_env=None):
         payload = {
@@ -5415,7 +5460,6 @@ def test_opd_spec_never_resolves_the_allocator_conf_that_kills_vllm(monkeypatch)
     for worker_env in (None, {"FLASH_OPD_BACKEND": "trl"}, {"FLASH_OPD_BACKEND": "bogus"}):
         spec = _spec(worker_env)
         assert spec.phase == "opd"
-        assert effective_backend(spec) == "verl", worker_env
         alloc = build_worker_env(spec, spec.seed)["PYTORCH_CUDA_ALLOC_CONF"]
         assert "expandable_segments" not in alloc, (worker_env, alloc)
 
@@ -5431,7 +5475,6 @@ def test_opd_spec_never_resolves_the_allocator_conf_that_kills_vllm(monkeypatch)
             "gpu": {"type": "B200", "count": 1, "provider": "runpod"},
         }
     )
-    assert effective_backend(sft) == "verl"
     assert build_worker_env(sft, sft.seed)["PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:True"
 
 

{freesolo_flash_dev-1.0.71 → freesolo_flash_dev-1.0.72}/uv.lock

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