diff --git a/tests/test_bad_mappings_redirect.py b/tests/test_bad_mappings_redirect.py index 49dab2d98b..7ae9ddbb4f 100644 --- a/tests/test_bad_mappings_redirect.py +++ b/tests/test_bad_mappings_redirect.py @@ -26,7 +26,7 @@ def _load_get_model_name(): namespace = dict(mapper_ns) namespace["SUPPORTS_FOURBIT"] = True namespace["_env_says_offline"] = lambda: True - namespace["_get_new_mapper"] = lambda: ({}, {}, {}) + namespace["_get_new_mapper"] = lambda: ({}, {}, {}, {}, {}) wanted = {"__get_model_name", "_resolve_with_mappers", "get_model_name"} for node in tree.body: diff --git a/tests/test_get_model_name.py b/tests/test_get_model_name.py index 33ad316d88..bbd5bfacb1 100644 --- a/tests/test_get_model_name.py +++ b/tests/test_get_model_name.py @@ -6,7 +6,8 @@ from unsloth.models.mapper import FLOAT_TO_INT_MAPPER, MAP_TO_UNSLOTH_16bit def _no_remote_mapper(): - return {}, {}, {} + # int_to_float, float_to_int, map_to_16bit, fp8_block, fp8_row + return {}, {}, {}, {}, {} class TestGetModelName(unittest.TestCase): diff --git a/tests/test_new_mapper_no_global_leak.py b/tests/test_new_mapper_no_global_leak.py new file mode 100644 index 0000000000..5ec00fba5a --- /dev/null +++ b/tests/test_new_mapper_no_global_leak.py @@ -0,0 +1,99 @@ +"""Regression test for ``_get_new_mapper`` leaking into ``loader_utils`` globals. + +``get_model_name`` calls ``_get_new_mapper()`` whenever a name misses the local +tables, purely to answer "would a newer Unsloth support this?". It fetches +``mapper.py`` from GitHub main, prefixes the three mappers it wants with +``NEW_``, and ``exec``s the result into ``globals()``. + +The slice starts at ``__INT_TO_FLOAT_MAPPER``, so it also carries +``FLOAT_TO_FP8_BLOCK_MAPPER``/``FLOAT_TO_FP8_ROW_MAPPER`` and the two +``_add_*`` helpers, and those names are NOT renamed. Exec'ing into +``globals()`` therefore rebinds the FP8 tables that ``loader_utils`` imported +from the installed ``mapper``, so every later ``get_model_name(..., +load_in_fp8 = ...)`` in the process resolves through GitHub main's table +instead of the installed one. The probe is supposed to read, not to swap the +installed mappings out from under the caller. + +``loader_utils`` imports torch, so ast-extract ``_get_new_mapper`` and run it +against a stubbed ``requests`` rather than importing unsloth (which needs a GPU). +""" + +import ast +import os +import sys +import types + +_MODELS = os.path.join(os.path.dirname(__file__), os.pardir, "unsloth", "models") + + +def _mapper_source(): + with open(os.path.join(_MODELS, "mapper.py"), encoding = "utf-8") as f: + return f.read() + + +def _extract_get_new_mapper(namespace): + 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.FunctionDef) and node.name == "_get_new_mapper": + exec(compile(ast.Module([node], []), node.name, "exec"), namespace) + return namespace["_get_new_mapper"] + raise AssertionError("_get_new_mapper not found in loader_utils.py") + + +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 test_get_new_mapper_does_not_rebind_the_installed_fp8_tables(monkeypatch): + _install_fake_requests(monkeypatch, _mapper_source()) + + installed = {} + exec(compile(_mapper_source(), "mapper.py", "exec"), installed) + block = installed["FLOAT_TO_FP8_BLOCK_MAPPER"] + row = installed["FLOAT_TO_FP8_ROW_MAPPER"] + assert block and row, "the installed FP8 tables should not be empty" + + # Stand in for loader_utils' module globals, which import the FP8 tables. + namespace = {"FLOAT_TO_FP8_BLOCK_MAPPER": block, "FLOAT_TO_FP8_ROW_MAPPER": row} + get_new_mapper = _extract_get_new_mapper(namespace) + + int_to_float, float_to_int, map_to_16bit, fp8_block, fp8_row = get_new_mapper() + + # _get_new_mapper swallows every exception and returns empty dicts, so assert + # it actually ran before trusting anything below. + assert int_to_float and float_to_int and map_to_16bit, "the fetch/exec path did not run" + + # the probe has to hand the FETCHED fp8 tables back, or a newly added fp8 repo would + # miss both the installed tables and the probe and skip the upgrade message + assert fp8_block and fp8_row + assert fp8_block is not block and fp8_row is not row + + assert namespace["FLOAT_TO_FP8_BLOCK_MAPPER"] is block + assert namespace["FLOAT_TO_FP8_ROW_MAPPER"] is row + + +def test_get_new_mapper_leaves_no_helpers_behind(monkeypatch): + _install_fake_requests(monkeypatch, _mapper_source()) + + namespace = {} + get_new_mapper = _extract_get_new_mapper(namespace) + before = set(namespace) + + assert all(get_new_mapper()), "the fetch/exec path did not run" + + leaked = set(namespace) - before + assert not leaked, f"_get_new_mapper leaked {sorted(leaked)} into its module globals" diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py index 7fd8cd66b4..8214adc0bf 100644 --- a/unsloth/models/loader_utils.py +++ b/unsloth/models/loader_utils.py @@ -191,19 +191,39 @@ def _get_new_mapper(): .replace("MAP_TO_UNSLOTH_16bit", "NEW_MAP_TO_UNSLOTH_16bit") ) - exec(new_mapper, globals()) + # Exec into a throwaway namespace, never globals(). The slice also carries + # FLOAT_TO_FP8_BLOCK_MAPPER / FLOAT_TO_FP8_ROW_MAPPER, the _add_* helpers + # and the builder's loop variables, so exec'ing into globals() would swap + # the FP8 tables this module imported from the installed mapper for the + # ones on GitHub main. This is only a probe for "would a newer Unsloth + # support this name?", so it must not change what the installed version + # resolves; the fetched FP8 tables are returned for the probe to use + # instead of being written over the installed ones. + namespace = {} + exec(new_mapper, namespace) return ( - NEW_INT_TO_FLOAT_MAPPER, - NEW_FLOAT_TO_INT_MAPPER, - NEW_MAP_TO_UNSLOTH_16bit, + 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"], ) except: - return {}, {}, {} + return {}, {}, {}, {}, {} def _resolve_with_mappers( - model_name, load_in_4bit, load_in_fp8, int_to_float, float_to_int, map_to_unsloth_16bit + model_name, + load_in_4bit, + load_in_fp8, + int_to_float, + float_to_int, + map_to_unsloth_16bit, + fp8_block = None, + fp8_row = None, ): + # fp8_block/fp8_row default to the installed tables; the newer-mapper probe passes the + # fetched ones so it can answer for new FP8 repos without rebinding the installed ones. return __get_model_name( model_name = model_name, load_in_4bit = load_in_4bit, @@ -211,8 +231,8 @@ def _resolve_with_mappers( FLOAT_TO_INT_MAPPER = float_to_int, MAP_TO_UNSLOTH_16bit = map_to_unsloth_16bit, load_in_fp8 = load_in_fp8, - FLOAT_TO_FP8_BLOCK_MAPPER = FLOAT_TO_FP8_BLOCK_MAPPER, - FLOAT_TO_FP8_ROW_MAPPER = FLOAT_TO_FP8_ROW_MAPPER, + FLOAT_TO_FP8_BLOCK_MAPPER = FLOAT_TO_FP8_BLOCK_MAPPER if fp8_block is None else fp8_block, + FLOAT_TO_FP8_ROW_MAPPER = FLOAT_TO_FP8_ROW_MAPPER if fp8_row is None else fp8_row, ) @@ -252,9 +272,13 @@ def get_model_name( and not _env_says_offline() # offline: skip the remote (raw GitHub) mapper refresh ): # Try checking if a new Unsloth version allows it! - NEW_INT_TO_FLOAT_MAPPER, NEW_FLOAT_TO_INT_MAPPER, NEW_MAP_TO_UNSLOTH_16bit = ( - _get_new_mapper() - ) + ( + NEW_INT_TO_FLOAT_MAPPER, + NEW_FLOAT_TO_INT_MAPPER, + NEW_MAP_TO_UNSLOTH_16bit, + NEW_FP8_BLOCK_MAPPER, + NEW_FP8_ROW_MAPPER, + ) = _get_new_mapper() upgraded_model_name = _resolve_with_mappers( model_name = model_name, load_in_4bit = load_in_4bit, @@ -262,6 +286,10 @@ def get_model_name( int_to_float = NEW_INT_TO_FLOAT_MAPPER, float_to_int = NEW_FLOAT_TO_INT_MAPPER, map_to_unsloth_16bit = NEW_MAP_TO_UNSLOTH_16bit, + # the fp8 probe has to look at the FETCHED tables too, or a new fp8 repo would + # miss both here and in the installed tables and skip the upgrade message + fp8_block = NEW_FP8_BLOCK_MAPPER, + fp8_row = NEW_FP8_ROW_MAPPER, ) if upgraded_model_name is not None: raise NotImplementedError(