Fix SFTTrainer chunked-CE patch crash when forward is not a bound method

[verma8076]

What does this PR do?

Fixes the crash in _patch_chunked_ce_lm_head reported in #6483. The issue title says it's a Qwen3.5 problem, but it's actually broader: the patch assumes model.forward is always a bound method and reads original_forward.__func__ to build the patched forward's signature. That breaks for any model whose forward has been wrapped by accelerate's device-dispatch hooks — which is exactly what device_map="auto" does when accelerate decides to dispatch/offload the model.

accelerate.hooks.add_hook_to_module replaces model.forward with:

module.forward = functools.update_wrapper(functools.partial(new_forward, module), old_forward)

A functools.partial has no __func__, hence the AttributeError in the issue.

Fix: use inspect.unwrap(original_forward) instead of original_forward.__func__. functools.update_wrapper always sets __wrapped__ on the wrapper, pointing back at old_forward, so inspect.unwrap recovers the original callable regardless of whether it went through an accelerate hook. It also happens to correctly handle another layer of wrapping that's unrelated to this bug: some transformers forwards are themselves decorated with @functools.wraps (e.g. via can_return_tuple), which also sets __wrapped__. I checked that inspect.unwrap produces a byte-identical signature to the old .__func__ access on a plain (non-hooked) model, so this isn't a behavior change for the common case — it just also works for the hooked one.

I reproduced the crash locally against a small model wrapped with accelerate.hooks.add_hook_to_module (no GPU or the actual Qwen3.5 checkpoint needed — the bug is in signature introspection, not model computation), confirmed the fix resolves it, and confirmed a real forward/backward pass through the patched model still produces correct results afterward.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline, Pull Request section?
• Was this discussed/approved via a GitHub issue? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

This wasn't pre-discussed with a maintainer — it's a small, self-contained bug fix (no public API or checkpoint format change) with a repro and regression test, so I went ahead and opened the PR directly rather than waiting on a design discussion. Happy to adjust if there's a reason to handle it differently.

AI writing disclosure

• No AI usage: the PR was written entirely by a human.
• AI-assisted: some parts were suggested or improved by AI, but the PR was written and reviewed by a human.
• AI-generated: the PR was mostly or fully generated by an AI tool.

To be upfront about this: the investigation, fix, and test were done by Claude Code under my direction — I reviewed the reasoning and the diff, ran the tests myself, and I'm able to discuss/defend the change, but I want to flag the AI involvement honestly rather than undersell it.

Who can review?

@qgallouedec — you commented on the original issue and might have context on why the patch does the .__func__ dance in the first place.

---

Note

Low Risk
Small, localized fix to signature copying during SFT chunked-CE patching, with a targeted regression test and no training API changes.

Overview
Fixes AttributeError when _patch_chunked_ce_lm_head runs on models whose forward is not a plain bound method (e.g. device_map="auto" via Accelerate functools.partial, or Transformers @wraps layers).

Signature introspection now uses inspect.signature(inspect.unwrap(original_forward)) instead of original_forward.__func__, so the patched forward keeps the real parameter list for generate / VLMs without changing behavior on unwrapped models.

Adds a CPU regression test that wraps forward with AlignDevicesHook, asserts patching succeeds, and checks loss matches the reference model.

^{Reviewed by Cursor Bugbot for commit 571c9cb. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

Want higher recall? High effort reviews run extra passes and find more bugs. A team admin can switch effort levels in the Cursor dashboard.

^{Reviewed by Cursor Bugbot for commit 571c9cb. Configure here.}

[cursor]

Unwrap drops self from signature

Medium Severity

inspect.unwrap(original_forward) does not always land on an unbound function with self. For a plain bound method with no __wrapped__ chain, unwrap returns the bound method, so inspect.signature already omits self. After MethodType binding, introspection strips the first real parameter, breaking the signature that generate's _validate_model_kwargs relies on. The old .__func__ path preserved self in that case.

 

^{Reviewed by Cursor Bugbot for commit 571c9cb. Configure here.}