freesolo-flash-dev 1.2.98__tar.gz → 1.2.100__tar.gz

{freesolo_flash_dev-1.2.98 → freesolo_flash_dev-1.2.100}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: freesolo-flash-dev
-Version: 1.2.98
+Version: 1.2.100
 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.2.98 → freesolo_flash_dev-1.2.100}/flash/engine/worker/train/opd/overrides.py

@@ -23,6 +23,7 @@ from flash.engine.worker.backend_common import (
     trainer_dtype_overrides,
 )
 from flash.engine.worker.sft_train import _build_verl_child_env, _hydra_val
+from flash.engine.worker.verl.capabilities import rollout_max_num_seqs
 from flash.teacher.limits import OPD_NO_SIGNAL_ATTEMPTS
 
 _REQUIRED_OVERRIDE_KEYS = (
@@ -214,6 +215,15 @@ def _actor_rollout_overrides(config: dict, *, max_tokens: int) -> list[str]:
             "actor_rollout_ref.rollout.agent.num_workers="
             f"{agent_loop_workers(int(config['train_batch_size']) * int(config['group_size']), cap=4 if config.get('multimodal') else 8)}"
         ),
+        # one opd step submits exactly train_batch_size * group_size student generations. verl's
+        # 1024 default makes vllm reserve cuda-graph capture sizes -- and, on a gdn/mamba hybrid,
+        # one recurrent state block per decode slot -- for a batch this run never submits. both are
+        # allocated before the first token, so the over-provisioning OOMs a card that stays mostly
+        # free. `rollout.max_num_seqs` is a real verl field, so this is a plain override.
+        (
+            "actor_rollout_ref.rollout.max_num_seqs="
+            f"{rollout_max_num_seqs(int(config['train_batch_size']) * int(config['group_size']))}"
+        ),
     ]
 
 

{freesolo_flash_dev-1.2.98 → freesolo_flash_dev-1.2.100}/flash/engine/worker/train/rl/verl_config.py

@@ -26,6 +26,7 @@ from flash.engine.worker.backend_common import (
 )
 from flash.engine.worker.runtime.pkg_proxy import W as _w
 from flash.engine.worker.sft_train import _hydra_val, _verl_image_message_content
+from flash.engine.worker.verl.capabilities import rollout_max_num_seqs
 from flash.engine.worker.verl.parallelism import (
     ULYSSES_SEQUENCE_PARALLEL_SIZE,
     resolve_reshard_after_forward,
@@ -321,6 +322,11 @@ def _rollout_overrides(cfg: dict) -> list[str]:
         # child-silence watchdog kills it 1200s later. capping the pool leaves group_size and
         # prompts_per_step exactly as authored -- only scheduling parallelism narrows.
         f"actor_rollout_ref.rollout.agent.num_workers={agent_loop_workers(int(cfg['prompts_per_step']) * int(cfg['group_size']), cap=4 if cfg.get('multimodal') else 8)}",
+        # one grpo step submits exactly prompts_per_step * group_size sequences. left unset verl
+        # keeps its 1024 default, and vllm sizes cuda-graph capture and (on a gdn/mamba hybrid)
+        # per-slot recurrent state from that number, both up front -- so a 32-sequence run paid a
+        # 1024-sequence reservation and OOMed during capture with most of the card still free.
+        f"actor_rollout_ref.rollout.max_num_seqs={rollout_max_num_seqs(int(cfg['prompts_per_step']) * int(cfg['group_size']))}",
         # multi-turn runs flash's own agent loop, registered in the child by flash_grpo_plugin. the
         # stock single_turn_agent generates once and returns; it has no notion of an environment
         # reply, so a multi-turn env on the default loop would train on first turns only.

{freesolo_flash_dev-1.2.98 → freesolo_flash_dev-1.2.100}/flash/engine/worker/verl/capabilities.py

@@ -773,6 +773,45 @@ def rollout_fp8_kv(cc_ok: bool, gdn_hybrid: bool, model_id: str) -> bool:
     return rollout_sleep_unsupported(model_id)
 
 
+_MIN_ROLLOUT_MAX_NUM_SEQS = 16
+
+
+def rollout_max_num_seqs(rollout_batch: int) -> int:
+    """concurrent sequence slots to provision the rollout engine for.
+
+    verl leaves ``rollout.max_num_seqs`` at its yaml default of 1024 (rollout.yaml:79) and passes
+    it straight to the vllm engine (vllm_async_server.py:261). vllm sizes two up-front, fixed
+    allocations from that number, neither of which shrinks to the batch actually submitted:
+
+    * cuda-graph capture. ``_set_cudagraph_sizes`` (vllm ``config/vllm.py``, 0.19.1 -- the version
+      ``Dockerfile.worker`` pins) derives a capture ladder from
+      ``max_cudagraph_capture_size = min(max_num_seqs * 2, 512, max_num_batched_tokens)``, taking
+      ``[1,2,4]``, then ``range(8, min(size+1, 256), 8)``, then ``range(256, size+1, 16)`` -- each
+      rung clamped by that ceiling. verl's ``max_num_batched_tokens`` default is 8192
+      (``rollout.yaml:73``), so it does not bind. at 1024 the ceiling is 512 and the ladder is 51
+      sizes; a lora run captures each size twice, since the dispatcher takes
+      ``product(cudagraph_capture_sizes, lora_cases)`` with ``lora_cases == [0, 1]`` when
+      ``cudagraph_specialize_lora`` is on (``v1/cudagraph_dispatcher.py:110-193``, default True) --
+      exactly the ``0/102`` seen in the failing console. the ladder saturates at 512, so 1024, 512
+      and 256 are indistinguishable; only dropping below 256 shortens it, and 32 yields a ceiling
+      of 64, 11 sizes, and 22 graphs. note the prose comment above that function omits the per-rung
+      clamp and predicts 34 sizes for any small value; the implementation is authoritative.
+    * recurrent state. a gdn/mamba hybrid reserves one state block per decode slot, a dense
+      allocation that -- unlike paged kv -- cannot be shared, paged out, or grown on demand.
+
+    both are paid before the first rollout token, inside the ``gpu_memory_utilization`` budget. a
+    grpo/opd step submits exactly ``prompts_per_step * group_size`` sequences, so a run with 32
+    concurrent generations was provisioning for 1024 and dying during graph capture with ~80% of
+    the card free -- on 234, 358, and 460 GB alike, since the cost is fixed rather than
+    proportional to the card.
+
+    the floor keeps small runs on a usable capture ladder rather than collapsing it to one size.
+    """
+    if rollout_batch <= 0:
+        raise ValueError("rollout_batch must be positive")
+    return max(_MIN_ROLLOUT_MAX_NUM_SEQS, int(rollout_batch))
+
+
 def rollout_resident_overrides(sleep_unsupported: bool) -> list[str]:
     """verl overrides that pin the rollout engine resident, or ``[]`` to keep verl's default.
 

{freesolo_flash_dev-1.2.98 → freesolo_flash_dev-1.2.100}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "freesolo-flash-dev"
-version = "1.2.98"
+version = "1.2.100"
 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).
@@ -137,7 +137,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.2.98"
+version = "1.2.100"
 
 [tool.pytest.ini_options]
 pythonpath = ["."]

{freesolo_flash_dev-1.2.98 → freesolo_flash_dev-1.2.100}/tests/test_backend_common.py

@@ -4779,3 +4779,41 @@ def test_the_lora_rollout_guard_does_not_wrap_itself_twice(monkeypatch):
 
     assert module.vLLMHttpServer.generate is first
     assert asyncio.run(module.vLLMHttpServer().generate([1], {}, "req-3"))[0] == "generated"
+
+
+@pytest.mark.parametrize(
+    ("rollout_batch", "expected"),
+    [(1, 16), (4, 16), (16, 16), (32, 32), (64, 64), (512, 512)],
+)
+def test_rollout_max_num_seqs_tracks_the_rollout_batch_above_the_floor(rollout_batch, expected):
+    from flash.engine.worker.verl.capabilities import rollout_max_num_seqs
+
+    assert rollout_max_num_seqs(rollout_batch) == expected
+
+
+def test_rollout_max_num_seqs_never_returns_verls_default_for_a_small_batch():
+    # the regression this exists for: verl leaves rollout.max_num_seqs at 1024, and vllm sizes
+    # cuda-graph capture (min(max_num_seqs * 2, 512) over a 51-rung ladder, doubled when lora
+    # specialization is on) and per-slot gdn/mamba recurrent state from it, both before the first
+    # rollout token. a 32-sequence grpo step paid a 1024-sequence reservation and OOMed during
+    # capture with ~80% of the card free.
+    from flash.engine.worker.verl.capabilities import rollout_max_num_seqs
+
+    for batch in (8, 32, 64, 128):
+        assert rollout_max_num_seqs(batch) < 1024
+
+
+def test_rollout_max_num_seqs_is_never_below_the_batch_it_must_serve():
+    # provisioning FEWER slots than the step submits would queue sequences instead of running them,
+    # turning one rollout into several scheduler passes.
+    from flash.engine.worker.verl.capabilities import rollout_max_num_seqs
+
+    for batch in range(1, 600):
+        assert rollout_max_num_seqs(batch) >= batch
+
+
+def test_rollout_max_num_seqs_rejects_a_nonpositive_batch():
+    from flash.engine.worker.verl.capabilities import rollout_max_num_seqs
+
+    with pytest.raises(ValueError, match="rollout_batch must be positive"):
+        rollout_max_num_seqs(0)

{freesolo_flash_dev-1.2.98 → freesolo_flash_dev-1.2.100}/tests/test_opd_train.py

@@ -9255,3 +9255,26 @@ def test_opd_accounting_gate_keeps_the_first_failure_reason():
     state.fail("bridge shutdown")
     with pytest.raises(RuntimeError, match="exited with code 1"):
         state.checkpoint_state(1, timeout_s=600.0)
+
+
+def test_overrides_size_max_num_seqs_to_the_opd_rollout_batch():
+    """OPD colocates the student engine with the trainer, so a wasted reservation is fatal.
+
+    Left unset, verl passes its 1024 default to vllm, which sizes cuda-graph capture and (on a
+    gdn/mamba hybrid) one recurrent state block per decode slot from it -- both before the first
+    student generation and inside the same gpu_memory_utilization budget the resident trainer
+    already shares.
+    """
+    overrides = dict(
+        value.split("=", 1)
+        for value in build_opd_overrides(_config(train_batch_size=8, group_size=4))
+    )
+    assert overrides["actor_rollout_ref.rollout.max_num_seqs"] == "32"
+
+
+def test_overrides_floor_max_num_seqs_for_a_tiny_opd_rollout_batch():
+    overrides = dict(
+        value.split("=", 1)
+        for value in build_opd_overrides(_config(train_batch_size=1, group_size=1))
+    )
+    assert overrides["actor_rollout_ref.rollout.max_num_seqs"] == "16"

{freesolo_flash_dev-1.2.98 → freesolo_flash_dev-1.2.100}/tests/test_rl_train.py

@@ -9273,3 +9273,26 @@ def test_resumed_grpo_seeds_the_dump_watermark_at_the_resume_boundary():
     ledger = RolloutIdentityLedger(1, 2)
     with pytest.raises(ValueError, match="no registered rollout identity set"):
         ledger.seal(2)
+
+
+def test_build_verl_overrides_sizes_max_num_seqs_to_the_rollout_batch():
+    """The rollout engine must be provisioned for the batch the step submits, not verl's default.
+
+    verl leaves ``rollout.max_num_seqs`` at 1024 (rollout.yaml:79). vllm derives
+    ``max_cudagraph_capture_size = min(max_num_seqs * 2, 512)`` over the ladder
+    ``[1,2,4] + range(8,256,8) + range(256,513,16)`` -- 51 sizes at the default, captured twice
+    when lora specialization is on -- and a gdn/mamba hybrid additionally reserves one recurrent
+    state block per decode slot. Both are paid up front, inside the gpu_memory_utilization budget,
+    so a 32-sequence run died at graph capture on 234, 358 and 460 GB alike.
+    """
+    o = rl_train.build_verl_overrides(_overrides_cfg(prompts_per_step=8, group_size=4))
+    assert "actor_rollout_ref.rollout.max_num_seqs=32" in o
+    # and never verl's default, which is what made the reservation fixed rather than proportional.
+    assert not any(override == "actor_rollout_ref.rollout.max_num_seqs=1024" for override in o)
+
+
+def test_build_verl_overrides_floors_max_num_seqs_for_a_tiny_rollout_batch():
+    # a batch of 1 would otherwise capture a single graph size and push every wider decode step
+    # onto the eager path.
+    o = rl_train.build_verl_overrides(_overrides_cfg(prompts_per_step=1, group_size=1))
+    assert "actor_rollout_ref.rollout.max_num_seqs=16" in o

{freesolo_flash_dev-1.2.98 → freesolo_flash_dev-1.2.100}/uv.lock

@@ -1052,7 +1052,7 @@ wheels = [
 
 [[package]]
 name = "freesolo-flash"
-version = "1.2.98"
+version = "1.2.100"
 source = { editable = "." }
 
 [package.optional-dependencies]