FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

Import print_dist in auto_tp by alanhuangyoo · Pull Request #8311 · deepspeedai/DeepSpeed · GitHub

Import print_dist in auto_tp - #8311

Merged
delock merged 1 commit into
deepspeedai:masterfrom
alanhuangyoo:fix/auto-tp-import-print-dist
Aug 25, 2026
Merged

Import print_dist in auto_tp#8311
delock merged 1 commit into
deepspeedai:masterfrom
alanhuangyoo:fix/auto-tp-import-print-dist

Conversation

Copy link
Copy Markdown
Contributor

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.

register_replicated_grad_hooks() calls print_dist() but auto_tp.py only
imports log_dist, so the call raises NameError whenever any replicated
parameter gets a grad all-reduce hook:

  NameError: name 'print_dist' is not defined

DeepSpeedEngine calls this during AutoTP setup (engine.py:806, :845).

Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>
delock added this pull request to the merge queue Aug 25, 2026
Merged via the queue into deepspeedai:master with commit b753aec Aug 25, 2026
13 of 15 checks passed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants


Back | FazBrowse Home | New Git URL