freesolo-flash-dev 1.2.37__tar.gz → 1.2.43__tar.gz

{freesolo_flash_dev-1.2.37 → freesolo_flash_dev-1.2.43}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: freesolo-flash-dev
-Version: 1.2.37
+Version: 1.2.43
 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.43/flash/adapters/fused_experts.py

@@ -0,0 +1,375 @@
+"""Fused-expert LoRA targeting, validation, and export normalization."""
+
+from __future__ import annotations
+
+import re
+from collections import Counter
+from collections.abc import Mapping
+from typing import Any
+
+from flash.adapters.lora_rank import (
+    _rank_for_module,
+    declared_lora_ranks,
+    lora_tensor_rank_disagrees,
+)
+from flash.core.catalog import get_model
+
+_QWEN36_MODEL_ID = "Qwen/Qwen3.6-35B-A3B"
+_QWEN36_EXPERT_TARGET_PARAMETERS = (
+    "mlp.experts.gate_up_proj",
+    "mlp.experts.down_proj",
+)
+_QWEN36_FUSED_TARGET_DIMENSIONS = {
+    "mlp.experts.gate_up_proj": (2048, 1024),
+    "mlp.experts.down_proj": (512, 2048),
+}
+_QWEN36_FUSED_TARGET_RUNGS = {
+    "mlp.experts.down_proj": "",
+    "mlp.experts.gate_up_proj": "base_layer",
+}
+_FUSED_EXPERT_SYNTHETIC_MODULES = frozenset({"experts", "base_layer"})
+_FUSED_EXPERT_WRAPPER_MODULES = ("mlp.experts", "mlp.experts.base_layer")
+_INVALID_FUSED_LOCATION = object()
+
+_LoraTensor = tuple[str, str, str, str, tuple[int, ...]]
+_LoraPair = tuple[tuple[int, ...], tuple[int, ...]]
+
+
+def lora_target_parameters(model_id: str | None) -> list[str] | None:
+    """Return direct parameter targets required by the model's fused expert layout."""
+    if model_id == _QWEN36_MODEL_ID:
+        return list(_QWEN36_EXPERT_TARGET_PARAMETERS)
+    return None
+
+
+def validate_fused_expert_adapter_config(config: Mapping[str, Any], model_id: str) -> None:
+    """Validate the exact current adapter config required by a fused-expert model."""
+    required = lora_target_parameters(model_id)
+    if not required:
+        return
+
+    targets = config.get("target_parameters")
+    if not isinstance(targets, list):
+        raise ValueError(
+            f"adapter for {model_id} omits required expert targets; "
+            "target_parameters must be a list of strings"
+        )
+    if any(not isinstance(target, str) for target in targets):
+        raise ValueError(
+            f"adapter for {model_id} must declare target_parameters as a list of strings"
+        )
+    if len(targets) != len(set(targets)):
+        raise ValueError(f"adapter for {model_id} must declare unique target_parameters")
+    if set(targets) != set(required):
+        raise ValueError(
+            f"adapter for {model_id} must declare exactly the fused expert targets {required}"
+        )
+
+    rank_pattern = config.get("rank_pattern")
+    _validate_rank_pattern(rank_pattern, model_id)
+    declared = declared_lora_ranks(config)
+    unresolved = [target for target in required if _rank_for_module(target, declared) is None]
+    if unresolved:
+        raise ValueError(
+            f"adapter for {model_id} has no resolved LoRA rank for fused targets {unresolved}"
+        )
+
+    modules = config.get("target_modules")
+    if isinstance(modules, str):
+        if modules != "all-linear":
+            raise ValueError(f"adapter for {model_id} string target_modules must be 'all-linear'")
+        # all-linear names no modules, so only a default rank can resolve every ordinary target.
+        if declared.default is None:
+            raise ValueError(
+                f"adapter for {model_id} has no resolved default LoRA rank for 'all-linear' "
+                "target_modules"
+            )
+        return
+    if (
+        not isinstance(modules, list)
+        or not modules
+        or any(not isinstance(module, str) or not module for module in modules)
+    ):
+        raise ValueError(
+            f"adapter for {model_id} target_modules must be 'all-linear' or a non-empty "
+            "list of non-empty strings"
+        )
+    synthetic = [module for module in modules if _targets_fused_expert_wrapper(module)]
+    if synthetic:
+        raise ValueError(
+            f"adapter for {model_id} contains invalid synthetic target_modules {sorted(synthetic)}"
+        )
+    unresolved = [module for module in modules if _rank_for_module(module, declared) is None]
+    if unresolved:
+        raise ValueError(
+            f"adapter for {model_id} has no resolved LoRA rank for ordinary target_modules "
+            f"{unresolved}"
+        )
+
+
+def _validate_rank_pattern(rank_pattern: Any, model_id: str) -> None:
+    """Validate positive ranks and PEFT's anchored regex pattern syntax."""
+    message = (
+        f"adapter for {model_id} rank_pattern must map valid non-empty string patterns "
+        "to positive integer ranks"
+    )
+    if rank_pattern is None:
+        return
+    if not isinstance(rank_pattern, Mapping):
+        raise ValueError(message)
+    for pattern, rank in rank_pattern.items():
+        if (
+            not isinstance(pattern, str)
+            or not pattern.strip()
+            or not isinstance(rank, int)
+            or isinstance(rank, bool)
+            or rank <= 0
+        ):
+            raise ValueError(message)
+        try:
+            re.compile(rf"(.*\.)?({pattern})$")
+        except re.error as exc:
+            raise ValueError(message) from exc
+
+
+def _targets_fused_expert_wrapper(module: str) -> bool:
+    """Return whether PEFT suffix matching binds this target to an expert wrapper."""
+    return any(
+        module == wrapper or module.endswith(f".{wrapper}") or wrapper.endswith(f".{module}")
+        for wrapper in _FUSED_EXPERT_WRAPPER_MODULES
+    )
+
+
+def normalize_verl_fused_expert_export(config: dict[str, Any], model_id: str) -> None:
+    """Write the canonical fused targets into a verl-exported adapter config."""
+    targets = lora_target_parameters(model_id)
+    if not targets:
+        return
+    config["target_parameters"] = targets
+    modules = config.get("target_modules")
+    if modules is None:
+        raise ValueError(
+            f"exported adapter for {model_id} must retain ordinary LoRA target_modules"
+        )
+    if isinstance(modules, list):
+        config["target_modules"] = [
+            module
+            for module in modules
+            if not isinstance(module, str) or module not in _FUSED_EXPERT_SYNTHETIC_MODULES
+        ]
+
+
+def has_complete_fused_expert_tensors(
+    tensors: Mapping[str, tuple[int, ...]], config: Mapping[str, Any], model_id: str
+) -> bool:
+    """Return whether fused and ordinary LoRA tensors match the declared PEFT topology."""
+    expected = _expected_fused_expert_rungs(config, model_id)
+    if expected is None:
+        return False
+    parsed = [
+        tensor
+        for key, shape in tensors.items()
+        if (tensor := _parse_lora_tensor(key, shape)) is not None
+    ]
+    if len({adapter_name for _, _, adapter_name, _, _ in parsed}) != 1:
+        return False
+    return _has_complete_fused_rungs(parsed, expected, model_id) and _has_ordinary_evidence(
+        parsed, config, expected
+    )
+
+
+def _parse_lora_tensor(key: str, shape: tuple[int, ...]) -> _LoraTensor | None:
+    """Parse one canonical PEFT LoRA tensor key."""
+    matches = [
+        (factor, infix) for factor, infix in (("A", ".lora_A."), ("B", ".lora_B.")) if infix in key
+    ]
+    if len(matches) != 1:
+        return None
+    factor, infix = matches[0]
+    module_path, _, leaf = key.partition(infix)
+    leaf_parts = leaf.split(".")
+    if not module_path or len(leaf_parts) != 2:
+        return None
+    adapter_name, parameter = leaf_parts
+    if not adapter_name or parameter != "weight":
+        return None
+    return module_path, factor, adapter_name, key, shape
+
+
+def _expected_fused_expert_rungs(
+    config: Mapping[str, Any], model_id: str
+) -> dict[str, dict[str, _LoraPair]] | None:
+    """Return the exact target-specific PEFT rung geometry for each fused owner."""
+    targets = lora_target_parameters(model_id)
+    if model_id != _QWEN36_MODEL_ID or not targets:
+        return None
+    model = get_model(model_id)
+    experts = model.lora_expert_count
+    if experts <= 0 or model.num_layers <= 0:
+        return None
+    fused_count = model.num_layers * experts
+    catalog_dimensions = Counter(
+        (input_dim, output_dim)
+        for input_dim, output_dim, count in model.lora_target_shapes
+        if count == fused_count
+    )
+    if catalog_dimensions != Counter(_QWEN36_FUSED_TARGET_DIMENSIONS.values()):
+        return None
+
+    declared = declared_lora_ranks(config)
+    expected: dict[str, dict[str, _LoraPair]] = {}
+    for target in targets:
+        rank = _rank_for_module(target, declared)
+        dimensions = _QWEN36_FUSED_TARGET_DIMENSIONS.get(target)
+        rung = _QWEN36_FUSED_TARGET_RUNGS.get(target)
+        owner, separator, _ = target.rpartition(".")
+        if rank is None or dimensions is None or rung is None or not separator or not owner:
+            return None
+        input_dim, output_dim = dimensions
+        stacked_rank = rank * experts
+        expected.setdefault(owner, {})[rung] = (
+            (stacked_rank, input_dim),
+            (output_dim, stacked_rank),
+        )
+    return expected
+
+
+def _has_complete_fused_rungs(
+    tensors: list[_LoraTensor], expected: Mapping[str, Mapping[str, _LoraPair]], model_id: str
+) -> bool:
+    """Validate every concrete fused owner, rung, namespace, and factor shape."""
+    model = get_model(model_id)
+    for owner, expected_rungs in expected.items():
+        factors: dict[str, dict[str, dict[str, dict[str, tuple[int, ...]]]]] = {}
+        for module_path, factor, adapter_name, _key, shape in tensors:
+            location = _fused_tensor_location(module_path, owner, frozenset(expected_rungs))
+            if location is _INVALID_FUSED_LOCATION:
+                return False
+            if location is None:
+                continue
+            instance, rung = location
+            factors.setdefault(instance, {}).setdefault(rung, {}).setdefault(adapter_name, {})[
+                factor
+            ] = shape
+        if not factors:
+            return False
+        for seen_rungs in factors.values():
+            if set(seen_rungs) != set(expected_rungs):
+                return False
+            for rung, expected_pair in expected_rungs.items():
+                namespaces = seen_rungs[rung]
+                if len(namespaces) != 1:
+                    return False
+                seen_factors = next(iter(namespaces.values()))
+                if _lora_pair_shapes(seen_factors) != expected_pair:
+                    return False
+        expert_layers = {
+            prefix
+            for instance in factors
+            for prefix in [_layer_prefix(instance)]
+            if prefix is not None
+        }
+        layer_roots = {prefix.rsplit(".", 1)[0] for prefix in expert_layers}
+        if len(layer_roots) != 1:
+            return False
+        layer_root = next(iter(layer_roots))
+        expected_layers = {f"{layer_root}.{index}" for index in range(model.num_layers)}
+        if expert_layers != expected_layers:
+            return False
+    return True
+
+
+def _fused_tensor_location(
+    module_path: str, owner: str, rungs: frozenset[str]
+) -> tuple[str, str] | object | None:
+    """Locate one tensor on an exact concrete fused owner rung."""
+    if module_path.endswith(f".{owner}"):
+        instance, rung = module_path[: -len(owner) - 1], ""
+    else:
+        marker = f".{owner}."
+        if marker not in module_path:
+            return None
+        instance, rung = module_path.rsplit(marker, 1)
+    if rung not in rungs:
+        return _INVALID_FUSED_LOCATION
+    layer_prefix = _layer_prefix(instance)
+    if layer_prefix is None or instance != layer_prefix:
+        return _INVALID_FUSED_LOCATION
+    return f"{instance}.{owner}", rung
+
+
+def _has_ordinary_evidence(
+    tensors: list[_LoraTensor],
+    config: Mapping[str, Any],
+    fused_rungs: Mapping[str, Mapping[str, _LoraPair]],
+) -> bool:
+    """Validate concrete ordinary target evidence with per-module ranks and namespaces."""
+    modules = config.get("target_modules")
+    all_linear = modules == "all-linear"
+    targets = () if all_linear else tuple(modules) if isinstance(modules, list) else ()
+    groups: dict[str, dict[str, dict[str, tuple[str, tuple[int, ...]]]]] = {}
+    evidence: set[str] = set()
+    for module_path, factor, adapter_name, key, shape in tensors:
+        if _is_fused_rung(module_path, fused_rungs):
+            continue
+        matched = tuple(target for target in targets if _anchored_suffix_match(module_path, target))
+        if not all_linear and not matched:
+            continue
+        evidence.update(matched)
+        groups.setdefault(module_path, {}).setdefault(adapter_name, {})[factor] = (key, shape)
+    if not groups or (not all_linear and evidence != set(targets)):
+        return False
+
+    declared = declared_lora_ranks(config)
+    for module_path, namespaces in groups.items():
+        if len(namespaces) != 1 or _rank_for_module(module_path, declared) is None:
+            return False
+        factors = next(iter(namespaces.values()))
+        if set(factors) != {"A", "B"}:
+            return False
+        for key, shape in factors.values():
+            if not _is_positive_2d(shape) or lora_tensor_rank_disagrees(key, shape, declared):
+                return False
+    return True
+
+
+def _is_fused_rung(module_path: str, fused_rungs: Mapping[str, Mapping[str, _LoraPair]]) -> bool:
+    """Return whether a module path is one of the exact fused wrapper rungs."""
+    return any(
+        _fused_tensor_location(module_path, owner, frozenset(rungs))
+        not in (None, _INVALID_FUSED_LOCATION)
+        for owner, rungs in fused_rungs.items()
+    )
+
+
+def _anchored_suffix_match(module_path: str, target: str) -> bool:
+    """Mirror PEFT list-target matching without accepting partial segment suffixes."""
+    return module_path == target or module_path.endswith(f".{target}")
+
+
+def _is_positive_2d(shape: tuple[int, ...]) -> bool:
+    """Return whether a tensor shape is positive, integral, and two-dimensional."""
+    return len(shape) == 2 and all(
+        isinstance(dimension, int) and not isinstance(dimension, bool) and dimension > 0
+        for dimension in shape
+    )
+
+
+def _lora_pair_shapes(factors: Mapping[str, tuple[int, ...]]) -> _LoraPair | None:
+    """Return one complete positive 2-D A/B shape pair."""
+    if set(factors) != {"A", "B"}:
+        return None
+    shape_a = factors["A"]
+    shape_b = factors["B"]
+    if not _is_positive_2d(shape_a) or not _is_positive_2d(shape_b):
+        return None
+    return shape_a, shape_b
+
+
+def _layer_prefix(path: str) -> str | None:
+    """Return the concrete ``...layers.N`` prefix in a tensor path, if present."""
+    segments = path.split(".")
+    for index in range(len(segments) - 2, -1, -1):
+        if segments[index] == "layers" and segments[index + 1].isdigit():
+            return ".".join(segments[: index + 2])
+    return None

{freesolo_flash_dev-1.2.37 → freesolo_flash_dev-1.2.43}/flash/core/catalog.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from dataclasses import asdict, dataclass, replace
+from math import ceil
 from typing import Any
 
 ALGORITHMS = ("sft", "grpo", "opd")
@@ -622,6 +623,10 @@ def resolve_vocab_size(model_id: str, revision: str = "") -> int:
     return vocab_size_for(model_id)
 
 
+def _with_merge_disk_floor(info: ModelInfo) -> ModelInfo:
+    return replace(info, min_disk_gb=max(info.min_disk_gb, ceil(info.params_b * 2) + 64))
+
+
 def resolve_model(model_id: str, algorithm: str, model_revision: str = "") -> ModelInfo:
     """Resolve a curated model, validated for ``algorithm``; anything uncataloged is rejected.
 
@@ -643,9 +648,8 @@ def resolve_model(model_id: str, algorithm: str, model_revision: str = "") -> Mo
             params_b=params_b,
             params=f"{params_b:.1f}B",
             vocab_size=vocab_size,
-            min_disk_gb=max(info.min_disk_gb, int(params_b * 2) + 64),
         )
-    return info
+    return _with_merge_disk_floor(info)
 
 
 def validate_model_for_algorithm(model_id: str, algorithm: str) -> ModelInfo:
@@ -660,4 +664,4 @@ def validate_model_for_algorithm(model_id: str, algorithm: str) -> ModelInfo:
 
 
 def public_model_rows() -> list[dict[str, Any]]:
-    return [m.to_dict() for m in list_models()]
+    return [_with_merge_disk_floor(model).to_dict() for model in list_models()]

{freesolo_flash_dev-1.2.37 → freesolo_flash_dev-1.2.43}/flash/engine/worker/__init__.py

@@ -17,6 +17,7 @@ import time
 import traceback
 
 from flash._internal.diagnostics import sanitize_diagnostic
+from flash.adapters.fused_experts import lora_target_parameters
 from flash.core.spec import FIXED_SEED, load_job_spec_from_env
 from flash.engine.result.accounting import RunMetrics
 from flash.engine.worker.backend_common import collect_ray_failure_logs
@@ -56,9 +57,8 @@ from flash.engine.worker.io.wandb_log import (
 )
 from flash.engine.worker.model.adapter import (
     _download_adapter,
-    lora_target_parameters,
     make_lora,
-    validate_lora_target_parameters,
+    validate_warmstart_adapter,
 )
 from flash.engine.worker.model.decoding import (
     graded_text,
@@ -503,7 +503,7 @@ __all__ = [
     "think_token_count",
     "thinking_text",
     "upload_resume_checkpoint",
-    "validate_lora_target_parameters",
+    "validate_warmstart_adapter",
     "wait_for_gpu",
     "wandb_run_name",
     "write_base_model_provenance",

{freesolo_flash_dev-1.2.37 → freesolo_flash_dev-1.2.43}/flash/engine/worker/backend_common.py

@@ -471,29 +471,14 @@ def ray_num_cpus(gpu_count: int = 1, *, cap: int = 16) -> int:
     return max(floor, min(affinity, cap))
 
 
-def run_verl_training(
+def _run_streaming_verl_subprocess(
     cmd: list[str],
     *,
     env: dict[str, str],
-    on_step: Callable[[int], None] | None = None,
-    on_line: Callable[[str], None] | None = None,
-    heartbeat: Callable[[], None] | None = None,
-    step_pattern: str = r"step:\s*(\d+)",
-    heartbeat_interval_s: float = 20.0,
-    tail: ChildOutputTail | None = None,
+    on_line: Callable[[str], None],
+    errors: str | None = None,
 ) -> int:
-    """run a verl trainer subprocess, streaming stdout and surfacing step progress.
-
-    returns the process exit code. stdout+stderr are merged and scanned line by line: ``on_line``
-    receives every line, ``on_step`` receives each parsed training step, and ``heartbeat`` is called
-    at most once per ``heartbeat_interval_s``. callback failures terminate the child before they are
-    re-raised so a failed required checkpoint upload cannot leave paid training running unattended.
-    """
-    step_re = re.compile(step_pattern)
-    child_tail = tail if tail is not None else ChildOutputTail()
-    # before the child exists, so any grandchild it orphans reparents here and can actually be
-    # reaped. this process is not pid 1 -- the runpod handler is -- so without it every wait below
-    # answers ChildProcessError for a zombie nobody will collect.
+    """stream a verl subprocess under the shared process-group lifecycle supervisor."""
     adopt_orphaned_descendants()
     proc = subprocess.Popen(
         cmd,
@@ -501,87 +486,88 @@ def run_verl_training(
         stdout=subprocess.PIPE,
         stderr=subprocess.STDOUT,
         text=True,
+        errors=errors,
         bufsize=1,
         start_new_session=True,
     )
-    # start_new_session makes the leader pid the stable group id even after that leader is reaped.
+    # start_new_session keeps the group addressable after the leader is reaped.
     process_group_id = proc.pid
-    last_hb = 0.0
     try:
-        # the child's exit is watched independently of pipe EOF. a grandchild holding the inherited
-        # pipe open after the trainer dies would otherwise keep the loop below running forever, on a
-        # paid gpu, having already lost the process whose output it is waiting for.
         with _ChildExitWatchdog(
             proc, process_group_id=process_group_id, grace_s=_ORPHANED_PIPE_GRACE_S
         ) as watchdog:
             assert proc.stdout is not None
             for line in proc.stdout:
-                # held across the callbacks, not just the arrival: on_step uploads a checkpoint, and
-                # a reader inside one is working, not stuck.
                 with watchdog.handling_line():
-                    print(line, end="", flush=True)
-                    child_tail.record(line)
-                    if on_line is not None:
-                        on_line(line)
-                    m = step_re.search(line)
-                    if m and on_step is not None:
-                        on_step(int(m.group(1)))
-                    if heartbeat is not None:
-                        now = time.monotonic()
-                        if now - last_hb >= heartbeat_interval_s:
-                            heartbeat()
-                            last_hb = now
+                    on_line(line)
     except BaseException:
         kill_process_group(proc, process_group_id=process_group_id)
         raise
     finally:
         if proc.poll() is None:
-            # BOUNDED. EOF is not the child's exit, so a child that outlives its own stdout would
-            # park the attempt here. one that has not exited within the grace after its stdout closed
-            # is not going to, so tear the group down and let the wait inside that call collect it.
             try:
                 proc.wait(timeout=_TEARDOWN_GRACE_S)
             except subprocess.TimeoutExpired:
                 kill_process_group(proc, process_group_id=process_group_id)
-        # every job boundary, not just the failing ones. a straggler an earlier teardown SIGKILLed
-        # but could not drain in time exits shortly after, and `kill_process_group` -- the only
-        # other caller of this -- runs on exceptions alone: a worker whose later jobs all succeed
-        # would hold that zombie for its whole life, as pid 1 with nothing else to reap it.
         reap_stragglers()
-    collected = proc.returncode
-    if collected is None:
-        # the bounded wait above can end with nothing collected, which the unbounded one it replaced
-        # could not: a group member wedged in uninterruptible io outlives even the SIGKILL inside
-        # teardown. `int(None)` would replace that diagnosis with a TypeError, so name the survivor
-        # instead -- callers already handle a RuntimeError from here.
+    if proc.returncode is None:
         raise RuntimeError(
             f"verl subprocess {proc.pid} did not exit after teardown; its process group is still "
             "holding the gpu"
         )
-    return_code = int(collected)
+    return_code = int(proc.returncode)
     if watchdog.tore_down and return_code == 0:
-        # the child exited 0 but a descendant held the pipe open past the grace, so the group was
-        # killed to release it. reporting success here would upload whatever partial artifacts exist
-        # as a completed run: the trainer's own exit status says nothing about the descendant that
-        # was still running, and that descendant is why the pipe never closed.
         raise RuntimeError(
             f"verl subprocess {proc.pid} exited 0 but a descendant held its output pipe open for "
             f"{_ORPHANED_PIPE_GRACE_S:.0f}s; the process group was torn down to release the gpu"
         )
-    try:
-        raise_for_classified_verl_exit(return_code, child_tail)
-    except BaseException:
-        kill_process_group(proc, process_group_id=process_group_id)
-        raise
     if return_code != 0:
-        # a nonzero exit that carried no recognized signature RETURNS from the classifier rather
-        # than raising, so this is the one failing path that reached neither teardown above. the
-        # direct child is gone but its group need not be: a surviving EngineCore keeps its cuda
-        # context, and a reusable worker then hands the next attempt an occupied gpu.
         kill_process_group(proc, process_group_id=process_group_id)
     return return_code
 
 
+def run_verl_training(
+    cmd: list[str],
+    *,
+    env: dict[str, str],
+    on_step: Callable[[int], None] | None = None,
+    on_line: Callable[[str], None] | None = None,
+    heartbeat: Callable[[], None] | None = None,
+    step_pattern: str = r"step:\s*(\d+)",
+    heartbeat_interval_s: float = 20.0,
+    tail: ChildOutputTail | None = None,
+) -> int:
+    """run a verl trainer subprocess, streaming stdout and surfacing step progress.
+
+    returns the process exit code. stdout+stderr are merged and scanned line by line: ``on_line``
+    receives every line, ``on_step`` receives each parsed training step, and ``heartbeat`` is called
+    at most once per ``heartbeat_interval_s``. callback failures terminate the process group before
+    they are re-raised.
+    """
+    step_re = re.compile(step_pattern)
+    child_tail = tail if tail is not None else ChildOutputTail()
+    last_hb = 0.0
+
+    def handle_line(line: str) -> None:
+        nonlocal last_hb
+        print(line, end="", flush=True)
+        child_tail.record(line)
+        if on_line is not None:
+            on_line(line)
+        match = step_re.search(line)
+        if match and on_step is not None:
+            on_step(int(match.group(1)))
+        if heartbeat is not None:
+            now = time.monotonic()
+            if now - last_hb >= heartbeat_interval_s:
+                heartbeat()
+                last_hb = now
+
+    return_code = _run_streaming_verl_subprocess(cmd, env=env, on_line=handle_line)
+    raise_for_classified_verl_exit(return_code, child_tail)
+    return return_code
+
+
 _TEARDOWN_GRACE_S = 10.0
 
 # grace for a descendant holding stdout open after the trainer exits. observe child exit separately

{freesolo_flash_dev-1.2.37 → freesolo_flash_dev-1.2.43}/flash/engine/worker/model/adapter.py

@@ -5,7 +5,14 @@ from __future__ import annotations
 import json
 import os
 import shutil
+from collections.abc import Mapping
+from typing import Any
 
+from flash.adapters.fused_experts import (
+    has_complete_fused_expert_tensors,
+    lora_target_parameters,
+    validate_fused_expert_adapter_config,
+)
 from flash.adapters.lora_rank import resolve_adapter_ref
 from flash.engine.plan.recipe import RECIPE
 from flash.engine.worker.io.hf import (
@@ -17,35 +24,23 @@ from flash.engine.worker.io.hf import (
 )
 from flash.engine.worker.model.lora import (
     _read_adapter_tensor_keys,
+    _read_adapter_tensor_metadata,
 )
 from flash.engine.worker.runtime.pkg_proxy import W as _w
 
 _ADAPTER_DOWNLOAD_RETRIES = 4
 _ADAPTER_DOWNLOAD_BACKOFF_S = 5.0
-_QWEN35_EXPERT_TARGET_PARAMETERS = (
-    "mlp.experts.gate_up_proj",
-    "mlp.experts.down_proj",
-)
-
-
-def lora_target_parameters(model_id: str | None) -> list[str] | None:
-    """return direct parameter targets required by the model's fused expert layout."""
-    if model_id == "Qwen/Qwen3.6-35B-A3B":
-        return list(_QWEN35_EXPERT_TARGET_PARAMETERS)
-    return None
 
 
-def validate_lora_target_parameters(config: dict, model_id: str) -> None:
-    """fail closed when a warm-start adapter omits required fused expert parameters."""
-    required = set(lora_target_parameters(model_id) or ())
-    if not required:
+def validate_warmstart_adapter(config: Mapping[str, Any], model_id: str, adapter_dir: str) -> None:
+    """Validate a downloaded warm-start adapter without changing its config or files."""
+    validate_fused_expert_adapter_config(config, model_id)
+    if not lora_target_parameters(model_id):
         return
-    actual = set(config.get("target_parameters") or ())
-    missing = sorted(required - actual)
-    if missing:
+    tensors = _read_adapter_tensor_metadata(adapter_dir) or {}
+    if not has_complete_fused_expert_tensors(tensors, config, model_id):
         raise ValueError(
-            f"warm-start adapter for {model_id} omits required expert targets {missing}; "
-            "retrain the source adapter with the current Flash version"
+            f"warm-start adapter for {model_id} does not contain complete fused expert LoRA weights"
         )