Fix ZeRO-3 hooks for attribute-delegating modules

[Functionhx]

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

[tohtana]

Hi @Functionhx,
Thank you for the fix! Overall this looks good to me. I would like to have a few changes before merge.

For the new test file, we now use this copyright (please refer to existing files):

# Copyright (c) DeepSpeed Team.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

The tests should assert that the wrapper and child own distinct pre/post backward hook classes and that both modules' applied_pre_backward_ref_cnt values return to zero after each backward. The current training loop can still pass if the counter lookup runs based on the existing hasattr behavior. (draft of the fix)

[Functionhx]

@tohtana Thanks for the review! I've addressed the requested changes. Could you please take another look?🩷