Fix compiler codegen for a forward hidden behind a bare decorator#949
Open
danielhanchen wants to merge 1 commit into
mainunslothai/unsloth-zoo:mainfrom
fix-compiler-decorated-forwardunslothai/unsloth-zoo:fix-compiler-decorated-forwardCopy head branch name to clipboard
Open
Fix compiler codegen for a forward hidden behind a bare decorator#949danielhanchen wants to merge 1 commit intomainunslothai/unsloth-zoo:mainfrom fix-compiler-decorated-forwardunslothai/unsloth-zoo:fix-compiler-decorated-forwardCopy head branch name to clipboard
danielhanchen wants to merge 1 commit into
mainunslothai/unsloth-zoo:mainfrom
fix-compiler-decorated-forwardunslothai/unsloth-zoo:fix-compiler-decorated-forwardCopy head branch name to clipboard
Conversation
`create_standalone_class` read the forward through the class attribute at
two places, `inspect.getsource(f.forward)` and
`inspect.signature(f.forward).parameters`, with no allowance for a
decorator that does not use `functools.wraps`.
transformers 5.13.0 decorates the linear attention / mixer forward of 11
model families with `@force_accelerate_hooks(...)`, and
`transformers/integrations/accelerate.py` returns a bare inner closure
`def wrapped(self, *args, **kwargs)` with no `functools.wraps`. The class
attribute is therefore
`force_accelerate_hooks.<locals>.decorator.<locals>.wrapped`, so
`__name__` is `wrapped`, `__wrapped__` is absent, `inspect.getsource`
returns the wrapper from accelerate.py and `inspect.signature` reports
`(self, *args, **kwargs)`.
Three failures chained from that:
1. The `re.search(r"def\s+(\w+)\s*\(", forward_source)` check saw the name
`wrapped` rather than `forward`, misclassified the method as an Unsloth
temporary patch rename and set `disable = None`, silently dropping the
caller's compile decision and with it the `DISABLE_COMPILE_MODULES`
listing.
2. The decorator's closure body was emitted at module scope, where its free
variables `child_module_name` and `forward_func` do not exist. The real
body was never emitted at all.
3. The forwarding call was built from the wrapper's `(self, *args,
**kwargs)` signature while the definition came from the class's real
`def forward(...)`, producing `return wrapped(self, *args, **kwargs)`
under a named signature.
The first forward then raised `NameError: name 'args' is not defined`,
which made Qwen3.5 training fail immediately with the module compiler on
(its default).
Add `_unwrap_undecorated_method`, which walks `__wrapped__` and then
single function closure cells and accepts a candidate only when its
`__qualname__` belongs to the class being compiled, returning the input
unchanged otherwise. Use it at both call sites.
`inspect.getsource` on the real method returns the class body text
including the decorator line, and that is the exact substring replaced in
`full_class`, so the decorator moves onto the generated standalone
function and still fires exactly once.
Verified by generating every nn.Module class of 30 transformers families
twice, once per compiler: 22 of 686 outputs change, and they are exactly
the 11 `force_accelerate_hooks` classes across bamba, falcon_h1,
falcon_mamba, granitemoehybrid, jamba, lfm2, mamba, nemotron_h, qwen3_5,
qwen3_next and zamba2, times the compile and disable variants. Every other
family, llama, qwen2, qwen3, mistral, gemma, gemma2, gemma3, phi3,
mixtral, falcon, gpt2, starcoder2, cohere, olmo, granite, siglip, clip,
qwen2_vl and qwen3_moe, is byte identical.
| if getattr(current, "__qualname__", "").startswith(prefix): | ||
| return current | ||
| return func | ||
| pass |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What breaks
Fine-tuning Qwen3.5 fails at the very first forward with the module compiler enabled, which is the default:
The same generated-source defect hits the linear attention / mixer module of 14 classes across 14 model families, not just Qwen3.5. See the measured list below.
This is a regression from transformers, not from unsloth_zoo. Qwen3.5 fine-tuning worked on every transformers release from v5.2.0 (when the model landed, undecorated) through v5.12.1, and breaks from v5.13.0 onward. Both compiler call sites below have read the forward off the class attribute since compiler.py was first committed in November 2024, and have never changed; nothing in transformers put a
functools.wraps-less decorator on aforwarduntil 2026-07-01.Root cause
unsloth_zoo/compiler.py:create_standalone_classreads the model's forward through the class attribute at two places:old_source = inspect.getsource(f.forward)parameters = inspect.signature(f.forward).parametersBoth assume the class attribute is the method. That holds for an undecorated forward, and for any decorator that uses
functools.wraps(which copies__qualname__and sets__wrapped__, and bothinspect.getsourceandinspect.signaturealready follow__wrapped__). It does not hold for a decorator that returns a bare closure.transformers commit 36778030cb (PR #46978, 2026-07-01, released in v5.13.0 on 2026-07-03) added
force_accelerate_hooksand applied it to the mixer forward of 11 classes in one change; v5.14.0 and v5.14.1 took it to 14 (git grep -l @force_accelerate_hooks v5.14.1 -- 'src/transformers/models/*/modeling_*.py'returns 14 files, against 0 at v5.12.1).transformers/models/qwen3_5/modeling_qwen3_5.py:440(v5.14.1; line 458 on current main) has:and
transformers/integrations/accelerate.py:931-951returns a bare inner closure:There is no
functools.wrapsanywhere in that file, andfunctoolsis not even imported. So forQwen3_5GatedDeltaNet.forwardthe class attribute isforce_accelerate_hooks.<locals>.decorator.<locals>.wrapped:__name__iswrapped,__wrapped__is absent,inspect.signaturereports(self, *args, **kwargs),inspect.getsourcereturns the wrapper from accelerate.py, andco_freevarsis('child_module_name', 'forward_func').Three failures chain from that:
re.search(r"def\s+(\w+)\s*\(", forward_source)check atcompiler.py:1490sees the namewrappedrather thanforward, misclassifies the method as an Unsloth temporary-patch rename, and setsdisable = None. That suppresses the@torch.compiler.disablethe caller asked for, and with it the fact that the class is listed inDISABLE_COMPILE_MODULES.child_module_nameandforward_funcdo not exist. The real mixer body is never emitted at all.(self, *args, **kwargs)signature while the definition comes from the class's realdef forward(...).The generated file before this change, verbatim:
Note there is no
@torch.compiler.disable(recursive = False)on that file even thoughdisable = Truewas requested. That is failure 1.The fix
Add
_unwrap_undecorated_method, and use it at both call sites.It walks
__wrapped__and then single-function closure cells, and accepts a candidate only when its__qualname__belongs to the class being compiled. Anything else returns the input unchanged, so no other decorator shape changes behaviour:__qualname__fast path;functools.wrapsdecorator also returns immediately, becausewrapscopies__qualname__;The load-bearing detail:
inspect.getsourceon the real method returns the class-body text including the decorator line, and that is the exact substring replaced infull_class. So the decorator moves onto the generated standalone function and is removed from the class method, which means it still fires exactly once. There is a test asserting exactly that, including a numerical comparison against eager.After this change the same file is generated as:
Measured before and after
End to end,
unsloth/Qwen3.5-2BLoRA fine-tune on a B200, torch 2.11.0+cu128, transformers 5.14.1, module compiler enabled:NameError: name 'args' is not definedNew regression test, run against the same tree with and without the fix:
The 2 that pass both ways are deliberate controls: the premise test (the class attribute really does report the wrapper's name, signature and source) and the negative case (an undecorated forward is untouched).
Why this cannot affect other models
I generated every
nn.Moduleclass of 30 transformers families throughcreate_standalone_class, in both thedisableandcompilevariants, once with the pristine compiler and once with the patched one, and diffed the two output trees byte for byte.22 of 686 generated files differ. They are exactly the 11
force_accelerate_hooksclasses times the 2 variants:BambaMixer,FalconH1Mixer,FalconMambaMixer,GraniteMoeHybridMambaLayer,JambaMambaMixer,Lfm2ShortConv,MambaMixer,NemotronHMamba2Mixer,Qwen3_5GatedDeltaNet,Qwen3NextGatedDeltaNet,Zamba2MambaMixer.The set of files that differ is identical to the set of files that contained the broken
return wrapped(before the change, and no generated file containsdef wrapped(orreturn wrapped(after it.Every other family is byte identical: llama, qwen2, qwen3, mistral, gemma, gemma2, gemma3, phi3, mixtral, falcon, gpt2, starcoder2, cohere, olmo, granite, siglip, clip, qwen2_vl, qwen3_moe.
Existing suites, same tree with and without the fix, identical results either way:
Known limitation
The recovered function must belong to the class being compiled. A class that inherits an already-decorated forward from a base class is therefore left alone rather than unwrapped. That is deliberately conservative: it keeps the change inert where the correct answer is not obvious. No class in the 30 families swept hits this case, all 14 affected classes define their own forward.
Tests added
tests/test_compiler_decorated_forward.py, GPU-free and transformers-version independent. It builds a fake modeling module whose decorator mirrorsforce_accelerate_hooksexactly (bare inner closure, nofunctools.wraps), writes it to disk soinspect.getsourcecan read it, and runs it through the realcreate_standalone_class.Covers: the premise; the real body being emitted with no unbound name; the
disableclassification surviving; the decorator firing exactly once with numerics matching eager; an undecorated forward being untouched; and six inertness cases for_unwrap_undecorated_method(plain method,functools.wraps, inherited, ambiguous closure, function from another class, non-function attribute, plus stacked bare decorators at depth 1/2/3).Why here and not upstream
force_accelerate_hooksintransformers/integrations/accelerate.pyreturns a bare inner closure with nofunctools.wraps, which is what makes the decoratedforwardunreadable. That is the trigger, but it is nota defect on their side: the decorator is correct code doing what it documents,
forwardon a mixer module isinternal, and nothing guarantees that a class attribute named
forwardis the undecorated function or thatits source text is recoverable. Reading a method's source with
inspect.getsourceand regenerating it is amuch deeper assumption than any library offers, and that assumption is ours. The fix belongs here.
Adding the affected classes to
DISABLE_COMPILE_MODULESdoes not work either.Qwen3_5GatedDeltaNetalready matches that list through its
GatedDeltaNetsuffix entry. Matching only setsdisable = True;the class still goes through
create_standalone_class, which is where the defect is, so it is generatedjust as wrongly.
Loud versus silent
Worth separating what transformers caused from what this repo owns. Against a decorated forward, the
codegen used to raise, and the caller catches that and skips the module with a warning. The
rename-detection branch was originally guarded by
if "@torch.compiler.disable" in forward_source:so itonly fired on Unsloth's own temporary patches. That guard was dropped, restored a day later, and then
replaced by a module-name check that applies to everything except one module. In the two unguarded states a
decorated forward is silently mis-generated instead of raising, which is why this surfaces as a
NameErrorat the first forward rather than as a skipped module. This PR fixes the generation; the guard history is
noted here only so the failure mode is not mistaken for the cause.