From da4aa35130f47b9c5228625c71448f6d47d69e16 Mon Sep 17 00:00:00 2001 From: Guokai Ma Date: Mon, 24 Aug 2026 23:43:21 +0800 Subject: [PATCH] Fix NameError: print_dist is called but never imported in auto_tp 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 --- deepspeed/module_inject/auto_tp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/module_inject/auto_tp.py b/deepspeed/module_inject/auto_tp.py index e47c5e3ddcc5..74598898f053 100755 --- a/deepspeed/module_inject/auto_tp.py +++ b/deepspeed/module_inject/auto_tp.py @@ -17,7 +17,7 @@ from .fusedqkv_utils import require_tp_fused_qkvw 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 from deepspeed.module_inject.layers import is_autotp_training_mode from deepspeed.module_inject.layers import _build_param_uc_restore_meta from deepspeed.checkpoint.constants import DS_AUTOTP_UC_META