unsloth/studio/backend/core/export
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
..
__init__.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
export.py Harden model fetching (#6391) 2026-06-18 05:39:52 -07:00
orchestrator.py Studio: persistent per-user trust_remote_code approval cache (#6551) 2026-06-22 05:12:49 -07:00
worker.py Generalize transformers tier selection by probing AutoConfig (#6550) 2026-06-22 08:20:06 -07:00