Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

Appearance settings

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#949
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

@danielhanchen

@danielhanchen danielhanchen commented Jul 26, 2026

Copy link
Copy Markdown
Member

What breaks

Fine-tuning Qwen3.5 fails at the very first forward with the module compiler enabled, which is the default:

  File "unsloth_compiled_cache/unsloth_compiled_module_qwen3_5.py", line 95, in forward
    return wrapped(self, *args, **kwargs)
                          ^^^^
NameError: name 'args' is not defined

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 a forward until 2026-07-01.

Root cause

unsloth_zoo/compiler.py:create_standalone_class reads the model's forward through the class attribute at two places:

  • old_source = inspect.getsource(f.forward)
  • parameters = inspect.signature(f.forward).parameters

Both 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 both inspect.getsource and inspect.signature already 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_hooks and 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:

@force_accelerate_hooks("conv1d")
def forward(
    self,
    hidden_states: torch.Tensor,
    cache_params: Cache | None = None,
    attention_mask: torch.Tensor | None = None,
    **kwargs: Unpack[TransformersKwargs],
):

and transformers/integrations/accelerate.py:931-951 returns a bare inner closure:

def decorator(forward_func: Callable) -> Callable:
    def wrapped(self, *args, **kwargs):
        hooked_module = getattr(self, child_module_name)
        ...
        output = forward_func(self, *args, **kwargs)
        ...
        return output

    return wrapped

There is no functools.wraps anywhere in that file, and functools is not even imported. So for Qwen3_5GatedDeltaNet.forward the class attribute is force_accelerate_hooks.<locals>.decorator.<locals>.wrapped: __name__ is wrapped, __wrapped__ is absent, inspect.signature reports (self, *args, **kwargs), inspect.getsource returns the wrapper from accelerate.py, and co_freevars is ('child_module_name', 'forward_func').

Three failures chain from that:

  1. The compile decision is silently dropped. The re.search(r"def\s+(\w+)\s*\(", forward_source) check at compiler.py:1490 sees the name wrapped rather than forward, misclassifies the method as an Unsloth temporary-patch rename, and sets disable = None. That suppresses the @torch.compiler.disable the caller asked for, and with it the fact that the class is listed in DISABLE_COMPILE_MODULES.
  2. The wrong body is emitted. The decorator's closure body goes out at module scope, where its free variables child_module_name and forward_func do not exist. The real mixer body is never emitted at all.
  3. The call does not match the definition. The forwarding call is built from the wrapper's (self, *args, **kwargs) signature while the definition comes from the class's real def forward(...).

The generated file before this change, verbatim:

def wrapped(self, *args, **kwargs):
    hooked_module = getattr(self, child_module_name)      # NameError waiting to happen
    ...
    output = forward_func(self, *args, **kwargs)
    ...

class Qwen3_5GatedDeltaNet(nn.Module):
    def forward(
        self,
        hidden_states: torch.Tensor,
        cache_params: Cache | None = None,
        attention_mask: torch.Tensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ):
        return wrapped(self, *args, **kwargs)             # NameError: name 'args' is not defined

Note there is no @torch.compiler.disable(recursive = False) on that file even though disable = True was 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:

  • a plain method returns immediately via the __qualname__ fast path;
  • a functools.wraps decorator also returns immediately, because wraps copies __qualname__;
  • an inherited forward, an ambiguous closure holding more than one function, a recovered function belonging to a different class, and a non-function attribute all return the input unchanged.

The load-bearing detail: 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 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:

@torch.compiler.disable(recursive = False)
@force_accelerate_hooks("conv1d")
def Qwen3_5GatedDeltaNet_forward(
    self,
    hidden_states: torch.Tensor,
    ...
):
    hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)   # the real body
    ...

class Qwen3_5GatedDeltaNet(nn.Module):
    def forward(self, hidden_states, cache_params=None, attention_mask=None, **kwargs):
        return Qwen3_5GatedDeltaNet_forward(self, hidden_states=hidden_states, cache_params=cache_params, attention_mask=attention_mask, **kwargs)

Measured before and after

End to end, unsloth/Qwen3.5-2B LoRA fine-tune on a B200, torch 2.11.0+cu128, transformers 5.14.1, module compiler enabled:

before after
10-step LoRA fine-tune NameError: name 'args' is not defined passes
16bit merge + export not reached 4.55 GB, bfloat16, all files present
generation from base / adapter / merged not reached all three coherent

New regression test, run against the same tree with and without the fix:

################ VARIANT: pristine ################
12 failed, 2 passed, 14 warnings in 0.28s
################ VARIANT: patched ################
14 passed, 14 warnings in 0.08s

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.Module class of 30 transformers families through create_standalone_class, in both the disable and compile variants, 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_hooks classes 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 contains def wrapped( or return 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:

tests/test_compiler_dynamic_exec.py tests/test_compiler_output_capture.py
tests/test_compiler_rewriter_exhaustive.py tests/test_compile_cache_security.py
tests/test_temporary_patches_exhaustive.py
-> 286 passed, 28 skipped   (both pristine and patched)

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 mirrors force_accelerate_hooks exactly (bare inner closure, no functools.wraps), writes it to disk so inspect.getsource can read it, and runs it through the real create_standalone_class.

Covers: the premise; the real body being emitted with no unbound name; the disable classification 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_hooks in transformers/integrations/accelerate.py returns a bare inner closure with no
functools.wraps, which is what makes the decorated forward unreadable. That is the trigger, but it is not
a defect on their side: the decorator is correct code doing what it documents, forward on a mixer module is
internal, and nothing guarantees that a class attribute named forward is the undecorated function or that
its source text is recoverable. Reading a method's source with inspect.getsource and regenerating it is a
much deeper assumption than any library offers, and that assumption is ours. The fix belongs here.

Adding the affected classes to DISABLE_COMPILE_MODULES does not work either. Qwen3_5GatedDeltaNet
already matches that list through its GatedDeltaNet suffix entry. Matching only sets disable = True;
the class still goes through create_standalone_class, which is where the defect is, so it is generated
just 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 it
only 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 NameError
at 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.

`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.
Comment thread unsloth_zoo/compiler.py
if getattr(current, "__qualname__", "").startswith(prefix):
return current
return func
pass
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant

Morty Proxy This is a proxified and sanitized view of the page, visit original site.