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

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

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

@@ -3,9 +3,7 @@
 from __future__ import annotations
 
 import argparse
-import difflib
 import math
-import re
 import shlex
 import sys
 from typing import NoReturn
@@ -57,6 +55,7 @@ from flash.cli.commands.traces import (
     RECORDS_FORMAT,
     cmd_traces_export,
 )
+from flash.cli.errors import friendly_error
 from flash.cli.ui import render
 from flash.client.config import shadowed_login_warning
 from flash.core.catalog import ALGORITHMS
@@ -170,20 +169,6 @@ def _gpu_count_override(value: str) -> str:
         raise argparse.ArgumentTypeError(f"expected a number of gpus, got {value!r}") from None
 
 
-def _friendly_message(message: str) -> str:
-    """Shorten argparse's verbose ``invalid choice: 'x' (choose from a, b, c, ...)`` into a concise
-    ``unknown command 'x' (did you mean 'y'?)`` — the single closest match instead of dumping the
-    whole list. Other messages pass through untouched. Styled path only; the machine path keeps
-    argparse's exact text (scripts and the error tests match on the literal `invalid choice`)."""
-    m = re.search(r"invalid choice: '([^']*)'(?: \(choose from (.*)\))?", message)
-    if not m:
-        return message
-    bad, raw_choices = m.group(1), m.group(2) or ""
-    choices = [c.strip().strip("'\"") for c in raw_choices.split(",") if c.strip()]
-    near = difflib.get_close_matches(bad, choices, n=1)
-    return f"unknown command '{bad}'" + (f" (did you mean '{near[0]}'?)" if near else "")
-
-
 class _ThemedParser(argparse.ArgumentParser):
     """Base parser whose usage errors match the rest of the themed CLI.
 
@@ -201,9 +186,13 @@ class _ThemedParser(argparse.ArgumentParser):
         if not render.styled():
             super().error(message)  # argparse's raw usage + `prog: error: msg`, then exit 2
         # themed twin: the red ✗ error line (same idiom as main()'s catch-all and `flash login`),
-        # then a dimmed pointer at this parser's own --help instead of the raw usage block. An
-        # "invalid choice" becomes a short "did you mean" suggestion (see _friendly_message).
-        print(render.error(_friendly_message(message)), file=sys.stderr)
+        # then a dimmed pointer at this parser's own --help instead of the raw usage block. unknown
+        # commands and flags get a short "did you mean" suggestion (see `friendly_error`). the argv
+        # stash is what tells it which subcommand was selected; only the root parser carries one.
+        print(
+            render.error(friendly_error(message, self, getattr(self, "_flash_argv", None))),
+            file=sys.stderr,
+        )
         # dimmed pointer at THIS parser's own --help (argparse sets prog per parser: `flash --help`
         # for the root, `flash <cmd> --help` for a subcommand) instead of the raw usage block.
         print(render.arrow(f"run `{self.prog} --help` for usage"), file=sys.stderr)
@@ -902,6 +891,12 @@ def _redacted_args(raw_args: list[str]) -> list[str]:
 def main(argv: list[str] | None = None) -> int:
     raw_args = list(argv) if argv is not None else sys.argv[1:]
     parser = _build_parser()
+    # the tokens this invocation was actually given, for the error hook: argparse reports a
+    # subcommand's unknown flags from the ROOT parser, so scoping the suggestion to the selected
+    # command needs the argv `main` was called with. reading sys.argv there instead would be right
+    # only for the console-script path and wrong for every embedded caller -- including the tests,
+    # where it silently reads pytest's own argv and restores the whole-tree pool.
+    parser._flash_argv = raw_args
     args = parser.parse_args(argv)
     configure_logging(verbosity=getattr(args, "verbose", 0))
     debug = getattr(args, "debug", False)

{freesolo_flash_dev-1.2.89 → freesolo_flash_dev-1.2.95}/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.95/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.95/flash/cli/errors.py

@@ -0,0 +1,277 @@
+"""Turning argparse's usage errors into the themed CLI's short suggestions.
+
+argparse reports a mistyped flag by printing the whole usage block, which buries the one token that
+was actually wrong. These helpers reduce that to a single line naming the token and, where one
+exists, the closest real alternative.
+
+The hard part is making the suggestion one the user can actually follow. A subcommand's unknown
+tokens are handed back to the root parser to report, so candidate options must be recovered from the
+parser in effect at the rejected token's argv position. Looking only at the final selected command
+crosses nested command boundaries, while looking through the whole tree offers unrelated flags.
+
+Root flags are also positional here (`flash --debug login`, never `flash login --debug`), and `--`
+stops option parsing. The suggestion context therefore carries both boundaries, plus the original
+spelling of a repeated short root flag that must be moved without changing its count.
+"""
+
+from __future__ import annotations
+
+import argparse
+import difflib
+import re
+from collections.abc import Iterable
+from dataclasses import dataclass
+
+from flash._internal.channel import CLI_NAME
+
+
+@dataclass(frozen=True)
+class _SuggestionContext:
+    options: tuple[str, ...]
+    root_only: frozenset[str]
+    reposition: str | None = None
+    after_terminator: bool = False
+
+
+def friendly_error(
+    message: str, parser: argparse.ArgumentParser, argv: Iterable[str] | None
+) -> str:
+    """Rewrite one argparse usage error for the themed CLI.
+
+    The whole suggestion policy lives behind this call: which parser's flags are candidates, which
+    of them need repositioning, and how the message reads. The caller is an ``error()`` override on
+    every parser in the tree and should stay a two-liner -- deciding a candidate pool inline there
+    would put policy in a hook that exists to print.
+    """
+    if argv is None:
+        # a subparser raising its own error is already the parser in effect at the rejected token.
+        context = _SuggestionContext(tuple(sorted(_direct_options(parser))), frozenset())
+    else:
+        context = _suggestion_context(message, parser, tuple(argv))
+    return _friendly_message(message, context)
+
+
+def _friendly_message(message: str, context: _SuggestionContext) -> str:
+    """Shorten argparse's verbose errors into concise suggestions on the styled path.
+
+    An ``invalid choice: 'x' (choose from a, b, c, ...)`` becomes ``unknown command 'x' (did you
+    mean 'y'?)`` instead of dumping the whole list. An ``unrecognized arguments: --flag value``
+    suggests the closest option registered on the parser that rejected it. Other messages pass
+    through untouched. The machine path keeps argparse's exact text for scripts and error tests.
+    """
+    m = re.search(r"invalid choice: '([^']*)'(?: \(choose from (.*)\))?", message)
+    if m:
+        bad, raw_choices = m.group(1), m.group(2) or ""
+        choices = [c.strip().strip("'\"") for c in raw_choices.split(",") if c.strip()]
+        near = difflib.get_close_matches(bad, choices, n=1)
+        return f"unknown command '{bad}'" + (f" (did you mean '{near[0]}'?)" if near else "")
+
+    prefix = "unrecognized arguments: "
+    if not message.startswith(prefix):
+        return message
+    # a bare `-` starts with a dash but names no option, so it is an ordinary positional argparse
+    # could not place. requiring a character after the dash keeps it out of the typo machinery,
+    # which would otherwise answer it with the nearest short flag.
+    bad = next(
+        (
+            token.split("=", 1)[0]
+            for token in message[len(prefix) :].split()
+            if token.startswith("-") and len(token.split("=", 1)[0]) > 1
+        ),
+        None,
+    )
+    if bad is None:
+        return message
+    if context.after_terminator:
+        return message
+    # a root flag in the wrong position is a placement error, not a typo. repeatable short flags
+    # retain the complete spelling so moving `-vv` never downgrades it to one `-v`.
+    if context.reposition is not None:
+        return _reposition_message(context.reposition)
+    if bad in context.root_only:
+        return _reposition_message(bad)
+    # the position-scoped pool can still include an exact spelling shared with the root. dropping it
+    # keeps a rejected token from being echoed back as its own correction.
+    candidates = [option for option in context.options if option != bad]
+    near = difflib.get_close_matches(bad, candidates, n=1)
+    if not near:
+        return message
+    suggestion = near[0]
+    if suggestion in context.root_only:
+        # a root flag only parses before the subcommand, so naming it alone hands back a correction
+        # that fails the same way. say where it goes, or the second attempt is the first error again.
+        return f"unrecognized argument '{bad}' (did you mean '{suggestion}'? {_placement_hint(suggestion)})"
+    return f"unrecognized argument '{bad}' (did you mean '{suggestion}'?)"
+
+
+def _placement_hint(flag: str) -> str:
+    """Say where a root flag belongs, and claim nothing about the rest of the line.
+
+    This hint knows one thing: the flag parses only ahead of the subcommand. It does not know what
+    removing the flag from where it sat exposes. `--wait` takes an optional value, so a flag lifted
+    from between `--wait` and a run id leaves the run id to be swallowed as the timeout -- which
+    `_wait_seconds` then catches and explains precisely. Stating the destination and stopping there
+    keeps this hint true for every line; the one case it cannot see already answers itself.
+
+    Both the repositioning message and the respelling-plus-repositioning message end in this same
+    advice, so they share it rather than each spelling it out and drifting apart.
+    """
+    return f"it goes before the command, as `{CLI_NAME} {flag} ...`"
+
+
+def _reposition_message(flag: str) -> str:
+    """Answer a correctly spelled root flag that landed after the subcommand."""
+    return f"unrecognized argument '{flag}' ({_placement_hint(flag)})"
+
+
+def _suggestion_context(
+    message: str, root: argparse.ArgumentParser, argv: tuple[str, ...]
+) -> _SuggestionContext:
+    root_only = _root_only_options(root)
+    bad = _unrecognized_option(message)
+    if bad is None:
+        return _SuggestionContext((), root_only)
+    location = _option_location(root, argv, bad)
+    if location is None:
+        return _SuggestionContext((), root_only)
+    current, after_terminator = location
+    reposition = bad if _is_repeated_root_short_option(root, bad, root_only) else None
+    options = tuple(sorted({*_direct_options(current), *root_only}))
+    return _SuggestionContext(options, root_only, reposition, after_terminator)
+
+
+def _unrecognized_option(message: str) -> str | None:
+    prefix = "unrecognized arguments: "
+    if not message.startswith(prefix):
+        return None
+    return next(
+        (
+            token.split("=", 1)[0]
+            for token in message[len(prefix) :].split()
+            if token.startswith("-")
+        ),
+        None,
+    )
+
+
+def _option_location(
+    root: argparse.ArgumentParser, argv: tuple[str, ...], bad: str
+) -> tuple[argparse.ArgumentParser, bool] | None:
+    """Return the parser and terminator state at the rejected token's argv position."""
+    current = root
+    path_open = True
+    after_terminator = False
+    fallback: tuple[argparse.ArgumentParser, bool] | None = None
+
+    for token in argv:
+        if token == "--":
+            after_terminator = True
+            path_open = False
+            continue
+        if token.split("=", 1)[0] == bad:
+            location = (current, after_terminator)
+            fallback = location
+            # repeated spellings can include one accepted root occurrence and one rejected
+            # subcommand occurrence. choose the occurrence argparse could not consume here.
+            if after_terminator or len(_option_tuples(current, token)) != 1:
+                return location
+        if after_terminator or token.startswith("-") or not path_open:
+            continue
+        subparsers = next(
+            (a for a in current._actions if isinstance(a, argparse._SubParsersAction)), None
+        )
+        if subparsers is None:
+            path_open = False
+            continue
+        nxt = subparsers.choices.get(token)
+        if nxt is None:
+            path_open = False
+            continue
+        current = nxt
+
+    return fallback
+
+
+def _option_tuples(parser: argparse.ArgumentParser, token: str) -> list:
+    """The prefix matches for a token, for tokens argparse's own helper cannot be asked about.
+
+    `_get_option_tuples` indexes the token's second character, so a bare ``-`` raises IndexError
+    there. A lone dash is an ordinary unrecognized positional and must reach argparse's own message
+    rather than crash the code that was trying to improve it.
+    """
+    if len(token) < 2:
+        return []
+    return parser._get_option_tuples(token)
+
+
+def _parse_optional_fields(
+    parser: argparse.ArgumentParser, token: str
+) -> tuple[object, str | None, str | None] | None:
+    """The action, option spelling and explicit remainder argparse reads out of a token.
+
+    `_parse_optional` has three shapes across the interpreters `requires-python` admits: a 3-tuple
+    before 3.11.9/3.12.3, a 4-tuple once a separator was inserted, and from 3.12.11 a LIST of those
+    tuples, one per prefix match. Normalizing here keeps that churn in one place, and keeps the
+    separator out of the caller: the exact-spelling reconstruction the caller does already settles
+    what the separator would have told it. An unfamiliar shape returns None rather than raising,
+    because this runs inside the handler whose whole job is to render someone else's error.
+    """
+    parsed = parser._parse_optional(token)
+    if isinstance(parsed, list):
+        # several prefix matches means the spelling is ambiguous, so it is not one repeated flag.
+        parsed = parsed[0] if len(parsed) == 1 else None
+    if not isinstance(parsed, tuple) or len(parsed) not in (3, 4):
+        return None
+    return parsed[0], parsed[1], parsed[-1]
+
+
+def _is_repeated_root_short_option(
+    root: argparse.ArgumentParser, token: str, root_only: frozenset[str]
+) -> bool:
+    """Whether token is a compact repeat of a count-style root option, such as ``-vv``."""
+    if not token.startswith("-") or token.startswith("--"):
+        return False
+    fields = _parse_optional_fields(root, token)
+    if fields is None:
+        return False
+    action, option, explicit = fields
+    return (
+        isinstance(action, argparse._CountAction)
+        and option in root_only
+        and isinstance(explicit, str)
+        and bool(explicit)
+        # the whole token must be that one short flag repeated, so `-vx` and `-v=2` are not repeats.
+        and token == option + option[-1] * len(explicit)
+    )
+
+
+def _direct_options(parser: argparse.ArgumentParser) -> Iterable[str]:
+    for action in parser._actions:
+        yield from action.option_strings
+
+
+def _parser_options(parser: argparse.ArgumentParser) -> Iterable[str]:
+    """Yield option strings registered on a parser and its subparsers."""
+    yield from _direct_options(parser)
+    for action in parser._actions:
+        if isinstance(action, argparse._SubParsersAction):
+            for subparser in action.choices.values():
+                yield from _parser_options(subparser)
+
+
+def _root_only_options(parser: argparse.ArgumentParser) -> frozenset[str]:
+    """Option strings accepted only BEFORE the subcommand.
+
+    This CLI's root flags are positional: `flash --debug login` parses, `flash login --debug` does
+    not. A suggestion drawn from the whole option pool can therefore name a real flag that still
+    fails in the position the user typed it, so the caller needs to know which ones to reposition.
+    """
+    root: set[str] = set()
+    sub: set[str] = set()
+    for action in parser._actions:
+        if isinstance(action, argparse._SubParsersAction):
+            for subparser in action.choices.values():
+                sub.update(_parser_options(subparser))
+        else:
+            root.update(action.option_strings)
+    return frozenset(root - sub)

{freesolo_flash_dev-1.2.89 → freesolo_flash_dev-1.2.95}/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.95}/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.95}/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.95}/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.95}/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,
     )