[danielhanchen]

GPT-OSS MXFP4 has been loading with dims 1 and 2 swapped on every transformers 5.x

The red CI is the alarm, not the fire. unsloth_zoo/temporary_patches/mxfp4.py replaces convert_moe_packed_tensors with a variant that deliberately returns the un-transposed [E, D, G*B*2] layout, and leaves it to whichever loader hook is live to restore GPT-OSS's [E, G*B*2, D]. Which hook that is depends on the transformers version:

• 4.x: module level mxfp4.dequantize, called from quantizer_mxfp4. Zoo patches it, and it applies the transpose. Correct.
• 5.0 and newer: the live path is Mxfp4Dequantize (a ConversionOps) into dequantize_convertops. Nothing calls dequantize any more.

Zoo only ever patched dequantize. So from transformers 5.0.0 onward the compensating transpose was silently dropped and the model loaded with dims 1 and 2 swapped. dequantize still existed on 5.0 through 5.15, so the patch appeared to apply and nothing complained.

Measured with the real patch_convert_moe_packed_tensors():

transformers	loader output, before this PR	stock transformers gives
5.16.1	(2, 4, 96)	(2, 96, 4)
5.4.0	(2, 4, 96)	(2, 96, 4)
4.57.6	correct (legacy hook live and compensating)	matches

What actually changed upstream

Both CI failures come from a single upstream PR, huggingface/transformers#47579 (the DTensor tensor parallel rewrite), first released in v5.16.0:

tag	mxfp4.dequantize	shard_and_distribute_module
v5.15.1	present	full 8 argument signature
v5.16.0	deleted	(*args, **kwargs) tombstone that raises RuntimeError

dequantize was not renamed or moved, it was deleted as dead code. The tlatest5-trl1latest drift lane installs unpinned transformers>=5,<6, so it picked 5.16.0 up the day it shipped, with no zoo commit to blame.

Worth stating for the next person who reads the drift output: shard_and_distribute_module reporting (*args, **kwargs) is not a decorator wrapper hiding a real signature. It has no __wrapped__, __closure__ is None, co_argcount is 0, and co_filename points at tensor_parallel.py:54, the tombstone's own code object. inspect.signature(..., follow_wrapped=True) returns the identical thing. The detector was reading the truth; the source was wrong.

Attribution: no zoo commit is to blame

Worth stating outright so nobody has to re-derive it. The regression was introduced by an upstream release, not by any commit in this repo. The tlatest5-trl1latest drift lane installs unpinned transformers>=5,<6, so it adopts each new transformers the moment it is published, with no zoo change involved.

The run history dates it to the second:

event	time (UTC)
Core drift job on 183a1ee49, green	started 12:28:11, completed 12:32:08
transformers 5.16.0 published to PyPI	12:32:10
transformers 5.16.1 published to PyPI	14:48:55
Core drift job on af37769fb, red	started 15:55:49, completed 15:59:52

The last green drift job finished two seconds before the breaking release existed on PyPI. The next one started three and a half hours later, after the queue, and installed 5.16.1.

af37769fb (#1108) is simply the commit that happened to be at main when the next run started. It touches compiler.py, empty_model.py, peft_utils.py, saving_utils.py, tokenizer_utils.py, vllm_utils.py, lint config and security tests, and no mxfp4, tensor_parallel or gpt_oss code at all. Blaming it would be wrong. The 5.4 hour queue delay between run creation and job start is what made an unpinned lane flip mid-day.

Was upstream right to do this?

Taking the other reading seriously rather than assuming the newer change is correct:

On dequantize, upstream is right and the removal is clean. It was genuinely dead code from 5.0.0 onward: get_weight_conversions() has returned Mxfp4Dequantize since v5.0.0, and nothing has called the module level dequantize since. Upstream also removed the matching re-export from transformers/integrations/__init__.py, so there is no dangling public name left behind. from transformers.integrations import dequantize worked on 5.4.0 and now raises ImportError cleanly rather than resolving to something broken.

One inconsistency worth noting, not worth an issue. dequantize was publicly re-exported, and it was hard deleted in a minor release with no deprecation cycle, while its sibling in the very same PR, shard_and_distribute_module, got a proper FutureWarning tombstone. Two public symbols removed by one PR, two different removal policies. The tombstone is the better treatment and is why that half of the breakage produced a legible error instead of an AttributeError. This does not change the fix on our side and I am not filing anything upstream, but if someone else hits the dequantize removal cold, that asymmetry is the reason it is harder to diagnose than it should be.

The fix

Patch dequantize_convertops on 5.x. It is a clean choke point: Mxfp4Dequantize.convert is its only caller, verified on both 5.4.0 and 5.16.1. Critically this leaves zoo's un-transposed convention intact, so saving_utils._mxfp4_base_returns_transposed and the export convention tests need no changes at all.

Rejected alternatives:

• Making zoo's convert_moe_packed_tensors self-transpose. Double-transposes on stock 4.55.x, and forces coherent edits to saving_utils, the CPU export branch, and both convention tests.
• Skipping the convert_moe_packed_tensors patch on 5.16.0+. Only fixes 5.16+, leaves 5.0 through 5.15 corrupt, and desyncs the _cpu-presence signal saving_utils keys off.

The legacy dequantize patch is gated to < 5.0.0 and both shard_and_distribute_module call sites to < 5.16.0.

Two more defects CI never reported

1. temporary_patches/gpt_oss.py:476 also calls shard_and_distribute_module with 8 positionals. Its patch target load_and_swizzle_mxfp4 was dropped in the same 5.16.0 release, so the patch is inert there, but the import bailed the whole patch_gpt_oss early and took replace_with_mxfp4_linear down with it. Now resolved leniently.
2. Zoo passed set_param=False, which only exists on 5.x. Now that the branch carrying it is 4.x only, that kwarg would hit a 4.57.6 signature with no **kwargs to absorb it and TypeError on every 4.x tensor parallel load. Dropped, matching upstream 4.57.6's own 8 positional call.

Test evidence

	tx 4.57.6	tx 5.16.1
-k mxfp4, before	8 passed, 1 skipped	2 failed, 7 passed
-k mxfp4, after	8 passed, 1 skipped	7 passed, 2 skipped

Remaining failures on both versions are a pre-existing local bitsandbytes artifact (libbitsandbytes_cuda128.so missing in my workspace), identical on pristine main and unrelated. CI installs a working bitsandbytes.

The drift tests do not catch the corruption. Reverting only the source fix from a byte snapshot, with the test changes left in place:

layout proof 5.16.1: (2, 4, 96)   FAIL - loader gets corrupted layout
drift tests  5.16.1: 7 passed, 2 skipped   <- STILL GREEN

That is the single most important result here. Version gating the two red tests would have turned CI green over a broken loader.

New regression test

tests/test_mxfp4_load_path_layout.py asserts the loaded layout against what pristine transformers produces for the same input. This gap is why the bug survived since 5.0.0: test_mxfp4_transpose_convention and test_mxfp4_convert_no_double_transpose both exercise the save path, which stayed self-consistent throughout.

Subprocess isolated because applying the patch mutates transformers.integrations.mxfp4 process-wide, CPU only, UNSLOTH_ALLOW_CPU=1 since conftest's GPU-free harness does not reach a subprocess. It captures its golden value before patching rather than hardcoding a shape, so it needs no version gating of its own.

Mutation tested in both directions. With the source reverted it fails on 5.16.1 and on 5.4.0, catching the pre-5.16 corruption that no existing test saw:

FAIL transformers 5.4.0: loader got (2, 4, 96), stock transformers gives (2, 96, 4).

It skips cleanly on 4.57.6, which has no dequantize_convertops.

Deliberate scope boundary

Tensor parallel MXFP4 loading on transformers 5.16.0+ now raises a clear RuntimeError pointing at from_pretrained(..., tp_plan=...) rather than silently corrupting weights. That matches what upstream itself now does. Genuinely restoring it means porting to apply_tensor_parallelism, which is a feature port and does not belong in a CI fix. Happy to be overruled if you would rather that land here.

Also fixed the drift message wording: a (*args, **kwargs) signature now reports that it accepts any call and forwards nothing, rather than "accepts only 0 positionals", which sent the reader looking for a signature change that never happened.

Robustifying the detector, not just the loader

The detector was right here and still let a silent corruption run from 5.0.0 to 5.16.0, because it only ever asked "does this symbol exist with this signature" and never "is the symbol we patch the one actually on the live path". mxfp4.dequantize existed with an unchanged signature the entire time. Every pin was green. The patch reached nothing.

This PR adds _referenced_global_names and _assert_on_live_path, which statically walk a caller's code objects and fail when the symbol zoo patches has no live caller. Nested code is included via co_consts, so a call made inside a comprehension, a nested def or a lambda is still seen, which a bare co_names check would miss. It is wired into the ConversionOps detector, and a 4.x equivalent asserts quantizer_mxfp4 still calls the legacy hook.

Verified non-vacuous rather than assumed:

dequantize_convertops referenced by Mxfp4Dequantize.convert : True
legacy 'dequantize'   referenced by Mxfp4Dequantize.convert : False   <- the bug condition
bogus name detected (must be False)                         : False
positive control (symbol with no caller)                    : correctly FAILS
nested comprehension name seen (must be True)               : True

The second line is the point: this mechanism can see the exact condition that made the old patch a no-op.

Known limitation, stated rather than implied. This is a static reference check, so it catches a symbol losing all its callers, which is the failure mode that actually bit us. It will not catch a caller that is present but unreachable at runtime (guarded by a config branch that is never taken, say), and it covers the mxfp4 detectors here rather than the whole suite. The 4.x check also skips rather than fails when transformers.quantizers cannot import, since it pulls optional heavy backends like torchao and a broken one of those says nothing about drift.

Unrelated

#1112 fixes the other main red (the DEVICE_TYPE collection error in the security lane). It is independent of this one.

[danielhanchen]

@codex review

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. Can't wait for the next one!

Reviewed commit: 058b711fc8

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".