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>