From 92c3e48529c8b7f96033f52f93845819b2ae53e3 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Wed, 8 Jul 2026 13:37:44 -0700 Subject: [PATCH] Fix BAD_MAPPINGS not redirecting the -unsloth-bnb-4bit dynamic quants (#6949) --------- Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- .github/workflows/consolidated-tests-ci.yml | 1 + tests/test_bad_mappings_redirect.py | 47 +++++++++++++++++++++ unsloth/models/loader_utils.py | 5 +++ 3 files changed, 53 insertions(+) create mode 100644 tests/test_bad_mappings_redirect.py diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index ae4b386589..6ff3d19ba2 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -364,6 +364,7 @@ jobs: tests/utils/test_attention_masks.py \ tests/utils/test_trunc_normal_patch.py \ tests/python/test_fast_language_model_text_only.py \ + tests/test_bad_mappings_redirect.py \ tests/test_prefetch_snapshot_scope.py \ --deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap' # The deselected test monkeypatches flash_attn_varlen_func, which is diff --git a/tests/test_bad_mappings_redirect.py b/tests/test_bad_mappings_redirect.py new file mode 100644 index 0000000000..49dab2d98b --- /dev/null +++ b/tests/test_bad_mappings_redirect.py @@ -0,0 +1,47 @@ +"""Regression test for BAD_MAPPINGS redirecting oversized dynamic quants. + +get_model_name previously applied BAD_MAPPINGS only to the resolver's output, +but several listed names (the `-unsloth-bnb-4bit` dynamic quants, plus any name +the resolver doesn't map) come back as None, so their BAD_MAPPINGS entries were +dead and the oversized model loaded. Asserting over every entry catches all of +them. The mapper table and the resolver have no heavy imports of their own, +so we exec the import-free mapper module and ast-extract the resolver functions +rather than importing unsloth (which needs a GPU). +""" + +import ast +import os + +_MODELS = os.path.join(os.path.dirname(__file__), os.pardir, "unsloth", "models") + + +def _load_get_model_name(): + mapper_ns = {} + with open(os.path.join(_MODELS, "mapper.py"), encoding = "utf-8") as f: + exec(compile(f.read(), "mapper.py", "exec"), mapper_ns) + + with open(os.path.join(_MODELS, "loader_utils.py"), encoding = "utf-8") as f: + tree = ast.parse(f.read()) + + namespace = dict(mapper_ns) + namespace["SUPPORTS_FOURBIT"] = True + namespace["_env_says_offline"] = lambda: True + namespace["_get_new_mapper"] = lambda: ({}, {}, {}) + + wanted = {"__get_model_name", "_resolve_with_mappers", "get_model_name"} + for node in tree.body: + if isinstance(node, ast.Assign) and any( + getattr(target, "id", None) == "BAD_MAPPINGS" for target in node.targets + ): + exec(compile(ast.Module([node], []), "", "exec"), namespace) + elif isinstance(node, ast.FunctionDef) and node.name in wanted: + exec(compile(ast.Module([node], []), node.name, "exec"), namespace) + + return namespace["get_model_name"], namespace["BAD_MAPPINGS"] + + +def test_bad_mappings_redirect_every_listed_name(): + get_model_name, bad_mappings = _load_get_model_name() + assert bad_mappings, "BAD_MAPPINGS should not be empty" + for name, expected in bad_mappings.items(): + assert get_model_name(name, load_in_4bit = True) == expected, name diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py index d6d6ce877e..fa6282bcf6 100644 --- a/unsloth/models/loader_utils.py +++ b/unsloth/models/loader_utils.py @@ -239,6 +239,11 @@ def get_model_name( and new_model_name.lower() in BAD_MAPPINGS ): new_model_name = BAD_MAPPINGS[new_model_name.lower()] + elif new_model_name is None and model_name.lower() in BAD_MAPPINGS: + # Some bad names (e.g. the `-unsloth-bnb-4bit` dynamic quants) are keys + # of the mappers, not values, so the resolver returns None for them and + # the remap above is skipped; remap the input name directly instead. + new_model_name = BAD_MAPPINGS[model_name.lower()] if ( new_model_name is None