Import print_dist in auto_tp

[alanhuangyoo]

AutoTP.register_replicated_grad_hooks() calls print_dist(), but auto_tp.py imports only log_dist:

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.