diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index c16515afb8..a230efd1c9 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -697,6 +697,18 @@ 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, 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 \ @@ -711,6 +723,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_mxfp4_load_path_layout.py b/tests/test_mxfp4_load_path_layout.py new file mode 100644 index 0000000000..70e57096d5 --- /dev/null +++ b/tests/test_mxfp4_load_path_layout.py @@ -0,0 +1,111 @@ +# 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. + +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 +import subprocess +import sys +import textwrap + +import pytest + + +# 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 + 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" + # 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, + ) + 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..e62d57a6c2 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 @@ -36,6 +37,10 @@ _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 +# 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") def _skip_if_transformers_5x(reason: str) -> None: @@ -234,6 +239,58 @@ def _assert_params_superset( ) +# 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", + "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. + + 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"): + return set() + found: set[str] = set() + stack = [code] + while stack: + current = stack.pop() + 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 + + +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 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 " + 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) @@ -1241,8 +1298,21 @@ 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 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( + 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 +1321,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 +1336,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 +1344,133 @@ def test_mxfp4_dequantize_signature(): ) +def test_mxfp4_dequantize_convertops_signature(): + """mxfp4.py patches ``mxfp4.dequantize_convertops`` (blocks, scales) on + transformers 5.x. + + 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: + 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", + ) + # 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( + "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, + 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 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_referenced_global_names_ignores_attribute_accesses(): + """The live-path detector must not be satisfied by a same-named ATTRIBUTE. + + ``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 + + 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: @@ -1288,24 +1485,35 @@ 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 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: 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 +1522,29 @@ 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 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() + ): + 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..82a7ed8bd4 100644 --- a/unsloth_zoo/temporary_patches/gpt_oss.py +++ b/unsloth_zoo/temporary_patches/gpt_oss.py @@ -457,10 +457,17 @@ 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 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 + 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 +480,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..14875e1f6a 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,73 @@ 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 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, 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 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 + 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, 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, + 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.