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