From f60b982a09ed258b9469ef1cbca54c466d9f0cb1 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Mon, 13 Jul 2026 17:34:25 -0300 Subject: [PATCH] Studio: Fix torch_dtype deprecation warning on startup and ASR load (#6999) --- studio/backend/core/inference/inference.py | 3 +- studio/backend/core/rag/embeddings.py | 5 +- .../backend/tests/test_transformers_dtype.py | 77 +++++++++++++++++++ studio/backend/utils/transformers_dtype.py | 51 ++++++++++++ 4 files changed, 132 insertions(+), 4 deletions(-) create mode 100644 studio/backend/tests/test_transformers_dtype.py create mode 100644 studio/backend/utils/transformers_dtype.py diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 7e69e05124..bb4be39ef6 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -15,6 +15,7 @@ from pathlib import Path from typing import Optional, Union, Generator, Tuple from utils.models import ModelConfig, get_base_model_from_lora from utils.paths import is_model_cached +from utils.transformers_dtype import dtype_kwargs from utils.utils import format_error_message from utils.hardware import ( get_device, @@ -440,7 +441,7 @@ class InferenceBackend: feature_extractor = tokenizer.feature_extractor, processor = tokenizer, return_language = True, - torch_dtype = torch.float16, + **dtype_kwargs(torch.float16), ) self.models[model_name]["model"] = model self.models[model_name]["tokenizer"] = tokenizer diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 47d26209b4..b0ecedd593 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -21,6 +21,7 @@ from functools import lru_cache from typing import Callable from utils.hardware.hardware import DeviceType, get_device +from utils.transformers_dtype import dtype_kwargs from . import config @@ -157,9 +158,7 @@ def _get(model_name: str | None = None): device = _device() logger.info("loading embedding model %s on %s", name, device) _guard_model_security(name) - _model = SentenceTransformer( - name, device = device, model_kwargs = {"torch_dtype": "float16"} - ) + _model = SentenceTransformer(name, device = device, model_kwargs = dtype_kwargs("float16")) _name = name return _model diff --git a/studio/backend/tests/test_transformers_dtype.py b/studio/backend/tests/test_transformers_dtype.py new file mode 100644 index 0000000000..28629e5620 --- /dev/null +++ b/studio/backend/tests/test_transformers_dtype.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the version-safe torch_dtype/dtype kwarg helper.""" + +import sys +import types + +import pytest + +from utils.transformers_dtype import _has_torch_dtype_kwarg, dtype_kwargs + + +@pytest.fixture(autouse = True) +def _clear_cache(): + _has_torch_dtype_kwarg.cache_clear() + yield + _has_torch_dtype_kwarg.cache_clear() + + +def _stub_transformers(monkeypatch, version): + stub = types.ModuleType("transformers") + stub.__version__ = version + monkeypatch.setitem(sys.modules, "transformers", stub) + + +def test_old_transformers_uses_torch_dtype(monkeypatch): + _stub_transformers(monkeypatch, "4.51.3") + assert _has_torch_dtype_kwarg() is True + assert dtype_kwargs("float16") == {"torch_dtype": "float16"} + + +def test_new_transformers_uses_dtype(monkeypatch): + _stub_transformers(monkeypatch, "4.57.6") + assert _has_torch_dtype_kwarg() is False + assert dtype_kwargs("float16") == {"dtype": "float16"} + + +def test_rename_boundary_uses_dtype(monkeypatch): + _stub_transformers(monkeypatch, "4.56.0") + assert _has_torch_dtype_kwarg() is False + + +def test_just_below_boundary_uses_torch_dtype(monkeypatch): + _stub_transformers(monkeypatch, "4.55.4") + assert _has_torch_dtype_kwarg() is True + + +@pytest.mark.parametrize("version", ["4.56.0.dev0", "4.56.0rc1"]) +def test_rename_prerelease_uses_dtype(monkeypatch, version): + """A pre-release of the rename version sorts *below* ``4.56.0`` but already + accepts (and prefers) ``dtype``; the release-tuple check must not fall back to + the legacy name there, or it re-emits the deprecation warning it suppresses.""" + _stub_transformers(monkeypatch, version) + assert _has_torch_dtype_kwarg() is False + + +def test_malformed_version_prefers_modern_name(monkeypatch): + """A non-PEP440 __version__ raises InvalidVersion; the except branch must + swallow it and default to the modern name rather than crash the embedder warm-up.""" + _stub_transformers(monkeypatch, "not-a-version") + assert _has_torch_dtype_kwarg() is False + assert dtype_kwargs("float16") == {"dtype": "float16"} + + +def test_missing_transformers_prefers_modern_name(monkeypatch): + monkeypatch.delitem(sys.modules, "transformers", raising = False) + real_import = __import__ + + def _raise(name, *args, **kwargs): + if name == "transformers": + raise ImportError("no transformers") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr("builtins.__import__", _raise) + assert _has_torch_dtype_kwarg() is False + assert dtype_kwargs("float16") == {"dtype": "float16"} diff --git a/studio/backend/utils/transformers_dtype.py b/studio/backend/utils/transformers_dtype.py new file mode 100644 index 0000000000..daeb6e2452 --- /dev/null +++ b/studio/backend/utils/transformers_dtype.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Version-safe fp-dtype kwarg for transformers/sentence-transformers loads. + +transformers renamed the ``torch_dtype`` kwarg to ``dtype`` in 4.56.0, and emits +``torch_dtype is deprecated! Use dtype instead!`` when the old name is passed. But our floor (``transformers>=4.51.3``) predates ``dtype`` and only +accepts ``torch_dtype``, so a bare rename would ``TypeError`` on the floor. Pick +the name the installed version accepts instead. + +Answers the same question as ``unsloth_zoo.hf_utils.HAS_TORCH_DTYPE`` but derives +it independently, for two reasons. It uses a ``packaging.version`` check rather +than that constant's ``"torch_dtype" in PretrainedConfig.__doc__`` sniffing, which +raises ``TypeError`` under ``python -OO`` / ``PYTHONOPTIMIZE=2`` (docstrings are +stripped to ``None``, and ``"torch_dtype" in None`` is a type error). And it avoids +importing the constant at all: the RAG +embedder warms here at startup in the lean main process, and reading it would run +``unsloth_zoo``'s package ``__init__`` (torch import, GPU/Pytorch checks, the +patching banner) as a side effect. The embedder is deliberately torch-optional (it +degrades to the ``llama-server`` GGUF backend), so it must not drag in that +heavyweight import just to read one bool. +""" + +from functools import lru_cache + + +@lru_cache(maxsize = 1) +def _has_torch_dtype_kwarg() -> bool: + """True if the installed transformers still expects the legacy ``torch_dtype`` + name (i.e. predates the ``dtype`` rename). False when ``dtype`` is the accepted + name, or when transformers is missing/broken (prefer the modern name).""" + try: + import transformers + from packaging.version import Version + + # Compare on the release tuple so a pre-release of the rename version + # (``4.56.0.dev0``/``rc1``, which sort *below* ``4.56.0``) still counts as + # new and picks ``dtype`` -- those builds already accept it, and picking + # ``torch_dtype`` there would re-emit the very warning this suppresses. + return Version(transformers.__version__).release < (4, 56, 0) + except Exception: + return False + + +def dtype_kwargs(value) -> dict: + """``{"torch_dtype": value}`` on old transformers, ``{"dtype": value}`` on new. + + Splat into a load call (``pipeline(..., **dtype_kwargs(torch.float16))``) or use + directly as ``model_kwargs`` (``model_kwargs = dtype_kwargs("float16")``). + """ + return {"torch_dtype" if _has_torch_dtype_kwarg() else "dtype": value}