Fix NameError: print_dist is called but never imported in auto_tp - #8310
Fix NameError: print_dist is called but never imported in auto_tp#8310delock wants to merge 1 commit into
Conversation
register_replicated_grad_hooks() logs its summary via print_dist, which exists in deepspeed/utils/logging.py but was not imported in deepspeed/module_inject/auto_tp.py. Any model whose HuggingFace tp_plan contains replicated_with_grad_allreduce entries (e.g. Qwen3's q_norm / k_norm with recent transformers) hits NameError at deepspeed.initialize() whenever tensor-parallel size > 1. Verified with the AutoTP equivalence check from deepspeedai/DeepSpeedExamples#1008 (Qwen3-0.6B, tp=3 uneven / tp=4 even, 500 steps each): tp=3 and tp=4 both crashed at initialize before this fix and complete with agreeing loss curves after it. Same missing import is fixed in passing on the #8241 branch; this is the minimal standalone hotfix so the crash is not blocked on that refactor. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: da4aa35130
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list | ||
| from deepspeed.utils import groups | ||
| from deepspeed.utils.logging import log_dist | ||
| from deepspeed.utils.logging import log_dist, print_dist |
There was a problem hiding this comment.
This is a non-merge commit, but its commit message has no Signed-off-by trailer, so it does not meet the repository's mandatory DCO requirement and may be rejected by commit/CI policy. Recreate the commit with git commit --signoff using the configured author identity.
AGENTS.md reference: AGENTS.md:L6-L8
Useful? React with 👍 / 👎.
|
close as #8311 will fix it. |
1 participant
Motivation
register_replicated_grad_hooks()indeepspeed/module_inject/auto_tp.pylogs itssummary through
print_dist, which exists indeepspeed/utils/logging.pybut wasnever imported in this module. As a result, any model whose HuggingFace tp_plan
contains
replicated_with_grad_allreduceentries crashes withNameErroratdeepspeed.initialize()whenever tensor-parallel size > 1 — before a singletraining step runs.
This regressed in #8185, which introduced the grad-allreduce hooks path. With recent
transformers, Qwen3's tp_plan marks
q_norm/k_normasreplicated_with_grad_allreduce, so it is hit by simply launching a TP run of Qwen3.Reproducer (before this fix)
Triggers e.g. via the AutoTP equivalence check from
deepspeedai/DeepSpeedExamples#1008:
Description of changes
Add
print_distto the existing logging import — one line, matching the fix alreadycarried in passing by the #8241 refactor branch, so the crash is not blocked on that
larger review:
print_dist(rather thanlog_dist) matches the call site's intent: its docstringsays to use it when log level should not decide whether the message is printed, which
fits this one-line setup diagnostic.
Validation
Ran the Examples#1008 equivalence suite (Qwen3-0.6B, fp32, 500 steps per width,
4x RTX 4080, transformers 5.14.0.dev0) against this branch:
NameErroratdeepspeed.initialize().pre-commit passes on the changed file (yapf, flake8, codespell, check-torchdist,
check-license, check-torchcuda, check-extraindexurl).