[Bugfix] Resolve decorated forwards without functools.wraps in autowrap_forward

[malaiwah]

SUMMARY:
Fixes #3057.

The sequential-pipeline calibration path (GPTQ, SparseGPT, AWQ via this pipeline) cannot trace any model whose module forward is wrapped by a decorator that does not apply functools.wraps. This blocks quantization of such models entirely.

Concretely, in transformers >= 5.x Qwen3_5GatedDeltaNet.forward is decorated with @force_accelerate_hooks("conv1d") (see transformers/integrations/accelerate.py), which returns an inner def wrapped(self, *args, **kwargs) without copying __wrapped__ or __name__.

Root cause

In src/llmcompressor/pipelines/sequential/ast_helpers.py, autowrap_forward does:

source = inspect.getsource(module.forward)   # returns the wrapper's source
...
exec(code, namespace)
new_forward = namespace["forward"].__get__(module)   # <-- KeyError: 'forward'

Because inspect.getsource(module.forward) returns the source of the wrapper (the inner def wrapped(...)), the re-exec'd code defines a function named wrapped, not forward, and namespace["forward"] raises KeyError: 'forward'. inspect.unwrap does not help because the wrapper has no __wrapped__ attribute to follow.

The fix

Resolve the decorated method to the original forward function before calling inspect.getsource, by searching the wrapper's closure cells for a function named forward:

target = inspect.unwrap(module.forward)
if getattr(target, "__name__", "") != "forward" and getattr(target, "__closure__", None):
    from types import FunctionType

    for cell in target.__closure__:
        try:
            contents = cell.cell_contents
        except ValueError:
            continue
        if isinstance(contents, FunctionType) and contents.__name__ == "forward":
            target = contents
            break
source = inspect.getsource(target)

Why decorator behaviour is preserved

inspect.getsource on the original forward includes its decorator line, so re-executing the source re-applies the decorator. The accelerate hook setup (and any other wrapper side effects) therefore still take effect on the autowrapped forward. This is important: we are not bypassing the decorator, we are recovering the decorated definition so that torch.fx autowrapping can be applied to the real forward body while the decorator is re-applied around it.

There is no behaviour change for undecorated forwards: inspect.unwrap returns the method unchanged and the closure search is skipped.

TEST PLAN:

• Added regression tests in tests/llmcompressor/pipelines/sequential/test_ast_helpers.py:
  • test_autowrap_forward_handles_unwrapped_decorator: a module whose forward is wrapped by a no-functools.wraps decorator (mirroring force_accelerate_hooks). Previously raised KeyError: 'forward'; now autowraps successfully, produces correct output, and the decorator's side effect still fires (proving the decorator is re-applied).
  • test_autowrap_forward_plain_module_unchanged: an undecorated forward, asserting behaviour is unchanged.
• Both tests fail with KeyError: 'forward' on main and pass with this change.
• ruff check and ruff format --check pass on the touched files.
• End-to-end validation: a full GPTQ run of Qwen/Qwen3.8-27B (64-layer hybrid: 16 full-attention + 48 Qwen3_5GatedDeltaNet layers) over 482 subgraphs completed in ~45 minutes on a single RTX 5090, where on main it could not trace at all. This confirms tracing/calibration now completes end-to-end for this model class. (Note: this PR claims only that tracing/calibration completes; it makes no claim about the resulting checkpoint quality.)

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 35e8e6de-08a8-414e-bfe2-58dba85c8f6d

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews

	Protection	Waiting on
🔴	Require one maintainer review	👀 reviews

🔴 Require one maintainer review

Waiting for any of

• approved-reviews-by=HDCharles
• approved-reviews-by=brian-dellabetta
• approved-reviews-by=dsikka
• approved-reviews-by=kylesayrs
• approved-reviews-by=yiliu30

This rule is failing.

All PRs must have at least one approving review from a maintainer before merging.

• any of:
  • approved-reviews-by=HDCharles
  • approved-reviews-by=brian-dellabetta
  • approved-reviews-by=dsikka
  • approved-reviews-by=kylesayrs
  • approved-reviews-by=yiliu30
• #changes-requested-reviews-by = 0

[gemini-code-assist]

Code Review

This pull request updates autowrap_forward in ast_helpers.py to handle cases where module.forward is wrapped by a decorator that does not use functools.wraps (such as force_accelerate_hooks from transformers) by searching the wrapper's closure cells. It also adds regression tests to verify this behavior. However, a critical issue was identified: because module.forward returns a bound method, it does not have a __closure__ attribute directly (only its underlying __func__ does), which causes the closure-unwrapping logic to be skipped. You should extract the underlying function from the bound method before checking the closure cells.

[gemini-code-assist]

When module.forward is accessed, Python's descriptor protocol binds the function to the module instance, returning a bound method object (types.MethodType).

inspect.unwrap only follows __wrapped__ chains and does not unwrap bound methods. Therefore, target remains a bound method object. Since bound method objects do not have a __closure__ attribute (only their underlying function __func__ does), getattr(target, "__closure__", None) will always return None. As a result, the entire closure-unwrapping block is skipped, and the original issue remains unresolved for actual module instances.

To fix this, you should first extract the underlying function from the bound method (using inspect.unwrap or accessing __func__ if it is a method) before performing the closure search.

    target = inspect.unwrap(module.forward)
    # If target is a bound method, extract the underlying function
    if inspect.ismethod(target):
        target = target.__func__

    # Some decorators (e.g. transformers' `force_accelerate_hooks`, see
    # transformers/integrations/accelerate.py) wrap `forward` with an inner
    # function but do not apply `functools.wraps`. inspect.unwrap therefore
    # cannot recover the original (there is no `__wrapped__` attribute) and
    # inspect.getsource returns the *wrapper's* source, so the re-exec'd code
    # defines a function named `wrapped` rather than `forward` and the
    # `namespace["forward"]` lookup below raises KeyError. Recover the original
    # `forward` by searching the wrapper's closure cells. inspect.getsource on
    # the original includes its decorator line, so exec re-applies the
    # decorator and preserves its behaviour (e.g. accelerate hook setup).
    if getattr(target, "__name__", "") != "forward" and getattr(
        target, "__closure__", None
    ):

[kylesayrs]

Thanks for the fix! Could you please

1. Break this functionality into a helper function (something like get_unwrapped_forward)
2. Make sure that imports are declared at the top module level

Thank you!