ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

use config.json model_type for tier detection, add unsloth/nvidia namespace guard

Browse source
This commit is contained in:
Roland Tannous 2026-04-06 19:40:11 +00:00
commit 567df493ab
4 changed files with 99 additions and 99 deletions
Download patch file Download diff file Expand all files Collapse all files

6
studio/backend/core/inference/worker.py
View file

@ -328,7 +328,11 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
trust_remote_code = config.get("trust_remote_code", False)
if not trust_remote_code:
model_name = config["model_name"]
if "nemotron" in model_name.lower():
_mn_lower = model_name.lower()
if (
"nemotron" in _mn_lower
and (_mn_lower.startswith("unsloth/") or _mn_lower.startswith("nvidia/"))
):
trust_remote_code = True
logger.info(
"Auto-enabled trust_remote_code for Nemotron model: %s",

6
studio/backend/core/training/worker.py
View file

@ -405,7 +405,11 @@ def run_training_process(
# (Qwen3.5, Gemma 4, etc.) are native and do NOT need it — enabling it
# bypasses the compiler (disabling fused CE).
_lowered = model_name.lower()
if "nemotron" in _lowered and not config.get("trust_remote_code", False):
if (
"nemotron" in _lowered
and (_lowered.startswith("unsloth/") or _lowered.startswith("nvidia/"))
and not config.get("trust_remote_code", False)
):
config["trust_remote_code"] = True
logger.info(
"Auto-enabled trust_remote_code for Nemotron model: %s",

25
studio/backend/tests/test_transformers_version.py
View file

@ -32,8 +32,9 @@ from utils.transformers_version import (
_resolve_base_model,
_check_tokenizer_config_needs_v5,
_check_config_needs_550,
_get_config_json,
_tokenizer_class_cache,
_config_needs_550_cache,
_config_json_cache,
needs_transformers_5,
get_transformers_tier,
)
@ -202,7 +203,7 @@ class TestCheckConfigNeeds550:
"""Tests for _check_config_needs_550() local config.json checks."""
def setup_method(self):
_config_needs_550_cache.clear()
_config_json_cache.clear()
def test_gemma4_architecture(self, tmp_path: Path):
"""config.json with Gemma4ForConditionalGeneration should return True."""
@ -242,8 +243,8 @@ class TestCheckConfigNeeds550:
key = str(tmp_path)
_check_config_needs_550(key)
assert key in _config_needs_550_cache
assert _config_needs_550_cache[key] is True
assert key in _config_json_cache
assert _config_json_cache[key] is not None
def test_local_file_skips_network(self, tmp_path: Path):
"""When local config.json exists, no network request should be made."""
@ -265,7 +266,7 @@ class TestGetTransformersTier:
def setup_method(self):
_tokenizer_class_cache.clear()
_config_needs_550_cache.clear()
_config_json_cache.clear()
def test_gemma4_substring_returns_550(self):
assert get_transformers_tier("google/gemma-4-E2B-it") == "550"
@ -317,6 +318,20 @@ class TestGetTransformersTier:
# This shouldn't happen in practice, but verifies priority
assert get_transformers_tier("gemma-4-model") == "550"
def test_config_json_model_type_530(self, tmp_path: Path):
"""Local checkpoint with qwen3_moe model_type → 530."""
cfg = {"model_type": "qwen3_moe", "architectures": ["Qwen3MoeForCausalLM"]}
(tmp_path / "config.json").write_text(json.dumps(cfg))
assert get_transformers_tier(str(tmp_path)) == "530"
def test_config_json_model_type_glm4_moe(self, tmp_path: Path):
"""Local checkpoint with glm4_moe model_type → 530."""
cfg = {"model_type": "glm4_moe", "architectures": ["Glm4MoeForCausalLM"]}
(tmp_path / "config.json").write_text(json.dumps(cfg))
assert get_transformers_tier(str(tmp_path)) == "530"
def test_needs_transformers_5_compat(self):
"""needs_transformers_5 should return True for both 530 and 550 models."""
assert needs_transformers_5("google/gemma-4-E2B-it") is True

161
studio/backend/utils/transformers_version.py
View file

@ -60,14 +60,25 @@ TRANSFORMERS_550_MODEL_SUBSTRINGS: tuple[str, ...] = (
"gemma4", # Gemma-4 alternate naming
)
# Architecture classes / model_type values that require transformers 5.5.0.
# Checked via config.json (local or HuggingFace).
# Architecture classes that require transformers 5.5.0.
_TRANSFORMERS_550_ARCHITECTURES: set[str] = {
"Gemma4ForConditionalGeneration",
}
# model_type values (from config.json) → tier mapping.
_TRANSFORMERS_550_MODEL_TYPES: set[str] = {
"gemma4",
}
_TRANSFORMERS_530_MODEL_TYPES: set[str] = {
"qwen3_moe",
"qwen3_5_moe",
"qwen3_vl_moe",
"qwen3_next",
"deepseek_v3_moe",
"glm4_moe",
"glm4_moe_lite",
"ministral",
}
# Tokenizer classes that only exist in transformers>=5.x
_TRANSFORMERS_5_TOKENIZER_CLASSES: set[str] = {
@ -77,15 +88,14 @@ _TRANSFORMERS_5_TOKENIZER_CLASSES: set[str] = {
# Cache for dynamic tokenizer_config.json lookups to avoid repeated fetches
_tokenizer_class_cache: dict[str, bool] = {}
# Cache for dynamic config.json lookups (architecture/model_type checks)
_config_needs_550_cache: dict[str, bool] = {}
# Cache for config.json lookups (returns the parsed dict or None)
_config_json_cache: dict[str, dict | None] = {}
# Versions
TRANSFORMERS_550_VERSION = "5.5.0"
TRANSFORMERS_530_VERSION = "5.3.0"
TRANSFORMERS_DEFAULT_VERSION = "4.57.6"
# Backwards-compat alias — points to 5.5.0 (the highest 5.x tier).
# Consumers should prefer TRANSFORMERS_530_VERSION / TRANSFORMERS_550_VERSION.
# Backwards-compat alias used by other modules
TRANSFORMERS_5_VERSION = TRANSFORMERS_550_VERSION
# Pre-installed directories — created by setup.sh / setup.ps1
@ -95,45 +105,6 @@ _VENV_T5_550_DIR = str(Path.home() / ".unsloth" / "studio" / ".venv_t5_550")
_VENV_T5_DIR = _VENV_T5_550_DIR
def activate_transformers_for_subprocess(model_name: str) -> None:
"""Activate the correct transformers version in a subprocess worker.
Call this BEFORE any ML imports. Resolves LoRA adapters to their base
model, determines the required tier, and prepends the appropriate
``.venv_t5_*`` directory to ``sys.path``. Also propagates the path
via ``PYTHONPATH`` for child processes (e.g. GGUF converter).
Used by training, inference, and export workers.
"""
resolved = _resolve_base_model(model_name)
tier = get_transformers_tier(resolved)
if tier == "550":
if not _ensure_venv_t5_550_exists():
raise RuntimeError(
f"Cannot activate transformers 5.5.0: "
f".venv_t5_550 missing at {_VENV_T5_550_DIR}"
)
if _VENV_T5_550_DIR not in sys.path:
sys.path.insert(0, _VENV_T5_550_DIR)
logger.info("Activated transformers 5.5.0 from %s", _VENV_T5_550_DIR)
_pp = os.environ.get("PYTHONPATH", "")
os.environ["PYTHONPATH"] = _VENV_T5_550_DIR + (os.pathsep + _pp if _pp else "")
elif tier == "530":
if not _ensure_venv_t5_530_exists():
raise RuntimeError(
f"Cannot activate transformers 5.3.0: "
f".venv_t5_530 missing at {_VENV_T5_530_DIR}"
)
if _VENV_T5_530_DIR not in sys.path:
sys.path.insert(0, _VENV_T5_530_DIR)
logger.info("Activated transformers 5.3.0 from %s", _VENV_T5_530_DIR)
_pp = os.environ.get("PYTHONPATH", "")
os.environ["PYTHONPATH"] = _VENV_T5_530_DIR + (os.pathsep + _pp if _pp else "")
else:
logger.info("Using default transformers (4.57.x) for %s", model_name)
def _resolve_base_model(model_name: str) -> str:
"""If *model_name* points to a LoRA adapter, return its base model.
@ -259,43 +230,28 @@ def _check_tokenizer_config_needs_v5(model_name: str) -> bool:
return False
def _check_config_needs_550(model_name: str) -> bool:
"""Check ``config.json`` for architectures or model_type that require
transformers 5.5.0 (e.g. Gemma 4).
_SENTINEL = object() # distinguishes "not cached" from "cached as None"
Checks locally first, then falls back to fetching from HuggingFace.
Results are cached in ``_config_needs_550_cache``.
Returns False on any error (fail-open to lower tier).
def _get_config_json(model_name: str) -> dict | None:
"""Read and cache ``config.json`` for *model_name*.
Checks local path first, then fetches from HuggingFace.
Returns the parsed dict, or ``None`` on any error (fail-open).
The result is cached in ``_config_json_cache``.
"""
if model_name in _config_needs_550_cache:
return _config_needs_550_cache[model_name]
def _check_cfg(cfg: dict) -> bool:
archs = cfg.get("architectures", [])
if any(a in _TRANSFORMERS_550_ARCHITECTURES for a in archs):
return True
if cfg.get("model_type") in _TRANSFORMERS_550_MODEL_TYPES:
return True
return False
cached = _config_json_cache.get(model_name, _SENTINEL)
if cached is not _SENTINEL:
return cached
# --- Check local config.json first ------------------------------------
local_path = Path(model_name)
local_cfg = local_path / "config.json"
local_cfg = Path(model_name) / "config.json"
if local_cfg.is_file():
try:
with open(local_cfg) as f:
cfg = json.load(f)
result = _check_cfg(cfg)
if result:
logger.info(
"Local config.json check: %s needs transformers 5.5.0 "
"(architectures=%s, model_type=%s)",
model_name,
cfg.get("architectures", []),
cfg.get("model_type"),
)
_config_needs_550_cache[model_name] = result
return result
_config_json_cache[model_name] = cfg
return cfg
except Exception as exc:
logger.debug("Could not read %s: %s", local_cfg, exc)
@ -307,21 +263,26 @@ def _check_config_needs_550(model_name: str) -> bool:
req = urllib.request.Request(url, headers = {"User-Agent": "unsloth-studio"})
with urllib.request.urlopen(req, timeout = 10) as resp:
cfg = json.loads(resp.read().decode())
result = _check_cfg(cfg)
if result:
logger.info(
"Dynamic config.json check: %s needs transformers 5.5.0 "
"(architectures=%s, model_type=%s)",
model_name,
cfg.get("architectures", []),
cfg.get("model_type"),
)
_config_needs_550_cache[model_name] = result
return result
_config_json_cache[model_name] = cfg
return cfg
except Exception as exc:
logger.debug("Could not fetch config.json for '%s': %s", model_name, exc)
_config_needs_550_cache[model_name] = False
logger.debug(
"Could not fetch config.json for '%s': %s", model_name, exc
)
_config_json_cache[model_name] = None
return None
def _check_config_needs_550(model_name: str) -> bool:
"""Check ``config.json`` for architectures or model_type that require
transformers 5.5.0. Uses the shared ``_get_config_json`` cache."""
cfg = _get_config_json(model_name)
if cfg is None:
return False
archs = cfg.get("architectures", [])
if any(a in _TRANSFORMERS_550_ARCHITECTURES for a in archs):
return True
return cfg.get("model_type") in _TRANSFORMERS_550_MODEL_TYPES
def get_transformers_tier(model_name: str) -> str:
@ -331,19 +292,35 @@ def get_transformers_tier(model_name: str) -> str:
``"530"`` for models needing transformers 5.3.0 (e.g. Ministral-3, Qwen3 MoE),
or ``"default"`` for everything else (4.57.x).
The 5.5.0 check runs first, then 5.3.0.
Fast path: substring checks (no I/O) for both tiers run first.
Slow path: single config.json fetch (cached) checks model_type for
both tiers, then tokenizer_config.json as final fallback.
"""
lowered = model_name.lower()
# --- Fast substring checks (no I/O) ------------------------------------
# --- Fast substring checks (no I/O) -----------------------------------
if any(sub in lowered for sub in TRANSFORMERS_550_MODEL_SUBSTRINGS):
return "550"
if any(sub in lowered for sub in TRANSFORMERS_5_MODEL_SUBSTRINGS):
return "530"
# --- Slow config fallbacks (local file first, then network) -----------
if _check_config_needs_550(model_name):
return "550"
# --- config.json model_type / architecture check (single fetch) -------
cfg = _get_config_json(model_name)
if cfg is not None:
model_type = cfg.get("model_type", "")
archs = cfg.get("architectures", [])
# Check 5.5.0 first
if model_type in _TRANSFORMERS_550_MODEL_TYPES:
return "550"
if any(a in _TRANSFORMERS_550_ARCHITECTURES for a in archs):
return "550"
# Check 5.3.0
if model_type in _TRANSFORMERS_530_MODEL_TYPES:
return "530"
# --- Final fallback: tokenizer_config.json for 5.3.0 ------------------
if _check_tokenizer_config_needs_v5(model_name):
return "530"