freesolo-flash-dev 1.2.130__tar.gz → 1.2.131__tar.gz

{freesolo_flash_dev-1.2.130 → freesolo_flash_dev-1.2.131}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: freesolo-flash-dev
-Version: 1.2.130
+Version: 1.2.131
 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.130 → freesolo_flash_dev-1.2.131}/flash/core/spec.py

@@ -20,6 +20,7 @@ detail. See their docstrings before changing either.
 from __future__ import annotations
 
 import json
+import math
 import os
 from dataclasses import asdict, dataclass, field, fields
 from typing import Any, Literal
@@ -365,6 +366,9 @@ class TrainSpec:
     hf_repo: str = ""
     # None -> worker's tuned recipe default.
     learning_rate: float | None = field(default=None, metadata={"introduced_in": "0.2.0"})
+    # sft only: none preserves a legacy persisted run's 16x optimizer. new configs freeze the
+    # recipe default at parse time; 1 uses equal adapter learning rates.
+    loraplus_ratio: float | None = field(default=None, metadata={"introduced_in": "1.2.131"})
     # sft only. the packaged-dataset estimate resolves this against the selected row count into
     # examples_per_update (packed) or pins the optimizer batch to 1 (unpacked). grpo/opd have no
     # profile, so they take prompts_per_step instead and reject this key: the two are not the same
@@ -399,6 +403,17 @@ class TrainSpec:
     )
 
     def __post_init__(self) -> None:
+        if self.loraplus_ratio is not None:
+            ratio = self.loraplus_ratio
+            if isinstance(ratio, bool) or not isinstance(ratio, (int, float)):
+                raise TypeError("train.loraplus_ratio must be a number")
+            try:
+                ratio = float(ratio)
+            except OverflowError as exc:
+                raise ValueError("train.loraplus_ratio must be a finite number >= 1") from exc
+            if not math.isfinite(ratio) or ratio < 1:
+                raise ValueError("train.loraplus_ratio must be a finite number >= 1")
+            object.__setattr__(self, "loraplus_ratio", ratio)
         if not self.lora_alpha:
             object.__setattr__(self, "lora_alpha", 2 * self.lora_rank)
         max_steps = parse_max_steps(self.max_steps)
@@ -566,6 +581,8 @@ class JobSpec:
     project: str = ""
 
     def __post_init__(self) -> None:
+        if self.train.loraplus_ratio is not None and self.algorithm != "sft":
+            raise ValueError("train.loraplus_ratio only applies to sft")
         object.__setattr__(self, "seed", parse_seed(self.seed))
         object.__setattr__(self, "model_revision", _model_revision(self.model_revision))
         for field_name in ("model_revision_auto", "model_revision_force_pin"):
@@ -681,6 +698,9 @@ class JobSpec:
         differently, so removing either pop would break existing digests.
         """
         data = asdict(self)
+        # preserve historical preparation digests when the new option was not authored.
+        if data["train"]["loraplus_ratio"] is None:
+            data["train"].pop("loraplus_ratio")
         # missing and unset are the same internal state. omitting the empty value preserves that state
         # without serializing it into an explicit empty preference, which every parser rejects.
         if not data["gpu"].get("providers"):
@@ -786,6 +806,7 @@ class JobSpec:
                 init_from_adapter_revision=str(train.get("init_from_adapter_revision") or ""),
                 hf_repo=str(train.get("hf_repo") or ""),
                 learning_rate=opt_float(train.get("learning_rate")),
+                loraplus_ratio=train.get("loraplus_ratio"),
                 batch_size=opt_int(train.get("batch_size")),
                 prompts_per_step=opt_int(train.get("prompts_per_step")),
                 max_context_tokens=opt_int(train.get("max_context_tokens")),

{freesolo_flash_dev-1.2.130 → freesolo_flash_dev-1.2.131}/flash/engine/plan/recipe.py

@@ -170,6 +170,7 @@ class SFTConfig:
     max_seq_len: int = 1024
     max_seq_len_thinking: int = 2048
     learning_rate: float = 1e-4
+    loraplus_ratio: float = 1.0
     warmup_frac: float = 0.03
     effective_batch: int = 32
     num_epochs: int = 2

{freesolo_flash_dev-1.2.130 → freesolo_flash_dev-1.2.131}/flash/engine/worker/train/entry/sft_train.py

@@ -28,7 +28,6 @@ from flash.engine.worker.verl.checkpoints import restore_verl_resume
 from flash.engine.worker.verl.parallelism import ULYSSES_SEQUENCE_PARALLEL_SIZE
 
 # todo: run the two-gpu sft smoke on the exact runpod image and command assembled below.
-_SFT_LORAPLUS_RATIO = 16.0
 # consecutive zero-grad-norm steps tolerated before the run is failed as untrainable (GRAD-001).
 # any nonzero grad norm is proof the backward graph is intact and resets the count. 2 is enough to
 # separate a one-off fully-masked batch from a severed graph, and keeps the wasted spend to a couple
@@ -488,6 +487,7 @@ def _write_sft_result(options, data, model, child, progress, verified, outputs)
             "device_peak_gpu_gb": device_peak_gpu_gb,
             "loraplus_optim": _VERL_OPTIMIZER_NAME,
             "loraplus_applied": progress.loraplus_applied,
+            "loraplus_ratio": options.loraplus_ratio,
             "verl_backend": "fsdp2",
             # sft shards by DATA -- see ULYSSES_SEQUENCE_PARALLEL_SIZE for why. fsdp splits the batch
             # across the ranks actually LAUNCHED, which is the allocated card count only when the
@@ -585,7 +585,7 @@ def run_sft_train(spec=None) -> None:
     child = _prepare_sft_child(
         options, data, model, capabilities, use_remove_padding, gdn_reset_arch
     )
-    child_progress = _prepare_sft_progress(data, model, child)
+    child_progress = _prepare_sft_progress(options, data, model, child)
     progress = child_progress.values
 
     def on_line(line: str) -> None:

{freesolo_flash_dev-1.2.130 → freesolo_flash_dev-1.2.131}/flash/engine/worker/train/entry/sft_train_runner.py

@@ -33,7 +33,6 @@ from flash.providers.core.base import rentable_gpu_counts
 
 RECIPE = _sft_train.RECIPE
 _MAX_ZERO_GRAD_STEPS = _sft_train._MAX_ZERO_GRAD_STEPS
-_SFT_LORAPLUS_RATIO = _sft_train._SFT_LORAPLUS_RATIO
 _LORAPLUS_READY_MARKER = _sft_train._LORAPLUS_READY_MARKER
 _VERL_OPTIMIZER_IMPL = _sft_train._VERL_OPTIMIZER_IMPL
 _VERL_OPTIMIZER_NAME = _sft_train._VERL_OPTIMIZER_NAME
@@ -70,6 +69,7 @@ class _SftOptions:
     model_revision: str
     epochs: int
     learning_rate: float
+    loraplus_ratio: float
     effective_batch: int
     max_steps: int
     save_at_steps: tuple[int, ...]
@@ -147,6 +147,7 @@ class _SftProgress:
     loss_curve: list[float]
     train_tokens: int
     loraplus_applied: bool
+    loraplus_ratio: float
     wandb_link: dict[str, str | None]
     # what the child plugin must prove applied; verified at the first optimizer step and again at
     # child exit.
@@ -187,6 +188,9 @@ def _resolve_sft_options(spec) -> _SftOptions:
     model_id = spec.model if spec else RECIPE.hf_model_id
     model_revision = getattr(spec, "model_revision", "") if spec else ""
     train_spec = spec.train if spec else None
+    # new authored sft specs persist their ratio. an unset stored value predates the option and
+    # must retain the 16x optimizer groups used by its checkpoint, including across further retries.
+    default_loraplus_ratio = 16.0 if train_spec is not None else RECIPE.sft.loraplus_ratio
 
     def train_opt(name, default):
         value = getattr(train_spec, name, None) if train_spec else None
@@ -215,6 +219,7 @@ def _resolve_sft_options(spec) -> _SftOptions:
         model_revision=model_revision,
         epochs=int(train_opt("epochs", RECIPE.sft.num_epochs)),
         learning_rate=float(train_opt("learning_rate", RECIPE.sft.learning_rate)),
+        loraplus_ratio=float(train_opt("loraplus_ratio", default_loraplus_ratio)),
         effective_batch=int(train_opt("batch_size", RECIPE.sft.effective_batch)),
         max_steps=int(train_opt("max_steps", 0) or 0),
         save_at_steps=tuple(getattr(train_spec, "save_at_steps", ()) or ()),
@@ -588,7 +593,7 @@ def _write_sft_child_shims(
         {
             "marker_file": shim_markers,
             "seed": int(seed),
-            "loraplus_ratio": float(_SFT_LORAPLUS_RATIO),
+            "loraplus_ratio": options.loraplus_ratio,
             "loraplus_ready_marker": _LORAPLUS_READY_MARKER,
             "save_at_steps": list(options.save_at_steps),
             "total_steps": int(model.update_horizon),
@@ -758,7 +763,9 @@ def _prepare_sft_child(
     )
 
 
-def _prepare_sft_progress(data: _SftData, model: _SftModelSetup, child: _SftChild) -> _SftProgress:
+def _prepare_sft_progress(
+    options: _SftOptions, data: _SftData, model: _SftModelSetup, child: _SftChild
+) -> _SftProgress:
     resume_step = child.resume_step
     progress = {"step": resume_step, "loss": None, "grad_norm": None, "lr": None}
     # consecutive steps seen with grad_norm == 0.0 at a nonzero lr. one step can legitimately be
@@ -785,7 +792,8 @@ def _prepare_sft_progress(data: _SftData, model: _SftModelSetup, child: _SftChil
         observed_grad_norms=observed_grad_norms,
         loss_curve=[],
         train_tokens=train_tokens,
-        loraplus_applied=resume_step >= model.update_horizon,
+        loraplus_applied=options.loraplus_ratio > 1 and resume_step >= model.update_horizon,
+        loraplus_ratio=options.loraplus_ratio,
         wandb_link={},
         shim_markers=child.shim_markers,
         expected_shims=child.expected_shims,
@@ -846,7 +854,7 @@ def _consume_sft_marker_line(progress: _SftProgress, line: str) -> bool:
     if progress.expected_shims and not progress.shims_verified:
         verify_applied_shim_markers(progress.shim_markers, progress.expected_shims)
         progress.shims_verified = True
-    if _SFT_LORAPLUS_RATIO > 1 and not progress.loraplus_applied:
+    if progress.loraplus_ratio > 1 and not progress.loraplus_applied:
         raise RuntimeError(
             "verl reached an optimizer step before the required lora+ shim succeeded"
         )
@@ -914,7 +922,7 @@ def _finish_sft_child(
     if progress.expected_shims and not progress.shims_verified:
         verify_applied_shim_markers(progress.shim_markers, progress.expected_shims)
         progress.shims_verified = True
-    if _SFT_LORAPLUS_RATIO > 1 and not progress.loraplus_applied:
+    if progress.loraplus_ratio > 1 and not progress.loraplus_applied:
         raise RuntimeError("required lora+ shim did not emit its success marker")
     return train_wall, device_peak_gpu_gb
 

{freesolo_flash_dev-1.2.130 → freesolo_flash_dev-1.2.131}/flash/schema/__init__.py

@@ -25,6 +25,7 @@ from flash.core.spec import (
     parse_seed,
     require_project_id,
 )
+from flash.engine.plan.recipe import RECIPE
 from flash.providers.core.base import (
     UnsupportedGpuError,
     authored_gpu_ceiling,
@@ -297,6 +298,8 @@ def validate_train_keys_for_algorithm(train_raw: Mapping[str, Any], algorithm: s
     every spec fail to reparse its own output and break resubmit, warm-start and server reparse.
     A null is the absence of an authored value, which is exactly what the user is being asked for.
     """
+    if train_raw.get("loraplus_ratio") is not None and algorithm != "sft":
+        raise ConfigError("train.loraplus_ratio only applies to sft")
     for key, allowed in _ALGORITHM_ONLY_TRAIN_KEYS.items():
         if train_raw.get(key) is None or algorithm in allowed:
             continue
@@ -678,6 +681,10 @@ def spec_from_dict(
     )
     wandb_spec = _validate_wandb_section(raw)
 
+    loraplus_ratio = _train_float(train_raw, "loraplus_ratio", minimum=1.0)
+    if algorithm == "sft" and loraplus_ratio is None:
+        # freeze the new default; an absent persisted value identifies the legacy 16x optimizer.
+        loraplus_ratio = RECIPE.sft.loraplus_ratio
     try:
         train_spec = TrainSpec(
             epochs=_train_int(train_raw, "epochs", minimum=1),
@@ -686,6 +693,7 @@ def spec_from_dict(
             init_from_adapter=init_from_adapter,
             hf_repo="",  # assigned server-side; see submit_job._assign_managed_hf_repo
             learning_rate=_train_float(train_raw, "learning_rate", minimum=0.0, exclusive=True),
+            loraplus_ratio=loraplus_ratio,
             batch_size=_train_int(train_raw, "batch_size", minimum=1),
             prompts_per_step=_train_int(train_raw, "prompts_per_step", minimum=1),
             max_context_tokens=_train_int(train_raw, "max_context_tokens", minimum=1),

{freesolo_flash_dev-1.2.130 → freesolo_flash_dev-1.2.131}/flash/schema/fields.py

@@ -65,7 +65,10 @@ def _train_float(
         return None
     if isinstance(v, bool) or not isinstance(v, (int, float)):
         raise ConfigError(f"train.{key} must be a number")
-    v = float(v)
+    try:
+        v = float(v)
+    except OverflowError as exc:
+        raise ConfigError(f"train.{key} must be a finite number") from exc
     # nan/inf slip past range checks (nan compares false, inf passes any minimum).
     if not math.isfinite(v):
         raise ConfigError(f"train.{key} must be a finite number")

{freesolo_flash_dev-1.2.130 → freesolo_flash_dev-1.2.131}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "freesolo-flash-dev"
-version = "1.2.130"
+version = "1.2.131"
 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).
@@ -196,7 +196,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.130"
+version = "1.2.131"
 
 [tool.pytest.ini_options]
 pythonpath = ["."]

{freesolo_flash_dev-1.2.130 → freesolo_flash_dev-1.2.131}/tests/test_lora_targeting.py

@@ -456,7 +456,7 @@ def test_sft_grpo_and_opd_plugins_install_text_targeting_only_for_text_jobs(tmp_
     shim_dir = tmp_path / "sft-text-shim"
     shim_dir.mkdir()
     _marker, expected, raw_sft = sft_train_runner._write_sft_child_shims(
-        SimpleNamespace(model_id="Qwen/Qwen3.5-9B", save_at_steps=()),
+        SimpleNamespace(model_id="Qwen/Qwen3.5-9B", save_at_steps=(), loraplus_ratio=16.0),
         SimpleNamespace(
             update_horizon=1,
             reentrant_gradient_checkpointing=False,
@@ -474,7 +474,7 @@ def test_sft_grpo_and_opd_plugins_install_text_targeting_only_for_text_jobs(tmp_
     multimodal_dir = tmp_path / "sft-multimodal-shim"
     multimodal_dir.mkdir()
     _marker, expected, raw_sft = sft_train_runner._write_sft_child_shims(
-        SimpleNamespace(model_id="Qwen/Qwen3.5-9B", save_at_steps=()),
+        SimpleNamespace(model_id="Qwen/Qwen3.5-9B", save_at_steps=(), loraplus_ratio=16.0),
         SimpleNamespace(
             update_horizon=1,
             reentrant_gradient_checkpointing=False,

freesolo_flash_dev-1.2.131/tests/test_loraplus_config.py

@@ -0,0 +1,97 @@
+"""user-authored lora+ ratios survive cli and worker spec boundaries."""
+
+import pytest
+
+from flash.core.spec import JobSpec, TrainSpec
+from flash.schema import ConfigError, spec_and_train_keys_from_file, spec_from_dict
+
+
+def _raw(train, algorithm="sft"):
+    return {
+        "model": "Qwen/Qwen3.5-9B",
+        "algorithm": algorithm,
+        "project": "11111111-1111-4111-8111-111111111111",
+        "environment": {"id": "owner/project/env"},
+        "train": train,
+    }
+
+
+@pytest.mark.parametrize("ratio", [1, 3.5, 16, 32])
+def test_loraplus_ratio_survives_public_and_worker_round_trips(ratio):
+    spec = spec_from_dict(_raw({"loraplus_ratio": ratio, "learning_rate": 1e-5}))
+    public = spec_from_dict(spec.to_dict())
+    worker = JobSpec.from_json(spec.to_json())
+    for result in (spec, public, worker):
+        assert result.train.loraplus_ratio == ratio
+        assert result.train.learning_rate == 1e-5
+
+
+@pytest.mark.parametrize("train", [{}, {"loraplus_ratio": None}])
+def test_new_unset_ratio_is_frozen_for_future_retries(train):
+    spec = spec_from_dict(_raw(train))
+    assert spec.train.loraplus_ratio == 1.0
+    assert spec.to_dict()["train"]["loraplus_ratio"] == 1.0
+    assert JobSpec.from_json(spec.to_json()).train.loraplus_ratio == 1.0
+
+
+def test_legacy_unset_ratio_preserves_historical_payload_shape():
+    spec = JobSpec.from_dict(_raw({}))
+    assert spec.train.loraplus_ratio is None
+    assert "loraplus_ratio" not in spec.to_dict()["train"]
+    assert "loraplus_ratio" not in spec.to_internal_dict()["train"]
+    assert JobSpec.from_json(spec.to_json()).train.loraplus_ratio is None
+
+
+@pytest.mark.parametrize(
+    "ratio",
+    [
+        True,
+        False,
+        "16",
+        [],
+        0,
+        -1,
+        0.5,
+        float("nan"),
+        float("inf"),
+        -float("inf"),
+        10**400,
+        -(10**400),
+    ],
+)
+def test_invalid_ratio_is_rejected_at_public_and_internal_boundaries(ratio):
+    with pytest.raises(ConfigError, match="loraplus_ratio"):
+        spec_from_dict(_raw({"loraplus_ratio": ratio}))
+    with pytest.raises((TypeError, ValueError), match="loraplus_ratio"):
+        JobSpec.from_dict(_raw({"loraplus_ratio": ratio}))
+    with pytest.raises((TypeError, ValueError), match="loraplus_ratio"):
+        TrainSpec(loraplus_ratio=ratio)
+
+
+@pytest.mark.parametrize("algorithm", ["grpo", "opd"])
+def test_ratio_is_rejected_for_algorithms_without_loraplus(algorithm):
+    raw = _raw({"loraplus_ratio": 4}, algorithm)
+    with pytest.raises(ConfigError, match="only applies to sft"):
+        spec_from_dict(raw)
+    with pytest.raises(ValueError, match="only applies to sft"):
+        JobSpec.from_dict(raw)
+
+
+def test_toml_ratio_and_cli_override_are_user_authored_keys(tmp_path):
+    config = tmp_path / "flash.toml"
+    config.write_text(
+        'model = "Qwen/Qwen3.5-9B"\n'
+        'algorithm = "sft"\n'
+        'project = "11111111-1111-4111-8111-111111111111"\n'
+        '[environment]\nid = "owner/project/env"\n'
+        "[train]\nloraplus_ratio = 4.0\nlearning_rate = 0.00001\n"
+    )
+    spec, keys = spec_and_train_keys_from_file(str(config), run_id="test")
+    assert spec.train.loraplus_ratio == 4.0
+    assert "loraplus_ratio" in keys
+    spec, keys = spec_and_train_keys_from_file(
+        str(config), run_id="test", overrides=["train.loraplus_ratio=1.0"]
+    )
+    assert spec.train.loraplus_ratio == 1.0
+    assert spec.train.learning_rate == 1e-5
+    assert "loraplus_ratio" in keys

{freesolo_flash_dev-1.2.130 → freesolo_flash_dev-1.2.131}/tests/test_server_api.py

@@ -809,6 +809,25 @@ def test_dry_run_rejects_invalid_structured_serving_constraints(
     assert message in response.json()["detail"]
 
 
+@pytest.mark.parametrize("ratio", [10**400, -(10**400)])
+def test_dry_run_rejects_overflowing_loraplus_ratio(api, monkeypatch, ratio) -> None:
+    import flash.server.routes.runs as runs_route
+
+    monkeypatch.setattr(
+        runs_route._app,
+        "prepare_job",
+        lambda *_a, **_k: pytest.fail("invalid ratio must fail before preparation"),
+    )
+    spec = {**SPEC, "algorithm": "sft", "train": {"loraplus_ratio": ratio}}
+    response = api.post(
+        "/v1/runs",
+        headers=_bearer(_login()),
+        json={"spec": spec, "dry_run": True},
+    )
+    assert response.status_code == 400, response.text
+    assert "train.loraplus_ratio must be a finite number" in response.json()["detail"]
+
+
 def test_warmstart_dry_run_preserves_serving_preflight_error(api) -> None:
     spec = {
         **SPEC,

{freesolo_flash_dev-1.2.130 → freesolo_flash_dev-1.2.131}/tests/test_sft_train.py

@@ -1750,8 +1750,9 @@ def test_seeded_dataloader_handles_missing_and_present_validation_sampler(
     assert exc_info.value is original_error
 
 
+@pytest.mark.parametrize("ratio", [1.0, 3.5, 16.0])
 def test_generated_sitecustomize_installs_linear_scheduler_and_required_loraplus(
-    monkeypatch, capsys
+    monkeypatch, capsys, ratio
 ):
     class FakeLoader:
         def __init__(self, *args, **kwargs):
@@ -1837,7 +1838,8 @@ def test_generated_sitecustomize_installs_linear_scheduler_and_required_loraplus
 
     sft_plugin._install_seeded_dataloader(43)
     sft_plugin._install_linear_scheduler()
-    sft_plugin._install_loraplus(16, "CUSTOM_LORAPLUS_READY")
+    original_builder = FakeEngine._build_optimizer
+    sft_plugin._install_loraplus(ratio, "CUSTOM_LORAPLUS_READY")
 
     engine = FakeEngine()
     engine.optimizer_config = SimpleNamespace(
@@ -1851,11 +1853,17 @@ def test_generated_sitecustomize_installs_linear_scheduler_and_required_loraplus
         lr_warmup_steps_ratio=0.1,
         total_training_steps=20,
     )
-    assert engine._build_optimizer(SimpleNamespace()) == "lora+"
+    assert engine._build_optimizer(SimpleNamespace()) == ("lora+" if ratio > 1 else "plain")
     output = capsys.readouterr().out
-    assert "CUSTOM_LORAPLUS_READY ratio=16 optimizer=AdamW" in output
     assert _LORAPLUS_READY_MARKER not in output
-    assert optimizer_calls[0]["optimizer_kwargs"]["eps"] == 1e-8
+    if ratio > 1:
+        assert f"CUSTOM_LORAPLUS_READY ratio={ratio:g} optimizer=AdamW" in output
+        assert optimizer_calls[0]["optimizer_kwargs"]["eps"] == 1e-8
+        assert optimizer_calls[0]["loraplus_lr_ratio"] == ratio
+    else:
+        assert FakeEngine._build_optimizer is original_builder
+        assert optimizer_calls == []
+        assert "CUSTOM_LORAPLUS_READY" not in output
     assert engine._build_lr_scheduler("optimizer") == "linear"
     assert scheduler_calls == [("optimizer", {"num_warmup_steps": 2, "num_training_steps": 20})]
 
@@ -1866,7 +1874,7 @@ def test_sft_plugin_config_carries_the_canonical_loraplus_marker(tmp_path):
     shim_dir = tmp_path / "shim"
     shim_dir.mkdir()
     _, _, raw_config = sft_train_runner._write_sft_child_shims(
-        SimpleNamespace(save_at_steps=(3,)),
+        SimpleNamespace(save_at_steps=(3,), loraplus_ratio=16.0),
         SimpleNamespace(update_horizon=7, reentrant_gradient_checkpointing=False),
         shim_dir=str(shim_dir),
         custom_dataset_path=str(shim_dir / "dataset.py"),
@@ -1879,6 +1887,34 @@ def test_sft_plugin_config_carries_the_canonical_loraplus_marker(tmp_path):
     assert json.loads(raw_config)["loraplus_ready_marker"] == _LORAPLUS_READY_MARKER
 
 
+@pytest.mark.parametrize("ratio", [1.0, 3.5, 16.0])
+@pytest.mark.parametrize("gate", ["step", "exit"])
+def test_loraplus_marker_is_required_only_for_the_runs_selected_ratio(ratio, gate):
+    from flash.engine.worker.train.entry import sft_train_runner
+
+    progress = sft_train_runner._prepare_sft_progress(
+        SimpleNamespace(loraplus_ratio=ratio),
+        SimpleNamespace(rows=[]),
+        SimpleNamespace(update_horizon=1),
+        SimpleNamespace(resume_step=0, shim_markers="", expected_shims=()),
+    )
+
+    def check():
+        if gate == "step":
+            return sft_train_runner._consume_sft_marker_line(progress, "step:1 - train/loss:1.0")
+        return sft_train_runner._finish_sft_child(
+            SimpleNamespace(stop_gb=lambda: 0.0), 0.0, 0, progress
+        )
+
+    if ratio > 1:
+        with pytest.raises(RuntimeError, match="required lora\\+ shim"):
+            check()
+        sft_train_runner._consume_sft_marker_line(
+            progress, f"{_LORAPLUS_READY_MARKER} ratio={ratio:g} optimizer=AdamW"
+        )
+    check()
+
+
 def _exec_dataloader_shim(monkeypatch):
     """Execute the real rendered dataloader shim against fake torch classes, and return them.
 
@@ -3036,18 +3072,80 @@ def test_sft_collapse_warning_stays_quiet_for_structured_singleturn_targets(monk
     assert "bare assistant target coerced" not in output
 
 
-def test_run_sft_train_orchestrates_exact_dataset_and_resume_accounting(monkeypatch):
+@pytest.mark.parametrize("legacy", [True, False])
+def test_sft_retries_preserve_the_prepared_optimizer_ratio(monkeypatch, legacy):
+    from flash.core.spec import JobSpec
+    from flash.engine.worker.train.entry import sft_train_runner
+    from flash.runner.lifecycle.state import RunStatus
+    from flash.runner.lifecycle.status import reallocation_spec_from_status
+    from flash.runner.lifecycle.submit import _effective_preparation_snapshot
+    from flash.schema import spec_from_dict
+
+    _stub_sft_run(monkeypatch)
+    raw = {
+        "model": "Qwen/Qwen3.5-9B",
+        "algorithm": "sft",
+        "project": "11111111-1111-4111-8111-111111111111",
+        "environment": {"id": "owner/project/env"},
+        "train": {"epochs": 1, "learning_rate": 1e-5},
+        "gpu": {"type": "H100", "count": 1},
+    }
+    spec = JobSpec.from_dict(raw) if legacy else spec_from_dict(raw)
+    spec = replace(spec, run_id="legacy-sft", model_revision="a" * 40, model_revision_auto=True)
+    snapshot = _effective_preparation_snapshot(spec, spec, None)
+    if legacy:
+        # frozen before the recovery fix: old snapshots must still pass integrity validation.
+        assert snapshot["preparation_digest"] == (
+            "d190288e79077c89f6b79353947a393e2a0463c141757d844bd889511f91022b"
+        )
+        assert "loraplus_ratio" not in snapshot["worker_spec"]["train"]
+    original_digest = snapshot["preparation_digest"]
+    status = RunStatus(
+        state="running",
+        run_id=spec.run_id,
+        spec=spec.to_dict(),
+        effective_preparation=json.loads(json.dumps(snapshot)),
+    )
+    for _attempt in range(2):
+        recovered = reallocation_spec_from_status(status)
+        worker = JobSpec.from_json(recovered.to_json())
+        options = sft_train_runner._resolve_sft_options(worker)
+        assert options.loraplus_ratio == (16.0 if legacy else 1.0)
+        status.effective_preparation = _effective_preparation_snapshot(
+            JobSpec.from_dict(status.spec), worker, None, stored_public=status.spec
+        )
+        assert status.effective_preparation["preparation_digest"] == original_digest
+
+
+@pytest.mark.parametrize("ratio", [None, 1.0, 3.5, 16.0])
+def test_run_sft_train_orchestrates_exact_dataset_and_resume_accounting(monkeypatch, ratio):
     from flash._internal.diagnostics import SECRET_ENV_KEYS_ENV
+    from flash.core.spec import JobSpec
     from flash.engine.worker.train.entry import sft_train
+    from flash.schema import spec_from_dict
 
     monkeypatch.setenv("PYTHONPATH", "synthetic-sft-parent-path")
     monkeypatch.setenv(SECRET_ENV_KEYS_ENV, "PYTHONPATH")
     spec, captured = _stub_sft_run(monkeypatch)
+    authored = spec_from_dict(
+        {
+            "model": spec.model,
+            "algorithm": "sft",
+            "project": "11111111-1111-4111-8111-111111111111",
+            "environment": {"id": "owner/project/env"},
+            "train": {} if ratio is None else {"loraplus_ratio": ratio},
+        }
+    )
+    spec.train.loraplus_ratio = JobSpec.from_json(authored.to_json()).train.loraplus_ratio
+    effective_ratio = 1.0 if ratio is None else ratio
 
     def fake_training(command, *, env, on_step, on_line, heartbeat):
         captured["command"] = command
         captured["child_env"] = env
-        on_line(f"{_LORAPLUS_READY_MARKER} ratio=16 optimizer=AdamW\n")
+        config = json.loads(env["FLASH_SFT_PLUGIN_CONFIG"])
+        assert config["loraplus_ratio"] == effective_ratio
+        if effective_ratio > 1:
+            on_line(f"{_LORAPLUS_READY_MARKER} ratio={effective_ratio:g} optimizer=AdamW\n")
         on_line("step:2 - train/loss:1.0 - train/global_tokens:8\n")
         on_step(2)
         heartbeat()
@@ -3080,7 +3178,8 @@ def test_run_sft_train_orchestrates_exact_dataset_and_resume_accounting(monkeypa
     assert captured["meta"]["step"] == 2
     assert captured["meta"]["train_tokens"] > 0
     assert captured["meta"]["notes"]["loss_curve"] == [1.0]
-    assert captured["meta"]["notes"]["loraplus_applied"] is True
+    assert captured["meta"]["notes"]["loraplus_applied"] is (effective_ratio > 1)
+    assert captured["meta"]["notes"]["loraplus_ratio"] == effective_ratio
     notes = captured["meta"]["notes"]
     realized_max_length = notes["realized_max_length"]
     assert 0 < realized_max_length < notes["configured_max_length"]
@@ -3287,7 +3386,8 @@ def test_the_sft_runner_seeds_the_watcher_with_the_step_it_resumed_from(monkeypa
     assert captured["published"][0][1] == 2
 
 
-def test_a_resume_at_the_horizon_still_publishes_the_final_deployable(monkeypatch):
+@pytest.mark.parametrize("ratio", [1.0, 3.5, 16.0])
+def test_a_resume_at_the_horizon_still_publishes_the_final_deployable(monkeypatch, ratio):
     """the seeded resume step must not suppress the final publish.
 
     the previous attempt's per-step deployable publish is best-effort (``required=False``) while
@@ -3296,6 +3396,7 @@ def test_a_resume_at_the_horizon_still_publishes_the_final_deployable(monkeypatc
     from flash.engine.worker.train.entry import sft_train
 
     spec, captured = _stub_sft_run(monkeypatch)
+    spec.train.loraplus_ratio = ratio
     # max_steps is 2, so resuming at 2 means the watcher never runs and finalization is the only
     # path left that can publish the step.
     monkeypatch.setattr(
@@ -3312,6 +3413,8 @@ def test_a_resume_at_the_horizon_still_publishes_the_final_deployable(monkeypatc
     sft_train.run_sft_train(spec)
 
     assert [step for _adapter, step in captured["published"]] == [2]
+    assert captured["meta"]["notes"]["loraplus_ratio"] == ratio
+    assert captured["meta"]["notes"]["loraplus_applied"] is (ratio > 1)
 
 
 def _sft_model_save_freq(monkeypatch, *, save_at_steps, save_every, horizon):
@@ -3332,6 +3435,7 @@ def _sft_model_save_freq(monkeypatch, *, save_at_steps, save_every, horizon):
         model_revision="revision",
         epochs=1,
         learning_rate=5e-5,
+        loraplus_ratio=16.0,
         effective_batch=64,
         max_steps=0,
         save_at_steps=save_at_steps,

{freesolo_flash_dev-1.2.130 → freesolo_flash_dev-1.2.131}/tests/test_spec_and_validation.py

@@ -193,6 +193,7 @@ def test_train_key_registry_is_derived_from_trainspec_metadata() -> None:
             "credit_assignment",
             "entropy_quantile",
             "lora_alpha",
+            "loraplus_ratio",
             "prompts_per_step",
         }
     } == {"0.2.0"}
@@ -340,6 +341,7 @@ def test_historical_train_schema_shapes_are_immutable_source_snapshots() -> None
         "save_at_steps",
         "entropy_quantile",
         "prompts_per_step",
+        "loraplus_ratio",
     }
     assert "opd_eos_loss_coef" not in TRAIN_SCHEMA_KEYS
     assert "advantage_clip" not in TRAIN_SCHEMA_KEYS

{freesolo_flash_dev-1.2.130 → freesolo_flash_dev-1.2.131}/uv.lock

@@ -1480,7 +1480,7 @@ wheels = [
 
 [[package]]
 name = "freesolo-flash"
-version = "1.2.130"
+version = "1.2.131"
 source = { editable = "." }
 
 [package.optional-dependencies]