3 commits

Author	SHA1	Message	Date
Daniel Han	7729c37fd0	Studio: restore Qwen causal-conv1d on demand, narrow fatal mamba-ssm, provision SSM kernels in export (#13971) ... Follow up to #13861 and #13808. An audit of the merged on demand SSM kernel change plus a T4 benchmark found a few gaps, fixed here. ## Why - On a Kaggle T4 (sm75, torch 2.11 cu130, fp16, Qwen3.5-0.8B), Studio chat decode was 4.4% slower without causal-conv1d (19.10 vs 19.95 tok/s, 6 alternating repeats, slower in all 9 pairs). Training speed, losses and peak memory were unchanged. The zoo Triton conv covers `causal_conv1d_fn` only, so single token decode fell back to the torch update path. Training with trainable conv weights also allocates an fp32 dW scratch in the Triton backward, and packed `seq_idx` is not accepted by the zoo dispatcher under transformers 5.20. - Mamba-family models that have a working transformers torch path (mamba, mamba2, falcon_mamba, jamba, zamba, zamba2, bamba) were made fatal when mamba-ssm could not be installed. - Export, local LoRA continuation, and inference model switching did not provision or pick up kernels correctly once setup stopped installing them. - Setup on pre-Ampere GPUs (T4) downloaded flash-attn even though Unsloth only enables FA2 at sm80+. ## Changes 1. Qwen3.5 / Qwen3.6 / Qwen3-Next request causal-conv1d on demand again in training, inference and Clef (same as before #13861, minus the setup time install). The zoo Triton platform probe is removed. 2. A missing mamba-ssm is fatal only for `nemotron_h`, `falcon_h1`, `granitemoehybrid`. Other Mamba families install best effort and continue on the torch path with a status message. 3. Local LoRA folders (`adapter_config.json`, no `config.json`) resolve kernel needs from `base_model_name_or_path`; the training `adapter_base_pin` wins when set. 4. The export worker provisions SSM kernels before importing Unsloth / the export backend, with the same required vs optional split. 5. Inference re-enables the modeling source scan by reading `transformers/models/<type>/modeling_*.py` as text, without importing transformers. 6. If a kernel is newly installed after the model's `transformers.models.<type>` module was already imported, the inference worker asks for a restart and the orchestrator respawns it once. 7. Torch repair removes causal-conv1d / mamba-ssm whose compiled extension no longer imports (never installs anything new). A hung import check keeps the package. 8. Setup and the long context runtime install skip flash-attn below sm80 (and remove a stale copy at setup). No visible GPU and ROCm keep current behavior. ## Tests New and updated tests fail on main and pass here (`test_ssm_runtime.py`, `test_ssm_restart_respawn.py`, `test_export_ssm_kernels.py`, `test_clef_fast_paths.py`, `test_kernel_install.py`, `test_training_worker_flash_attn.py`, `test_ssm_kernels_install_python_stack.py`, `test_flash_attn_install_python_stack.py`, `test_windows_torch_flavor_invariant.py`). With `test_install_python_stack.py` and `test_cuda_repair.py`: 966 passed after merging current main. ## Hardening (second commit) - Stale SSM kernel removal after torch repair needs proof: torch must import first, and only an ImportError / OSError from the extension itself removes the package. Signals, hangs, launch errors and a broken torch keep it. - `unsloth install-kernels --without-gpu` (or `UNSLOTH_INSTALL_KERNELS_WITHOUT_GPU=1`) installs GPU gated kernels on a GPU-less image builder. An already working kernel counts as success before any skip, skipped kernels are listed, and the exit code is unchanged (0) unless `--strict` is passed. The CLI respects `UV_OFFLINE` / HF offline variables and never falls back from uv to pip offline. - Inference respawn refuses to spawn while the old worker survives shutdown. SSM restart and the Xet fallback each get their own one-use retry (3 spawns max). - Restart detection only counts native `modeling_*` modules and the selected checkpoint's remote code modules from its `auto_map`, matched on the exact `transformers_modules` namespace. - Export provisions SSM kernels inside the existing transformers activation window (no extra Hub reachability probe), and skips provisioning on macOS and for the MLX / CPU export backends. - `UV_OFFLINE` counts as offline for on-demand installs (no URL probes, source builds, or uv to pip fallback). - Long context flash-attn install returns early when flash_attn already imports, reuses the worker's own torch for the capability check, and never applies the sm80 floor on ROCm. ## Install robustness (platform matrix review) - A wheel is downloaded whenever its HEAD check is not a definite 404 (proxies that refuse HEAD no longer push required hybrids into a failure), then import-verified; fallback only after a real failure. - A failed install is remembered for the worker process (keyed by kernel, torch, toolchain), so preflight and load share one attempt. Optional Mamba families get mamba-ssm only as a prebuilt wheel, never a source build. ROCm source builds get an 1800 s deadline and CUDA source builds 7200 s (`UNSLOTH_SOURCE_BUILD_TIMEOUT` overrides); a timeout kills the whole build tree. - A runtime install whose extension does not import is uninstalled, so transformers never finds a broken package. - One cross-process lock per kernel and environment covers check, install, verify and cleanup (training and runtime), with a recheck after acquiring it. - Only a CUDA or ROCm torch provisions anything; CPU, XPU, MPS and unknown torch builds skip (export included). - Optional mamba-ssm follows setup's sm80 floor (sm75 with Triton 3.4+); the required hybrids keep the previous behavior. - The resolver only returns URLs for cells that are actually published (Unsloth release: torch 2.13/2.14 cp313, win_amd64 torch 2.10 to 2.12 cp311 to cp313; upstream causal-conv1d / mamba-ssm release cells), and there is no source build of the pinned 1.6.1 / 2.3.1 sources on torch 2.13+, which need C++20. - Training stays best effort for mamba-ssm, as before #13861. - Export runs the malware and remote code consent gates before any kernel provisioning. - The extension is probed again after the install lock is acquired, so a concurrent replacement is not treated as broken. - The automatic HF offline window (Hub unreachable) no longer blocks package installs; only user set offline variables do. - `unsloth install-kernels --dry-run` prints the resolved URL while offline. - A Mamba model whose code needs mamba-ssm fails with a clear error on a torch build without CUDA or ROCm. ## GPU checks (merge base vs head 4a6429a65) Kaggle T4, torch 2.11.0+cu130, Qwen3.5-0.8B fp16: | | base | head | |---|---|---| | causal-conv1d on demand | not installed | installed, `causal_conv1d_update_kernel` engaged | | decode tok/s (6 alternating repeats) | 18.70 | 19.41 (+3.8%, 6/6 pairs) | | decode peak alloc | 1.748 GB | 1.748 GB | | LoRA step time (3 repeats) | 1975.5 ms | 1970.7 ms (inside spread) | | LoRA peak alloc | 2.338 GB | 2.338 GB | | flash-attn at setup on sm75 | downloaded | skipped | AMD gfx1151 (torch 2.11.0+rocm7.13.0): Qwen3.5 LoRA losses within 1.2e-4 relative (head vs merge commit bit identical), peak 2.389 GB both, backend suites 262/262 and 275/275 passing. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>	2026-10-19 04:32:13 -07:00
Daniel Han	90ac3b7734	Studio: install causal-conv1d and mamba-ssm only when a model needs them (#13861) ... Studio setup stops downloading causal-conv1d and mamba-ssm. They are installed when a model that needs them is first trained or loaded, and Qwen3.5 / Qwen3.6 / Qwen3-Next no longer request causal-conv1d at all. ## Why #13396 added both kernels to every NVIDIA Linux / Windows setup next to flash-attn: about 555 MB of downloads (~620 MB installed) for every user, while only a few model families use them. Studio already installs them on demand when such a model is trained (`core/training/worker.py`) or loaded for chat (`core/inference/worker.py` through `utils/ssm_runtime.ensure_ssm_runtime`), so the setup step only duplicated that. Qwen3.5 / 3.6 / Qwen3-Next do not need causal-conv1d with Unsloth: unsloth-zoo runs their GatedDeltaNet with its vendored FLA kernels and its own Triton causal conv (`unsloth_zoo/temporary_patches/gdn_causal_conv1d.py`, model types `qwen3_5`, `qwen3_5_moe`, `qwen3_next`). Studio still downloaded the CUDA package for them because their transformers modeling files mention `causal_conv1d`. ## What changes - `studio/install_python_stack.py`: the setup-time SSM install is gone from all three places (initial setup, Linux torch repair, Windows torch flavor repair), with its progress step. flash-attn stays where it was. `unsloth install-kernels` still installs every kernel, for offline images. - `studio/backend/utils/ssm_runtime.py`: `ZOO_GDN_MODEL_TYPES` (equal to zoo's list; a test reads zoo's source to keep them equal). Names and configs of those models, including `_text` / `_vision` sub-configs, resolve to no causal-conv1d request, checked before the modeling-source scan. A config that mixes in another model type that needs it still asks for it. Zoo's own kill switch `UNSLOTH_DISABLE_TRITON_CAUSAL_CONV1D=1` brings the old request back. - `core/training/worker.py` and `core/systemone/clef_runtime.py` use the same predicate, so the eager install, the `is_causal_conv1d_available` hook and Clef (a Qwen3.5 backbone) all agree. - Only where zoo's Triton conv can run: a CUDA torch (not ROCm, not CPU), Triton importable, not macOS. This is read from torch's distribution version (`+rocm` / `+cpu` / `+cu`), so no load imports torch to decide. On ROCm, macOS or Windows without triton-windows, Qwen keeps getting causal-conv1d as before. - With setup no longer covering for it, the on-demand install now decides from the model's config, not only its name: `resolved_model_is_ssm` maps `nemotron_h`, `falcon_h1`, `granitemoehybrid`, `bamba`, `jamba`, `zamba`, `zamba2`, `mamba`, `mamba2`, `falcon_mamba` (and their `_text` sub-configs) to mamba-ssm + causal-conv1d, for training and for the inference load of the resolved base. Previously `NVIDIA-Nemotron-Nano-9B-v2`, `granite-4.0-tiny-preview`, renamed local checkpoints, Jamba, Zamba and Bamba never got mamba-ssm on demand. - After an install inside an inference worker that already imported transformers, the cached `is_causal_conv1d_available` / `is_mamba_ssm_available` / `is_mamba_2_ssm_available` and cached hub-kernel misses are cleared, so switching from Qwen3.5 to Nemotron-H in one session gets the fast path. - The mamba-ssm error raised when it cannot be installed (e.g. offline) says to run `unsloth install-kernels` once while online. - `install_kernel(install_without_gpu=)` is gone; setup was its only caller. - Clef still applies the FLA TileLang guard when it skips the install. Unchanged: LFM2, Mamba / Jamba / Zamba / Bamba still get causal-conv1d on demand; Nemotron-H, Falcon-H1 and Granite-4.0-H still get causal-conv1d then mamba-ssm on demand; nothing already installed is removed, and a usable causal_conv1d still wins inside zoo. ## Tests - `tests/python/test_ssm_kernels_install_python_stack.py`: setup has no SSM step; flash-attn is still installed at all three sites; the CLI kernels are still there. - `test_windows_torch_flavor_invariant.py`: the torch progress totals drop by one on Windows and Linux. - `studio/backend/tests/test_ssm_runtime.py`: Qwen3.5 / 3.6 / Next names and configs install nothing, the kill switch restores the request, a mixed config still asks, the Mamba2 hybrids still install both kernels in order, and the model-type list equals zoo's. - 21 of the new cases fail against main, and the follow-up cases (config-based SSM detection, platform check, cache refresh, offline message) fail against this PR's first commit; all pass here. The touched test files and the setup / resolver suites pass (test_ssm_runtime 84, test_training_worker_flash_attn 67, test_install_python_stack 295, test_cuda_repair 284, test_windows_torch_flavor_invariant 54, among others). ## Accuracy and speed without causal-conv1d (Qwen3.5) Tiny Qwen3.5 / Qwen3-Next through Unsloth, bf16: zoo's Triton conv is bit-identical to torch's reference conv (logits, loss, LoRA grads, 5 optimizer steps); the CUDA package differs from both at bf16 rounding (5.9e-3 logits, grads cos 0.99993). Prefill plus cached token-by-token decode matches a full forward equally with and without the package (4.0e-3 vs 3.9e-3). The profiler shows which kernels ran in every arm. Qwen3.5-0.8B bf16 on a Colab L4, torch 2.14.1+cu130, 3 alternating repeats per arm: | | Without causal-conv1d (zoo Triton) | With causal-conv1d (CUDA) | |---|---|---| | LoRA training, batch 4 x 1024, 30 timed steps | 911.5 ms/step (spread 11.2) | 917.5 ms/step (spread 15.9) | | Greedy decode, 256 new tokens, batch 1 | 19.3 tok/s (spread 0.6) | 19.0 tok/s (spread 0.2) | Both differences are inside the run-to-run spread. Without the package, single-token decode uses torch's `conv_depthwise2d`; with it, `causal_conv1d_update`. A tiny NemotronH also runs identically with a pure-Python mamba-ssm (the Mamba2 path never calls `selective_scan_cuda`), and a merged 16-bit export of a NemotronH LoRA works with neither kernel installed. ## Not covered here - `test_decision_native.py` / `test_decision_serving.py` were not collected locally (missing `fastmcp`).	2026-10-15 08:20:03 -07:00
Daniel Han	9f4a868d43	Studio setup: install causal-conv1d and mamba-ssm prebuilt wheels next to flash-attn (#13396) ... * Studio setup: install causal-conv1d and mamba-ssm prebuilt wheels next to flash-attn Setup already installed flash-attn from prebuilt-wheels-cu13 but left causal-conv1d and mamba-ssm to the first Mamba / hybrid model load. Install both at setup through the same wheel-only install_kernel the CLI and SSM runtime use: never a source build, never a torch change, skipped where no wheel is published, GPU-gated for mamba-ssm, and resynced wherever flash-attn is after a torch repair. UNSLOTH_STUDIO_SKIP_SSM_KERNELS_INSTALL=1 opts out. * SSM kernel setup step: local check first, bounded probes, fresh runner, no-GPU install, stale removal Skip the wheel lookup when the resident kernel already loads, so an offline re-run does not wait on two HEAD requests. Each call gets its own runner, which bounds the import probes and resets install_kernel's per-runner GPU and Triton cache after a torch repair. Setup installs mamba-ssm where no GPU is visible yet (an image build); a visible pre-Ampere GPU still skips it. A resident build that no longer loads and has nothing replacing it is removed, as flash-attn's resync does. * SSM kernel setup step: honour UV_OFFLINE No wheel lookup and no install when uv is told to stay offline: the pip fallback in install_wheel does not read UV_OFFLINE. A kernel that already loads is still reported. * SSM kernel setup step: offline still removes a stale build; a probe that cannot launch is inconclusive * SSM kernel setup step: bound the torch probe on every platform and skip on a launch failure * SSM kernel setup step: an uninstall that cannot launch warns instead of aborting * Trim comments * SSM kernel setup step: redact captured install output before printing it install_kernel only replaces the wheel URL; uv / pip failure text or an exception can name another index carrying credentials. * Tests: pin the no-GPU mamba-ssm case to Linux so Windows' Triton gate does not answer first	2026-10-12 06:23:20 -07:00