Commit graph

3 commits

Author SHA1 Message Date
Daniel Han
ce0323263e
Fix test isolation: restore sys.modules after the pre-import gate test (#6578)
* Restore sys.modules in test_pre_import_gate_is_transformers_free

The test pops transformers and utils.models.model_config from sys.modules to
assert the pre-import security gate does not re-import them, but never put them
back. A later importer then rebound a fresh utils.models.model_config, so tests
that had captured the original instance missed their patches and hit the real
path: test_vision_cache patches _is_vision_model_uncached on the original
module, but is_vision_model (still bound to that original) ran the real network
lookup instead. This produced 17 spurious failures whenever test_ssm_runtime
ran before test_vision_cache in the same process.

Snapshot the removed modules and restore the original objects in a finally, so
the assertions still run against a clean slate while later tests see the same
module instances they captured at import time.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-22 08:45:34 -07:00
Daniel Han
007a21235c
Generalize transformers tier selection by probing AutoConfig (#6550)
* Resolve the transformers tier by probing AutoConfig instead of guessing

When the only signal is a 5.x tokenizer class, get_transformers_tier guessed the
lowest 5.x sidecar (530). That misroutes models whose built-in config parser needs
a higher tier: dense NemotronH ships a 5.x tokenizer but its '-' (MLP) layer only
transformers 5.10 can parse, so 5.3/5.5 raise KeyError '-'. The config.json
transformers_version field records the saving version, not the minimum to load, so
it cannot drive routing either.

Replace the weak tokenizer->530 guesses (local and remote) with a probe: parse
config.json with the built-in parser (trust_remote_code=False) in each sidecar,
escalating 530->550->510, and pick the first that succeeds. This generalizes to any
architecture without hardcoded lists. Strong signals stay fast paths (no subprocess);
the probe runs only when the tier is otherwise ambiguous and is cached by (model,
commit sha). It never executes repo code, never downloads weights, never raises, and
falls back to the legacy 530 guess on a transient/auth/offline failure or when no
sidecar is available. UNSLOTH_DISABLE_TIER_PROBE restores the old behavior.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address review: tier probe fallbacks and cross-platform robustness

Codex:
- Never escalate to 510 on uncertainty. When every sidecar was probed and none
  parsed with the built-in parser, the model is a remote-code / custom model_type
  that loads via its own code; keep the legacy 530 route instead of jumping to
  510 (which would change the behavior of models that worked on the 5.3 stack).
- Only cache the 530 fallback when the result is conclusive (every tier actually
  probed). If a sidecar was missing/uninstallable the environment is incomplete,
  so return 530 uncached and retry on the next call.
- Do not pin the tier cache under an unknown revision: _resolve_commit_sha no
  longer memoizes a None sha (a transient Hub failure is retried), and _probe_tier
  only caches a tier when the commit sha is known.

Gemini:
- Wrap Path.exists() in the sha resolver in try/except OSError (a remote repo id
  can raise WinError 123 on Windows).
- Probe script writes the error to sys.stderr.buffer as UTF-8 bytes so a non-ASCII
  message cannot itself raise UnicodeEncodeError under cp1252.
- subprocess.run decodes stderr with errors="replace" to avoid UnicodeDecodeError
  on non-UTF-8 consoles.

Tests: 72 passed (added partial-sidecar uncached, sha-unresolved not cached,
all-failed stays 530 + cached, sha resolver retries None / handles OSError).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address review round 2: authenticate tier checks, stop memoizing local sigs

Codex:
- Thread hf_token through _check_config_needs_510/550 and
  _check_tokenizer_config_needs_v5 (and the underlying raw fetches). Previously a
  gated/private model whose only 5.x signal is tokenizer_config.json never reached
  the authenticated probe: the unauthenticated raw fetch failed and cached False,
  so the model fell through to the default 4.x tier. The per-check caches are now
  keyed by (model, token) so an unauthenticated miss cannot poison a later authed
  read, mirroring _load_config_json.
- _resolve_commit_sha no longer memoizes a local directory signature. A local
  signature is mutable (size/mtime of config/tokenizer), so a reused/overwritten
  checkpoint path would otherwise keep selecting the previous tier; it is now
  recomputed every call. Only the immutable remote commit sha is memoized.

Tests: 75 passed (added token-cache isolation + auth header, local signature not
memoized, token threaded into all checks/probe).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address review round 3: reach activation with the token, drop SHA tier cache

Codex round 3:
- Thread hf_token into the activation path that actually selects a sidecar. The
  token-aware tier checks added last round were unreachable:
  activate_transformers_for_subprocess called get_transformers_tier without a
  token, and the inference/training/export workers passed only the model name even
  though they hold a request-scoped hf_token. activate_transformers_for_subprocess
  now takes hf_token and the three workers forward config["hf_token"], so a
  gated/private model whose only 5.x signal is an authenticated config/tokenizer is
  routed to the right sidecar instead of falling to default 4.x.
- Stop importing huggingface_hub during tier detection. _probe_tier no longer
  resolves a commit sha, so it never pulls huggingface_hub into the worker before
  the sidecar venv is prepended to sys.path (activation only prepends, never
  purges), which would otherwise pin the default-env hub over the sidecar's
  pinned huggingface_hub==1.8.0.
- The tier cache is now keyed by model_name for the process lifetime (a model's
  required tier is a property of its architecture; cleared on restart). This drops
  the mutable-SHA memo that masked remote revision changes and the mutable
  local-signature memo, removing _resolve_commit_sha / _local_dir_signature /
  _probe_sha_cache entirely.
- Do not cache a probe success that depended on a skipped lower tier: if a lower
  sidecar was unavailable, the lowest valid tier may change once it installs, so
  the result is returned uncached and re-probed next call.

Tests: 73 passed (probe imports no hub; success uncached when a lower tier is
skipped; activation forwards the token).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Trim comments to be more succinct

* Re-probe overwritten local checkpoints and authenticate the probe child

The AutoConfig tier probe cached its result under the bare model_name, so a
local checkpoint overwritten in place (same path, new config.json) kept serving
the stale sidecar. Fold a cheap config.json signature (size + mtime) into the
cache key for local paths; remote ids stay name-keyed so no huggingface_hub
import lands before the sidecar is activated.

The probe relies on the implicit HF_TOKEN env, so an inherited
HF_HUB_DISABLE_IMPLICIT_TOKEN=1 left it unauthenticated and a gated repo 401ed
into the 530 fail-safe. Clear that flag in the child env when a token is set.

* Keep tier probes off the log-only path and probe new 5.x archs default-first

- get_transformers_tier gains probe=True/False. needs_transformers_5 (a coarse
  4-vs-5 boolean used only for a spawn log and a vision-check branch) now passes
  probe=False, so a parent/log-only caller never spawns sidecar probes. The real
  activation path keeps probe=True and resolves the exact tier in the worker.
- A config.json saved by transformers 5.x but matched by no fast path is now probed
  default-first: _probe_tier gains include_default + floor, prepending the ambient
  4.57.x tier to the escalation. A model that still parses on the default is left on
  it (no mis-route onto a sidecar); only a config the default parser cannot read
  escalates to the lowest 5.x tier that parses. The transformers_version field is a
  cheap 'worth probing' hint only, read from the already-fetched config (no extra
  network); ordinary 4.x configs never probe.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Separate probe cache by mode and keep version-field 5.x visible to needs_transformers_5

- _probe_tier cache was keyed only by config.json signature, so a default-first probe
  that returned 'default' could be handed back to a later tokenizer/known-5.x caller
  (floor=530), leaving a model with a 5.x-only tokenizer on transformers 4.x. Key the
  cache by probe mode (floor + include_default); the legacy 530 mode keeps the bare key.
- The version-field 5.x detection is a cheap config read, not a probe, so run it even
  when probe=False: a standard-tokenizer model whose only signal is transformers_version
  >= 5 now classifies as 5.x via needs_transformers_5 (returns '530' without spawning a
  probe), so the vision-routing fallback uses the 5.x subprocess instead of failing the
  default parser and marking it non-vision. The real activation path still probes
  default-first and may resolve 'default'.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Don't treat local checkpoints as Hub ids, and fix stale activation test double

- _load_config_json / _check_tokenizer_config_needs_v5: a local checkpoint dir whose
  config.json / tokenizer_config.json is not yet present was being fetched from the Hub
  as if the path were a repo id, and the 404 miss was cached. A later call after the
  file is written (in-progress checkpoint) then served the stale miss, so a
  TokenizersBackend checkpoint fell through to the default tier. Skip the Hub fetch for
  local dirs and do not cache the miss, so the file is read once it appears.
- test_activate_transformers_version_or_warn_*: the worker now threads hf_token into
  _activate_transformers_version (model_name, hf_token); update the one-arg test doubles
  to the real two-arg signature so the silent-success path stays silent.

* Tighten comments in the AutoConfig probe and tier-selection paths

* Address review: canonical probe cache key and reuse _token_cache_key

- _probe_cache_key resolves config.json to its absolute realpath before
  keying, so a relative path or a changed cwd can't collide with or miss a
  prior probe result. Remote ids still fall back to the name (stat raises,
  caught).
- _cached_config_json reuses _token_cache_key instead of re-hashing the
  token inline, keeping the (model, token) key derivation in one place.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-22 08:20:06 -07:00
Daniel Han
2a05426adb
Auto-install SSM kernels (causal-conv1d, mamba-ssm) for inference loads (#6535)
* Auto-install SSM kernels (causal-conv1d, mamba-ssm) for inference loads

Mamba/SSM hybrids (Nemotron-H/Nano, Falcon-H1, Granite-4.0-H, ...) lazily import
mamba_ssm / causal_conv1d during from_pretrained, so loading them for chat failed
with 'mamba-ssm is required by the Mamba model but cannot be imported'. The training
worker already wheel-first installs these before a fine-tune; the inference worker
did not. Add utils/ssm_runtime.ensure_ssm_runtime and call it from the inference load
path so the same models load for inference. Training worker is untouched; a drift
test keeps the shared detection and pinned versions in lockstep.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* ssm_runtime: invalidate import caches, skip MLX, cover LoRA base

- Invalidate importlib finder caches in _is_importable and after a successful
  wheel install, so a kernel installed earlier in this same process is actually
  importable when the modeling code lazy-imports it during from_pretrained.
- Skip the SSM kernel install entirely on the MLX (Apple Silicon) load path:
  these are CUDA/ROCm Torch kernels with no MLX use and no macOS prebuilt wheel,
  so the source build would fail before the MLX backend loads the model.
- For LoRA loads, also run detection over the resolved base model, since an
  adapter id like 'me/my-lora' won't match the SSM heuristics but its SSM base
  (Nemotron-H, ...) is what needs the kernels.

Adds tests for cache invalidation and the MLX-skip / LoRA-base worker wiring.

* Tighten SSM autoinstall comments

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* ssm_runtime: verify wheel imports, HIP-aware source build, build heartbeat

Address review feedback:
- Verify a prebuilt wheel actually imports before trusting it; a CUDA/ABI-mismatched
  wheel now falls back to a source build instead of returning success and failing later
  with the cryptic lazy-import error.
- HIP-aware source build: require hipcc on ROCm, inject clang --gcc-install-dir, and use
  the 1800s timeout, mirroring the training worker (ROCm has no prebuilt wheel).
- Emit a status heartbeat every 60s during the source build so a long (ROCm) build does
  not trip the orchestrator's 300s inactivity timeout.

Tests cover the wheel-not-importable fallback and the missing-hipcc ROCm bail.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Make causal-conv1d best-effort and harden the SSM source build

- causal-conv1d is a fast path: models that merely want it (Qwen3-Next, LFM2)
  fall back to torch, so a failed install must not reject an otherwise loadable
  chat model on Windows/CPU/macOS or an ABI without a wheel. Only a true SSM
  model's mamba-ssm requirement stays fatal, matching the training worker which
  treats causal-conv1d as best-effort.
- The source build is reached only when not importable, including a wheel that
  installed but failed to import; add --reinstall/--force-reinstall so it
  replaces the broken install instead of no-opping as already satisfied.
- Add --no-cache to the ROCm uv source build to avoid reusing stale artifacts
  from a partial HIP build, mirroring the training worker.

* Address review: install SSM kernels before transformers, harden import + Windows

Codex:
- Install the SSM kernels before importing transformers. run_inference_process
  imported core.inference.inference (which imports unsloth/transformers) before the
  load, and a sidecar transformers can evaluate its optional-backend gates against
  the import state; installing causal_conv1d/mamba_ssm afterwards left those gates
  unsatisfied and a Nemotron/Falcon/Granite load still failed with "mamba-ssm is
  required". The initial model's kernels are now installed in run_inference_process
  before the ML import, via a shared _ensure_ssm_kernels helper; _handle_load keeps
  calling it (idempotent) for a LoRA's base and for later in-process loads.
- _is_importable now treats any import failure as "not importable", not only
  ImportError. An ABI-incompatible native kernel (undefined symbol after a torch/CUDA
  upgrade) raises OSError/RuntimeError; letting those escape reported
  ssm_runtime_install_failed instead of falling back to reinstall/source build.
- Skip causal-conv1d on Windows (no prebuilt wheel), mirroring the training worker.
  A causal-conv1d-only model (Qwen3-Next/LFM2) no longer drops a chat load into a
  multi-minute untimed source build; it uses the torch fallback. mamba-ssm is still
  attempted for true SSM hybrids.

Tests: test_ssm_runtime.py +5 (broken-kernel exceptions read as not-importable;
causal-conv1d skipped on win32 while mamba-ssm still installs). 36 passed.

* Trim comments to be more succinct

* Run security gates before installing SSM kernels

The SSM kernel auto-install is name-based (model_is_ssm is a substring match, no
config fetch), so a model id merely containing an SSM substring triggered a
native-package install (possibly a slow source build) before the malware and
remote-code consent gates ran. Extract those gates into _run_security_gates and
call it before the kernel install in both the pre-import path of
run_inference_process and in _handle_load, so a blocked or nonexistent model is
refused before any build. The gates are metadata-only and do not import
transformers, so they are safe to run before the pre-import install.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Resolve remote LoRA bases before importing transformers

_resolve_base_model only reads a local adapter_config.json, so a remote LoRA
adapter whose own id has no SSM substring but whose base is a Nemotron/Falcon/
Granite model had its base discovered only by ModelConfig in _handle_load, after
transformers was imported and its optional-backend availability snapshotted, so
the SSM kernel install there was too late. Add _remote_lora_base, a metadata-only
adapter_config.json fetch (no huggingface_hub / transformers import), and use it
in the pre-import path so the base is gated and its kernels pre-installed.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Gate only loaded roots, tier on the resolved base, read offline LoRA cache

Three follow-ups to the pre-import resolution:

- The security gate reused the SSM target list, which for a local full fine-tune
  includes the config.json-recorded base. That base is never loaded, so scanning
  it could falsely block a safe local checkpoint. Gate only the model plus a
  genuine LoRA base (matching _handle_load's mc.is_lora), separate from the
  broader SSM-install list.

- Tier activation ran on the raw adapter id, so a remote LoRA whose base needs a
  sidecar transformers version imported the default and failed. Resolve the base
  once up front and activate on it.

- _remote_lora_base bailed on offline before checking the hub cache, missing a
  cached adapter's base. Read the cached adapter_config.json when offline or when
  the fetch fails.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Keep the pre-import gate transformers-free; harden remote LoRA resolution

The pre-import security gate called security_load_subdirs, which imports
model_config and thus transformers, snapshotting optional-backend availability
before the SSM kernels are installed and defeating the ordering. Add
compute_subdirs to _run_security_gates and pass False in the preflight so it scans
from the root only (transformers-free); _handle_load still runs the authoritative
gate with full subdir scoping after the import.

_remote_lora_base now skips existing local relative paths (is_local_path) so a
checkpoint like outputs/run1 is never treated as a Hub repo, and distinguishes a
definitive 404 (not a LoRA -> None) from transient/offline failures (read the
cache), so a repo that is now a full model no longer resolves a stale cached base.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Probe a real model id for SSM kernels; respect HF_ENDPOINT

model_is_ssm is a substring match, so an arbitrary name could false-match and
force a mamba-ssm install that fails the load for a non-SSM model:
- a LoRA adapter id like user/falcon-h1-lora (the SSM-relevant code is the base's);
- a local checkpoint under an SSM-named parent dir, e.g. /runs/falcon-h1/llama-ckpt.

Add ssm_probe_identifier, which resolves the base (or a bare local checkpoint's
basename) and feed that to ensure_ssm_runtime from both the pre-import path and
_handle_load, so detection runs against a real model id, never an adapter id or
parent folders.

_remote_lora_base now honors HF_ENDPOINT so enterprise/mirror deployments resolve
the adapter base instead of always hitting huggingface.co.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Tighten comments in the pre-import SSM gate/install path

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <michaelhan2050@gmail.com>
2026-06-22 04:48:29 -07:00