Fix ZeRO-3 hooks for attribute-delegating modules (#8144)
## Summary
- make ZeRO-3 backward-hook state checks instance-local for
attribute-delegating modules
- add a regression test that exercises two training steps through a
delegating wrapper
## Root cause
ZeRO-3 registers child modules before their parents. Wrappers such as
PEFT's `ModulesToSaveWrapper` delegate unknown attributes to a wrapped
child, so `hasattr(wrapper, "pre_bwd_fn")` and `hasattr(wrapper,
"post_bwd_fn")` can incorrectly find the child's hook classes. DeepSpeed
then skips creating wrapper-local classes even though those classes
capture module-specific closures and counters. During forward, the
delegated post-backward class can access the child before its
`ds_grads_remaining` counter is initialized, raising `AttributeError`.
Check the module's own `__dict__` for these internal attributes so each
module receives hook classes and counters bound to itself.
## User impact
This restores ZeRO-3 training with attribute-delegating wrappers,
including PEFT LoRA configurations that use
`modules_to_save=["lm_head"]`.
Fixes #7615.
## Testing
- `pre-commit run --files deepspeed/runtime/zero/parameter_offload.py
tests/unit/runtime/zero/test_zero_hook_attribute_delegation.py`
- `PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 pytest -q
tests/unit/runtime/zero/test_zero_hook_attribute_delegation.py
tests/unit/runtime/zero/test_zero_dynamic_class.py -s` (`3 passed`)
- two ZeRO-3 training steps with PEFT's real `ModulesToSaveWrapper`
---------
Signed-off-by: Yuchen Fan <functionhx@gmail.com>