316 lines
12 KiB
Python
316 lines
12 KiB
Python
"""
|
|
Automatic transformers version switching.
|
|
|
|
Some newer model architectures (Ministral-3, GLM-4.7-Flash, Qwen3-30B-A3B MoE,
|
|
tiny_qwen3_moe) require transformers>=5.1.0, while everything else needs the
|
|
default 4.57.x that ships with Unsloth.
|
|
|
|
When loading a LoRA adapter with a custom name, we resolve the base model from
|
|
``adapter_config.json`` and check *that* against the model list.
|
|
|
|
Strategy (sys.path overlay):
|
|
• The default transformers (4.57.x) always lives in site-packages.
|
|
• When 5.x is needed, we ``pip install --target <dir> --no-deps`` to a
|
|
separate directory and **prepend** it to ``sys.path``.
|
|
• To revert, we simply **remove** that directory from ``sys.path``.
|
|
• After either change we purge cached modules so the next import picks
|
|
up the correct version.
|
|
"""
|
|
|
|
import importlib
|
|
import importlib.metadata
|
|
import json
|
|
import logging
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Ensure our logger is visible even if root logger isn't configured for INFO.
|
|
if not logger.handlers:
|
|
_handler = logging.StreamHandler()
|
|
_handler.setLevel(logging.INFO)
|
|
_handler.setFormatter(
|
|
logging.Formatter("[%(name)s|%(levelname)s]%(message)s")
|
|
)
|
|
logger.addHandler(_handler)
|
|
logger.setLevel(logging.INFO)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Detection
|
|
# ---------------------------------------------------------------------------
|
|
|
|
# Lowercase substrings — if ANY appears anywhere in the lowered model name,
|
|
# we need transformers 5.x.
|
|
TRANSFORMERS_5_MODEL_SUBSTRINGS: tuple[str, ...] = (
|
|
"ministral-3-", # Ministral-3-{3,8,14}B-{Instruct,Reasoning,Base}-2512
|
|
"glm-4.7-flash", # GLM-4.7-Flash
|
|
"qwen3-30b-a3b", # Qwen3-30B-A3B-Instruct-2507 and variants
|
|
"tiny_qwen3_moe", # imdatta0/tiny_qwen3_moe_2.8B_0.7B
|
|
)
|
|
|
|
# Versions
|
|
TRANSFORMERS_5_VERSION = "5.1.0"
|
|
TRANSFORMERS_DEFAULT_VERSION = "4.57.1"
|
|
|
|
# Persistent directory for the transformers 5.x overlay — lives next to .venv/
|
|
_PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent.parent # studio/backend/utils/ → project root
|
|
_OVERLAY_DIR = str(_PROJECT_ROOT / ".venv_overlay")
|
|
|
|
|
|
def _resolve_base_model(model_name: str) -> str:
|
|
"""If *model_name* points to a LoRA adapter, return its base model.
|
|
|
|
Checks for ``adapter_config.json`` locally first. Only calls the heavier
|
|
``get_base_model_from_lora`` for paths that are actual local directories
|
|
(avoids noisy warnings for plain HF model IDs).
|
|
|
|
Returns the original *model_name* unchanged if it is not a LoRA adapter.
|
|
"""
|
|
# --- Fast local check ---------------------------------------------------
|
|
local_path = Path(model_name)
|
|
adapter_cfg_path = local_path / "adapter_config.json"
|
|
if adapter_cfg_path.is_file():
|
|
try:
|
|
with open(adapter_cfg_path) as f:
|
|
cfg = json.load(f)
|
|
base = cfg.get("base_model_name_or_path")
|
|
if base:
|
|
logger.info(
|
|
"Resolved LoRA adapter '%s' → base model '%s'",
|
|
model_name, base,
|
|
)
|
|
return base
|
|
except Exception as exc:
|
|
logger.debug("Could not read %s: %s", adapter_cfg_path, exc)
|
|
|
|
# --- Only try the heavier fallback for local directories ----------------
|
|
if local_path.is_dir():
|
|
try:
|
|
from utils.models import get_base_model_from_lora
|
|
base = get_base_model_from_lora(model_name)
|
|
if base:
|
|
logger.info(
|
|
"Resolved LoRA adapter '%s' → base model '%s' "
|
|
"(via get_base_model_from_lora)",
|
|
model_name, base,
|
|
)
|
|
return base
|
|
except Exception as exc:
|
|
logger.debug(
|
|
"get_base_model_from_lora failed for '%s': %s",
|
|
model_name, exc,
|
|
)
|
|
|
|
return model_name
|
|
|
|
|
|
def needs_transformers_5(model_name: str) -> bool:
|
|
"""Return True if *model_name* belongs to an architecture that requires
|
|
``transformers>=5.1.0``."""
|
|
lowered = model_name.lower()
|
|
return any(sub in lowered for sub in TRANSFORMERS_5_MODEL_SUBSTRINGS)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Version switching
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _get_in_memory_version() -> str | None:
|
|
"""Return the transformers version currently loaded in this process."""
|
|
tf = sys.modules.get("transformers")
|
|
if tf is not None:
|
|
return getattr(tf, "__version__", None)
|
|
return None
|
|
|
|
|
|
# All top-level prefixes that hold references to transformers internals.
|
|
_PURGE_PREFIXES = (
|
|
"transformers",
|
|
"huggingface_hub",
|
|
"unsloth",
|
|
"unsloth_zoo",
|
|
"peft",
|
|
"trl",
|
|
"accelerate",
|
|
"auto_gptq",
|
|
# NOTE: bitsandbytes is intentionally EXCLUDED — it registers torch custom
|
|
# operators at import time via torch.library.define(). Those registrations
|
|
# live in torch's global operator registry which survives module purge.
|
|
# Re-importing bitsandbytes after purge → duplicate registration → crash.
|
|
# Our own modules that import from transformers at module level
|
|
# (e.g. model_config.py: `from transformers import AutoConfig`)
|
|
"utils.models",
|
|
"core.training",
|
|
"core.inference",
|
|
"core.export",
|
|
)
|
|
|
|
|
|
def _purge_modules() -> int:
|
|
"""Remove all cached modules for transformers and its dependents.
|
|
|
|
Returns the number of modules purged.
|
|
"""
|
|
importlib.invalidate_caches()
|
|
to_remove = [
|
|
k for k in list(sys.modules.keys())
|
|
if any(k == p or k.startswith(p + ".") for p in _PURGE_PREFIXES)
|
|
]
|
|
for key in to_remove:
|
|
del sys.modules[key]
|
|
return len(to_remove)
|
|
|
|
|
|
# Packages to install into the overlay (each with --no-deps).
|
|
_OVERLAY_PACKAGES = (
|
|
f"transformers=={TRANSFORMERS_5_VERSION}",
|
|
"huggingface_hub>=1.3.0,<2.0",
|
|
)
|
|
|
|
|
|
def _install_overlay() -> None:
|
|
"""Install transformers 5.x into BOTH site-packages and the overlay
|
|
directory, then prepend the overlay to ``sys.path``.
|
|
|
|
We install into site-packages so that:
|
|
- ``importlib.metadata.version()`` returns the correct version
|
|
- sub-package resolution (``transformers.models.*``) uses 5.x code
|
|
The overlay is kept as a safety net for sys.path-based resolution.
|
|
"""
|
|
# --- 1. Install into site-packages (updates code + metadata) -----------
|
|
for pkg in _OVERLAY_PACKAGES:
|
|
cmd = [sys.executable, "-m", "pip", "install", pkg]
|
|
logger.info("Installing %s into site-packages: %s", pkg, " ".join(cmd))
|
|
result = subprocess.run(
|
|
cmd,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
)
|
|
if result.returncode != 0:
|
|
logger.error("pip install failed:\n%s", result.stdout)
|
|
raise RuntimeError(
|
|
f"Failed to install {pkg}.\npip output:\n{result.stdout}"
|
|
)
|
|
logger.info("Site-packages install succeeded")
|
|
|
|
# --- 2. Install into overlay (safety net for sys.path resolution) ------
|
|
needs_overlay = (
|
|
not os.path.isdir(_OVERLAY_DIR)
|
|
or not os.listdir(_OVERLAY_DIR)
|
|
or not os.path.isdir(os.path.join(_OVERLAY_DIR, "huggingface_hub"))
|
|
)
|
|
if needs_overlay:
|
|
if os.path.isdir(_OVERLAY_DIR):
|
|
shutil.rmtree(_OVERLAY_DIR)
|
|
os.makedirs(_OVERLAY_DIR, exist_ok=True)
|
|
for pkg in _OVERLAY_PACKAGES:
|
|
cmd = [
|
|
sys.executable, "-m", "pip", "install",
|
|
"--target", _OVERLAY_DIR,
|
|
"--no-deps",
|
|
pkg,
|
|
]
|
|
logger.info("Installing %s to overlay: %s", pkg, " ".join(cmd))
|
|
subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
|
|
logger.info("Overlay install succeeded")
|
|
|
|
# Prepend to sys.path (if not already there)
|
|
if _OVERLAY_DIR not in sys.path:
|
|
sys.path.insert(0, _OVERLAY_DIR)
|
|
logger.info("Prepended %s to sys.path", _OVERLAY_DIR)
|
|
|
|
# --- 3. Purge old modules and force fresh import -----------------------
|
|
count = _purge_modules()
|
|
logger.info("Purged %d cached modules", count)
|
|
|
|
import transformers
|
|
logger.info("Loaded transformers %s", transformers.__version__)
|
|
|
|
|
|
def _remove_overlay() -> None:
|
|
"""Revert to the default transformers (4.57.x) by restoring site-packages
|
|
and removing the overlay from ``sys.path``."""
|
|
# --- 1. Restore default versions in site-packages ----------------------
|
|
default_packages = (
|
|
f"transformers=={TRANSFORMERS_DEFAULT_VERSION}",
|
|
"huggingface_hub==0.36.0",
|
|
)
|
|
for pkg in default_packages:
|
|
cmd = [sys.executable, "-m", "pip", "install", pkg]
|
|
logger.info("Restoring %s in site-packages: %s", pkg, " ".join(cmd))
|
|
result = subprocess.run(
|
|
cmd,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
)
|
|
if result.returncode != 0:
|
|
logger.error("pip install failed:\n%s", result.stdout)
|
|
|
|
# --- 2. Remove overlay from sys.path -----------------------------------
|
|
while _OVERLAY_DIR in sys.path:
|
|
sys.path.remove(_OVERLAY_DIR)
|
|
logger.info("Removed %s from sys.path", _OVERLAY_DIR)
|
|
|
|
# --- 3. Purge old modules and force fresh import -----------------------
|
|
count = _purge_modules()
|
|
logger.info("Purged %d cached modules", count)
|
|
|
|
import transformers
|
|
logger.info("Reverted to transformers %s", transformers.__version__)
|
|
|
|
|
|
def ensure_transformers_version(model_name: str) -> None:
|
|
"""Ensure the correct ``transformers`` version is active for *model_name*.
|
|
|
|
Uses sys.path overlay:
|
|
• Need 5.x → install to separate dir, prepend sys.path, purge modules.
|
|
• Need 4.x → remove overlay from sys.path, purge modules.
|
|
|
|
For LoRA adapters with custom names, the base model is resolved from
|
|
``adapter_config.json`` before checking.
|
|
|
|
Call this at the top of every model-loading code path (training ``/start``,
|
|
inference ``/load``, export ``/load-checkpoint``).
|
|
"""
|
|
# Resolve LoRA adapters to their base model for accurate detection
|
|
resolved = _resolve_base_model(model_name)
|
|
want_5 = needs_transformers_5(resolved)
|
|
target_version = TRANSFORMERS_5_VERSION if want_5 else TRANSFORMERS_DEFAULT_VERSION
|
|
target_major = int(target_version.split(".")[0])
|
|
|
|
# Check what's actually loaded in memory
|
|
in_memory = _get_in_memory_version()
|
|
overlay_active = _OVERLAY_DIR in sys.path
|
|
|
|
logger.info(
|
|
"Version check for '%s' (resolved: '%s'): need=%s, "
|
|
"in_memory=%s, overlay_active=%s",
|
|
model_name, resolved, target_version, in_memory, overlay_active,
|
|
)
|
|
|
|
# --- Already correct? ---------------------------------------------------
|
|
if in_memory is not None:
|
|
in_memory_major = int(in_memory.split(".")[0])
|
|
if in_memory_major == target_major:
|
|
logger.info(
|
|
"transformers %s already loaded — correct for '%s'",
|
|
in_memory, model_name,
|
|
)
|
|
return
|
|
|
|
# --- Switch version -----------------------------------------------------
|
|
if want_5:
|
|
logger.info("Activating transformers %s overlay…", TRANSFORMERS_5_VERSION)
|
|
_install_overlay()
|
|
else:
|
|
logger.info("Reverting to default transformers %s…", TRANSFORMERS_DEFAULT_VERSION)
|
|
_remove_overlay()
|
|
|
|
final = _get_in_memory_version()
|
|
logger.info("✓ transformers version is now %s", final)
|