Keep the newer-mapper probe from replacing the installed FP8 mappers (#7478)

* Keep the newer-mapper probe from replacing the installed FP8 mappers

get_model_name calls _get_new_mapper() whenever a name misses the local
tables, only to answer whether a newer Unsloth would support it. That
helper fetches mapper.py from main, prefixes INT_TO_FLOAT_MAPPER,
FLOAT_TO_INT_MAPPER and MAP_TO_UNSLOTH_16bit with NEW_, and execs 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, the _add_* helpers
and the builder's loop variables, and none of those are renamed.
Exec'ing into globals() therefore rebinds the two FP8 tables that
loader_utils imported from the installed mapper, so every later
get_model_name(..., load_in_fp8 = ...) in the process resolves through
main's table instead of the installed one. The probe deliberately does
not adopt the new 4bit mappers (it raises NotImplementedError asking the
user to upgrade), so silently adopting the new FP8 ones is inconsistent,
and it also leaves loader_utils and mapper disagreeing about the same
tables. Reaching it needs nothing unusual: any org/model name absent
from the tables triggers the fetch.

Exec into a throwaway namespace and read the three mappers out of it, so
the probe stays a read and the installed mappings are left alone.

Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>

* Hand the fetched FP8 tables back from the probe instead of dropping them

Isolating the exec stopped the probe corrupting the installed FP8 tables, but
it also removed the only reason the probe ever saw the fetched ones: the
_resolve_with_mappers call still read FLOAT_TO_FP8_BLOCK_MAPPER and
FLOAT_TO_FP8_ROW_MAPPER off the module globals. A newly added FP8 repo would
then miss both the installed tables and the probe, so an older install would
stop raising the upgrade NotImplementedError for it.

Return the two fetched tables and let _resolve_with_mappers take them as
optional arguments, defaulting to the installed ones. The probe now answers
for new FP8 repos without writing over what the installed version resolves.

_get_new_mapper returns five tables now, so the two existing stubs in
test_get_model_name.py and test_bad_mappings_redirect.py are updated to match.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

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:

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):

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"

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(