DeepSpeed
75f395fd - Fix ZeRO-3 hooks for attribute-delegating modules (#8144)

Commit
21 days ago
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>
Author
Parents
Loading