diff --git a/tests/test_new_mapper_fetched_fp8.py b/tests/test_new_mapper_fetched_fp8.py new file mode 100644 index 0000000000..2835aadb59 --- /dev/null +++ b/tests/test_new_mapper_fetched_fp8.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Regression tests for what ``_get_new_mapper`` hands back to the upgrade probe. + +``test_new_mapper_no_global_leak.py`` serves the repo's own ``mapper.py`` as both installed +and fetched source, so it cannot tell a fetched table from a fresh copy of the installed one. +Two gaps it misses: + +1. The probe must answer for an fp8 repo only the FETCHED mapper knows, so an extra ``"8"`` + entry is spliced into the fetched source only. Isolating the exec without returning the + fetched fp8 tables would silently drop the fp8 half of the upgrade check. +2. The probe must survive a fetched ``mapper.py`` with no fp8 tables (anything older, or a + future rename): reading them with ``[]`` raises ``KeyError`` into the bare ``except``, + taking the 4bit half, the probe's whole purpose, down with it. + +``loader_utils`` imports torch, so ast-extract the resolvers and run them against a stubbed +``requests``, as in ``tests/test_bad_mappings_redirect.py``. +""" + +import ast +import os +import sys +import types + +_MODELS = os.path.join(os.path.dirname(__file__), os.pardir, "unsloth", "models") + +_WANTED = {"__get_model_name", "_resolve_with_mappers", "_get_new_mapper", "get_model_name"} + +# An fp8 ("8") model, spliced into the FETCHED mapper only. +_NEW_KEY = "unsloth/Zeta-9B-Only-On-Main" +_NEW_OFFICIAL = "zeta-org/Zeta-9B-Only-On-Main-FP8" +_NEW_BLOCK = "unsloth/Zeta-9B-Only-On-Main-FP8-Block" +_NEW_ROW = "unsloth/Zeta-9B-Only-On-Main-FP8-Row" +_ANCHOR = ' "unsloth/Kimi-K2-Instruct-BF16" : (' + + +def _mapper_source(): + with open(os.path.join(_MODELS, "mapper.py"), encoding = "utf-8") as f: + return f.read() + + +def _with_extra_fp8_model(source): + assert _ANCHOR in source, "anchor moved; update this test" + entry = ( + f' "{_NEW_KEY}" : {{\n' + f' "16" : ("{_NEW_KEY}", "zeta-org/Zeta-9B-Only-On-Main"),\n' + f' "8" : ("{_NEW_OFFICIAL}", "{_NEW_BLOCK}", "{_NEW_ROW}"),\n' + f" }},\n" + ) + return source.replace(_ANCHOR, entry + _ANCHOR, 1) + + +def _without_fp8_tables(source): + """A mapper.py from before the fp8 tables existed.""" + return source.replace("FLOAT_TO_FP8_BLOCK_MAPPER", "SOME_OTHER_BLOCK_TABLE").replace( + "FLOAT_TO_FP8_ROW_MAPPER", "SOME_OTHER_ROW_TABLE" + ) + + +class _FakeResponse: + def __init__(self, text): + self.text = text + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + +def _install_fake_requests(monkeypatch, text): + module = types.ModuleType("requests") + module.get = lambda url, timeout = None: _FakeResponse(text) + monkeypatch.setitem(sys.modules, "requests", module) + + +def _install_fake_vllm_absent(monkeypatch, namespace): + """vllm >= 0.12.0 returns early from __get_model_name, leaving the probe unreachable.""" + monkeypatch.delitem(sys.modules, "vllm", raising = False) + fake = types.ModuleType("importlib") + fake.util = types.SimpleNamespace(find_spec = lambda name: None) + namespace["importlib"] = fake + + +def _load_resolver(installed_source): + """Stand-in for loader_utils' module globals, built from `installed_source`.""" + from unsloth_zoo.utils import Version + + mapper_ns = {} + exec(compile(installed_source, "mapper.py", "exec"), mapper_ns) + + namespace = { + "INT_TO_FLOAT_MAPPER": mapper_ns["INT_TO_FLOAT_MAPPER"], + "FLOAT_TO_INT_MAPPER": mapper_ns["FLOAT_TO_INT_MAPPER"], + "MAP_TO_UNSLOTH_16bit": mapper_ns["MAP_TO_UNSLOTH_16bit"], + "FLOAT_TO_FP8_BLOCK_MAPPER": mapper_ns["FLOAT_TO_FP8_BLOCK_MAPPER"], + "FLOAT_TO_FP8_ROW_MAPPER": mapper_ns["FLOAT_TO_FP8_ROW_MAPPER"], + "SUPPORTS_FOURBIT": True, + "transformers_version": Version("4.57.6"), + "Version": Version, + "os": os, + } + with open(os.path.join(_MODELS, "loader_utils.py"), encoding = "utf-8") as f: + tree = ast.parse(f.read()) + for node in tree.body: + if isinstance(node, ast.Assign) and any( + getattr(t, "id", None) in ("BAD_MAPPINGS", "_OFFLINE_ENV_VALUES", "_OFFLINE_ENV_KEYS") + for t in node.targets + ): + exec(compile(ast.Module([node], []), "", "exec"), namespace) + elif isinstance(node, ast.FunctionDef) and ( + node.name in _WANTED or node.name == "_env_says_offline" + ): + exec(compile(ast.Module([node], []), node.name, "exec"), namespace) + return namespace + + +def test_probe_answers_for_an_fp8_repo_only_the_fetched_mapper_knows(monkeypatch): + installed = _mapper_source() + namespace = _load_resolver(installed) + installed_block = namespace["FLOAT_TO_FP8_BLOCK_MAPPER"] + installed_row = namespace["FLOAT_TO_FP8_ROW_MAPPER"] + assert _NEW_OFFICIAL.lower() not in installed_block, "the installed table must not know it" + + _install_fake_requests(monkeypatch, _with_extra_fp8_model(installed)) + _install_fake_vllm_absent(monkeypatch, namespace) + + try: + resolved = namespace["get_model_name"]( + _NEW_OFFICIAL, load_in_4bit = False, load_in_fp8 = "block" + ) + except NotImplementedError as error: + assert "not supported in your current Unsloth version" in str(error) + else: + raise AssertionError( + f"a fetched-only fp8 repo must raise the upgrade error, got {resolved!r}" + ) + + # Answering must not have adopted the fetched tables. + assert namespace["FLOAT_TO_FP8_BLOCK_MAPPER"] is installed_block + assert namespace["FLOAT_TO_FP8_ROW_MAPPER"] is installed_row + assert _NEW_OFFICIAL.lower() not in namespace["FLOAT_TO_FP8_BLOCK_MAPPER"] + + +def test_probe_survives_a_fetched_mapper_without_the_fp8_tables(monkeypatch): + installed = _mapper_source() + namespace = _load_resolver(installed) + _install_fake_requests(monkeypatch, _without_fp8_tables(installed)) + + int_to_float, float_to_int, map_to_16bit = namespace["_get_new_mapper"]()[:3] + + assert ( + int_to_float and float_to_int and map_to_16bit + ), "a fetched mapper.py without the fp8 tables must not take the 4bit upgrade check down" diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py index 8214adc0bf..7aae75fe4f 100644 --- a/unsloth/models/loader_utils.py +++ b/unsloth/models/loader_utils.py @@ -205,8 +205,12 @@ def _get_new_mapper(): namespace["NEW_INT_TO_FLOAT_MAPPER"], namespace["NEW_FLOAT_TO_INT_MAPPER"], namespace["NEW_MAP_TO_UNSLOTH_16bit"], - namespace["FLOAT_TO_FP8_BLOCK_MAPPER"], - namespace["FLOAT_TO_FP8_ROW_MAPPER"], + # .get, not []: these two come from the fetched file under its own names (unlike + # the NEW_ names above, renamed here), so an older or renamed mapper.py would + # KeyError into the bare except and take the 4bit half of the probe down too. + # {} is safe: the probe runs only after the installed tables already missed. + namespace.get("FLOAT_TO_FP8_BLOCK_MAPPER", {}), + namespace.get("FLOAT_TO_FP8_ROW_MAPPER", {}), ) except: return {}, {}, {}, {}, {}