freesolo-flash-dev 1.2.89__tar.gz → 1.2.94__tar.gz

{freesolo_flash_dev-1.2.89 → freesolo_flash_dev-1.2.94}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: freesolo-flash-dev
-Version: 1.2.89
+Version: 1.2.94
 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.89 → freesolo_flash_dev-1.2.94}/flash/cli/commands/__init__.py

@@ -470,16 +470,10 @@ def _log_follow_progress(status: dict | None, fallback_state: str) -> tuple[str,
     parts = [state]
     heartbeat = status.get("last_heartbeat") if isinstance(status, dict) else None
     # retries rewind steps while state remains running, so surface the 0-based attempt identity.
-    # prefer live `remote.attempt`; the heartbeat may belong to the superseded worker. fall back only
-    # when `remote` is absent, not explicitly null: null marks the allocation window after teardown,
-    # when the new attempt is unknown and the attached heartbeat is stale.
-    from flash.providers._lifecycle.poll import _attempt_int
-
-    remote = status.get("remote")
-    remote_cleared = "remote" in status and remote is None
-    attempt = _attempt_int(remote.get("attempt")) if isinstance(remote, dict) else None
-    if attempt is None and not remote_cleared and isinstance(heartbeat, dict):
-        attempt = _attempt_int(heartbeat.get("attempt"))
+    # `live_attempt` owns the provenance order (live `remote.attempt` first, heartbeat only when
+    # `remote` is absent rather than cleared at teardown) and is shared with the worker-artifact
+    # labelling below, so the spinner and the appended sections name the same current attempt.
+    attempt = render.live_attempt(status)
     if isinstance(heartbeat, dict):
         heartbeat_age_seconds = render._heartbeat_age_seconds(heartbeat.get("ts"))
         # stage and step come from the heartbeat, attempt from `remote` below. during the relaunch
@@ -712,20 +706,6 @@ def _follow_status(
         return 130
 
 
-def _print_worker_output(client: ApiClient, run_id: str, *, printed_any: bool = False) -> bool:
-    for name, text in (client.get_worker_output(run_id) or {}).items():
-        if not text:
-            continue
-        sep = "\n" if printed_any else ""
-        if render.styled():
-            print(f"{sep}{render.log_section(name)}")
-        else:
-            print(f"{sep}----- {name} -----")
-        print(text, end="" if text.endswith("\n") else "\n")
-        printed_any = True
-    return printed_any
-
-
 def cmd_log(args) -> int:
     client = client_from_config()
     if getattr(args, "follow", False):
@@ -980,3 +960,10 @@ from flash.cli.commands.train_cost import (  # noqa: E402,F401
     _sft_cost_rows,
     _warn_if_wandb_requested_without_key,
 )
+
+# re-exported so `commands._print_worker_output(...)` keeps resolving for `cmd_log` and the CLI tests.
+from flash.cli.commands.worker_output import (  # noqa: E402,F401
+    _artifact_attempt,
+    _print_worker_output,
+    _worker_section_name,
+)

freesolo_flash_dev-1.2.94/flash/cli/commands/worker_output.py

@@ -0,0 +1,67 @@
+"""Printing the worker-uploaded text artifacts that `flash runs log` appends to a run log.
+
+These sections are NOT part of the chronological log. `cmd_log` prints the control-plane log first
+and appends these afterwards, so they always land at the tail no matter how old they are -- and the
+plane returns the highest UPLOADED attempt, which during a retry is the attempt that just died, not
+the one now running. Read without a label, the last thing on screen is a traceback from a worker
+that has already been replaced.
+
+So the job here is provenance: say which attempt produced each section, and say when that attempt is
+not the live one. The old output is kept rather than hidden because it carries the failure that
+explains why a retry exists at all.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import TYPE_CHECKING
+
+from flash.cli.ui import render
+
+if TYPE_CHECKING:
+    from flash.client import ApiClient
+
+
+def _artifact_attempt(name: str) -> int | None:
+    """Return the bounded attempt encoded by a worker artifact name."""
+    from flash.providers._lifecycle.poll import _attempt_int
+
+    match = re.search(r"_attempt(\d+)\.txt$", name)
+    return _attempt_int(int(match.group(1))) if match else None
+
+
+def _worker_section_name(name: str, current_attempt: int | None) -> str:
+    """Label an attempt-scoped artifact and identify superseded output."""
+    attempt = _artifact_attempt(name)
+    if attempt is None:
+        return name
+    if current_attempt is None:
+        return f"{name} (attempt={attempt})"
+    if attempt != current_attempt:
+        return f"{name} (attempt={attempt}, previous attempt; current attempt={current_attempt})"
+    return f"{name} (attempt={attempt}, current attempt)"
+
+
+def _print_worker_output(client: ApiClient, run_id: str, *, printed_any: bool = False) -> bool:
+    worker_output = client.get_worker_output(run_id) or {}
+    if not worker_output:
+        return printed_any
+    try:
+        current_attempt = render.live_attempt(client.get_run(run_id) or {})
+    except Exception:
+        # artifact text is best-effort diagnostic output; a status refresh failure must not hide it.
+        current_attempt = None
+    for name, text in worker_output.items():
+        if not text:
+            continue
+        # label provenance from the filename rather than letting the final section, which is merely
+        # the last one appended, read as the live attempt's output.
+        section_name = _worker_section_name(name, current_attempt)
+        sep = "\n" if printed_any else ""
+        if render.styled():
+            print(f"{sep}{render.log_section(section_name)}")
+        else:
+            print(f"{sep}----- {section_name} -----")
+        print(text, end="" if text.endswith("\n") else "\n")
+        printed_any = True
+    return printed_any

{freesolo_flash_dev-1.2.89 → freesolo_flash_dev-1.2.94}/flash/cli/ui/heartbeat.py

@@ -43,6 +43,33 @@ _WARMUP_HEARTBEAT_FRESH_FOR_S = 1200.0
 _WARMUP_STAGES = frozenset({"rl_train_start", "rl_initializing"})
 
 
+def live_attempt(obj: dict) -> int | None:
+    """The attempt a status payload says is live, or None when it cannot be established.
+
+    `remote.attempt` is the plane's live attempt. It is preferred over `last_heartbeat.attempt`,
+    which is whichever worker produced the ping and may already be superseded. The heartbeat is a
+    fallback only when `remote` is ABSENT: an explicitly null `remote` is the teardown window, where
+    the attached ping belongs to a worker that is already gone, so falling back there would report a
+    dead attempt as the live one.
+
+    Shared so the status line, the log-follow spinner, and worker-artifact labelling cannot disagree
+    about which attempt is current -- three surfaces that a user reads within one screen of each
+    other, where a disagreement reads as a run that is on two attempts at once.
+    """
+    from flash.providers._lifecycle.poll import _attempt_int
+
+    remote = obj.get("remote")
+    if isinstance(remote, dict):
+        attempt = _attempt_int(remote.get("attempt"))
+        if attempt is not None:
+            return attempt
+    elif "remote" in obj:
+        # explicitly null: the teardown window. the ping is the dead worker's, so there is no answer.
+        return None
+    heartbeat = obj.get("last_heartbeat")
+    return _attempt_int(heartbeat.get("attempt")) if isinstance(heartbeat, dict) else None
+
+
 def heartbeat_is_current_attempt(obj: dict, heartbeat: dict) -> bool:
     """False only when the heartbeat provably belongs to a superseded retry attempt.
 

{freesolo_flash_dev-1.2.89 → freesolo_flash_dev-1.2.94}/flash/cli/ui/render.py

@@ -808,6 +808,7 @@ from flash.cli.ui.heartbeat import (  # noqa: E402,F401
     _superseded_hint,
     heartbeat_is_current_attempt,
     heartbeat_is_superseded,
+    live_attempt,
     warmup_message,
 )
 from flash.cli.ui.tables import (  # noqa: E402,F401

{freesolo_flash_dev-1.2.89 → freesolo_flash_dev-1.2.94}/flash/engine/plan/kv_sizing.py

@@ -19,6 +19,10 @@ _KV_BLOCK_TOKENS = 16
 # floor remains authoritative at short context; longer contexts retain 1.5 gb plus 25% fragmentation.
 _KV_PROFILE_OVERHEAD_GB = 1.5
 _KV_FRAGMENTATION_MARGIN = 1.25
+# the smallest share of a card vllm is given. below this the engine has no room for its own runtime
+# once the kv pool is carved out, so a budget that prices under it is not a smaller engine but a
+# broken one. the legacy equation and the multi-rank shard path both bottom out here.
+_MIN_ENGINE_UTIL = 0.10
 
 
 def _round_up(value: int, multiple: int) -> int:
@@ -138,15 +142,18 @@ def colocate_kv_util(
     fp8_kv: bool = False,
     model_info=None,
     preserve_legacy_floor: bool = False,
+    tensor_parallel: int = 1,
 ) -> float:
-    """vllm_gpu_memory_utilization for the colocated GRPO rollout engine.
+    """vllm_gpu_memory_utilization for one colocated rollout tensor-parallel rank.
 
-    Budget vLLM's weight copy plus KV cache. Catalog models use geometry with measured overhead
-    and an 8 GB floor; uncataloged models retain the conservative legacy equation.
+    Budget the rank's weight shard plus a conservative full KV cache. Catalog models use geometry
+    with measured overhead and an 8 GB floor; uncataloged models retain the legacy KV equation.
     """
-    weights_gb = (
-        max(0.5, float(params_b or 1.0)) * 2.0
-    )  # vLLM's bf16 weight copy lives in the budget
+    total_weights_gb = max(0.5, float(params_b or 1.0)) * 2.0
+    tp_size = max(1, int(tensor_parallel))
+    # tensor parallelism shards vllm's bf16 weights, but kv heads can replicate when tp is wider than
+    # the model's kv-head count. shard only the weight term and keep the full kv budget on every rank.
+    weights_gb = total_weights_gb if tp_size == 1 else total_weights_gb / tp_size
     # catalog geometry is authoritative. active params remain only for the uncataloged fallback, while
     # the weight copy always uses total parameters.
     kv_params_b = float(active_params_b) if active_params_b else params_b
@@ -169,7 +176,7 @@ def colocate_kv_util(
                 preserve_legacy_floor=preserve_legacy_floor,
             ),
         )
-        return max(0.10, min(_util_cap, (weights_gb + kv_gb) / max(1.0, total_vram_gb)))
+        return max(_MIN_ENGINE_UTIL, min(_util_cap, (weights_gb + kv_gb) / max(1.0, total_vram_gb)))
     # Sleep mode keeps a larger pool (1.5x margin): the engine is offloaded during the backward, so a
     # bigger rollout-phase KV does not compete with the training peak.
     kv_pool_gb = max(
@@ -184,4 +191,9 @@ def colocate_kv_util(
             preserve_legacy_floor=preserve_legacy_floor,
         ),
     )
-    return min(_util_cap, (weights_gb + kv_pool_gb) / max(1.0, total_vram_gb))
+    sized = min(_util_cap, (weights_gb + kv_pool_gb) / max(1.0, total_vram_gb))
+    # sharding the weight term can drive a small model's rank budget low enough that vllm has no
+    # room left for its own runtime once the kv pool is carved out -- a 4B model on 2 cards prices
+    # under this floor. single-rank sizing carries the whole weight copy and never gets that small,
+    # so the floor exists only where sharding created the problem.
+    return max(_MIN_ENGINE_UTIL, sized) if tp_size > 1 else sized

{freesolo_flash_dev-1.2.89 → freesolo_flash_dev-1.2.94}/flash/engine/worker/opd_train.py

@@ -374,7 +374,7 @@ def run_opd_train(spec=None) -> None:
                 # and on an image run that fan-out exhausted the grpo worker container's threads.
                 "multimodal": bool(multimodal),
                 "gpu_mem_util": _resolve_opd_gpu_mem_util(
-                    request, prompt_state, runtime, model_id, fp8_kv
+                    request, prompt_state, workload, runtime, model_id, fp8_kv
                 ),
                 "loggers": runtime.loggers,
                 # resolved from the out-of-process capability probe, never by opening cuda in this parent -- see fused_ce_backend.

{freesolo_flash_dev-1.2.89 → freesolo_flash_dev-1.2.94}/flash/engine/worker/opd_train_runner.py

@@ -535,6 +535,7 @@ def _spec_gpu_type(spec: Any) -> str:
 def _resolve_opd_gpu_mem_util(
     request: _OpdRequest,
     prompt_state: _PromptState,
+    workload: _WorkloadState,
     runtime: _RuntimeState,
     model_id: str,
     fp8_kv: bool,
@@ -559,12 +560,15 @@ def _resolve_opd_gpu_mem_util(
             "model_id": model_id,
             "model_revision": str(getattr(request, "model_revision", "") or ""),
             "engine_len": int(prompt_state.max_model_len),
-            "group_size": int(request.knobs.group_size),
+            # opd submits prompts_per_step * group_size generations to the rollout engine together.
+            # grpo's concurrency is group_size because its resolver receives one prompt group here.
+            "group_size": int(workload.prompts_per_step) * int(request.knobs.group_size),
         },
         gpu_type=_spec_gpu_type(getattr(request, "spec", None)),
         n_gpus=int(runtime.gpu_count),
         fp8_kv=bool(fp8_kv),
         sleep_unsupported=rollout_sleep_unsupported(model_id),
+        preserve_legacy_floor=True,
     )
 
 

{freesolo_flash_dev-1.2.89 → freesolo_flash_dev-1.2.94}/flash/engine/worker/train/rl/verl_config.py

@@ -499,24 +499,28 @@ def resolve_gpu_mem_util(
     n_gpus: int,
     fp8_kv: bool,
     sleep_unsupported: bool,
+    preserve_legacy_floor: bool = False,
 ) -> float:
     """size vllm's colocated executor budget from this run's geometry.
 
     ``gpu_memory_utilization`` covers the second bf16 weight copy plus kv pool. using
     ``colocate_kv_util`` keeps worker allocation aligned with preflight admission.
 
-    the flat constant is kept where the model does NOT apply, because a wrong number is worse than
+    for multi-gpu runs, vllm shards its bf16 weight copy across tensor-parallel ranks. the kv term
+    remains unsharded here because vllm can replicate kv heads when tensor parallelism is wider than
+    the model's kv-head count. the result is explicitly capped at the old 0.5 constant, so this path
+    can only lower a rank's reservation and cannot create a new over-reservation.
+
+    the flat constant is kept where the model does not apply, because a wrong number is worse than
     a conservative one:
 
-    - UNKNOWN CARD (empty/unmanaged ``gpu.type``): the budget is a fraction of the card, so without
+    - unknown card (empty/unmanaged ``gpu.type``): the budget is a fraction of the card, so without
       its size there is nothing to take a fraction of.
-    - MULTI-GPU (``n_gpus > 1``): the rollout is tensor-parallel, so vLLM's weight copy is sharded
-      ACROSS cards while ``colocate_kv_util`` sizes one whole copy against one card. a single-card
-      number would over-reserve per rank on exactly the shapes already tight.
-    - PINNED REVISION with no resolvable parameter count: the weight term dominates, so a guessed
+    - pinned revision with no resolvable parameter count: the weight term dominates, so a guessed
       size is not worth acting on.
+    - any sizing exception: launch falls back to the previous constant instead of failing.
     """
-    if n_gpus > 1 or not (gpu_type or "").strip():
+    if not (gpu_type or "").strip():
         return _DEFAULT_GPU_MEM_UTIL
     try:
         from flash.core.catalog import MODELS
@@ -545,7 +549,7 @@ def resolve_gpu_mem_util(
         )
         if params_b <= 0:
             return _DEFAULT_GPU_MEM_UTIL
-        return colocate_kv_util(
+        sized = colocate_kv_util(
             params_b,
             int(inp["engine_len"]),
             total_vram_gb,
@@ -558,7 +562,12 @@ def resolve_gpu_mem_util(
             active_params_b=float(getattr(info, "active_params_b", 0.0) or 0.0) or None,
             fp8_kv=fp8_kv,
             model_info=info,
+            preserve_legacy_floor=preserve_legacy_floor,
+            tensor_parallel=n_gpus,
         )
+        # multi-rank sizing is strictly one-directional against the previous worker contract: it may
+        # free memory but can never claim more than the flat reservation a working run already used.
+        return min(_DEFAULT_GPU_MEM_UTIL, sized) if n_gpus > 1 else sized
     except Exception as e:  # sizing must never be what stops a run from launching
         print(
             f"[rl-verl] gpu_memory_utilization sizing failed ({e}); using {_DEFAULT_GPU_MEM_UTIL}"

{freesolo_flash_dev-1.2.89 → freesolo_flash_dev-1.2.94}/flash/providers/_lifecycle/poll.py

@@ -224,6 +224,9 @@ def _format_child_tail(tail: object, silent_ticks: object = None) -> str:
 
 def _format_heartbeat(hb: dict) -> str:
     msg = f"worker: stage={hb.get('stage')}"
+    attempt = _attempt_int(hb.get("attempt"))
+    if attempt is not None:
+        msg += f" attempt={attempt}"
     if hb.get("liveness"):
         msg += " liveness=true"
     for key, digits in (

{freesolo_flash_dev-1.2.89 → freesolo_flash_dev-1.2.94}/flash/server/platform/runtime.py

@@ -428,9 +428,10 @@ def _deferred_resubmit_loop(spec) -> None:
 def _latest_worker_artifact_name(repo: str, prefix: str, phase: str, kind: str) -> str:
     """Return the newest worker artifact under a prefix.
 
-    Console uses one overwritten name; errors and raylogs use the highest attempt. Fall back when
-    listing fails, and derive attempt-scoped names from the writer's definition to prevent silent
-    reader/writer drift.
+    Attempt-scoped console and error artifacts use the highest uploaded attempt, which can still be
+    older than the live attempt while its replacement has not uploaded yet. Console falls back to its
+    canonical final name and errors to attempt 0 when listing fails; derive scoped names from the
+    writer's definition to prevent silent reader/writer drift.
     """
     import re
 

{freesolo_flash_dev-1.2.89 → freesolo_flash_dev-1.2.94}/pyproject.toml

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

{freesolo_flash_dev-1.2.89 → freesolo_flash_dev-1.2.94}/tests/test_cli_commands.py

@@ -1429,6 +1429,60 @@ def test_status_runs_and_log_command(fake_client, capsys, monkeypatch) -> None:
     assert "cost_usd" not in out
 
 
+def test_log_labels_previous_attempt_artifacts_after_the_live_attempt_log(
+    fake_client, capsys, monkeypatch
+) -> None:
+    live_heartbeat = {
+        "stage": "rl_step",
+        "step": 4,
+        "ts": 456.0,
+        "attempt": 1,
+        "gpu": {"device_name": "NVIDIA B200"},
+    }
+    fake_client.log_text = _format_heartbeat(live_heartbeat) + "\n"
+    monkeypatch.setattr(
+        fake_client,
+        "get_run",
+        lambda _run_id: {
+            "run_id": "flash-1",
+            "state": "running",
+            "remote": {"attempt": 1},
+            "last_heartbeat": live_heartbeat,
+        },
+    )
+    monkeypatch.setattr(
+        fake_client,
+        "get_worker_output",
+        lambda _run_id: {
+            "console_rl_attempt0.txt": (
+                'HEARTBEAT {"stage":"rl_step","step":0,"attempt":0,'
+                '"gpu":{"device_name":"NVIDIA H200"}}\n'
+            ),
+            "error_rl_attempt0.txt": "torch.OutOfMemoryError: CUDA OOM\n",
+            "raylogs_rl_attempt0.txt": "raylet exited\n",
+        },
+    )
+
+    assert _run(["runs", "log", "flash-1"]) == 0
+    out = capsys.readouterr().out
+
+    live_line = "worker: stage=rl_step attempt=1 step=4"
+    previous_header = (
+        "----- console_rl_attempt0.txt (attempt=0, previous attempt; current attempt=1) -----"
+    )
+    assert live_line in out
+    assert previous_header in out
+    assert (
+        "----- error_rl_attempt0.txt (attempt=0, previous attempt; current attempt=1) -----" in out
+    )
+    assert (
+        "----- raylogs_rl_attempt0.txt (attempt=0, previous attempt; current attempt=1) -----"
+        in out
+    )
+    assert out.index(live_line) < out.index(previous_header)
+    assert 'HEARTBEAT {"stage":"rl_step","step":0,"attempt":0' in out
+
+
 def test_log_prints_partial_log_line_with_newline(fake_client, capsys) -> None:
     fake_client.log_text = "partial log line"
     fake_client.get_worker_output = lambda run_id: {}

{freesolo_flash_dev-1.2.89 → freesolo_flash_dev-1.2.94}/tests/test_cli_render_theme.py

@@ -671,6 +671,30 @@ def test_heartbeat_is_current_attempt_rejects_malformed_identities() -> None:
     assert is_current({"remote": {"attempt": 2}}, {"attempt": 1}) is False
 
 
+def test_live_attempt_is_one_rule_for_every_surface_that_names_an_attempt() -> None:
+    """The spinner, the status line, and the worker-artifact labels must not disagree.
+
+    They are read within one screen of each other, so two provenance rules read as a run that is on
+    two attempts at once -- and the artifact labels would call the live attempt's console "previous".
+    """
+    live = render.live_attempt
+
+    # the plane's live attempt wins over the ping, which may be the superseded worker's.
+    assert live({"remote": {"attempt": 2}, "last_heartbeat": {"attempt": 1}}) == 2
+    # `remote` absent entirely (a managed payload that omits it): the ping is all there is.
+    assert live({"last_heartbeat": {"attempt": 3}}) == 3
+    # a `remote` that carries no usable attempt is not an answer either, so the ping still applies.
+    assert live({"remote": {}, "last_heartbeat": {"attempt": 1}}) == 1
+    for bad in ["1", True, 2.7, float("inf"), -1]:
+        assert live({"remote": {"attempt": bad}, "last_heartbeat": {"attempt": 4}}) == 4
+    # explicitly null `remote` is the teardown window: the attached ping belongs to a dead worker,
+    # so there is no live attempt to report rather than that worker's.
+    assert live({"remote": None, "last_heartbeat": {"attempt": 1}}) is None
+    # nothing usable anywhere reads as unknown, never as attempt 0.
+    assert live({}) is None
+    assert live({"last_heartbeat": {"attempt": "1"}}) is None
+
+
 def test_progress_age_always_adds_to_heartbeat_age(monkeypatch):
     from flash.cli.ui import heartbeat as heartbeat_ui
     from flash.cli.ui import render

{freesolo_flash_dev-1.2.89 → freesolo_flash_dev-1.2.94}/tests/test_jobs.py

@@ -479,7 +479,7 @@ def test_surface_heartbeat_logs_gpu_status(monkeypatch):
     assert stage == "sft_step"
     assert len(lines) == 1
     line = lines[0]
-    assert "worker: stage=sft_step step=12 loss=1.2346" in line
+    assert "worker: stage=sft_step attempt=0 step=12 loss=1.2346" in line
     assert "gpu[RTX 5090" in line
     assert "util=97%" in line
     assert "mem=22.2GB/31.8GB" in line

{freesolo_flash_dev-1.2.89 → freesolo_flash_dev-1.2.94}/tests/test_opd_train.py

@@ -9164,6 +9164,38 @@ def test_normalize_initial_prompt_rejects_a_state_with_no_prompt():
         normalize_initial_prompt(_Prompt(), state, None)
 
 
+def test_opd_multigpu_gpu_mem_util_matches_the_shared_tp_aware_resolver():
+    request = SimpleNamespace(
+        model_revision="",
+        knobs=SimpleNamespace(group_size=4),
+        spec=SimpleNamespace(gpu=SimpleNamespace(type="H200")),
+    )
+    prompt_state = SimpleNamespace(max_model_len=2048)
+    workload = SimpleNamespace(prompts_per_step=2)
+    runtime = SimpleNamespace(gpu_count=2)
+    model_id = "Qwen/Qwen3.6-35B-A3B"
+
+    got = opd_train._resolve_opd_gpu_mem_util(
+        request, prompt_state, workload, runtime, model_id, fp8_kv=False
+    )
+    want = rl_train.resolve_gpu_mem_util(
+        {
+            "model_id": model_id,
+            "model_revision": "",
+            "engine_len": 2048,
+            "group_size": 8,
+        },
+        gpu_type="H200",
+        n_gpus=2,
+        fp8_kv=False,
+        sleep_unsupported=True,
+        preserve_legacy_floor=True,
+    )
+
+    assert got == want
+    assert got < rl_train._DEFAULT_GPU_MEM_UTIL
+
+
 def test_build_opd_overrides_sizes_the_rollout_memory_budget():
     # opd emitted no gpu_memory_utilization at all, so verl substituted its own default of 0.5 and
     # the engine claimed half the CARD on every wake regardless of what the trainer held. this is

{freesolo_flash_dev-1.2.89 → freesolo_flash_dev-1.2.94}/tests/test_resume_on_retry.py

@@ -1132,6 +1132,7 @@ def test_a_retry_marks_where_the_previous_attempt_ends_in_the_log(orch, monkeypa
         if attempt == 0:
             print("Traceback (most recent call last):\ntorch.OutOfMemoryError: CUDA OOM", file=log)
             return PollResult(False, failure="stalled", detail="infra")
+        print("worker: stage=rl_step attempt=1 step=1", file=log)
         return PollResult(True, metrics={"train_tokens": 4096})
 
     monkeypatch.setattr(rp_jobs, "submit_run", fake_submit)
@@ -1147,6 +1148,9 @@ def test_a_retry_marks_where_the_previous_attempt_ends_in_the_log(orch, monkeypa
     assert text.index("CUDA OOM") < text.index(marker), (
         "the marker must sit after the failure it disowns, or it cannot separate the two attempts"
     )
+    assert text.index(marker) < text.index("worker: stage=rl_step attempt=1 step=1"), (
+        "the replacement attempt's heartbeat must follow the boundary that assigns its provenance"
+    )
 
 
 def test_the_marker_does_not_claim_one_previous_attempt_after_two_failures(orch, monkeypatch):

{freesolo_flash_dev-1.2.89 → freesolo_flash_dev-1.2.94}/tests/test_rl_train.py

@@ -1730,7 +1730,19 @@ def test_gpu_mem_util_is_the_sized_budget_not_a_constant():
     got = rl_train.resolve_gpu_mem_util(
         _mem_util_inp(), gpu_type="H100", n_gpus=1, fp8_kv=False, sleep_unsupported=False
     )
-    assert got == want
+    explicit_tp_one = colocate_kv_util(
+        float(info.params_b),
+        2048,
+        float(get_gpu_info("H100").vram_gb),
+        sleep_mode=True,
+        num_generations=8,
+        active_params_b=None,
+        fp8_kv=False,
+        model_info=info,
+        tensor_parallel=1,
+    )
+    assert got == want == explicit_tp_one
+    assert got.hex() == "0x1.11eb851eb851ep-2"
     # and it is genuinely NOT the old constant, so the test cannot pass on an unwired build.
     assert got != rl_train._DEFAULT_GPU_MEM_UTIL
     assert got < rl_train._DEFAULT_GPU_MEM_UTIL
@@ -1791,22 +1803,49 @@ def test_gpu_mem_util_sizing_reaches_the_launch_config():
     )
 
 
-def test_gpu_mem_util_keeps_the_constant_where_the_model_does_not_apply():
-    """The three shapes ``colocate_kv_util`` cannot size are left on the conservative constant.
-
-    Multi-gpu is the load-bearing one: the rollout is tensor-parallel, so vLLM's weight copy is
-    sharded ACROSS cards while the model sizes one whole copy against ONE card. Sizing it here would
-    over-reserve per rank on exactly the shapes that are already tight.
-    """
+def test_multigpu_gpu_mem_util_shards_only_weights_and_frees_the_observed_shortfall():
     default = rl_train._DEFAULT_GPU_MEM_UTIL
-    # multi-gpu: sharded rollout, unmodelled.
-    assert (
-        rl_train.resolve_gpu_mem_util(
-            _mem_util_inp(), gpu_type="H100", n_gpus=2, fp8_kv=False, sleep_unsupported=False
+    inp = _mem_util_inp(model_id="Qwen/Qwen3.6-35B-A3B", group_size=4)
+    for gpu_type, vram_gb in (("H200", 141.0), ("B200", 180.0)):
+        got = rl_train.resolve_gpu_mem_util(
+            inp, gpu_type=gpu_type, n_gpus=2, fp8_kv=False, sleep_unsupported=True
         )
-        == default
-    )
-    # unknown card: the budget is a FRACTION of the card, so there is nothing to take a fraction of.
+        assert got < default
+        assert (default - got) * vram_gb > 1.57
+
+
+def test_multigpu_gpu_mem_util_never_exceeds_the_previous_constant():
+    default = rl_train._DEFAULT_GPU_MEM_UTIL
+    model_ids = ("Qwen/Qwen3.5-0.8B", "Qwen/Qwen3.5-4B", "Qwen/Qwen3.6-35B-A3B")
+    for model_id in model_ids:
+        for gpu_type in ("H100", "H200", "B200"):
+            for tensor_parallel in (2, 4, 8):
+                for engine_len in (1024, 2048, 8192, 32768):
+                    for group_size in (2, 4, 8):
+                        got = rl_train.resolve_gpu_mem_util(
+                            _mem_util_inp(
+                                model_id=model_id,
+                                engine_len=engine_len,
+                                group_size=group_size,
+                            ),
+                            gpu_type=gpu_type,
+                            n_gpus=tensor_parallel,
+                            fp8_kv=False,
+                            sleep_unsupported=True,
+                        )
+                        assert 0.10 <= got <= default, (
+                            model_id,
+                            gpu_type,
+                            tensor_parallel,
+                            engine_len,
+                            group_size,
+                            got,
+                        )
+
+
+def test_gpu_mem_util_keeps_the_constant_where_the_model_does_not_apply(monkeypatch):
+    default = rl_train._DEFAULT_GPU_MEM_UTIL
+    # unknown card: the budget is a fraction of the card, so there is nothing to take a fraction of.
     for unknown in ("", "   ", "Nonexistent9000"):
         assert (
             rl_train.resolve_gpu_mem_util(
@@ -1815,9 +1854,13 @@ def test_gpu_mem_util_keeps_the_constant_where_the_model_does_not_apply():
             == default
         )
     # unresolvable model size: the weight copy is the dominant term.
+    monkeypatch.setattr("flash.engine.plan.vram.resolve_params_b", lambda *_args, **_kwargs: None)
     assert (
         rl_train.resolve_gpu_mem_util(
-            _mem_util_inp(model_id="not-a-real-org/not-a-real-model-xyz"),
+            _mem_util_inp(
+                model_id="Qwen/Qwen3.5-4B",
+                model_revision="a" * 40,
+            ),
             gpu_type="H100",
             n_gpus=1,
             fp8_kv=False,
@@ -1825,6 +1868,14 @@ def test_gpu_mem_util_keeps_the_constant_where_the_model_does_not_apply():
         )
         == default
     )
+    # any sizing exception keeps launch on the prior constant rather than making sizing fatal.
+    monkeypatch.setattr("flash.providers.base.get_gpu_info", lambda _gpu_type: 1 / 0)
+    assert (
+        rl_train.resolve_gpu_mem_util(
+            _mem_util_inp(), gpu_type="H100", n_gpus=2, fp8_kv=False, sleep_unsupported=False
+        )
+        == default
+    )
 
 
 def test_gpu_mem_util_does_not_credit_an_offload_a_resident_engine_never_takes():