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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: freesolo-flash-dev
-Version: 1.2.97
+Version: 1.2.98
 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.97 → freesolo_flash_dev-1.2.98}/flash/cli/commands/train_cost.py

@@ -61,21 +61,30 @@ def _cmd_train_cost(args) -> int:
 
 def _cmd_train_cost_offline(spec) -> int:
     """Catalog-only quote for the algorithms that do not need workload evidence yet (grpo, opd)."""
+    from dataclasses import replace
+
     from flash.cost import estimate_cost
+    from flash.schema import MIN_LORA_RANK
 
+    config = runconfig_from_spec(spec)
     if spec.train.init_from_adapter:
-        # --cost is offline/catalog-only and cannot read the source adapter, so the rank stays at the
-        # local default. Warm starts train and are priced at the SOURCE adapter's authoritative rank
-        # (resolved server-side at submit/dry-run), which can be higher — so this estimate may
-        # under-quote. stderr keeps stdout clean for machine-readable callers.
+        # --cost is offline/catalog-only and cannot read the source adapter's rank, and the parser
+        # rejects an authored `lora_rank` alongside `init_from_adapter`, so there is never an
+        # authored rank to quote at. stderr keeps stdout clean for machine-readable callers.
+        #
+        # quote at rank 1 rather than at the serialization placeholder. the placeholder is not a
+        # measurement of anything, and `estimate_cost` reads the rank for both the exact-card fit and
+        # cost-ranked hardware selection. rank 1 is a vram lower bound, but not a dollar lower bound:
+        # a higher source rank can move the run onto a cheaper multi-card shape.
+        config = replace(config, lora_rank=MIN_LORA_RANK)
         print(
-            "warning: warm-start (train.init_from_adapter) cost uses the default LoRA rank; the "
-            "source adapter's rank is authoritative and resolved at submit, so a higher-rank source "
-            f"may cost more than this estimate. Run `{_commands().CLI_NAME} train --dry-run` "
-            "for a source-rank quote.",
+            "warning: warm-start (train.init_from_adapter) cost is a provisional rank-1 estimate. "
+            "The authoritative source adapter rank is resolved server-side and can change the "
+            "selected hardware and cost in either direction. Run "
+            f"`{_commands().CLI_NAME} train --dry-run` for a quote using the resolved source rank.",
             file=sys.stderr,
         )
-    estimate = estimate_cost(runconfig_from_spec(spec))
+    estimate = estimate_cost(config)
     if render.styled():
         print(render.cost_panel(estimate))
     else:

{freesolo_flash_dev-1.2.97 → freesolo_flash_dev-1.2.98}/flash/engine/plan/kv_sizing.py

@@ -132,6 +132,26 @@ def _colocate_util_cap(weights_gb: float, total_vram_gb: float) -> float:
     return 0.55 if (big_weight_copy and leaves_room_for_trainer_copy) else 0.45
 
 
+def _capped_engine_util(
+    weights_gb: float,
+    adapter_gb: float,
+    requested_kv_gb: float,
+    total_vram_gb: float,
+    util_cap: float,
+) -> float:
+    """Fit weights, adapter, then KV into vllm's capped share of one card."""
+    card_gb = max(1.0, total_vram_gb)
+    max_engine_gb = util_cap * card_gb
+    # gpu_memory_utilization is vllm's whole allocation, and its lora weights live inside it. reserve
+    # weights plus adapter first; only a cap collision may reduce the requested kv pool. spelling out
+    # fitted_kv_gb makes that loss visible instead of incidentally subtracting adapter bytes from a
+    # weights-plus-kv fraction. if weights plus adapter alone exceed the cap, kv reaches zero and the
+    # returned cap truthfully exposes that this engine shape has no viable cache budget.
+    fitted_kv_gb = max(0.0, min(requested_kv_gb, max_engine_gb - weights_gb - adapter_gb))
+    fitted_engine_gb = weights_gb + adapter_gb + fitted_kv_gb
+    return min(util_cap, fitted_engine_gb / card_gb)
+
+
 def colocate_kv_util(
     params_b: float | None,
     vllm_max_len: int,
@@ -143,14 +163,19 @@ def colocate_kv_util(
     model_info=None,
     preserve_legacy_floor: bool = False,
     tensor_parallel: int = 1,
+    lora_rank: int = 0,
 ) -> float:
     """vllm_gpu_memory_utilization for one colocated rollout tensor-parallel rank.
 
-    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.
+    budget the rank's weight shard, rank-local bf16 lora footprint, and conservative full kv cache.
+    catalog models use geometry with measured overhead and an 8 gb floor; uncataloged models retain
+    the legacy kv equation and no unverified adapter margin.
     """
+    from flash.engine.plan.vram import _lora_weight_memory_gb
+
     total_weights_gb = max(0.5, float(params_b or 1.0)) * 2.0
     tp_size = max(1, int(tensor_parallel))
+    adapter_gb = _lora_weight_memory_gb(lora_rank, model_info, tensor_parallel=tp_size) or 0.0
     # 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
@@ -162,7 +187,7 @@ def colocate_kv_util(
     # but not H200. `_NvidiaSmiPeakSampler` verifies runtime headroom.
     _util_cap = _colocate_util_cap(weights_gb, total_vram_gb)
     if not sleep_mode:
-        # `gpu_memory_utilization` covers weights plus KV, so budget both. scale KV with
+        # `gpu_memory_utilization` covers weights, adapter, and kv, so budget all three. scale kv with
         # context and group, floor at `_KV_CAP`, and keep the 0.45 cap aligned with
         # `estimate_vram_gb(..., sleep_offload=False)`.
         kv_gb = max(
@@ -176,7 +201,8 @@ def colocate_kv_util(
                 preserve_legacy_floor=preserve_legacy_floor,
             ),
         )
-        return max(_MIN_ENGINE_UTIL, min(_util_cap, (weights_gb + kv_gb) / max(1.0, total_vram_gb)))
+        sized = _capped_engine_util(weights_gb, adapter_gb, kv_gb, total_vram_gb, _util_cap)
+        return max(_MIN_ENGINE_UTIL, sized)
     # 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(
@@ -191,7 +217,7 @@ def colocate_kv_util(
             preserve_legacy_floor=preserve_legacy_floor,
         ),
     )
-    sized = min(_util_cap, (weights_gb + kv_pool_gb) / max(1.0, total_vram_gb))
+    sized = _capped_engine_util(weights_gb, adapter_gb, kv_pool_gb, total_vram_gb, _util_cap)
     # 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,

{freesolo_flash_dev-1.2.97 → freesolo_flash_dev-1.2.98}/flash/engine/plan/vram.py

@@ -129,6 +129,33 @@ def _legacy_lora_floor_gb(lora_rank: int, effective_params_b: float) -> float:
     return (lora_rank / 16.0) * (0.3 + 0.04 * effective_params_b)
 
 
+def _lora_parameter_count(
+    lora_rank: int, model_info=None, *, tensor_parallel: int = 1
+) -> int | None:
+    """lora parameters resident on one tensor-parallel rank from concrete target geometry."""
+    shapes = tuple(getattr(model_info, "lora_target_shapes", ()) or ())
+    rank = int(lora_rank)
+    if not shapes or rank <= 0:
+        return None
+    tp_size = max(1, int(tensor_parallel))
+    if tp_size == 1:
+        return rank * sum((int(i) + int(o)) * int(count) for i, o, count in shapes)
+    # vllm shards one factor, except the current min-dim-1 gate keeps that factor whole. integer shard
+    # sizing covers that catalog shape exactly; this is not a general replicated-layer classifier.
+    return sum(
+        rank * int(count) * (max(int(i), int(o)) + math.ceil(min(int(i), int(o)) / tp_size))
+        for i, o, count in shapes
+    )
+
+
+def _lora_weight_memory_gb(
+    lora_rank: int, model_info=None, *, tensor_parallel: int = 1
+) -> float | None:
+    """rank-local resident bf16 lora weights, or none without catalog target geometry."""
+    params = _lora_parameter_count(lora_rank, model_info, tensor_parallel=tensor_parallel)
+    return None if params is None else params * _BYTES_PER_PARAM["bf16"] / 1e9
+
+
 def _lora_memory_gb(
     lora_rank: int,
     effective_params_b: float,
@@ -137,14 +164,9 @@ def _lora_memory_gb(
 ) -> float:
     """Adapter, gradient, and optimizer memory for the actual PEFT all-linear targets."""
     floor = _legacy_lora_floor_gb(lora_rank, effective_params_b)
-    shapes = tuple(getattr(model_info, "lora_target_shapes", ()) or ())
-    if not shapes:
+    trainable_params = _lora_parameter_count(lora_rank, model_info)
+    if trainable_params is None:
         return floor
-    target_dims = sum(
-        (int(in_features) + int(out_features)) * int(count)
-        for in_features, out_features, count in shapes
-    )
-    trainable_params = max(1, int(lora_rank)) * target_dims
     bytes_per_param = (
         _LORA_PAGED_BYTES_PER_PARAM
         if (algorithm or "").lower() in ("grpo", "rl")
@@ -225,15 +247,15 @@ def grpo_kv_floor_gb(
     fp8_kv: bool = False,
     model_info=None,
     preserve_legacy_floor: bool = False,
+    lora_rank: int = 0,
+    tensor_parallel: int = 1,
 ) -> int:
-    """Smallest card (GB) whose colocated vLLM executor budget still leaves a viable KV pool.
-
-    The 0.45/0.55 cap must hold the weight copy plus half the concurrent group's KV; otherwise
-    vLLM can fail with "No available memory for the cache blocks".
-    """
+    """Smallest card whose cap holds rank-local weights, adapter, and half the KV pool."""
     kv_params_b = float(active_params_b) if active_params_b else float(params_b)
-    weights_gb = max(0.5, float(params_b)) * 2.0
-    need = weights_gb + 0.5 * _resident_kv_gb(
+    tp_size = max(1, int(tensor_parallel))
+    weights_gb = max(0.5, float(params_b)) * 2.0 / tp_size
+    adapter_gb = _lora_weight_memory_gb(lora_rank, model_info, tensor_parallel=tp_size) or 0.0
+    kv_gb = _resident_kv_gb(
         kv_params_b,
         vllm_max_len,
         group_size,
@@ -241,6 +263,7 @@ def grpo_kv_floor_gb(
         model_info=model_info,
         preserve_legacy_floor=preserve_legacy_floor,
     )
+    need = weights_gb + adapter_gb + 0.5 * kv_gb
     lower = math.ceil(need / 0.55)
     upper = math.ceil(need / 0.45)
     for total_vram_gb in range(lower, upper + 1):
@@ -295,6 +318,7 @@ def _rollout_kv_floor_gb(
     model_info=None,
     model_id: str = "",
     preserve_legacy_floor: bool = False,
+    lora_rank: int = 0,
 ) -> int:
     floor = grpo_kv_floor_gb(
         params_b,
@@ -303,6 +327,7 @@ def _rollout_kv_floor_gb(
         active_params_b=active_params_b,
         model_info=model_info,
         preserve_legacy_floor=preserve_legacy_floor,
+        lora_rank=lora_rank,
     )
     from flash.providers.base import max_non_fp8_kv_vram_gb
 
@@ -323,6 +348,7 @@ def _rollout_kv_floor_gb(
         fp8_kv=True,
         model_info=model_info,
         preserve_legacy_floor=preserve_legacy_floor,
+        lora_rank=lora_rank,
     )
     return max(fp8_floor, ceiling + 1)
 
@@ -838,6 +864,7 @@ def _catalog_model_required_vram_gb(
                 model_info=sizing_info,
                 model_id=model_id,
                 preserve_legacy_floor=is_opd,
+                lora_rank=lora_rank,
             ),
         )
     return need

{freesolo_flash_dev-1.2.97 → freesolo_flash_dev-1.2.98}/flash/engine/worker/opd_train_runner.py

@@ -560,6 +560,7 @@ 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),
+            "lora_rank": int(workload.lora_rank),
             # 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),

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

@@ -503,13 +503,16 @@ def resolve_gpu_mem_util(
 ) -> 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
+    ``gpu_memory_utilization`` covers the second bf16 weight copy, lora adapter, and kv pool. using
     ``colocate_kv_util`` keeps worker allocation aligned with preflight admission.
 
-    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.
+    for multi-gpu runs, vllm shards its bf16 weight copy across tensor-parallel ranks. its default
+    lora layout shards only one factor of each projection's a/b pair, so the rank-local adapter uses
+    the larger safe orientation from catalog geometry rather than dividing the full adapter evenly.
+    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 adapter lives inside vllm's reservation, and 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:
@@ -564,6 +567,7 @@ def resolve_gpu_mem_util(
             model_info=info,
             preserve_legacy_floor=preserve_legacy_floor,
             tensor_parallel=n_gpus,
+            lora_rank=int(inp["lora_rank"]),
         )
         # 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.

{freesolo_flash_dev-1.2.97 → freesolo_flash_dev-1.2.98}/flash/schema/__init__.py

@@ -50,6 +50,12 @@ from flash.schema.fields import (
     _wandb_spec,
 )
 
+# the smallest rank the parser accepts, and so the smallest a source adapter can turn out to have.
+# unresolved warm starts use it instead of the serialization default for permissive client-side
+# sizing. it is a vram lower bound, not a dollar lower bound because cost-ranked hardware selection
+# is non-monotonic across shapes. keep it aligned with the `lora_rank` parser floor.
+MIN_LORA_RANK = 1
+
 _OWNER_REPO_RE = r"[A-Za-z0-9][A-Za-z0-9._-]*"
 _RUN_ID_RE = r"[A-Za-z0-9][A-Za-z0-9._-]{0,127}"
 # canonical short checkpoint references name a run alias or a saved checkpoint. immutable adapter
@@ -437,6 +443,7 @@ def _validate_gpu_section(
     model_revision: str,
     algorithm: str,
     train_raw: dict[str, Any],
+    init_from_adapter: str,
     thinking: bool,
 ) -> tuple[GpuSpec, bool]:
     """Validate the GPU section and return its canonical spec plus auto-sizing provenance."""
@@ -483,10 +490,18 @@ def _validate_gpu_section(
         # provider preferences are soft; configured unnamed providers remain eligible.
 
     requested_gpu_count = authored_gpu_ceiling(gpu_spec.type, gpu_count)
+    unresolved_warmstart_rank = bool(init_from_adapter) and "lora_rank" not in train_raw
+    # the client cannot resolve a run id to the source adapter metadata, so rank 32 is only a
+    # serialization placeholder here. keep geometry checks and reject shapes that cannot fit even at
+    # the minimum valid rank; the server replaces this value from adapter_config.json before its
+    # authoritative allocation-time vram check.
+    preflight_train = (
+        {**train_raw, "lora_rank": MIN_LORA_RANK} if unresolved_warmstart_rank else train_raw
+    )
     preflight_gpu_count = provisional_gpu_count(
         model,
         algorithm,
-        train=train_raw,
+        train=preflight_train,
         thinking=thinking,
         geometry_model_revision=model_revision,
         gpu_count=requested_gpu_count,
@@ -495,10 +510,16 @@ def _validate_gpu_section(
         if gpu_types and preflight_gpu_count <= 1 and not model_revision:
             from flash.providers.allocator import required_vram_gb
 
+            # sized from `preflight_train`, so an unresolved warm start is measured at rank 1 rather
+            # than at the placeholder. that is a true vram lower bound: no source adapter can need
+            # less. a card rejected here therefore cannot fit at any rank the source could turn out to
+            # have, which keeps an impossible pin (an 80 GB A100 for a run needing 180 GB at rank 1)
+            # rejected at parse time. relaxing the rank is not the same as dropping the check --
+            # dropping it would let the authored `gpu.type` go unvalidated entirely.
             required_vram = required_vram_gb(
                 model,
                 algorithm,
-                train=train_raw,
+                train=preflight_train,
                 thinking=thinking,
             )
             # every authored class is checked, not just the head: a fallback too small to hold the
@@ -517,7 +538,7 @@ def _validate_gpu_section(
         provisional_gpu(
             model,
             algorithm=algorithm,
-            train=train_raw,
+            train=preflight_train,
             thinking=thinking,
             geometry_model_revision=model_revision,
             gpu_count=preflight_gpu_count,
@@ -529,7 +550,12 @@ def _validate_gpu_section(
 
 
 def _validate_algorithm_model_consistency(
-    model: str, algorithm: str, thinking: bool, train_raw: dict[str, Any]
+    model: str,
+    algorithm: str,
+    thinking: bool,
+    train_raw: dict[str, Any],
+    *,
+    init_from_adapter: str,
 ) -> tuple[str, int, int]:
     """Validate algorithm and model-info consistency."""
     try:
@@ -547,7 +573,6 @@ def _validate_algorithm_model_consistency(
             f"{model} always emits <think> reasoning and cannot run with thinking "
             f"disabled; set thinking = true"
         )
-    init_from_adapter = _init_from_adapter_ref(train_raw)
     if init_from_adapter and "lora_rank" in train_raw:
         raise ConfigError(
             "train.lora_rank cannot be set with train.init_from_adapter because source adapter "
@@ -558,7 +583,7 @@ def _validate_algorithm_model_consistency(
             "train.lora_alpha cannot be set with train.init_from_adapter because source adapter "
             "alpha metadata is authoritative"
         )
-    lora_rank = _train_int(train_raw, "lora_rank", minimum=1) or 32
+    lora_rank = _train_int(train_raw, "lora_rank", minimum=MIN_LORA_RANK) or 32
     # unset -> the tuned 2 x rank default. authored -> the user's value, which need not be 2 x rank.
     lora_alpha = _train_int(train_raw, "lora_alpha", minimum=1) or 2 * lora_rank
     max_lora_rank = serving_lora_rank_cap(info)
@@ -590,16 +615,22 @@ def spec_from_dict(
             )
         except (TypeError, ValueError) as exc:
             raise ConfigError(str(exc)) from exc
+    init_from_adapter = _init_from_adapter_ref(train_raw)
     gpu_spec, gpu_count_auto = _validate_gpu_section(
         raw,
         model=model,
         model_revision=model_revision,
         algorithm=algorithm,
         train_raw=train_raw,
+        init_from_adapter=init_from_adapter,
         thinking=thinking,
     )
     init_from_adapter, lora_rank, lora_alpha = _validate_algorithm_model_consistency(
-        model, algorithm, thinking, train_raw
+        model,
+        algorithm,
+        thinking,
+        train_raw,
+        init_from_adapter=init_from_adapter,
     )
     wandb_spec = _validate_wandb_section(raw)
 

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

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

{freesolo_flash_dev-1.2.97 → freesolo_flash_dev-1.2.98}/tests/test_cli_estimate.py

@@ -1021,8 +1021,8 @@ def test_sft_cost_warns_when_an_env_key_shadows_the_saved_login(tmp_path, monkey
 
 
 @pytest.mark.parametrize("algorithm", ["grpo", "opd"])
-def test_non_sft_cost_stays_offline(tmp_path, monkeypatch, capsys, algorithm):
-    """grpo/opd keep the local analytical quote until PR2 profiles their rollouts.
+def test_warm_start_non_sft_cost_stays_offline(tmp_path, monkeypatch, capsys, algorithm):
+    """warm-start grpo/opd keep the local analytical quote until PR2 profiles their rollouts.
 
     Asserted by making the client constructor itself fatal: a passing quote then proves no
     authenticated request was possible, not merely that none was observed.
@@ -1037,7 +1037,8 @@ def test_non_sft_cost_stays_offline(tmp_path, monkeypatch, capsys, algorithm):
     monkeypatch.setenv("FLASH_STYLE", "0")
     # group_size 2 is grpo's floor (advantages are group-relative) and is valid for opd too.
     body = SFT_TOML.replace('algorithm = "sft"', f'algorithm = "{algorithm}"').replace(
-        "batch_size = 8\n", "prompts_per_step = 8\nmax_examples = 40\ngroup_size = 2\n"
+        "batch_size = 8\n",
+        'prompts_per_step = 8\nmax_examples = 40\ngroup_size = 2\ninit_from_adapter = "source-run"\n',
     )
 
     assert cmd_train(_sft_args(tmp_path, body)) == 0
@@ -1047,6 +1048,121 @@ def test_non_sft_cost_stays_offline(tmp_path, monkeypatch, capsys, algorithm):
     assert "source rows in published copy" not in captured.out
     assert "selected from" not in captured.out
     assert "shadowed!" not in captured.err
+    assert (
+        "warning: warm-start (train.init_from_adapter) cost is a provisional rank-1 estimate. "
+        "The authoritative source adapter rank is resolved server-side and can change the selected "
+        "hardware and cost in either direction. Run `flash train --dry-run` for a quote using the "
+        "resolved source rank.\n" in captured.err
+    )
+
+
+@pytest.mark.parametrize("algorithm", ["grpo", "opd"])
+def test_plain_non_sft_cost_stays_offline(tmp_path, monkeypatch, capsys, algorithm):
+    """Ordinary grpo/opd quotes stay local independently of the warm-start branch."""
+    from flash.cli import commands
+
+    monkeypatch.setattr(
+        commands,
+        "client_from_config",
+        lambda *_a, **_kw: pytest.fail(f"{algorithm} --cost must not contact the control plane"),
+    )
+    monkeypatch.setenv("FLASH_STYLE", "0")
+    body = SFT_TOML.replace('algorithm = "sft"', f'algorithm = "{algorithm}"').replace(
+        "batch_size = 8\n", "prompts_per_step = 8\nmax_examples = 40\ngroup_size = 2\n"
+    )
+
+    assert cmd_train(_sft_args(tmp_path, body)) == 0
+    captured = capsys.readouterr()
+    assert "TOTAL" in captured.out
+    assert "warm-start" not in captured.err
+
+
+def test_warm_start_cost_quotes_the_card_the_parser_accepts(tmp_path, capsys, monkeypatch):
+    """The quote must not refuse the exact configuration the parser now admits.
+
+    `estimate_cost` reads the rank for both the price and the exact-card fit, so quoting at the
+    serialization placeholder rejected a single B200 as needing 199 GB and advised `--gpus 2`, a
+    shortfall the real source rank may not have. Sized at the rank-1 vram lower bound it fits in 180
+    GB, while the warning accurately labels the dollars as provisional.
+    """
+    from flash.cli import commands
+
+    monkeypatch.setattr(
+        commands,
+        "client_from_config",
+        lambda *_a, **_kw: pytest.fail("offline quote must not contact the plane"),
+    )
+    monkeypatch.setenv("FLASH_STYLE", "0")
+    body = SFT_TOML.replace('algorithm = "sft"', 'algorithm = "grpo"').replace(
+        "batch_size = 8\n",
+        'prompts_per_step = 8\nmax_examples = 40\ngroup_size = 2\ninit_from_adapter = "source-run"\n',
+    )
+    # SFT_TOML ends with an empty `[gpu]`, so the pin is appended into that section.
+    body = body.replace('model = "Qwen/Qwen3.5-4B"', 'model = "Qwen/Qwen3.6-35B-A3B"')
+    body += 'type = "B200"\ncount = 1\n'
+
+    assert cmd_train(_sft_args(tmp_path, body)) == 0
+    captured = capsys.readouterr()
+    assert "TOTAL" in captured.out
+    assert "--gpus 2" not in captured.err
+    assert (
+        "warning: warm-start (train.init_from_adapter) cost is a provisional rank-1 estimate. "
+        "The authoritative source adapter rank is resolved server-side and can change the selected "
+        "hardware and cost in either direction. Run `flash train --dry-run` for a quote using the "
+        "resolved source rank.\n" in captured.err
+    )
+
+
+def _warm_start_rank_boundary_config():
+    from flash.cost.spec import runconfig_from_spec
+
+    spec = spec_from_dict(
+        {
+            "project": "00000000-0000-0000-0000-000000000001",
+            "model": "Qwen/Qwen3.6-35B-A3B",
+            "algorithm": "grpo",
+            "environment": {"id": "my-org/my-proj/my-env"},
+            "train": {"prompts_per_step": 8, "max_examples": 40, "group_size": 2},
+        }
+    )
+    return runconfig_from_spec(spec)
+
+
+def test_higher_warm_start_rank_crosses_hardware_shape_boundary():
+    """Catalog geometry, not mutable prices, owns the rank-dependent hardware boundary."""
+    from dataclasses import replace
+
+    from flash.cost import estimate_cost
+    from flash.providers.base import GPU_INFO
+
+    config = _warm_start_rank_boundary_config()
+    rank_1 = estimate_cost(replace(config, lora_rank=1))
+    rank_8 = estimate_cost(replace(config, lora_rank=8))
+
+    assert GPU_INFO["B200"].vram_gb == 180
+    assert GPU_INFO["H200"].vram_gb == 141
+    assert (rank_1.required_vram_gb, rank_1.gpu_count, rank_1.gpu) == (180, 1, "B200")
+    assert (rank_8.required_vram_gb, rank_8.gpu_count, rank_8.gpu) == (182, 2, "H200")
+
+
+def test_higher_warm_start_rank_can_select_cheaper_hardware(monkeypatch):
+    """Controlled prices prove that crossing the shape boundary can lower total cost."""
+    from dataclasses import replace
+
+    from flash.cost import estimate_cost
+    from flash.providers import base
+
+    b200 = replace(base.GPU_INFO["B200"], hourly_usd=5.50)
+    h200 = replace(base.GPU_INFO["H200"], hourly_usd=4.00)
+    monkeypatch.setattr(base, "GPU_INFO", {"B200": b200, "H200": h200})
+
+    config = _warm_start_rank_boundary_config()
+    rank_1 = estimate_cost(replace(config, lora_rank=1))
+    rank_8 = estimate_cost(replace(config, lora_rank=8))
+
+    assert (rank_1.gpu, rank_1.gpu_count, rank_1.gpu_hourly_usd) == ("B200", 1, 5.50)
+    assert (rank_8.gpu, rank_8.gpu_count, rank_8.gpu_hourly_usd) == ("H200", 2, 4.00)
+    assert rank_8.total_usd < rank_1.total_usd
 
 
 def test_cmd_train_cost_rejects_context_above_serving_cap(tmp_path):