Import print_dist in auto_tp (#8311)
`AutoTP.register_replicated_grad_hooks()` calls `print_dist()`, but
`auto_tp.py` imports only `log_dist`:
```python
from deepspeed.utils.logging import log_dist
...
if registered:
print_dist(
f"AutoTP: registered tensor-parallel grad all-reduce for {len(registered)} replicated "
f"parameters, e.g. {registered[0]!r}",
ranks=[0])
```
So the call raises as soon as `registered` is non-empty:
```
NameError: name 'print_dist' is not defined
```
`DeepSpeedEngine` calls this during AutoTP setup
(`deepspeed/runtime/engine.py:806` and `:845`), so any AutoTP run that
registers a hook on a replicated parameter hits it.
`print_dist` and `log_dist` are both in `deepspeed/utils/logging.py` and
are not interchangeable — `print_dist` exists specifically for messages
that should appear regardless of log level — so this adds the import
rather than switching the call to `log_dist`.
Before, on master:
```
tests/unit/model_parallelism/test_tp_plan_real_models.py -> 2 failed, 4 passed
FAILED TestQwen3UnevenTPPlan::test_qwen3_tp3_keeps_attention_heads_aligned
FAILED TestQwen3UnevenTPPlan::test_autotp_size_above_kv_head_count_leaves_trailing_ranks_empty
```
After:
```
tests/unit/model_parallelism/test_tp_plan_real_models.py -> 6 passed
```
`pre-commit run --files deepspeed/module_inject/auto_tp.py` is clean.
The call came in with #8185 on 2026-08-20. It is not caught by the live
CI: `modal-torch-latest` runs `tests/unit/v1/` only, and the self-hosted
GPU workflows that would run `unit/` have not produced a run in a long
time (`nv-a6000` last ran 2025-08-01, `nv-nightly` 2026-01-15,
`nv-torch-latest-v100` and `nv-inference` have no runs listed).
Env: torch 2.13.0+cu130, H20, single node.
Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>