Motivation

register_replicated_grad_hooks() in deepspeed/module_inject/auto_tp.py logs its
summary through print_dist, which exists in deepspeed/utils/logging.py but was
never imported in this module. As a result, any model whose HuggingFace tp_plan
contains replicated_with_grad_allreduce entries crashes with NameError at
deepspeed.initialize() whenever tensor-parallel size > 1 — before a single
training step runs.

This regressed in #8185, which introduced the grad-allreduce hooks path. With recent
transformers, Qwen3's tp_plan marks q_norm / k_norm as
replicated_with_grad_allreduce, so it is hit by simply launching a TP run of Qwen3.

Reproducer (before this fix)

[rank0]:   File "/root/DeepSpeed/deepspeed/module_inject/auto_tp.py", line 620, in register_replicated_grad_hooks
[rank0]:     print_dist(
[rank0]: NameError: name 'print_dist' is not defined

Triggers e.g. via the AutoTP equivalence check from
deepspeedai/DeepSpeedExamples#1008:

cd training/autotp_equivalence
bash run_gpu.sh 500 0,1,2,3    # tp=3 and tp=4 both die at initialize

Description of changes

Add print_dist to the existing logging import — one line, matching the fix already
carried in passing by the #8241 refactor branch, so the crash is not blocked on that
larger review:

-from deepspeed.utils.logging import log_dist
+from deepspeed.utils.logging import log_dist, print_dist

print_dist (rather than log_dist) matches the call site's intent: its docstring
says 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:

• Before: tp=3 / tp=4 both NameError at deepspeed.initialize().
• After: both widths run to completion and agree with the AutoTP=1 baseline —

===== AutoTP=3 (uneven 6/6/4) vs AutoTP=1 =====
steps=500 first_rel=6.74e-08 mean_rel=7.45e-07 worst_rel=2.28e-04 at step 291
OK: 500 steps agree within 1e-02
===== AutoTP=4 (even 4/4/4/4, control) vs AutoTP=1 =====
steps=500 first_rel=0.00e+00 mean_rel=2.38e-06 worst_rel=7.68e-04 at step 291
OK: 500 steps agree within 1e-02

pre-commit passes on the changed file (yapf, flake8, codespell, check-torchdist,
check-license, check-torchcuda, check-extraindexurl).