freesolo-flash-dev 1.4.23__tar.gz → 1.4.24__tar.gz

{freesolo_flash_dev-1.4.23 → freesolo_flash_dev-1.4.24}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: freesolo-flash-dev
-Version: 1.4.23
+Version: 1.4.24
 Summary: Flash — managed LoRA post-training (SFT/GRPO) 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.4.23 → freesolo_flash_dev-1.4.24}/flash/core/catalog.py

@@ -23,6 +23,20 @@ _IMAGE_TRAINING_MODELS = frozenset(
 )
 
 
+# Native templates qualified against the shared Qwen3 coder grammar and replay contract.
+# New model families must pass the tool-template matrix before joining this set.
+TOOL_CALLING_MODELS = frozenset(
+    {
+        "Qwen/Qwen3.5-0.8B",
+        "Qwen/Qwen3.5-4B",
+        "Qwen/Qwen3.5-9B",
+        "Qwen/Qwen3.5-9B-Base",
+        "Qwen/Qwen3.8-27B",
+        "Qwen/Qwen3.6-35B-A3B",
+    }
+)
+
+
 def supports_image_training(model: str | ModelInfo | None) -> bool:
     """Return whether a curated model supports image-bearing SFT and GRPO."""
     model_id = model.id if isinstance(model, ModelInfo) else model

{freesolo_flash_dev-1.4.23 → freesolo_flash_dev-1.4.24}/flash/custom/backend.py

@@ -1,4 +1,4 @@
-"""Reference single-GPU LoRA engine for explicit custom-loop operations.
+"""GPU-resident LoRA engine for explicit custom-loop operations.
 
 No server-side user code: the client supplies only tokens and dLoss/dLogprob weights.
 """
@@ -56,7 +56,9 @@ class TorchBackend:
     def _tensor(self, tokens):
         import torch
 
-        return torch.tensor([tokens], device=self.parameters[0].device, dtype=torch.long)
+        return torch.tensor(
+            [tokens], device=self.model.get_input_embeddings().weight.device, dtype=torch.long
+        )
 
     def _logprobs(self, tokens, target_tokens=None):
         """Logprobs of the targets: the next token ([T]) or a [T, K] target matrix."""
@@ -65,7 +67,7 @@ class TorchBackend:
         ids = self._tensor(tokens)
         logits = self.model(input_ids=ids[:, :-1], use_cache=False).logits[0].float()
         if target_tokens is None:
-            targets = ids[0, 1:, None]
+            targets = ids[0, 1:, None].to(logits.device)
         else:
             targets = torch.tensor(target_tokens, device=logits.device, dtype=torch.long)
         # gather targets and one normalizer per position; never copy a vocabulary row each.
@@ -259,7 +261,9 @@ class TorchBackend:
         self.optimizer.load_state_dict(state["optimizer"])
         torch.set_rng_state(state["rng"])
         if torch.cuda.is_available():
-            torch.cuda.set_rng_state_all(state["cuda_rng"])
+            from flash.custom.model_parallel import restore_cuda_rng
+
+            restore_cuda_rng(state["cuda_rng"])
         self.algorithm_state = (
             json.loads((source / "algorithm.json").read_text())
             if (source / "algorithm.json").exists()
@@ -292,8 +296,10 @@ def _ensure_cuda_headroom(min_free_bytes=CUDA_HEADROOM_BYTES):
 
         if not torch.cuda.is_available():
             return False
-        free, _ = torch.cuda.mem_get_info()
-        if free >= min_free_bytes:
+        if all(
+            torch.cuda.mem_get_info(device)[0] >= min_free_bytes
+            for device in range(torch.cuda.device_count())
+        ):
             return False
         torch.cuda.empty_cache()
         return True
@@ -339,6 +345,7 @@ def load_backend(
     checkpoint_dir,
     max_tokens=2048,
     init_adapter=None,
+    gpu_count=1,
 ):
     import torch
     from peft import LoraConfig, get_peft_model
@@ -372,13 +379,19 @@ def load_backend(
         from transformers import AutoModelForImageTextToText
 
         model_class = AutoModelForImageTextToText
+    from flash.custom.model_parallel import load_placement, validate_placement
+
+    placement = load_placement(model_id, lora_rank, max_tokens, gpu_count)
     model = model_class.from_pretrained(
         model_id,
         revision=revision,
         trust_remote_code=False,
         torch_dtype=torch.bfloat16,
         attn_implementation="sdpa",
-    ).to("cuda")
+        **placement,
+    )
+    if gpu_count == 1:
+        model = model.to("cuda")
     if init_adapter:
         model = _warm_start(model, init_adapter, lora_rank, lora_alpha)
     else:
@@ -394,6 +407,7 @@ def load_backend(
                 **targeting,
             ),
         )
+    validate_placement(model, gpu_count)
     model.config._commit_hash = revision
     return TorchBackend(model, tokenizer, checkpoint_dir=checkpoint_dir, max_tokens=max_tokens)
 

{freesolo_flash_dev-1.4.23 → freesolo_flash_dev-1.4.24}/flash/custom/billing.py

@@ -8,6 +8,7 @@ from pathlib import Path
 def estimate(spec, *, allocation=None):
     from flash.cost.types import CostEstimate
     from flash.providers.core.allocator import allocate
+    from flash.providers.core.base import get_gpu_info
 
     rented = allocation or allocate(
         spec.model,
@@ -28,8 +29,9 @@ def estimate(spec, *, allocation=None):
         method="custom",
         steps=spec.train.max_steps,
         gpu=rented.gpu,
+        gpu_count=rented.gpu_count,
         provider=rented.provider,
-        gpu_vram_gb=rented.min_vram_gb,
+        gpu_vram_gb=get_gpu_info(rented.gpu).vram_gb,
         required_vram_gb=rented.min_vram_gb,
         gpu_hourly_usd=rented.hourly_usd,
         setup_seconds=0,
@@ -37,7 +39,7 @@ def estimate(spec, *, allocation=None):
         train_seconds=seconds,
         wall_clock_seconds=seconds,
         wall_capped=True,
-        total_usd=seconds * rented.hourly_usd / 3600,
+        total_usd=seconds * rented.hourly_usd * rented.gpu_count / 3600,
         notes=(
             "Upper bound for a dedicated custom session, including time waiting for the local script.",
         ),

freesolo_flash_dev-1.4.24/flash/custom/model_parallel.py

@@ -0,0 +1,47 @@
+"""Keep a custom model's whole decoder layers on one or two local CUDA devices."""
+
+
+def load_placement(model_id, rank, context, gpu_count):
+    import torch
+
+    from flash.custom.placement import MAX_GPUS, per_device_vram_gb
+
+    if type(gpu_count) is not int or gpu_count not in (1, MAX_GPUS):
+        raise ValueError("custom training supports one or two GPUs")
+    if torch.cuda.device_count() < gpu_count:
+        raise ValueError(f"custom training requires {gpu_count} visible GPUs")
+    if gpu_count == 1:
+        return {}
+    need = per_device_vram_gb(
+        model_id, {"lora_rank": rank, "max_context_tokens": context}, gpu_count
+    )
+    memory = {}
+    for device in range(gpu_count):
+        total = torch.cuda.get_device_properties(device).total_memory
+        if total < need * 1e9:
+            raise ValueError(f"custom training requires at least {need} GB on each GPU")
+        memory[device] = int(total * 0.8)
+    # Transformers/Accelerate put complete decoder blocks on devices and retain
+    # autograd across their transfers. No CPU/disk offload is permitted below.
+    return {"device_map": "balanced", "max_memory": memory}
+
+
+def validate_placement(model, gpu_count):
+    devices = {p.device for p in model.parameters()}
+    if any(device.type != "cuda" for device in devices):
+        raise ValueError("custom training requires GPU-resident weights; offload is unsupported")
+    if len(devices) != gpu_count:
+        raise ValueError("custom training did not place model weights on every allocated GPU")
+
+
+def restore_cuda_rng(states):
+    """Restore matching devices; seed additional shards from the saved primary RNG.
+
+    Checkpoints can move between GPU shapes. Never address a saved device that
+    does not exist on the new worker. Exact RNG replay applies to unchanged shapes.
+    """
+    import torch
+
+    if states:
+        for device in range(torch.cuda.device_count()):
+            torch.cuda.set_rng_state(states[device] if device < len(states) else states[0], device)

freesolo_flash_dev-1.4.24/flash/custom/placement.py

@@ -0,0 +1,77 @@
+"""Memory admission for one-process custom training with GPU-resident layer sharding."""
+
+import math
+
+MAX_GPUS = 2
+
+
+def per_device_vram_gb(model_id, train, count, *, total_need=None):
+    """Shard resident state, but reserve full activations/logits on every card.
+
+    Unlike FSDP, layer placement does not distribute a vocabulary projection or
+    optimizer tensors within one layer. Allow an extra layer and both vocabulary
+    matrices for partition imbalance; never admit using pooled memory alone.
+    """
+    from flash.core.catalog import get_model
+    from flash.engine.plan.vram import _lora_memory_gb, _positive_int_or_default, _sizing_value
+    from flash.providers.core.allocator import required_vram_gb, vram_headroom
+
+    need = (
+        total_need if total_need is not None else required_vram_gb(model_id, "custom", train=train)
+    )
+    if count == 1:
+        return need
+    if count != MAX_GPUS:
+        raise ValueError("custom training supports one or two Modal GPUs")
+    info = get_model(model_id)
+    rank = _positive_int_or_default(_sizing_value(train, "lora_rank"), 16)
+    resident = info.params_b * 2 + _lora_memory_gb(
+        rank, info.active_params_b or info.params_b, "custom", info
+    )
+    imbalance = resident / info.num_layers + 4 * info.vocab_size * info.hidden_size / 1e9
+    return math.ceil(
+        max(0, need - resident * vram_headroom()) + (resident / count + imbalance) * vram_headroom()
+    )
+
+
+def fitting_candidates(candidates, model_id, train, *, total_need):
+    """Keep only allocations that fit the custom worker on every device."""
+    return [
+        candidate
+        for candidate in candidates
+        if candidate.vram_gb
+        >= per_device_vram_gb(model_id, train, candidate.gpu_count, total_need=total_need)
+    ]
+
+
+def preview_gpu(model_id, train, max_count, *, total_need):
+    from flash.providers.core.base import GPU_INFO, UnsupportedGpuError
+    from flash.providers.core.sharding import combined_vram_gb
+    from flash.providers.modal.pricing import hourly_rate
+
+    shapes = [
+        (hourly_rate(gpu.name) * count, count, gpu.name)
+        for count in (1, 2)
+        if count <= max_count
+        for gpu in GPU_INFO.values()
+        if gpu.validated
+        and count <= gpu.max_count
+        and combined_vram_gb(gpu.vram_gb, count) >= total_need
+        and gpu.vram_gb >= per_device_vram_gb(model_id, train, count, total_need=total_need)
+    ]
+    if not shapes:
+        raise UnsupportedGpuError("custom training cannot fit the per-device memory budget")
+    return min(shapes)[2]
+
+
+def allocation_shape(model_id, provider, providers, gpu_type, fallbacks, max_count):
+    """Validate custom provider preferences before resolving the allocation shape."""
+    from flash.providers.core.base import UnsupportedGpuError, custom_gpu_shape
+    from flash.providers.core.registry import validated_provider_preferences
+
+    if provider and providers:
+        raise UnsupportedGpuError("provider and providers cannot both be set")
+    validated_provider_preferences(providers, allow_empty=isinstance(providers, tuple))
+    return custom_gpu_shape(
+        model_id, provider, (gpu_type, *fallbacks) if gpu_type else (), max_count
+    )

{freesolo_flash_dev-1.4.23 → freesolo_flash_dev-1.4.24}/flash/custom/server.py

@@ -1,4 +1,4 @@
-"""Authenticated, single-session worker. Run one process per dedicated GPU."""
+"""Authenticated, single-session worker for one dedicated GPU allocation."""
 
 import asyncio
 import hmac

{freesolo_flash_dev-1.4.23 → freesolo_flash_dev-1.4.24}/flash/custom/worker.py

@@ -59,6 +59,7 @@ def run(*, context):
             seed=spec.seed,
             checkpoint_dir="/tmp/flash-custom-checkpoints",
             max_tokens=spec.train.max_context_tokens,
+            gpu_count=spec.gpu.count,
             init_adapter=init_adapter,
         )
     backend.checkpoints = checkpoints

{freesolo_flash_dev-1.4.23 → freesolo_flash_dev-1.4.24}/flash/providers/core/allocator.py

@@ -7,6 +7,7 @@ from dataclasses import replace
 from typing import NoReturn
 
 from flash._internal.logging import get_logger
+from flash.custom.placement import allocation_shape, fitting_candidates
 from flash.providers.core.base import (
     GPU_INFO,
     Allocation,
@@ -18,7 +19,6 @@ from flash.providers.core.base import (
     _run_cost_key,
     authored_gpu_ceiling,
     canonical_gpu,
-    custom_gpu_shape,
     largest_rentable_count,
     providers_for,
     rentable_gpu_counts,
@@ -186,7 +186,7 @@ def _executed_gpu_count(algorithm: str, train, overrides, gpu_count: int) -> int
     if algo in ("grpo", "rl"):
         return _executed_rl_gpu_count(algo, train, gpu_count)
     if algo == "custom":
-        return 1
+        return gpu_count  # layer sharding uses every card even for a single sequence
     if algo != "sft":
         return gpu_count
     batch = _sizing_int(train, "batch_size", 0)
@@ -674,7 +674,6 @@ def _resolved_gpu_count(
 def _gather_candidates(
     available: tuple[str, ...],
     *,
-    per_card_need: float,
     constraints: AllocationConstraints,
     acceptable: tuple[str, ...],
     provider: str,
@@ -686,6 +685,17 @@ def _gather_candidates(
 
     ``acceptable`` is the authoritative class filter; empty means unpinned.
     """
+    # with multiple cards allowed, a card can contribute as one of N; query providers at the reduced
+    # per-card floor so classes too small to fit alone still surface, then keep only the shapes that
+    # actually fit below.
+    per_card_need = constraints.required_vram_gb
+    if constraints.max_gpu_count > 1:
+        per_card_need = max(
+            1,
+            math.ceil(
+                constraints.required_vram_gb / (constraints.max_gpu_count * SHARD_VRAM_EFFICIENCY)
+            ),
+        )
     allowed = set(acceptable)
     queries = tuple(replace(constraints, gpu_type=gpu) for gpu in acceptable) or (constraints,)
     candidates: list[Candidate] = []
@@ -813,11 +823,8 @@ def allocate(
     # off the authored batch over-reserves (4 GB on a 4B at batch 8 / 4096) and can reject a card the
     # run would have fit on. `_fits` reads it too, so the requirement and the fit cannot disagree.
     if algorithm == "custom":
-        if provider and providers:
-            raise UnsupportedGpuError("provider and providers cannot both be set")
-        validated_provider_preferences(providers, allow_empty=isinstance(providers, tuple))
-        provider, gpu_type, gpu_type_fallbacks, max_gpu_count = custom_gpu_shape(
-            model_id, provider, (gpu_type, *gpu_type_fallbacks) if gpu_type else (), max_gpu_count
+        provider, gpu_type, gpu_type_fallbacks, max_gpu_count = allocation_shape(
+            model_id, provider, providers, gpu_type, gpu_type_fallbacks, max_gpu_count
         )
         providers = ()  # only Modal has a validated custom worker transport today.
     sized_train = _overridden_train(train, overrides)
@@ -898,20 +905,15 @@ def allocate(
         required_vram_gb=need,
         max_gpu_count=cap,
     )
-    # with multiple cards allowed, a card can contribute as one of N; query providers at the reduced
-    # per-card floor so classes too small to fit alone still surface, then keep only the shapes that
-    # actually fit below.
-    per_card_need = need
-    if cap > 1:
-        per_card_need = max(1, math.ceil(need / (cap * SHARD_VRAM_EFFICIENCY)))
     candidates, lookup_failed, structurally_unsupported = _gather_candidates(
         available,
-        per_card_need=per_card_need,
         constraints=constraints,
         acceptable=acceptable,
         provider=provider,
     )
     candidates = _fitting_candidates(candidates, need, executed_width)
+    if algorithm == "custom":
+        candidates = fitting_candidates(candidates, model_id, sized_train, total_need=need)
     supported_available = tuple(name for name in available if name not in structurally_unsupported)
     if not candidates:
         _raise_no_candidate_error(

{freesolo_flash_dev-1.4.23 → freesolo_flash_dev-1.4.24}/flash/providers/core/base.py

@@ -309,16 +309,16 @@ def _run_cost_key(
 
 
 def custom_gpu_shape(model_id, provider, gpu_types, count):
-    """Keep the custom transport single-GPU; use shared sizing and hardware eligibility."""
+    """Bound custom layer sharding to validated one/two-GPU Modal shapes."""
     from flash.core.catalog import validate_model_for_algorithm
 
     validate_model_for_algorithm(model_id, "custom")
     if provider not in {"", "modal"} or (
-        count is not None and (type(count) is not int or count != 1)
+        count is not None and (type(count) is not int or count not in (1, 2))
     ):
-        raise ValueError("custom training currently requires one managed Modal GPU")
+        raise ValueError("custom training supports one or two managed Modal GPUs")
     allowed = tuple(gpu_types)
-    return "modal", allowed[0] if allowed else "", allowed[1:], 1
+    return "modal", allowed[0] if allowed else "", allowed[1:], count or 1
 
 
 def authored_gpu_ceiling(gpu_type: str, gpu_count: int | None) -> int | None:
@@ -465,7 +465,7 @@ def provisional_gpu_count(
         required_vram_gb,
     )
 
-    if algorithm == "custom":
+    if algorithm == "custom" and gpu_count is None:
         gpu_count = 1
     ceiling = MAX_COMBINATION_CARDS if gpu_count is None else gpu_count
     geometry_revision = (
@@ -546,6 +546,10 @@ def provisional_gpu(
         geometry_model_revision=geometry_model_revision,
         gpu_count=gpu_count,
     )
+    if algorithm == "custom":
+        from flash.custom.placement import preview_gpu
+
+        return preview_gpu(model_id, train, effective_count, total_need=min_vram)
     geometry_revision = (
         model_revision if geometry_model_revision is None else geometry_model_revision
     )

{freesolo_flash_dev-1.4.23 → freesolo_flash_dev-1.4.24}/flash/schema/__init__.py

@@ -534,13 +534,9 @@ def _validate_gpu_section(
         raise ConfigError(str(exc)) from exc
 
     if algorithm == "custom":
-        from flash.providers.core.base import custom_gpu_shape
+        from flash.schema.custom import validate_custom_gpu_count
 
-        try:
-            custom_gpu_shape(model, gpu_spec.provider, gpu_spec.acceptable_types, gpu_count)
-        except ValueError as exc:
-            raise ConfigError(str(exc)) from exc
-        gpu_count = 1
+        gpu_count = validate_custom_gpu_count(model, gpu_spec, gpu_count)
 
     gpu_types = gpu_spec.acceptable_types
     for candidate in gpu_types:
@@ -568,7 +564,7 @@ def _validate_gpu_section(
             geometry_model_revision=model_revision,
             gpu_count=requested_gpu_count,
         )
-        if gpu_types and preflight_gpu_count <= 1 and not model_revision:
+        if gpu_types and (preflight_gpu_count <= 1 or algorithm == "custom") and not model_revision:
             from flash.providers.core.allocator import required_vram_gb
 
             # sized from `preflight_train`, so an unresolved warm start is measured at rank 1 rather
@@ -587,6 +583,12 @@ def _validate_gpu_section(
             # run is one allocation would never rent, so accepting it here would advertise failover
             # the run does not actually have and surface the shortfall only after the head ran out
             # of capacity.
+            if algorithm == "custom":
+                from flash.custom.placement import per_device_vram_gb
+
+                required_vram = per_device_vram_gb(
+                    model, preflight_train, preflight_gpu_count, total_need=required_vram
+                )
             for candidate in gpu_types:
                 if get_gpu_info(candidate).vram_gb < required_vram:
                     raise ConfigError(

{freesolo_flash_dev-1.4.23 → freesolo_flash_dev-1.4.24}/flash/schema/custom.py

@@ -1,7 +1,7 @@
 """Custom-only fields and constraints; common sections use the normal JobSpec parser."""
 
 from flash.custom.types import integer
-from flash.schema.fields import ConfigError
+from flash.schema.fields import ConfigError, _train_int
 
 TRAIN_KEYS = {
     "max_steps",
@@ -65,9 +65,14 @@ def validate_custom_spec(spec):
         )
     try:
         integer(spec.train.max_steps, "train.max_steps", 1, 1_000_000)
-        integer(spec.train.max_context_tokens, "train.max_context_tokens", 2, 4096)
-        integer(spec.train.lora_rank, "train.lora_rank", 1, 64)
-        integer(spec.train.lora_alpha, "train.lora_alpha", 1, 256)
+        # Rank/alpha use the common parser and model's serving rank cap. Context
+        # uses the shared submission preflight and GPU memory-fit checks; custom
+        # next-token losses only add a two-token minimum, not a separate ceiling.
+        _train_int(
+            {"max_context_tokens": spec.train.max_context_tokens},
+            "max_context_tokens",
+            minimum=2,
+        )
         integer(spec.train.save_every, "train.save_every", 1, 1_000_000)
     except ValueError as exc:
         raise ConfigError(str(exc)) from exc
@@ -76,3 +81,14 @@ def validate_custom_spec(spec):
             "train.init_from_adapter starts fresh optimizer state from any adapter; "
             "custom.resume_from continues a custom checkpoint. Set only one"
         )
+
+
+def validate_custom_gpu_count(model, gpu_spec, gpu_count):
+    """Validate custom GPU placement and apply its single-device default."""
+    from flash.providers.core.base import custom_gpu_shape
+
+    try:
+        custom_gpu_shape(model, gpu_spec.provider, gpu_spec.acceptable_types, gpu_count)
+    except ValueError as exc:
+        raise ConfigError(str(exc)) from exc
+    return gpu_count or 1

{freesolo_flash_dev-1.4.23 → freesolo_flash_dev-1.4.24}/flash/serve/request/tool_calls.py

@@ -10,6 +10,7 @@ from dataclasses import dataclass
 from decimal import Decimal, DecimalException
 from typing import Any
 
+from flash.core.catalog import TOOL_CALLING_MODELS
 from flash.serve.contract.protocol import MAX_CHAT_REQUEST_BYTES, TEXT_TYPES
 from flash.serve.request import text_scan
 from flash.serve.request.schema_types import _decimal_is_integral, _matches_type
@@ -57,8 +58,7 @@ class FunctionTool:
 
 
 def qualified_tool_parser(base_model: str) -> str | None:
-    supported = {"Qwen/Qwen3.5-9B", "Qwen/Qwen3.5-9B-Base"}
-    return TOOL_PARSER_QWEN3_CODER if base_model in supported else None
+    return TOOL_PARSER_QWEN3_CODER if base_model in TOOL_CALLING_MODELS else None
 
 
 def normalize_tools(

{freesolo_flash_dev-1.4.23 → freesolo_flash_dev-1.4.24}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "freesolo-flash-dev"
-version = "1.4.23"
+version = "1.4.24"
 description ="Flash — managed LoRA post-training (SFT/GRPO) for Freesolo environments, driven by the `flash` CLI"
 # The training and serving stack is validated on Python 3.11 and 3.12.
 requires-python = ">=3.11,<3.13"
@@ -210,7 +210,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.4.23"
+version = "1.4.24"
 
 [tool.pytest.ini_options]
 pythonpath = ["."]