freesolo-flash-dev 0.2.37__tar.gz → 0.2.38__tar.gz

{freesolo_flash_dev-0.2.37 → freesolo_flash_dev-0.2.38}/Dockerfile.worker

@@ -42,7 +42,7 @@ ENV TRITON_CACHE_DIR=/opt/flash/kernelcache/triton \
 RUN mkdir -p /opt/flash/kernelcache/triton /opt/flash/kernelcache/inductor \
     /opt/flash/kernelcache/flashinfer_cubin /opt/flash/kernelcache/flashinfer
 
-# build-essential for Triton/Inductor (Liger) + any source builds; git for pip VCS installs.
+# build-essential for Triton/Inductor/chalk + any source builds; git for pip VCS installs.
 RUN apt-get update && apt-get install -y --no-install-recommends build-essential git curl \
     && rm -rf /var/lib/apt/lists/*
 
@@ -58,7 +58,6 @@ RUN pip install --no-cache-dir \
     "huggingface_hub>=0.25" \
     "accelerate>=1.4" \
     "wandb>=0.17" \
-    "liger-kernel>=0.5" \
     "flash-linear-attention @ git+https://github.com/fla-org/flash-linear-attention.git@f0e213dbd8b5fb90c3c7eca869ac1706d5377139" \
     "tilelang==0.1.11" \
     "apache-tvm-ffi==0.1.11" \
@@ -148,7 +147,7 @@ RUN TORCH_CUDA_ARCH_LIST="8.0 8.6 8.9 9.0 10.0 12.0" CAUSAL_CONV1D_FORCE_BUILD=T
 # affect another tenant). The redirect can't be a static image ENV: the volume mounts at a fixed path
 # only at runtime and only when attached, so a baked ENV would break cache-less (no-volume) cold runs.
 #
-# Compiled-kernel cache: OPT-IN bake, JIT fallback. The fused Triton (Liger/fla/chalk),
+# Compiled-kernel cache: OPT-IN bake, JIT fallback. The fused Triton (fla/chalk),
 # TorchInductor, and Hopper tilelang kernels the trainer builds on first use are content-addressed
 # by GPU ARCH + toolchain version, so they can only be compiled on a real GPU of each catalog arch
 # (8.0/8.6/8.9/9.0/12.0) — which the CI image builder normally lacks (the FA2 *source* path above
@@ -187,7 +186,7 @@ RUN if [ "${BUILD_KERNEL_CACHE}" = "true" ] && [ -f /opt/flash/kernelcache/mega_
 # Sanity: the worker imports must succeed (catches a broken stack at build time, not run time).
 # Import EVERY package the worker runtime uses (derived from a full AST scan of the engine/envs/
 # bootstrap code + the WORKER_DEPS list above): the core stack, the optimizer/accelerator libs
-# (liger, flash-linear-attention=fla, flash_attn, bitsandbytes, triton), and the HF/env libs. A
+# (flash-linear-attention=fla, flash_attn, bitsandbytes, triton), and the HF/env libs. A
 # missing or ABI-broken package fails the build HERE instead of silently shipping (the FA wheel in
 # particular can install yet fail to import against this torch). cuDNN comes from the base image.
 # flash_attn is BEST-EFFORT (the install above tolerates a build failure -> SDPA fallback), so it
@@ -195,7 +194,7 @@ RUN if [ "${BUILD_KERNEL_CACHE}" = "true" ] && [ -f /opt/flash/kernelcache/mega_
 # could never be produced. The rest of the stack is mandatory.
 RUN python -c "import importlib.metadata as M; \
 import torch, transformers, trl, peft, vllm, datasets, huggingface_hub, accelerate, wandb, \
-liger_kernel, fla, bitsandbytes, hf_transfer, triton, freesolo, runpod; \
+fla, bitsandbytes, hf_transfer, triton, freesolo, runpod; \
 import importlib; \
 fa_ok = False; \
 exec('try:\n import flash_attn\n fa_ok = True\nexcept Exception:\n fa_ok = False'); \
@@ -207,7 +206,7 @@ cudnn = torch.backends.cudnn.version(); assert cudnn, 'cuDNN not compiled into t
 ver = lambda n: M.version(n); \
 print('worker stack OK | torch', torch.__version__, '| cudnn', cudnn, '| triton', triton.__version__, '| flash_attn', 'present' if fa_ok else 'ABSENT (SDPA fallback)', '| flash_attn_3', 'present (Hopper FA3)' if fa3_ok else 'ABSENT (SDPA on Hopper)', '| flashinfer', 'present (B200 MoE/attn)' if fi_ok else 'ABSENT (Triton fused-MoE fallback)'); \
 [print('  ', d, ver(d)) for d in ['transformers','trl','peft','vllm','flashinfer-python','datasets','huggingface-hub', \
-'accelerate','wandb','liger-kernel','flash-linear-attention','bitsandbytes', \
+'accelerate','wandb','flash-linear-attention','bitsandbytes', \
 'hf-transfer','freesolo','runpod']]"
 
 # RunPod Serverless worker entrypoint (Pattern D — Flash custom image). Generate /rp_handler.py

{freesolo_flash_dev-0.2.37 → freesolo_flash_dev-0.2.38}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash-dev
-Version: 0.2.37
+Version: 0.2.38
 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-0.2.37 → freesolo_flash_dev-0.2.38}/docker/Dockerfile.kernelcache

@@ -1,7 +1,7 @@
 # Thin per-arch kernel-cache layer.
 #
 # Bakes a warmup-produced mega-cache ON TOP of the EXACT base image the warmup ran inside, so the
-# cache's toolchain (torch/triton/fla/liger) is guaranteed to match the image it ships in. Rebuilding
+# cache's toolchain (torch/triton/fla/chalk) is guaranteed to match the image it ships in. Rebuilding
 # the full Dockerfile.worker from the checkout instead could drift from the published base and bake a
 # cache that silently won't load (cold-JIT fallback) -- this avoids that by construction.
 #

{freesolo_flash_dev-0.2.37 → freesolo_flash_dev-0.2.38}/docker/bake_kernel_cache.py

@@ -100,16 +100,16 @@ def main() -> int:
     entry_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "bake_pod_entry.py")
     allowed_cuda = [v.strip() for v in args.allowed_cuda.split(",") if v.strip()] or None
 
-    # the chalk install spec the production worker uses (so the bake warms chalk's kernels too);
-    # FLASH_CHALK_SPEC overrides, else the source-of-truth DEFAULT_CHALK_SPEC, else a safe literal.
+    # The chalk install spec the production worker uses, so the bake warms the same kernels. Keep the
+    # default import fail-loud: a stale hidden fallback would silently bake against the wrong chalk SHA.
     chalk_spec = os.environ.get("FLASH_CHALK_SPEC", "").strip()
     if not chalk_spec:
         try:
             from flash.providers.runpod.train.deps import DEFAULT_CHALK_SPEC
 
             chalk_spec = DEFAULT_CHALK_SPEC
-        except Exception:
-            chalk_spec = "freesolo-chalk>=0.1.0,<0.2.0"
+        except Exception as exc:
+            raise RuntimeError("could not import DEFAULT_CHALK_SPEC for kernel-cache bake") from exc
 
     _upload_flash_code(api, repo, token)
 

{freesolo_flash_dev-0.2.37 → freesolo_flash_dev-0.2.38}/docker/kernel_fingerprint.py

@@ -11,7 +11,7 @@ Two fingerprints, because "out of date" has two flavors with very different cost
   * fp_cache = hash of the inputs whose kernels live in the baked mega-cache
     (torch.compiler.save_cache_artifacts: Triton/Inductor/torch.compile). Changing one of these
     INVALIDATES the cache, so the per-arch image needs a real GPU re-warm. These are: the base
-    FROM image (torch+triton), the fla git pin, liger, tilelang, apache-tvm-ffi, the default chalk
+    FROM image (torch+triton), the fla git pin, tilelang, apache-tvm-ffi, the default chalk
     spec, and the warmup script itself (it decides which kernels get compiled).
 
   * fp_base = hash of everything else baked into :cu128 that is NOT in the cache (FA2/FA3 wheels,
@@ -27,7 +27,7 @@ time, which is why every parse below FAILS LOUD rather than hashing a None.
 Known limitations (deliberately scoped -- each only ever costs a recoverable cold-JIT, never
 correctness, and the alternatives over-fire the paid GPU bake):
   * fp_cache hashes the Dockerfile dep PINS, not pip-resolved versions. A cache-affecting range
-    (liger-kernel>=, the chalk spec) could resolve a newer build on a later worker-image rebuild with
+    (the chalk spec) could resolve a newer build on a later worker-image rebuild with
     no text change, leaving fp_cache unmoved. Pin those exactly for airtight coverage. (fp_base DOES
     hash the whole Dockerfile.worker, so arbitrary base edits -- apt/ENV/CMD/cache-dir -- still
     trigger a free re-layer; only a cache-affecting change that isn't a parsed pin slips through.)
@@ -43,6 +43,7 @@ stdlib only, no flash/torch import, so it runs under a bare python3 in CI (no uv
 from __future__ import annotations
 
 import argparse
+import ast
 import hashlib
 import json
 import re
@@ -56,7 +57,7 @@ LABEL_REVISION = "org.opencontainers.image.revision"
 
 # pip specs in Dockerfile.worker's main stack whose kernels land in the mega-cache. everything else
 # in that block is a base-layer dep (fp_base). fla is matched separately (it carries the git sha).
-_CACHE_PKGS = {"liger-kernel", "tilelang", "apache-tvm-ffi"}
+_CACHE_PKGS = {"tilelang", "apache-tvm-ffi"}
 
 
 def _search(pattern: str, text: str, what: str, *, flags: int = 0) -> str:
@@ -71,6 +72,44 @@ def _sha256_file(path: Path) -> str:
     return hashlib.sha256(path.read_bytes()).hexdigest()
 
 
+def _python_string_constant(text: str, name: str, what: str) -> str:
+    """Resolve a simple module-level string constant, including f-strings using prior constants."""
+    values: dict[str, str] = {}
+
+    def _eval(node: ast.AST) -> str | None:
+        if isinstance(node, ast.Constant) and isinstance(node.value, str):
+            return node.value
+        if isinstance(node, ast.JoinedStr):
+            parts: list[str] = []
+            for value in node.values:
+                if isinstance(value, ast.Constant) and isinstance(value.value, str):
+                    parts.append(value.value)
+                elif (
+                    isinstance(value, ast.FormattedValue)
+                    and isinstance(value.value, ast.Name)
+                    and value.value.id in values
+                ):
+                    parts.append(values[value.value.id])
+                else:
+                    return None
+            return "".join(parts)
+        return None
+
+    tree = ast.parse(text)
+    for node in tree.body:
+        if not isinstance(node, ast.Assign):
+            continue
+        value = _eval(node.value)
+        if value is None:
+            continue
+        for target in node.targets:
+            if isinstance(target, ast.Name):
+                values[target.id] = value
+    if name not in values:
+        raise ValueError(f"kernel_fingerprint: could not parse {what}")
+    return values[name]
+
+
 def _pip_stack_specs(dockerfile: str) -> list[str]:
     """The quoted specs from Dockerfile.worker's first `pip install` block (the pinned worker stack)."""
     specs: list[str] = []
@@ -129,12 +168,11 @@ def collect_inputs(
         raise ValueError(
             "kernel_fingerprint: fla spec missing or not pinned to a 40-char commit sha"
         )
-    chalk = _search(r'DEFAULT_CHALK_SPEC\s*=\s*"([^"]+)"', deps, "deps.py DEFAULT_CHALK_SPEC")
+    chalk = _python_string_constant(deps, "DEFAULT_CHALK_SPEC", "deps.py DEFAULT_CHALK_SPEC")
 
     cache_inputs = {
         "from_image": from_image,
         "fla": fla,
-        "liger": _need("liger-kernel"),
         "tilelang": _need("tilelang"),
         "tvm_ffi": _need("apache-tvm-ffi"),
         "chalk": chalk,

{freesolo_flash_dev-0.2.37 → freesolo_flash_dev-0.2.38}/flash/catalog.py

@@ -115,9 +115,16 @@ class ModelInfo:
 
 DEFAULT_MODEL = "Qwen/Qwen3.5-4B"
 
+# The pre-quantized FP8 checkpoint each base model's serving engine LOADS (mirrors serving's
+# ``src.prequant_config``). Every dense model now serves a Freesolo-OWNED FP8_DYNAMIC checkpoint (no
+# community-repo dependence); the 35B VL MoE serves the OFFICIAL Qwen FP8 (it preserves the full
+# vision-language model). Informational for the catalog mirror — deploy gating reads only max_lora_rank.
 SERVING_FP8_MODEL_REPOS: dict[str, str] = {
-    "Qwen/Qwen3.5-4B": "lovedheart/Qwen3.5-4B-FP8",
-    "Qwen/Qwen3.5-9B": "lovedheart/Qwen3.5-9B-FP8",
+    "openbmb/MiniCPM5-1B": "Freesolo-Co/MiniCPM5-1B-FP8",
+    "Qwen/Qwen3.5-0.8B": "Freesolo-Co/Qwen3.5-0.8B-FP8",
+    "Qwen/Qwen3.5-2B": "Freesolo-Co/Qwen3.5-2B-FP8",
+    "Qwen/Qwen3.5-4B": "Freesolo-Co/Qwen3.5-4B-FP8",
+    "Qwen/Qwen3.5-9B": "Freesolo-Co/Qwen3.5-9B-FP8",
     "Qwen/Qwen3.6-35B-A3B": "Qwen/Qwen3.6-35B-A3B-FP8",
 }
 
@@ -133,9 +140,10 @@ MODELS: dict[str, ModelInfo] = {
         recommended_gpu="RTX 4090",
         serving=ServingCapacity(
             gpu="L4",
+            serve_model_id=SERVING_FP8_MODEL_REPOS["openbmb/MiniCPM5-1B"],
             max_loras=16,
-            max_lora_rank=64,
-            max_model_len=32768,
+            max_lora_rank=128,
+            max_model_len=8192,
         ),
         thinking="hybrid",
         notes="On-device class SLM (131k ctx); standard Llama architecture.",
@@ -151,9 +159,10 @@ MODELS: dict[str, ModelInfo] = {
         recommended_gpu="RTX 4090",
         serving=ServingCapacity(
             gpu="L4",
+            serve_model_id=SERVING_FP8_MODEL_REPOS["Qwen/Qwen3.5-0.8B"],
             max_loras=16,
-            max_lora_rank=64,
-            max_model_len=32768,
+            max_lora_rank=128,
+            max_model_len=8192,
         ),
         thinking="hybrid",
         notes="Smallest Qwen3.5; cheap smoke/dev runs with the modern arch.",
@@ -169,9 +178,10 @@ MODELS: dict[str, ModelInfo] = {
         recommended_gpu="RTX 4090",
         serving=ServingCapacity(
             gpu="L4",
+            serve_model_id=SERVING_FP8_MODEL_REPOS["Qwen/Qwen3.5-2B"],
             max_loras=16,
-            max_lora_rank=64,
-            max_model_len=32768,
+            max_lora_rank=128,
+            max_model_len=8192,
         ),
         thinking="hybrid",
     ),
@@ -187,8 +197,8 @@ MODELS: dict[str, ModelInfo] = {
         serving=ServingCapacity(
             gpu="L4",
             serve_model_id=SERVING_FP8_MODEL_REPOS["Qwen/Qwen3.5-4B"],
-            max_loras=64,
-            max_lora_rank=32,
+            max_loras=16,
+            max_lora_rank=64,
             max_model_len=8192,
             max_num_seqs=8,
             gpu_memory_utilization=0.98,
@@ -212,8 +222,8 @@ MODELS: dict[str, ModelInfo] = {
         serving=ServingCapacity(
             gpu="L4",
             serve_model_id=SERVING_FP8_MODEL_REPOS["Qwen/Qwen3.5-9B"],
-            max_loras=44,
-            max_lora_rank=32,
+            max_loras=16,
+            max_lora_rank=64,
             max_model_len=8192,
             max_num_seqs=8,
             gpu_memory_utilization=0.98,
@@ -250,8 +260,11 @@ MODELS: dict[str, ModelInfo] = {
         serving=ServingCapacity(
             gpu="A100-80GB",
             serve_model_id=SERVING_FP8_MODEL_REPOS["Qwen/Qwen3.6-35B-A3B"],
-            max_loras=12,
-            max_lora_rank=32,
+            # rank-64 at only 6 hot slots: the fused-MoE LoRA buffer scales with
+            # max_loras x rank x num_experts, so the A100-80GB ceiling is ~max_loras x rank = 384
+            # (6 x 64 fits at 99.3% util; 16 x 64 OOMs on every single/multi GPU). Serving-validated.
+            max_loras=6,
+            max_lora_rank=64,
             max_model_len=8192,
             max_num_seqs=8,
             max_num_batched_tokens=4096,

{freesolo_flash_dev-0.2.37 → freesolo_flash_dev-0.2.38}/flash/cli/__init__.py

@@ -253,8 +253,7 @@ def _build_parser() -> argparse.ArgumentParser:
     )
     env_pull.add_argument(
         "env_id",
-        help='the Freesolo environment id: a managed slug "your-name/your-env", a '
-        '"github:owner/repo@ref:path" ref, or a github.com URL',
+        help='the managed Freesolo environment slug "your-name/your-env"',
     )
     env_pull.add_argument(
         "path",

{freesolo_flash_dev-0.2.37 → freesolo_flash_dev-0.2.38}/flash/cli/envpush.py

@@ -65,38 +65,32 @@ def _normalize_managed_hub_id(raw: object) -> tuple[str | None, _ManagedHubIdErr
 def cmd_env_pull(args) -> int:
     """Download a published Freesolo environment (or a single file from it) to local disk.
 
-    Managed hub slugs pull package tarballs through the authenticated Flash control plane. Explicit
-    GitHub refs still pull via GitHub's tarball / raw media type, so files larger than 1 MB
-    (e.g. ``datasets/train.jsonl``) come back intact instead of empty.
+    Environments are addressed by their managed hub slug ``namespace/name`` and pulled as package
+    tarballs through the authenticated Flash control plane.
     """
     from flash.client import ClientError, client_from_config
-    from flash.envs.adapter import is_freesolo_environment_id
     from flash.envs.pull import (
-        download_environment_file,
         download_environment_file_from_archive,
         ensure_environment_pull_destination_available,
         environment_local_dirname,
-        pull_environment_package,
         pull_environment_package_from_archive,
     )
 
     env_id = str(args.env_id or "").strip()
-    if not is_freesolo_environment_id(env_id):
+    managed_env_id, managed_error = _normalize_managed_hub_id(env_id)
+    if managed_error == "not-canonical":
+        return _err(
+            f'env id must be lowercase "namespace/name" with no spaces (got {args.env_id!r})'
+        )
+    if managed_env_id is None:
         print(
-            'env id must be a Freesolo environment id — a managed slug "your-name/your-env", '
-            f'a "github:owner/repo@ref:path" ref, or a github.com URL (got {args.env_id!r})',
+            'env id must be a managed Freesolo hub slug "your-name/your-env" '
+            f"(got {args.env_id!r})",
             file=sys.stderr,
         )
         return 1
+    env_id = managed_env_id
     try:
-        managed_env_id, managed_error = _normalize_managed_hub_id(env_id)
-        if managed_error == "not-canonical":
-            return _err(
-                f'env id must be lowercase "namespace/name" with no spaces (got {args.env_id!r})'
-            )
-        managed = managed_env_id is not None
-        if managed:
-            env_id = managed_env_id
         if args.path:
             default_name = Path(args.path.replace("\\", "/")).name
             out = Path(args.output) if args.output else Path(default_name)
@@ -110,11 +104,8 @@ def cmd_env_pull(args) -> int:
             if (out.exists() or out.is_symlink()) and not args.force:
                 print(f"refusing to overwrite {out} (pass --force)", file=sys.stderr)
                 return 1
-            if managed:
-                package = client_from_config().download_env_package(env_id)
-                data = download_environment_file_from_archive(package, args.path)
-            else:
-                data = download_environment_file(env_id, args.path)
+            package = client_from_config().download_env_package(env_id)
+            data = download_environment_file_from_archive(package, args.path)
             out.parent.mkdir(parents=True, exist_ok=True)
             _atomic_write_bytes(out, data)
             if render.styled():
@@ -123,12 +114,9 @@ def cmd_env_pull(args) -> int:
                 print(f"pulled {args.path} from {env_id} -> {out} ({len(data):,} bytes)")
         else:
             out = Path(args.output) if args.output else Path(environment_local_dirname(env_id))
-            if managed:
-                ensure_environment_pull_destination_available(out, overwrite=args.force)
-                package = client_from_config().download_env_package(env_id)
-                pull_environment_package_from_archive(package, out, overwrite=args.force)
-            else:
-                pull_environment_package(env_id, out, overwrite=args.force)
+            ensure_environment_pull_destination_available(out, overwrite=args.force)
+            package = client_from_config().download_env_package(env_id)
+            pull_environment_package_from_archive(package, out, overwrite=args.force)
             if render.styled():
                 print(render.env_pulled(f"{out}/", env_id))
             else:
@@ -146,12 +134,6 @@ def cmd_env_pull(args) -> int:
         return 1
     except (ValueError, FileNotFoundError, RuntimeError, OSError) as exc:
         print(f"env pull failed: {exc}", file=sys.stderr)
-        if "GitHub environment request failed" in str(exc) and not os.environ.get("GITHUB_TOKEN"):
-            print(
-                "hint: private/internal environments need a GitHub token; set GITHUB_TOKEN "
-                "(e.g. `export GITHUB_TOKEN=$(gh auth token)`).",
-                file=sys.stderr,
-            )
         return 1
 
 

{freesolo_flash_dev-0.2.37 → freesolo_flash_dev-0.2.38}/flash/engine/chalk_kernels.py

@@ -1,7 +1,8 @@
 """Optional chalk GPU kernels (the ``freesolo-chalk`` package).
 
-Gap-filling Triton kernels for qwen3.5 applied on top of Liger (liger=False; TRL owns Liger).
-Degrades to a no-op if freesolo-chalk is not installed.
+Standalone Triton kernels for qwen3.5/3.6. Flash does not ask TRL to apply Liger; chalk owns the
+RMSNorm, SwiGLU, fused-linear-CE, RoPE, LoRA-delta, trainable attention epilogue, embedding, and GDN
+patches. Degrades to a no-op if freesolo-chalk is not installed.
 """
 
 from __future__ import annotations
@@ -14,11 +15,16 @@ log = get_logger(__name__)
 
 _KERNELS: list[tuple[str, bool]] = [
     ("rope", True),
+    ("rmsnorm", True),
+    ("swiglu", True),
+    ("fused_linear_cross_entropy", True),
     ("fused_lora_delta", True),
+    ("trainable_attn_epilogue", True),
     ("fused_embedding", True),
-    ("fused_mlp", False),  # off (Liger owns MLP/SwiGLU)
+    ("gdn", True),
+    ("fused_mlp", False),  # off: eval-only bf16 MLP forward, not the training activation path
     ("attn_epilogue", False),  # off (eval-only; needs q/k/v out of LoRA)
-    ("fp8_frozen_base", False),  # off (Hopper sm_90+ only)
+    ("fp8_frozen_base", False),  # off by default: speed/memory tradeoff, enable only after run A/B
 ]
 
 
@@ -31,8 +37,25 @@ def active_kernels(report: Mapping[str, object] | None) -> list[str]:
     )
 
 
+def chalk_supports_fused_ce_model(model_id: str | None) -> bool:
+    """True when chalk's standalone fused-CE patch supports this model family."""
+    mid = (model_id or "").lower()
+    return any(token in mid for token in ("qwen3.5", "qwen3_5", "qwen3.6", "qwen3_6"))
+
+
+def chalk_fused_ce_available(model_id: str | None = None) -> bool:
+    """Best-effort preflight for SFT fused-CE sizing before the trainer exists."""
+    if model_id is not None and not chalk_supports_fused_ce_model(model_id):
+        return False
+    try:
+        from chalk.transformers import apply_chalk_kernel_to_qwen35 as _apply
+    except Exception:
+        return False
+    return callable(_apply)
+
+
 def install_chalk_kernels(model=None) -> dict:
-    """Apply chalk's gap-filling kernels to ``model``; call AFTER TRL builds the trainer.
+    """Apply chalk standalone kernels to ``model``; call AFTER TRL builds the trainer.
 
     Returns chalk's per-kernel report, or ``{}`` when freesolo-chalk isn't installed.
     """
@@ -45,7 +68,7 @@ def install_chalk_kernels(model=None) -> dict:
     except ImportError:
         log.info(
             "freesolo-chalk is not installed on this worker (set FLASH_CHALK_SPEC to an installable "
-            "spec, or check the default PyPI install); chalk kernels off, using eager/Liger."
+            "spec, or check the default install); chalk kernels off, using eager kernels."
         )
         return {}
     except Exception as e:
@@ -53,7 +76,7 @@ def install_chalk_kernels(model=None) -> dict:
         return {}
 
     try:
-        # liger=False: TRL already applied Liger; chalk composes on top of the live Liger modules.
+        # liger=False: flash uses chalk standalone. TRL must not have applied Liger first.
         report = apply_chalk_kernel_to_qwen35(model, liger=False, **kwargs)
     except Exception as e:  # never block training on the optional kernel stack
         log.warning("chalk apply failed (ignored, kernels disabled): %s", e)

{freesolo_flash_dev-0.2.37 → freesolo_flash_dev-0.2.38}/flash/engine/vram.py

@@ -162,7 +162,8 @@ _VOCAB_DEFAULT = 248_320
 _LOGITS_BUDGET_GB = 6.0
 # 16 B/elem: fp32 logits+grad + bf16 logits+grad + CE temp. 8 B/elem under-counts (live OOM confirmed).
 _SFT_LOGITS_BYTES_PER_ELEM = 16.0
-# Single source of truth for Liger gate — engine.worker.perf imports these.
+# Single source of truth for the SFT fused-CE gate. Keep the historical constant names because
+# engine.worker.perf and tests import them.
 _LIGER_MIN_PARAMS_B = 3.0
 _LIGER_LONG_CTX_TOKENS = 2048
 
@@ -264,6 +265,7 @@ def estimate_vram_gb(
     sleep_offload: bool = True,
     active_params_b: float | None = None,
     fp8_kv: bool = False,
+    sft_fused_ce: bool | None = None,
 ) -> float:
     """Estimated peak VRAM (GB) for a LoRA job on one GPU.
 
@@ -292,7 +294,7 @@ def estimate_vram_gb(
         logits = min(completion * vocab * 4 / 1e9, _LOGITS_BUDGET_GB)
         train = activations + logits
         return base + (max(rollout, train) if sleep_offload else rollout + train)
-    fused = sft_logits_fused(params_b, seq_len)
+    fused = sft_logits_fused(params_b, seq_len) if sft_fused_ce is None else bool(sft_fused_ce)
     pd = sft_per_device(batch_size, seq_len=seq_len, vocab=vocab, fused=fused)
     activations = _ACT_COEF * pd * (seq_len / 1024.0) * width
     # Don't clamp to budget: pd=1 is irreducible and the logits can exceed the budget at near-2048 ctx.
@@ -372,7 +374,7 @@ def sft_gc_off_peak_gb(
 ) -> float:
     """Estimated peak VRAM (GB) for a FUSED-CE LoRA SFT step with gradient checkpointing OFF: the
     resident weights + optimizer/base + the no-recompute activations held across ALL ``num_layers``.
-    Fused CE (Liger FLCE) is assumed, so there is no ``[B, T, vocab]`` logits term (the thing that
+    Fused CE (chalk FLCE) is assumed, so there is no ``[B, T, vocab]`` logits term (the thing that
     made GC-off impossible at a 248k vocab). Unknown architecture dims -> ``inf`` (caller keeps GC on).
 
     MoE: the activation backbone scales with the model's real ``hidden`` x ``num_layers`` (geometry),
@@ -482,14 +484,22 @@ def model_required_vram_gb(
             use_vllm=use_vllm,
             vocab=vocab,
             active_params_b=active_params_b,
+            sft_fused_ce=sft_fused_ce,
         )
         return math.ceil(est * headroom)
 
     from flash.catalog import MODELS, vocab_size_for
+    from flash.engine.chalk_kernels import chalk_supports_fused_ce_model
 
     info = MODELS.get(model_id)
     model_vocab = vocab_size_for(model_id)
     is_grpo = (algorithm or "").lower() in ("grpo", "rl")
+    sft_fused_ce = (
+        None
+        if is_grpo
+        else sft_logits_fused(resolve_params_b(model_id), seq_len)
+        and chalk_supports_fused_ce_model(model_id)
+    )
     if info is not None:
         params_b = info.params_b  # curated, authoritative (required field) — no string parsing
         quant = getattr(info, "quant", "bf16") or "bf16"