From 6e410798cdda28121dce64092a9326ac70cf84b0 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Thu, 27 Aug 2026 05:05:58 +0000 Subject: [PATCH 1/5] Restore the MXFP4 dequantize transpose on the transformers 5.x load path Zoo's convert_moe_packed_tensors replacement returns the un-transposed [E, D, G*B*2] layout on purpose and leaves transpose(1, 2) to whichever loader hook is live. That hook is mxfp4.dequantize only on 4.x. From transformers 5.0 the live path is Mxfp4Dequantize (a ConversionOps) into dequantize_convertops, so the transpose was never applied and GPT-OSS loaded with dims 1 and 2 swapped. 5.16.0 then deleted dequantize and reduced shard_and_distribute_module to a raising tombstone, which is what turned the drift detector red. Patch dequantize_convertops on 5.x, gate the legacy dequantize patch to < 5.0.0, and gate both shard_and_distribute_module call sites to < 5.16.0. Also drop set_param=False: it only exists on 5.x, and the branch that passed it is now 4.x only, where the signature has no **kwargs to absorb it and every tensor parallel load would TypeError. Add tests/test_mxfp4_load_path_layout.py, which asserts the loaded layout against pristine transformers. The existing mxfp4 tests only cover the save path, which stayed self-consistent and missed this entirely. --- tests/test_mxfp4_load_path_layout.py | 108 ++++++++++++++++ tests/test_temporary_patches_exhaustive.py | 138 ++++++++++++++++----- unsloth_zoo/temporary_patches/gpt_oss.py | 22 +++- unsloth_zoo/temporary_patches/mxfp4.py | 118 +++++++++++------- 4 files changed, 310 insertions(+), 76 deletions(-) create mode 100644 tests/test_mxfp4_load_path_layout.py diff --git a/tests/test_mxfp4_load_path_layout.py b/tests/test_mxfp4_load_path_layout.py new file mode 100644 index 0000000000..680ff9c4d4 --- /dev/null +++ b/tests/test_mxfp4_load_path_layout.py @@ -0,0 +1,108 @@ +# Unsloth Zoo - Utilities for Unsloth +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +"""The mxfp4 LOAD path must hand the model the same layout stock transformers would. + +This is the gap that let a silent weight corruption live across the whole +transformers 5.x line. Zoo's `convert_moe_packed_tensors` replacement returns the +UN-transposed [E, D, G*B*2] layout on purpose, and something downstream has to +restore GPT-OSS's [E, G*B*2, D]. Which hook does that depends on the version: + + * 4.x -> module level `mxfp4.dequantize`, called by quantizer_mxfp4 + * 5.0 and newer -> `Mxfp4Dequantize` (a ConversionOps) -> `dequantize_convertops` + +Zoo only ever patched `dequantize`. From 5.0 nothing calls it, so the transpose was +dropped and GPT-OSS loaded with dims 1 and 2 swapped; 5.16.0 then deleted the +function outright, which is what finally turned the drift detector red. + +The existing mxfp4 tests did not catch any of that, because +test_mxfp4_transpose_convention and test_mxfp4_convert_no_double_transpose both +exercise the SAVE path, which stayed self-consistent throughout. So this file +asserts the LOADED layout specifically: patch, then compare against the value +pristine transformers produces for the same input. + +Runs in a subprocess because applying the patch mutates +`transformers.integrations.mxfp4` process-wide, and CPU-only because the +comparison is pure layout arithmetic with no accelerator needed. +""" + +import os +import subprocess +import sys +import textwrap + +import pytest + + +# Golden is captured BEFORE the patch, in the same process and from the same input, +# so this compares zoo against whatever the installed transformers actually does +# rather than against a hardcoded shape that would itself need version gating. +_PROBE = textwrap.dedent( + """ + import torch + import transformers + import transformers.integrations.mxfp4 as m + + if not hasattr(m, "dequantize_convertops"): + print("SKIP no dequantize_convertops (pre-5.0 ConversionOps path)") + raise SystemExit(0) + + E, D, G, B = 2, 4, 3, 16 + torch.manual_seed(0) + blocks = torch.randint(0, 255, (E, D, G, B), dtype=torch.uint8) + scales = torch.full((E, D, G), 127, dtype=torch.uint8) + + golden = m.dequantize_convertops(blocks.clone(), scales.clone()).clone() + + from unsloth_zoo.temporary_patches.mxfp4 import patch_convert_moe_packed_tensors + patch_convert_moe_packed_tensors() + + after = m.dequantize_convertops(blocks.clone(), scales.clone()) + + # .cpu(): zoo's convert_moe_packed_tensors moves inputs to the accelerator when + # one is present, which is intended and irrelevant to the layout being checked. + got, want = tuple(after.shape), tuple(golden.shape) + if got != want: + print( + "FAIL transformers %s: loader got %s, stock transformers gives %s. " + "Zoo's un-transposed convert_moe_packed_tensors reached the loader with " + "nothing restoring transpose(1, 2), so GPT-OSS loads with dims 1 and 2 " + "swapped." % (transformers.__version__, got, want) + ) + raise SystemExit(1) + if not torch.equal(after.cpu(), golden.cpu()): + print("FAIL transformers %s: shape %s matches but values differ" + % (transformers.__version__, got)) + raise SystemExit(1) + print("OK transformers %s: loader layout %s matches stock" + % (transformers.__version__, got)) + """ +) + + +def test_mxfp4_load_path_layout_matches_stock_transformers(): + pytest.importorskip("transformers.integrations.mxfp4") + env = dict(os.environ) + # conftest's GPU-free harness does not reach a subprocess, and importing + # unsloth_zoo calls get_device_type() at import time. + env["UNSLOTH_ALLOW_CPU"] = "1" + r = subprocess.run( + [sys.executable, "-c", _PROBE], capture_output=True, text=True, env=env, + ) + if "SKIP" in r.stdout: + pytest.skip(r.stdout.strip().split("SKIP", 1)[1].strip()) + assert r.returncode == 0, f"stdout={r.stdout}\nstderr={r.stderr}" + assert "OK" in r.stdout, f"stdout={r.stdout}\nstderr={r.stderr}" diff --git a/tests/test_temporary_patches_exhaustive.py b/tests/test_temporary_patches_exhaustive.py index b276b935f1..2dc7dc2943 100644 --- a/tests/test_temporary_patches_exhaustive.py +++ b/tests/test_temporary_patches_exhaustive.py @@ -36,6 +36,11 @@ _TX_VERSION = getattr(transformers, "__version__", "0.0.0") _TX_IS_5X = Version(_TX_VERSION) >= Version("5.0.0") +# 5.16.0 (upstream PR #47579, the DTensor tensor parallel rewrite) deleted +# transformers.integrations.mxfp4.dequantize and reduced +# transformers.integrations.tensor_parallel.shard_and_distribute_module to a +# tombstone that raises RuntimeError when called. +_TX_GE_5_16 = Version(_TX_VERSION) >= Version("5.16.0") def _skip_if_transformers_5x(reason: str) -> None: @@ -1241,8 +1246,22 @@ def test_mxfp4_convert_moe_packed_tensors_signature(): def test_mxfp4_dequantize_signature(): - """mxfp4.py:220 patches ``mxfp4.dequantize`` (module, param_name, - param_value, target_device, dq_param_name, **kwargs).""" + """mxfp4.py patches ``mxfp4.dequantize`` (module, param_name, param_value, + target_device, dq_param_name, **kwargs) on transformers 4.x ONLY. + + ``dequantize`` is the live per-parameter loader hook only on 4.x, where + quantizer_mxfp4 calls it. From 5.0 the live path is + ``Mxfp4Dequantize`` (a ConversionOps) -> ``dequantize_convertops``, leaving + ``dequantize`` defined but unreferenced, and 5.16.0 deleted it outright. Zoo + gates its patch to < 5.0.0 to match, so pinning the symbol above that range + would report drift against something zoo no longer touches. The 5.x + replacement is pinned by test_mxfp4_dequantize_convertops_signature below.""" + if _TX_IS_5X: + pytest.skip( + f"transformers {_TX_VERSION}: mxfp4.dequantize is not the live " + "loader hook on 5.x and zoo gates its patch to < 5.0.0; the " + "ConversionOps replacement is pinned separately" + ) mod_name = "transformers.integrations.mxfp4" try: mod = importlib.import_module(mod_name) @@ -1251,7 +1270,7 @@ def test_mxfp4_dequantize_signature(): fn = getattr(mod, "dequantize", None) if fn is None: pytest.fail( - "DRIFT DETECTED: zoo temporary_patches/mxfp4.py:220 expects " + "DRIFT DETECTED: zoo temporary_patches/mxfp4.py expects " "transformers.integrations.mxfp4.dequantize but it is missing on " f"transformers {_TX_VERSION}" ) @@ -1266,7 +1285,7 @@ def test_mxfp4_dequantize_signature(): ) if not _has_var_keyword(fn): pytest.fail( - f"DRIFT DETECTED: zoo temporary_patches/mxfp4.py:185 forwards " + f"DRIFT DETECTED: zoo temporary_patches/mxfp4.py forwards " f"by name (model=..., empty_param=..., casting_dtype=..., " f"to_contiguous=..., rank=..., device_mesh=...) via **kwargs but " f"upstream transformers.integrations.mxfp4.dequantize lost its " @@ -1274,6 +1293,51 @@ def test_mxfp4_dequantize_signature(): ) +def test_mxfp4_dequantize_convertops_signature(): + """mxfp4.py patches ``mxfp4.dequantize_convertops`` (blocks, scales) on + transformers 5.x. + + This is the 5.x replacement for the deleted module-level ``dequantize``: + ``Mxfp4Dequantize.convert`` is its only caller, and zoo rebinds it to + re-apply the ``transpose(1, 2)`` that its own un-transposed + ``convert_moe_packed_tensors`` leaves off. Absent on 4.x, where the legacy + ``dequantize`` hook carries the transpose instead.""" + try: + mod = importlib.import_module("transformers.integrations.mxfp4") + except Exception as exc: + pytest.skip(f"mxfp4 integrations unavailable: {exc}") + fn = getattr(mod, "dequantize_convertops", None) + if not _TX_IS_5X: + if fn is None: + pytest.skip( + f"transformers {_TX_VERSION} predates the ConversionOps mxfp4 " + "loading path; zoo patches mxfp4.dequantize instead" + ) + return + if fn is None: + pytest.fail( + "DRIFT DETECTED: zoo temporary_patches/mxfp4.py expects " + "transformers.integrations.mxfp4.dequantize_convertops but it is " + f"missing on transformers {_TX_VERSION}. Without it the loader " + "keeps zoo's un-transposed [E, D, G*B*2] layout and GPT-OSS loads " + "with dims 1 and 2 swapped." + ) + _assert_params_superset( + fn, + required=["blocks", "scales"], + zoo_file="mxfp4.py", + label="dequantize_convertops", + ) + # Zoo's replacement is what Mxfp4Dequantize.convert calls, so the class and + # its call site must both still be there for the patch to reach the loader. + if getattr(mod, "Mxfp4Dequantize", None) is None: + pytest.fail( + "DRIFT DETECTED: zoo temporary_patches/mxfp4.py relies on " + "transformers.integrations.mxfp4.Mxfp4Dequantize calling " + f"dequantize_convertops, but the class is missing on {_TX_VERSION}" + ) + + def test_mxfp4_fp4_values_constant_present(): """mxfp4.py:113 / :227 imports ``FP4_VALUES`` constant.""" try: @@ -1288,24 +1352,36 @@ def test_mxfp4_fp4_values_constant_present(): def test_mxfp4_shard_and_distribute_module_present(): - """mxfp4.py:181 imports - ``transformers.integrations.tensor_parallel.shard_and_distribute_module``; - patched dequantize delegates when device_mesh != None. Positional - arity must be 8 for zoo's call to land.""" + """mxfp4.py (4.x gate) and gpt_oss.py (< 5.16.0 gate) import + ``transformers.integrations.tensor_parallel.shard_and_distribute_module`` + and delegate to it when device_mesh != None. Positional arity must be 8 + for zoo's call to land. + + 5.16.0 replaced the function with a ``(*args, **kwargs)`` tombstone that + raises RuntimeError, so there is nothing left to pin above that version -- + zoo gates both call sites off instead. Note the tombstone is a genuine + upstream stub, NOT a decorator wrapper: it has no ``__wrapped__``, so + ``follow_wrapped`` would not recover a real signature.""" try: mod = importlib.import_module("transformers.integrations.tensor_parallel") except Exception as exc: pytest.fail( - "DRIFT DETECTED: zoo temporary_patches/mxfp4.py:181 imports " + "DRIFT DETECTED: zoo temporary_patches/mxfp4.py imports " f"transformers.integrations.tensor_parallel but it is missing: {exc}" ) fn = getattr(mod, "shard_and_distribute_module", None) if fn is None: pytest.fail( - "DRIFT DETECTED: zoo temporary_patches/mxfp4.py:181 expects " + "DRIFT DETECTED: zoo temporary_patches/mxfp4.py expects " "shard_and_distribute_module but it is missing on transformers " f"{_TX_VERSION}" ) + if _TX_GE_5_16: + pytest.skip( + f"transformers {_TX_VERSION}: shard_and_distribute_module is a " + "deprecation tombstone that raises RuntimeError; zoo gates both " + "call sites to < 5.16.0 and uses tp_plan= above it" + ) # 8 positionals: model, param_value, empty_param, dq_param_name, # casting_dtype, to_contiguous, rank, device_mesh. params = [ @@ -1314,30 +1390,32 @@ def test_mxfp4_shard_and_distribute_module_present(): inspect.Parameter.POSITIONAL_ONLY) ] if len(params) < 8: + sig = inspect.signature(fn) + # A bare (*args, **kwargs) is not a function that stopped taking + # arguments, it is a passthrough: upstream's 5.16.0 tombstone accepts + # any call and unconditionally raises RuntimeError. Reporting that as + # "accepts only 0 positionals" sends the reader looking for a signature + # change that never happened, so name what it actually is. Worth + # stating explicitly that it is NOT a decorator wrapper either: it has + # no __wrapped__, so follow_wrapped recovers nothing. + if _has_var_keyword(fn) and any( + p.kind is inspect.Parameter.VAR_POSITIONAL for p in sig.parameters.values() + ): + detail = ( + f"accepts any call and forwards nothing: {sig}. That is a " + "deprecation tombstone or passthrough, not a narrowed " + "signature, so the call site cannot be repaired by dropping " + "arguments and must be gated off or ported" + ) + else: + detail = f"accepts only {len(params)}: {sig}" pytest.fail( - f"DRIFT DETECTED: zoo temporary_patches/mxfp4.py:196 calls " - f"shard_and_distribute_module with 8 positionals but installed " - f"signature on transformers {_TX_VERSION} accepts only " - f"{len(params)}: {inspect.signature(fn)}" + f"DRIFT DETECTED: zoo temporary_patches/mxfp4.py calls " + f"shard_and_distribute_module with 8 positionals but the installed " + f"signature on transformers {_TX_VERSION} {detail}" ) -def test_mxfp4_shard_and_distribute_set_param_kwarg_or_4x_compat(): - """mxfp4.py:196 passes ``set_param=False`` (5.x kwarg). 4.x without - **kwargs catch-all TypeErrors on the TP path; surface as skip.""" - mod = importlib.import_module("transformers.integrations.tensor_parallel") - fn = mod.shard_and_distribute_module - if "set_param" in _param_names(fn): - return # 5.x - if _has_var_keyword(fn): - return # **kwargs swallows set_param - pytest.skip( - f"transformers {_TX_VERSION} predates set_param kwarg on " - "shard_and_distribute_module; zoo's TP path (device_mesh != None) " - "requires 5.x. Non-TP users unaffected." - ) - - def test_mxfp4_mxfp4_config_top_level_class(): """mxfp4.py:93 imports top-level ``transformers.Mxfp4Config``.""" if not hasattr(transformers, "Mxfp4Config"): diff --git a/unsloth_zoo/temporary_patches/gpt_oss.py b/unsloth_zoo/temporary_patches/gpt_oss.py index 94ff6ee56d..049ba795ed 100644 --- a/unsloth_zoo/temporary_patches/gpt_oss.py +++ b/unsloth_zoo/temporary_patches/gpt_oss.py @@ -457,10 +457,18 @@ def mlp_forward(self, hidden_states): except Exception as e: return raise_error("triton_kernels.matmul_ogs", e) - try: - from transformers.integrations.tensor_parallel import shard_and_distribute_module - except Exception as e: - return raise_error("transformers.integrations.tensor_parallel.shard_and_distribute_module", e) + # Legacy per-parameter TP loader hook. transformers 5.16.0 (upstream PR #47579, + # the DTensor tensor parallel rewrite) reduced it to a tombstone that raises + # RuntimeError, and dropped load_and_swizzle_mxfp4 in the same release, so the + # patch below is inert there. Resolve it leniently instead of bailing out, so + # replace_with_mxfp4_linear further down still gets patched on 5.16.0+. + if transformers_version < Version("5.16.0"): + try: + from transformers.integrations.tensor_parallel import shard_and_distribute_module + except Exception as e: + return raise_error("transformers.integrations.tensor_parallel.shard_and_distribute_module", e) + else: + shard_and_distribute_module = None def load_and_swizzle_mxfp4(module, param_name, param_value, target_device, *args, **kwargs): model = kwargs.get("model", None) @@ -473,6 +481,12 @@ def load_and_swizzle_mxfp4(module, param_name, param_value, target_device, *args for proj in ["gate_up_proj", "down_proj"]: if proj in param_name: if device_mesh is not None: + if shard_and_distribute_module is None: + raise RuntimeError( + "Unsloth: tensor parallel MXFP4 loading needs transformers < 5.16.0, " + "which still provides shard_and_distribute_module. On 5.16.0 and newer " + "load with from_pretrained(..., tp_plan=...) instead." + ) shard_and_distribute_module(model, param_value, empty_param, param_name, casting_dtype, to_contiguous, rank, device_mesh) else: setattr(module, param_name.rsplit(".", 1)[1], torch.nn.Parameter(param_value, requires_grad=False)) diff --git a/unsloth_zoo/temporary_patches/mxfp4.py b/unsloth_zoo/temporary_patches/mxfp4.py index 1769a53c28..624506ef12 100644 --- a/unsloth_zoo/temporary_patches/mxfp4.py +++ b/unsloth_zoo/temporary_patches/mxfp4.py @@ -21,9 +21,13 @@ import torch.nn as nn import os import math +from importlib.metadata import version as importlib_version +from unsloth_zoo.utils import Version from .common import TEMPORARY_PATCHES, UNSLOTH_ENABLE_LOGGING, logger from .utils import patch_function, raise_error +transformers_version = Version(importlib_version("transformers")) + # UNSLOTH_MXFP4_NO_DEQUANTIZE=1 keeps MXFP4 quantized (needs triton_kernels); # otherwise weights dequantize to bf16 for LoRA training. UNSLOTH_MXFP4_NO_DEQUANTIZE = os.environ.get("UNSLOTH_MXFP4_NO_DEQUANTIZE", "0") == "1" @@ -152,48 +156,78 @@ def convert_moe_packed_tensors( Transformers 4.55.4 did dequantized.transpose(1, 2).contiguous().to(target_device) but new versions > 4.56.0 removed the transpose(1, 2) and moved it into patch_convert_moe_packed_tensors """ - try: - import transformers.integrations.mxfp4 - from transformers.integrations.tensor_parallel import shard_and_distribute_module - except Exception as e: - return raise_error("transformers.integrations.mxfp4.dequantize", e) - - def dequantize(module, param_name, param_value, target_device, dq_param_name, **kwargs): - model = kwargs.get("model", None) - empty_param = kwargs.get("empty_param", None) - casting_dtype = kwargs.get("casting_dtype", None) - to_contiguous = kwargs.get("to_contiguous", None) - rank = kwargs.get("rank", None) - device_mesh = kwargs.get("device_mesh", None) - - for proj in ["gate_up_proj", "down_proj"]: - if proj in param_name: - if device_mesh is not None: - param_value = shard_and_distribute_module( - model, - param_value, - empty_param, - dq_param_name, - casting_dtype, - to_contiguous, - rank, - device_mesh, - set_param=False, - ) - blocks_attr = f"{proj}_blocks" - scales_attr = f"{proj}_scales" - setattr(module, param_name.rsplit(".", 1)[1], param_value) - if hasattr(module, blocks_attr) and hasattr(module, scales_attr): - dequantized = convert_moe_packed_tensors(getattr(module, blocks_attr), getattr(module, scales_attr)) - # [HERE] we must do transpose(1, 2) - dequantized = dequantized.transpose(1, 2).contiguous().to(target_device) - # TODO: this is perhaps necessary since if target_device is cpu, and the param was on gpu - if target_device == "cpu" and torch.cuda.is_available(): - torch.cuda.empty_cache() - setattr(module, proj, torch.nn.Parameter(dequantized)) - delattr(module, blocks_attr) - delattr(module, scales_attr) - patch_function(transformers.integrations.mxfp4, "dequantize", dequantize) + # convert_moe_packed_tensors above deliberately returns the UN-transposed + # [E, D, G*B*2] layout -- saving_utils._mxfp4_base_returns_transposed keys the + # export path off exactly that convention -- so whichever loader hook is live + # has to restore GPT-OSS's [E, G*B*2, D] layout itself. + # + # Which hook is live depends on the transformers version: + # 4.x -> module level mxfp4.dequantize, called by quantizer_mxfp4 + # 5.0 and newer -> Mxfp4Dequantize (a ConversionOps) -> dequantize_convertops + # On 5.0 to 5.15 mxfp4.dequantize still exists but nothing calls it, and 5.16.0 + # (upstream PR #47579, the DTensor tensor parallel rewrite) deleted it outright + # together with shard_and_distribute_module. Patching only dequantize therefore + # dropped the transpose across the whole 5.x line and loaded GPT-OSS with dims 1 + # and 2 silently swapped. + + if hasattr(transformers.integrations.mxfp4, "dequantize_convertops"): + # 5.x path. Its only caller is Mxfp4Dequantize.convert. This closes over the + # un-transposed convert_moe_packed_tensors defined above instead of re-reading + # the module attribute, so the transpose stays correct even if the + # patch_function call above did not take and the attribute is still upstream's + # self-transposing version. + def dequantize_convertops(blocks, scales): + dequantized = convert_moe_packed_tensors(blocks, scales) + return torch.nn.Parameter(dequantized.transpose(1, 2).contiguous()) + patch_function(transformers.integrations.mxfp4, "dequantize_convertops", dequantize_convertops) + + if transformers_version < Version("5.0.0"): + # 4.x path. shard_and_distribute_module is only reachable from here, and only + # when a device_mesh is present, so it is imported inside the gate: on 5.16.0+ + # the name survives as a tombstone that raises RuntimeError when called. + try: + import transformers.integrations.mxfp4 + from transformers.integrations.tensor_parallel import shard_and_distribute_module + except Exception as e: + return raise_error("transformers.integrations.mxfp4.dequantize", e) + + def dequantize(module, param_name, param_value, target_device, dq_param_name, **kwargs): + model = kwargs.get("model", None) + empty_param = kwargs.get("empty_param", None) + casting_dtype = kwargs.get("casting_dtype", None) + to_contiguous = kwargs.get("to_contiguous", None) + rank = kwargs.get("rank", None) + device_mesh = kwargs.get("device_mesh", None) + + for proj in ["gate_up_proj", "down_proj"]: + if proj in param_name: + if device_mesh is not None: + # 8 positionals and no set_param: that kwarg only exists on 5.x, + # and this branch is now 4.x only, where passing it TypeErrors. + param_value = shard_and_distribute_module( + model, + param_value, + empty_param, + dq_param_name, + casting_dtype, + to_contiguous, + rank, + device_mesh, + ) + blocks_attr = f"{proj}_blocks" + scales_attr = f"{proj}_scales" + setattr(module, param_name.rsplit(".", 1)[1], param_value) + if hasattr(module, blocks_attr) and hasattr(module, scales_attr): + dequantized = convert_moe_packed_tensors(getattr(module, blocks_attr), getattr(module, scales_attr)) + # [HERE] we must do transpose(1, 2) + dequantized = dequantized.transpose(1, 2).contiguous().to(target_device) + # TODO: this is perhaps necessary since if target_device is cpu, and the param was on gpu + if target_device == "cpu" and torch.cuda.is_available(): + torch.cuda.empty_cache() + setattr(module, proj, torch.nn.Parameter(dequantized)) + delattr(module, blocks_attr) + delattr(module, scales_attr) + patch_function(transformers.integrations.mxfp4, "dequantize", dequantize) """ Add a new CPU-optimized version of convert_moe_packed_tensors with smaller default chunk size. From 39acbdc0c1bbbb57ff83c03811e0abea6acb1b1c Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Thu, 27 Aug 2026 05:26:03 +0000 Subject: [PATCH 2/5] Assert the patched mxfp4 symbols are actually on the live load path Existence and signature pinning alone is what let this bug run from transformers 5.0.0 to 5.16.0. Zoo patched mxfp4.dequantize while every 5.x release loaded through Mxfp4Dequantize into dequantize_convertops. The symbol existed with an unchanged signature the whole time, so every detector stayed green while the patch reached nothing. Add _referenced_global_names and _assert_on_live_path, which walk a caller's code objects (nested code included, so a call inside a comprehension or nested def is still seen) and fail when the symbol zoo patches has no live caller. Wire it into the ConversionOps detector, and add the 4.x equivalent asserting quantizer_mxfp4 still calls the legacy hook. The 4.x check skips rather than fails when transformers.quantizers cannot import, since a broken optional backend says nothing about drift. --- tests/test_temporary_patches_exhaustive.py | 103 ++++++++++++++++++++- 1 file changed, 102 insertions(+), 1 deletion(-) diff --git a/tests/test_temporary_patches_exhaustive.py b/tests/test_temporary_patches_exhaustive.py index 2dc7dc2943..2518322d91 100644 --- a/tests/test_temporary_patches_exhaustive.py +++ b/tests/test_temporary_patches_exhaustive.py @@ -239,6 +239,47 @@ def _assert_params_superset( ) +def _referenced_global_names(obj) -> set[str]: + """Every global name reachable from ``obj``'s code, nested code included. + + ``co_names`` on the top-level code object alone is not enough: a call made + inside a comprehension, a nested def or a lambda lives in a child code + object hanging off ``co_consts``, so the name would be invisible. Walk the + whole tree.""" + code = getattr(obj, "__code__", None) or getattr(obj, "__func__", None) + code = getattr(code, "__code__", code) + if code is None or not hasattr(code, "co_names"): + return set() + found: set[str] = set() + stack = [code] + while stack: + current = stack.pop() + found |= set(current.co_names) + stack += [c for c in current.co_consts if hasattr(c, "co_names")] + return found + + +def _assert_on_live_path(callee_name: str, caller, caller_label: str, zoo_file: str): + """DRIFT-fail when ``caller`` no longer references ``callee_name``. + + Existence-and-signature pinning is not sufficient on its own, and this suite + has the scar to prove it: zoo patched ``mxfp4.dequantize`` while every + transformers 5.x actually loaded through ``Mxfp4Dequantize`` -> + ``dequantize_convertops``. The symbol existed with the right signature all + the way to 5.15, so every detector stayed green while the patch reached + nothing and GPT-OSS loaded with dims 1 and 2 swapped. Asking "is the thing + we patch still CALLED by the thing that loads" is what closes that class of + gap, cheaply and statically.""" + if callee_name not in _referenced_global_names(caller): + pytest.fail( + f"DRIFT DETECTED: zoo temporary_patches/{zoo_file} patches " + f"{callee_name}, but {caller_label} no longer references it on " + f"transformers {_TX_VERSION}. The patch would still apply cleanly " + f"and silently reach nothing, which is exactly how the un-transposed " + f"MXFP4 layout survived from 5.0.0 to 5.16.0." + ) + + def _has_var_keyword(func) -> bool: try: sig = inspect.signature(func) @@ -1330,12 +1371,72 @@ def test_mxfp4_dequantize_convertops_signature(): ) # Zoo's replacement is what Mxfp4Dequantize.convert calls, so the class and # its call site must both still be there for the patch to reach the loader. - if getattr(mod, "Mxfp4Dequantize", None) is None: + cls = getattr(mod, "Mxfp4Dequantize", None) + if cls is None: pytest.fail( "DRIFT DETECTED: zoo temporary_patches/mxfp4.py relies on " "transformers.integrations.mxfp4.Mxfp4Dequantize calling " f"dequantize_convertops, but the class is missing on {_TX_VERSION}" ) + convert = getattr(cls, "convert", None) + if convert is None: + pytest.fail( + "DRIFT DETECTED: transformers.integrations.mxfp4.Mxfp4Dequantize " + f"lost its .convert method on {_TX_VERSION}; zoo's " + "dequantize_convertops patch has no live caller" + ) + _assert_on_live_path( + "dequantize_convertops", convert, "Mxfp4Dequantize.convert", "mxfp4.py", + ) + + +def test_mxfp4_legacy_dequantize_is_on_the_live_path_on_4x(): + """On 4.x, ``mxfp4.dequantize`` must still be CALLED by quantizer_mxfp4. + + Zoo patches it there, and a patch on a function nothing calls is a silent + no-op. That is precisely what happened on 5.x, where the symbol survived + until 5.16.0 with an unchanged signature while the loader had already moved + to the ConversionOps path. + + Best effort: transformers.quantizers pulls in optional heavy backends + (torchao and friends), and a broken one of those says nothing about drift, + so an unimportable caller module skips rather than fails.""" + if _TX_IS_5X: + pytest.skip( + f"transformers {_TX_VERSION}: the legacy dequantize hook is not on " + "the live path on 5.x; test_mxfp4_dequantize_convertops_signature " + "pins the ConversionOps replacement instead" + ) + try: + mod = importlib.import_module("transformers.integrations.mxfp4") + quantizer = importlib.import_module( + "transformers.quantizers.quantizer_mxfp4" + ) + except Exception as exc: + pytest.skip(_unimportable("transformers.quantizers.quantizer_mxfp4", exc)) + if getattr(mod, "dequantize", None) is None: + pytest.skip( + f"transformers {_TX_VERSION} has no mxfp4.dequantize; covered by " + "test_mxfp4_dequantize_signature" + ) + cls = getattr(quantizer, "Mxfp4HfQuantizer", None) + if cls is None: + pytest.fail( + "DRIFT DETECTED: transformers.quantizers.quantizer_mxfp4." + f"Mxfp4HfQuantizer is missing on {_TX_VERSION}" + ) + callers = [ + name for name, fn in vars(cls).items() + if isinstance(fn, types.FunctionType) + and "dequantize" in _referenced_global_names(fn) + ] + if not callers: + pytest.fail( + "DRIFT DETECTED: zoo temporary_patches/mxfp4.py patches " + "transformers.integrations.mxfp4.dequantize, but no method on " + f"Mxfp4HfQuantizer calls it on transformers {_TX_VERSION}. The " + "patch would apply cleanly and reach nothing." + ) def test_mxfp4_fp4_values_constant_present(): From 669055c6cef712489834e8dc72c9eef18f5c2f1e Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Thu, 27 Aug 2026 06:07:29 +0000 Subject: [PATCH 3/5] Run the layout test in CI, and stop the live-path check passing vacuously Two gates that did not gate. The layout regression test was in no workflow. Zoo CI runs explicitly named files, so it was only ever collected by `pytest tests/ --collect-only`, which proves it imports and nothing more. It now runs in the Core drift job, whose three lanes resolve to 4.57.6, 5.5.0 and 5.16.x: two affected versions gate it and one exercises the skip. That job sets UNSLOTH_IS_PRESENT in env, which the probe subprocess needs and conftest cannot reach, and it has no continue-on-error and propagates its lane status. The reachability check read co_names, which mixes attribute operands in with globals, so every self.quantization_config.dequantize config read counted as a call site. On transformers 5.16.1, where upstream genuinely removed the call, it still reported three callers against zero real ones, so it would have stayed green through the entire window it exists to prevent. Filtering on the opcode narrows 4.x detection to create_quantized_param alone. --- .github/workflows/consolidated-tests-ci.yml | 9 ++++ tests/test_temporary_patches_exhaustive.py | 56 +++++++++++++++++++-- 2 files changed, 60 insertions(+), 5 deletions(-) diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index c16515afb8..48c4a414df 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -697,6 +697,14 @@ jobs: if [ "$rc" -eq 0 ]; then trc=0 + # test_mxfp4_load_path_layout runs here, not in repo-tests-cpu, for + # two reasons. It is version-aware by construction -- red on 5.5.0 + # and 5.16.x, skipped on 4.57.6 -- so it only means anything across + # more than one pin, and these lanes are the only ones that provide + # them. And its probe runs in a subprocess, which conftest does not + # reach: importing unsloth_zoo there needs UNSLOTH_IS_PRESENT, which + # this job sets in `env` rather than leaving to whether conftest's + # best-effort `import unsloth` happened to succeed. "$venv/bin/python" -m pytest -v --tb=short -rs \ tests/test_upstream_pinned_symbols_transformers.py \ tests/test_torchvision_video_removed_message.py \ @@ -711,6 +719,7 @@ jobs: tests/test_compiler_rewriter_exhaustive.py \ tests/test_compiler_dynamic_exec.py \ tests/test_temporary_patches_exhaustive.py \ + tests/test_mxfp4_load_path_layout.py \ tests/test_missing_parent_package_is_drift.py \ tests/test_unsloth_zoo_lora_merge.py \ tests/test_peft_paramwrapper_layout_drift.py \ diff --git a/tests/test_temporary_patches_exhaustive.py b/tests/test_temporary_patches_exhaustive.py index 2518322d91..2fbeac8a02 100644 --- a/tests/test_temporary_patches_exhaustive.py +++ b/tests/test_temporary_patches_exhaustive.py @@ -18,6 +18,7 @@ from __future__ import annotations +import dis import importlib import importlib.util import types @@ -239,13 +240,30 @@ def _assert_params_superset( ) +# Opcodes whose name operand is a module-level binding (a global or an import), +# i.e. exactly the bindings a ``setattr(module, name, ...)`` patch can reach. +# ``LOAD_ATTR``/``STORE_ATTR`` are deliberately NOT here: co_names is documented +# as "names other than arguments and function locals", so it mixes attribute +# names in with globals, and matching on it would let an unrelated +# ``self.quantization_config.dequantize`` masquerade as a call to the +# module-level ``dequantize``. +_GLOBAL_NAME_OPS = frozenset(( + "LOAD_GLOBAL", "STORE_GLOBAL", "DELETE_GLOBAL", + "LOAD_NAME", "STORE_NAME", "DELETE_NAME", + "IMPORT_NAME", "IMPORT_FROM", +)) + + def _referenced_global_names(obj) -> set[str]: """Every global name reachable from ``obj``'s code, nested code included. - ``co_names`` on the top-level code object alone is not enough: a call made - inside a comprehension, a nested def or a lambda lives in a child code - object hanging off ``co_consts``, so the name would be invisible. Walk the - whole tree.""" + Scanning the top-level code object alone is not enough: a call made inside a + comprehension, a nested def or a lambda lives in a child code object hanging + off ``co_consts``, so the name would be invisible. Walk the whole tree. + + Filter on the opcode rather than reading ``co_names`` wholesale, otherwise + attribute accesses land in the result and the drift check passes + vacuously.""" code = getattr(obj, "__code__", None) or getattr(obj, "__func__", None) code = getattr(code, "__code__", code) if code is None or not hasattr(code, "co_names"): @@ -254,7 +272,9 @@ def _referenced_global_names(obj) -> set[str]: stack = [code] while stack: current = stack.pop() - found |= set(current.co_names) + for instruction in dis.get_instructions(current): + if instruction.opname in _GLOBAL_NAME_OPS and isinstance(instruction.argval, str): + found.add(instruction.argval) stack += [c for c in current.co_consts if hasattr(c, "co_names")] return found @@ -1439,6 +1459,32 @@ def test_mxfp4_legacy_dequantize_is_on_the_live_path_on_4x(): ) +def test_referenced_global_names_ignores_attribute_accesses(): + """The live-path detector must not be satisfied by a same-named ATTRIBUTE. + + ``co_names`` is "names other than arguments and function locals", so it + carries ``LOAD_ATTR``/``STORE_ATTR`` operands alongside globals. Reading it + wholesale made the 4.x gate vacuous: ``Mxfp4HfQuantizer`` reads + ``self.quantization_config.dequantize`` in half a dozen methods, so the + name is present whether or not anything still calls the module-level + ``dequantize`` that zoo patches. Real transformers 5.x is the proof - the + call site is gone there while every config read remains.""" + def attribute_only(cfg): + cfg.dequantize = True + return cfg.dequantize + + def real_global_call(x): + return dequantize(x) # noqa: F821 + + def real_import_from(x): + from transformers.integrations.mxfp4 import dequantize + return dequantize(x) + + assert "dequantize" not in _referenced_global_names(attribute_only) + assert "dequantize" in _referenced_global_names(real_global_call) + assert "dequantize" in _referenced_global_names(real_import_from) + + def test_mxfp4_fp4_values_constant_present(): """mxfp4.py:113 / :227 imports ``FP4_VALUES`` constant.""" try: From ac22c1d65556696faae29b71548d16ff410dfff6 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Thu, 27 Aug 2026 06:11:11 +0000 Subject: [PATCH 4/5] Tighten mxfp4 comments and correct upstream facts Compress the comments and docstrings added by this PR, and fix three statements about upstream transformers: - set_param was removed in 4.57.0 and is absent from 5.x, not a 5.x only kwarg - shard_and_distribute_module was not deleted in 5.16.0; it still imports and is a tombstone that raises only when called, so the gates are defensive - the dropped transpose affects 5.1.0 and newer, since 5.0.0 dequantize_convertops is still 3-arg --- tests/test_mxfp4_load_path_layout.py | 41 ++++---- tests/test_temporary_patches_exhaustive.py | 108 +++++++++------------ unsloth_zoo/temporary_patches/gpt_oss.py | 7 +- unsloth_zoo/temporary_patches/mxfp4.py | 41 ++++---- 4 files changed, 83 insertions(+), 114 deletions(-) diff --git a/tests/test_mxfp4_load_path_layout.py b/tests/test_mxfp4_load_path_layout.py index 680ff9c4d4..73ea38aab8 100644 --- a/tests/test_mxfp4_load_path_layout.py +++ b/tests/test_mxfp4_load_path_layout.py @@ -16,27 +16,21 @@ """The mxfp4 LOAD path must hand the model the same layout stock transformers would. -This is the gap that let a silent weight corruption live across the whole -transformers 5.x line. Zoo's `convert_moe_packed_tensors` replacement returns the -UN-transposed [E, D, G*B*2] layout on purpose, and something downstream has to -restore GPT-OSS's [E, G*B*2, D]. Which hook does that depends on the version: - - * 4.x -> module level `mxfp4.dequantize`, called by quantizer_mxfp4 - * 5.0 and newer -> `Mxfp4Dequantize` (a ConversionOps) -> `dequantize_convertops` - -Zoo only ever patched `dequantize`. From 5.0 nothing calls it, so the transpose was -dropped and GPT-OSS loaded with dims 1 and 2 swapped; 5.16.0 then deleted the -function outright, which is what finally turned the drift detector red. - -The existing mxfp4 tests did not catch any of that, because -test_mxfp4_transpose_convention and test_mxfp4_convert_no_double_transpose both -exercise the SAVE path, which stayed self-consistent throughout. So this file -asserts the LOADED layout specifically: patch, then compare against the value -pristine transformers produces for the same input. - -Runs in a subprocess because applying the patch mutates -`transformers.integrations.mxfp4` process-wide, and CPU-only because the -comparison is pure layout arithmetic with no accelerator needed. +Zoo's `convert_moe_packed_tensors` replacement returns the UN-transposed +[E, D, G*B*2] layout on purpose, so the live loader hook has to restore GPT-OSS's +[E, G*B*2, D]: + + * 4.x -> module level `mxfp4.dequantize`, called by quantizer_mxfp4 + * 5.1.0 and newer -> `Mxfp4Dequantize` (a ConversionOps) -> `dequantize_convertops` + +Zoo only ever patched `dequantize`, so from 5.1.0 the transpose was dropped and +GPT-OSS loaded with dims 1 and 2 swapped; 5.16.0 deleting the function is what +finally turned the drift detector red. The existing mxfp4 tests missed it because +they only exercise the SAVE path, which stayed self-consistent. So assert the +LOADED layout: patch, then compare against pristine transformers on the same input. + +Subprocess because patching mutates `transformers.integrations.mxfp4` process-wide; +CPU-only because the comparison is pure layout arithmetic. """ import os @@ -47,9 +41,8 @@ import pytest -# Golden is captured BEFORE the patch, in the same process and from the same input, -# so this compares zoo against whatever the installed transformers actually does -# rather than against a hardcoded shape that would itself need version gating. +# Golden is captured BEFORE the patch, same process and input, so this compares zoo +# against the installed transformers rather than a hardcoded, version-gated shape. _PROBE = textwrap.dedent( """ import torch diff --git a/tests/test_temporary_patches_exhaustive.py b/tests/test_temporary_patches_exhaustive.py index 2fbeac8a02..e62d57a6c2 100644 --- a/tests/test_temporary_patches_exhaustive.py +++ b/tests/test_temporary_patches_exhaustive.py @@ -38,9 +38,8 @@ _TX_VERSION = getattr(transformers, "__version__", "0.0.0") _TX_IS_5X = Version(_TX_VERSION) >= Version("5.0.0") # 5.16.0 (upstream PR #47579, the DTensor tensor parallel rewrite) deleted -# transformers.integrations.mxfp4.dequantize and reduced -# transformers.integrations.tensor_parallel.shard_and_distribute_module to a -# tombstone that raises RuntimeError when called. +# mxfp4.dequantize and reduced tensor_parallel.shard_and_distribute_module to a +# tombstone that still imports but raises RuntimeError when called. _TX_GE_5_16 = Version(_TX_VERSION) >= Version("5.16.0") @@ -240,12 +239,10 @@ def _assert_params_superset( ) -# Opcodes whose name operand is a module-level binding (a global or an import), -# i.e. exactly the bindings a ``setattr(module, name, ...)`` patch can reach. -# ``LOAD_ATTR``/``STORE_ATTR`` are deliberately NOT here: co_names is documented -# as "names other than arguments and function locals", so it mixes attribute -# names in with globals, and matching on it would let an unrelated -# ``self.quantization_config.dequantize`` masquerade as a call to the +# Opcodes whose name operand is a module-level binding, i.e. what a +# ``setattr(module, name, ...)`` patch can reach. ``LOAD_ATTR``/``STORE_ATTR`` are +# deliberately excluded: co_names mixes attribute names in with globals, so they +# would let ``self.quantization_config.dequantize`` masquerade as a call to the # module-level ``dequantize``. _GLOBAL_NAME_OPS = frozenset(( "LOAD_GLOBAL", "STORE_GLOBAL", "DELETE_GLOBAL", @@ -257,13 +254,10 @@ def _assert_params_superset( def _referenced_global_names(obj) -> set[str]: """Every global name reachable from ``obj``'s code, nested code included. - Scanning the top-level code object alone is not enough: a call made inside a - comprehension, a nested def or a lambda lives in a child code object hanging - off ``co_consts``, so the name would be invisible. Walk the whole tree. - - Filter on the opcode rather than reading ``co_names`` wholesale, otherwise - attribute accesses land in the result and the drift check passes - vacuously.""" + Walks the whole tree because a call inside a comprehension, nested def or + lambda lives in a child code object off ``co_consts``. Filters on opcode + rather than reading ``co_names`` wholesale, otherwise attribute accesses land + in the result and the drift check passes vacuously.""" code = getattr(obj, "__code__", None) or getattr(obj, "__func__", None) code = getattr(code, "__code__", code) if code is None or not hasattr(code, "co_names"): @@ -282,14 +276,11 @@ def _referenced_global_names(obj) -> set[str]: def _assert_on_live_path(callee_name: str, caller, caller_label: str, zoo_file: str): """DRIFT-fail when ``caller`` no longer references ``callee_name``. - Existence-and-signature pinning is not sufficient on its own, and this suite - has the scar to prove it: zoo patched ``mxfp4.dequantize`` while every - transformers 5.x actually loaded through ``Mxfp4Dequantize`` -> - ``dequantize_convertops``. The symbol existed with the right signature all - the way to 5.15, so every detector stayed green while the patch reached - nothing and GPT-OSS loaded with dims 1 and 2 swapped. Asking "is the thing - we patch still CALLED by the thing that loads" is what closes that class of - gap, cheaply and statically.""" + Existence-and-signature pinning is not enough: zoo patched ``mxfp4.dequantize`` + while 5.x loaded through ``Mxfp4Dequantize`` -> ``dequantize_convertops``. The + symbol kept its signature to 5.15, so detectors stayed green while the patch + reached nothing. Asking whether the patched symbol is still CALLED closes that + gap statically.""" if callee_name not in _referenced_global_names(caller): pytest.fail( f"DRIFT DETECTED: zoo temporary_patches/{zoo_file} patches " @@ -1310,12 +1301,11 @@ def test_mxfp4_dequantize_signature(): """mxfp4.py patches ``mxfp4.dequantize`` (module, param_name, param_value, target_device, dq_param_name, **kwargs) on transformers 4.x ONLY. - ``dequantize`` is the live per-parameter loader hook only on 4.x, where - quantizer_mxfp4 calls it. From 5.0 the live path is - ``Mxfp4Dequantize`` (a ConversionOps) -> ``dequantize_convertops``, leaving - ``dequantize`` defined but unreferenced, and 5.16.0 deleted it outright. Zoo - gates its patch to < 5.0.0 to match, so pinning the symbol above that range - would report drift against something zoo no longer touches. The 5.x + ``dequantize`` is the live loader hook only on 4.x, where quantizer_mxfp4 calls + it. From 5.1.0 the live path is ``Mxfp4Dequantize`` (a ConversionOps) -> + ``dequantize_convertops``, leaving ``dequantize`` defined but unreferenced until + 5.16.0 deleted it. Zoo gates its patch to < 5.0.0, so pinning the symbol above + that range would report drift against something zoo no longer touches; the 5.x replacement is pinned by test_mxfp4_dequantize_convertops_signature below.""" if _TX_IS_5X: pytest.skip( @@ -1358,11 +1348,10 @@ def test_mxfp4_dequantize_convertops_signature(): """mxfp4.py patches ``mxfp4.dequantize_convertops`` (blocks, scales) on transformers 5.x. - This is the 5.x replacement for the deleted module-level ``dequantize``: - ``Mxfp4Dequantize.convert`` is its only caller, and zoo rebinds it to - re-apply the ``transpose(1, 2)`` that its own un-transposed - ``convert_moe_packed_tensors`` leaves off. Absent on 4.x, where the legacy - ``dequantize`` hook carries the transpose instead.""" + The 5.x replacement for module-level ``dequantize``: ``Mxfp4Dequantize.convert`` + is its only caller, and zoo rebinds it to re-apply the ``transpose(1, 2)`` its + own un-transposed ``convert_moe_packed_tensors`` leaves off. Absent on 4.x, where + the legacy ``dequantize`` hook carries the transpose instead.""" try: mod = importlib.import_module("transformers.integrations.mxfp4") except Exception as exc: @@ -1389,8 +1378,8 @@ def test_mxfp4_dequantize_convertops_signature(): zoo_file="mxfp4.py", label="dequantize_convertops", ) - # Zoo's replacement is what Mxfp4Dequantize.convert calls, so the class and - # its call site must both still be there for the patch to reach the loader. + # The class and its call site must both still exist for the patch to reach + # the loader. cls = getattr(mod, "Mxfp4Dequantize", None) if cls is None: pytest.fail( @@ -1413,14 +1402,13 @@ def test_mxfp4_dequantize_convertops_signature(): def test_mxfp4_legacy_dequantize_is_on_the_live_path_on_4x(): """On 4.x, ``mxfp4.dequantize`` must still be CALLED by quantizer_mxfp4. - Zoo patches it there, and a patch on a function nothing calls is a silent - no-op. That is precisely what happened on 5.x, where the symbol survived - until 5.16.0 with an unchanged signature while the loader had already moved - to the ConversionOps path. + Zoo patches it there, and a patch on a function nothing calls is a silent no-op, + which is what happened on 5.x: the symbol survived to 5.16.0 with an unchanged + signature while the loader had already moved to the ConversionOps path. - Best effort: transformers.quantizers pulls in optional heavy backends - (torchao and friends), and a broken one of those says nothing about drift, - so an unimportable caller module skips rather than fails.""" + Best effort: transformers.quantizers pulls in optional heavy backends (torchao + and friends), and a broken one says nothing about drift, so an unimportable + caller module skips rather than fails.""" if _TX_IS_5X: pytest.skip( f"transformers {_TX_VERSION}: the legacy dequantize hook is not on " @@ -1462,13 +1450,11 @@ def test_mxfp4_legacy_dequantize_is_on_the_live_path_on_4x(): def test_referenced_global_names_ignores_attribute_accesses(): """The live-path detector must not be satisfied by a same-named ATTRIBUTE. - ``co_names`` is "names other than arguments and function locals", so it - carries ``LOAD_ATTR``/``STORE_ATTR`` operands alongside globals. Reading it - wholesale made the 4.x gate vacuous: ``Mxfp4HfQuantizer`` reads - ``self.quantization_config.dequantize`` in half a dozen methods, so the - name is present whether or not anything still calls the module-level - ``dequantize`` that zoo patches. Real transformers 5.x is the proof - the - call site is gone there while every config read remains.""" + ``co_names`` carries ``LOAD_ATTR``/``STORE_ATTR`` operands alongside globals, so + reading it wholesale made the 4.x gate vacuous: ``Mxfp4HfQuantizer`` reads + ``self.quantization_config.dequantize`` in several methods, so the name is + present whether or not anything still calls the module-level ``dequantize``. On + real 5.x the call site is gone while every config read remains.""" def attribute_only(cfg): cfg.dequantize = True return cfg.dequantize @@ -1504,11 +1490,10 @@ def test_mxfp4_shard_and_distribute_module_present(): and delegate to it when device_mesh != None. Positional arity must be 8 for zoo's call to land. - 5.16.0 replaced the function with a ``(*args, **kwargs)`` tombstone that - raises RuntimeError, so there is nothing left to pin above that version -- - zoo gates both call sites off instead. Note the tombstone is a genuine - upstream stub, NOT a decorator wrapper: it has no ``__wrapped__``, so - ``follow_wrapped`` would not recover a real signature.""" + 5.16.0 replaced the function with a ``(*args, **kwargs)`` tombstone that raises + RuntimeError, so there is nothing left to pin above that version and zoo gates + both call sites off instead. It is a genuine upstream stub, NOT a decorator + wrapper: no ``__wrapped__``, so ``follow_wrapped`` recovers nothing.""" try: mod = importlib.import_module("transformers.integrations.tensor_parallel") except Exception as exc: @@ -1538,13 +1523,10 @@ def test_mxfp4_shard_and_distribute_module_present(): ] if len(params) < 8: sig = inspect.signature(fn) - # A bare (*args, **kwargs) is not a function that stopped taking - # arguments, it is a passthrough: upstream's 5.16.0 tombstone accepts - # any call and unconditionally raises RuntimeError. Reporting that as - # "accepts only 0 positionals" sends the reader looking for a signature - # change that never happened, so name what it actually is. Worth - # stating explicitly that it is NOT a decorator wrapper either: it has - # no __wrapped__, so follow_wrapped recovers nothing. + # A bare (*args, **kwargs) is a passthrough, not a narrowed signature: + # upstream's 5.16.0 tombstone accepts any call and raises RuntimeError. + # Reporting "accepts only 0 positionals" would send the reader looking for + # a signature change that never happened, so name what it actually is. if _has_var_keyword(fn) and any( p.kind is inspect.Parameter.VAR_POSITIONAL for p in sig.parameters.values() ): diff --git a/unsloth_zoo/temporary_patches/gpt_oss.py b/unsloth_zoo/temporary_patches/gpt_oss.py index 049ba795ed..82a7ed8bd4 100644 --- a/unsloth_zoo/temporary_patches/gpt_oss.py +++ b/unsloth_zoo/temporary_patches/gpt_oss.py @@ -458,10 +458,9 @@ def mlp_forward(self, hidden_states): return raise_error("triton_kernels.matmul_ogs", e) # Legacy per-parameter TP loader hook. transformers 5.16.0 (upstream PR #47579, - # the DTensor tensor parallel rewrite) reduced it to a tombstone that raises - # RuntimeError, and dropped load_and_swizzle_mxfp4 in the same release, so the - # patch below is inert there. Resolve it leniently instead of bailing out, so - # replace_with_mxfp4_linear further down still gets patched on 5.16.0+. + # the DTensor tensor parallel rewrite) reduced it to a tombstone that raises when + # called, and dropped load_and_swizzle_mxfp4, so the patch below is inert there. + # The import itself still succeeds; skipping it is defensive. if transformers_version < Version("5.16.0"): try: from transformers.integrations.tensor_parallel import shard_and_distribute_module diff --git a/unsloth_zoo/temporary_patches/mxfp4.py b/unsloth_zoo/temporary_patches/mxfp4.py index 624506ef12..14875e1f6a 100644 --- a/unsloth_zoo/temporary_patches/mxfp4.py +++ b/unsloth_zoo/temporary_patches/mxfp4.py @@ -156,35 +156,30 @@ def convert_moe_packed_tensors( Transformers 4.55.4 did dequantized.transpose(1, 2).contiguous().to(target_device) but new versions > 4.56.0 removed the transpose(1, 2) and moved it into patch_convert_moe_packed_tensors """ - # convert_moe_packed_tensors above deliberately returns the UN-transposed - # [E, D, G*B*2] layout -- saving_utils._mxfp4_base_returns_transposed keys the - # export path off exactly that convention -- so whichever loader hook is live - # has to restore GPT-OSS's [E, G*B*2, D] layout itself. - # - # Which hook is live depends on the transformers version: - # 4.x -> module level mxfp4.dequantize, called by quantizer_mxfp4 - # 5.0 and newer -> Mxfp4Dequantize (a ConversionOps) -> dequantize_convertops - # On 5.0 to 5.15 mxfp4.dequantize still exists but nothing calls it, and 5.16.0 - # (upstream PR #47579, the DTensor tensor parallel rewrite) deleted it outright - # together with shard_and_distribute_module. Patching only dequantize therefore - # dropped the transpose across the whole 5.x line and loaded GPT-OSS with dims 1 - # and 2 silently swapped. + # convert_moe_packed_tensors above returns the UN-transposed [E, D, G*B*2] layout + # on purpose (saving_utils._mxfp4_base_returns_transposed keys the export path off + # that convention), so the live loader hook must restore GPT-OSS's [E, G*B*2, D]. + # Which hook is live: + # 4.x -> module level mxfp4.dequantize, called by quantizer_mxfp4 + # 5.1.0 and newer -> Mxfp4Dequantize (a ConversionOps) -> dequantize_convertops + # (5.0.0's dequantize_convertops is 3-arg and excluded by pyproject.) mxfp4.dequantize + # survives unreferenced until 5.16.0 (upstream PR #47579, the DTensor TP rewrite) + # deletes it, so patching only dequantize dropped the transpose from 5.1.0 on and + # loaded GPT-OSS with dims 1 and 2 silently swapped. if hasattr(transformers.integrations.mxfp4, "dequantize_convertops"): - # 5.x path. Its only caller is Mxfp4Dequantize.convert. This closes over the - # un-transposed convert_moe_packed_tensors defined above instead of re-reading - # the module attribute, so the transpose stays correct even if the - # patch_function call above did not take and the attribute is still upstream's - # self-transposing version. + # 5.x path, called only by Mxfp4Dequantize.convert. Closes over the local + # un-transposed convert_moe_packed_tensors rather than re-reading the module + # attribute, so the transpose stays correct even if patch_function above did + # not take and upstream's self-transposing version is still installed. def dequantize_convertops(blocks, scales): dequantized = convert_moe_packed_tensors(blocks, scales) return torch.nn.Parameter(dequantized.transpose(1, 2).contiguous()) patch_function(transformers.integrations.mxfp4, "dequantize_convertops", dequantize_convertops) if transformers_version < Version("5.0.0"): - # 4.x path. shard_and_distribute_module is only reachable from here, and only - # when a device_mesh is present, so it is imported inside the gate: on 5.16.0+ - # the name survives as a tombstone that raises RuntimeError when called. + # 4.x path. shard_and_distribute_module is imported inside the gate because on + # 5.16.0+ it still imports fine but is a tombstone that raises when called. try: import transformers.integrations.mxfp4 from transformers.integrations.tensor_parallel import shard_and_distribute_module @@ -202,8 +197,8 @@ def dequantize(module, param_name, param_value, target_device, dq_param_name, ** for proj in ["gate_up_proj", "down_proj"]: if proj in param_name: if device_mesh is not None: - # 8 positionals and no set_param: that kwarg only exists on 5.x, - # and this branch is now 4.x only, where passing it TypeErrors. + # 8 positionals, no set_param: that kwarg was removed in 4.57.0 + # (and is absent from 5.x), so passing it TypeErrors there. param_value = shard_and_distribute_module( model, param_value, From 058b711fc8d57fb2cbbb56b111f0bf5465a7b793 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Thu, 27 Aug 2026 06:37:51 +0000 Subject: [PATCH 5/5] Stop a failed optional unsloth install turning every 5.x lane red The core-upstream lane installs unsloth with `|| true`, so the checkout is allowed to fail. The probe subprocess then imports unsloth_zoo, which raises at __init__.py:250 when find_spec("unsloth") is None, and the failure names unsloth rather than the layout this test exists to police. Neither variable already in play covers it. UNSLOTH_ALLOW_CPU only bypasses accelerator detection, and UNSLOTH_IS_PRESENT guards a separate later check that the import never reaches, so the comment above the invocation named the wrong mechanism as well. The probe now sets UNSLOTH_ZOO_DISABLE_GPU_INIT. Verified with transformers 5.16.1 and unsloth genuinely absent: before, the test fails with the unsloth ImportError; after, it passes. Reverting the source still fails with the layout diagnostic, so the gate is intact. --- .github/workflows/consolidated-tests-ci.yml | 10 +++++++--- tests/test_mxfp4_load_path_layout.py | 10 ++++++++++ 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index 48c4a414df..a230efd1c9 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -702,9 +702,13 @@ jobs: # and 5.16.x, skipped on 4.57.6 -- so it only means anything across # more than one pin, and these lanes are the only ones that provide # them. And its probe runs in a subprocess, which conftest does not - # reach: importing unsloth_zoo there needs UNSLOTH_IS_PRESENT, which - # this job sets in `env` rather than leaving to whether conftest's - # best-effort `import unsloth` happened to succeed. + # reach, so it sets its own environment rather than relying on this + # job's `env`. The binding constraint there is find_spec("unsloth"), + # not the UNSLOTH_IS_PRESENT set below: unsloth is installed with + # `|| true` above, and unsloth_zoo raises on the missing spec before + # it ever reaches the UNSLOTH_IS_PRESENT check. The probe therefore + # sets UNSLOTH_ZOO_DISABLE_GPU_INIT so a failed optional checkout + # cannot turn these lanes red. "$venv/bin/python" -m pytest -v --tb=short -rs \ tests/test_upstream_pinned_symbols_transformers.py \ tests/test_torchvision_video_removed_message.py \ diff --git a/tests/test_mxfp4_load_path_layout.py b/tests/test_mxfp4_load_path_layout.py index 73ea38aab8..70e57096d5 100644 --- a/tests/test_mxfp4_load_path_layout.py +++ b/tests/test_mxfp4_load_path_layout.py @@ -92,6 +92,16 @@ def test_mxfp4_load_path_layout_matches_stock_transformers(): # conftest's GPU-free harness does not reach a subprocess, and importing # unsloth_zoo calls get_device_type() at import time. env["UNSLOTH_ALLOW_CPU"] = "1" + # The core-upstream lane installs unsloth with `|| true`, so the checkout is + # allowed to fail, and unsloth_zoo/__init__.py raises + # ImportError("Please install Unsloth") when find_spec("unsloth") is None. + # UNSLOTH_IS_PRESENT does NOT cover that: it guards a separate, later check + # that this import never reaches. Without the line below, one transient git + # failure turns every 5.x lane red and blames unsloth rather than the layout + # this test exists to police. Both variables are needed and neither implies + # the other: this one skips the unsloth requirement, UNSLOTH_ALLOW_CPU keeps + # get_device_type() from raising on a driverless runner. + env["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = "1" r = subprocess.run( [sys.executable, "-c", _PROBE], capture_output=True, text=True, env=env, )