DeepSpeed
b753aec8 - Import print_dist in auto_tp (#8311)

Commit
2 days ago
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>
Author
Parents
Loading