From 52d24ab43d7d837cdf26df2efd6cc6bac51559c6 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Mon, 6 Apr 2026 19:40:11 +0000 Subject: [PATCH] use config.json model_type for tier detection, add unsloth/nvidia namespace guard --- studio/backend/core/inference/worker.py | 6 +- studio/backend/core/training/worker.py | 6 +- .../tests/test_transformers_version.py | 25 ++- studio/backend/utils/transformers_version.py | 161 ++++++++---------- 4 files changed, 99 insertions(+), 99 deletions(-) diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 1010e56dac..506e631c51 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -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", diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 2c4672231a..edd88cc10f 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -506,7 +506,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", diff --git a/studio/backend/tests/test_transformers_version.py b/studio/backend/tests/test_transformers_version.py index c031c2fea3..609a154a9a 100644 --- a/studio/backend/tests/test_transformers_version.py +++ b/studio/backend/tests/test_transformers_version.py @@ -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 diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index 36c3a4c22d..5ae92f2df1 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -61,14 +61,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] = { @@ -78,15 +89,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 @@ -96,45 +106,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. @@ -260,43 +231,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) @@ -308,21 +264,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: @@ -332,19 +293,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"