DeepSpeed
e6324af9 - fix(comm): Expose GradBucket in deepspeed.comm API (#7400)

Commit
183 days ago
fix(comm): Expose GradBucket in deepspeed.comm API (#7400) This PR fixes an omission in the `deepspeed.comm` API where `GradBucket` was not exposed, despite the package aiming for full compatibility with `torch.distributed`. ##The Problem As reported in issue #7393, when a user replaces `torch.distributed` with `deepspeed.comm`, they expect all public APIs to be available. However, attempting to access `deepspeed.comm.GradBucket` (for example, when using it as a type hint for DDP communication hooks) results in an `AttributeError`. ##The Solution This change resolves the issue by importing `GradBucket` directly from `torch.distributed` into the `deepspeed/comm/comm.py` file, making it part of the public `deepspeed.comm` namespace. A `# noqa: F401` comment has been added to the import line. This is necessary to bypass the `flake8` linter's "imported but unused" check, as the specific purpose of this import is to expose the symbol to the end-user, not for it to be used within the `comm.py` file itself. ##How This Was Tested The fix was verified with a local test script that confirms `deepspeed.comm.GradBucket` can now be accessed correctly and is identical to `torch.distributed.GradBucket`. The pre-commit hooks now pass successfully. ##Related run test Screenshout <img width="1250" alt="Screenshot 2025-06-30 at 22 41 10" src="https://github.com/user-attachments/assets/cadf18e1-9d1a-4164-a5ff-0b3e6804ac48" /> ##Related Issue Fixes #7393 Signed-off-by: Vensenmu <vensenmu@gmail.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>
Author
Parents
Loading