(1/2) Implementing Compiler Pass for AutoTP by therealnaveenkamal · Pull Request #8204 · deepspeedai/DeepSpeed

therealnaveenkamal · 2026-08-02T05:36:26Z

Working on #8104

Added support for AutoTP. Added two primitives: copy_to_tp and reduce_from_tp - the f and g nodes.
Identifies column/row-parallel matmuls by the injected layer type in nn_module_stack, reading back.
Wrote a test to verify correctness of module injection and compiler pass.

cc @tohtana

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: a3a2bb41dd

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

tohtana

Thank you @therealnaveenkamal, this is amazing!
I don't see any critical issue in the code. The approach that reuses existing AutoTP's patterns is great. It gives consistent results with non-DeepCompile AutoTP.

One remaining work is validating the correctness in a more realistic setting. I think it would be good to compare loss values from existing AutoTP and this one. I did similar work for AutoEP. The harness for the verification might be useful if you don't have such a script. We should check different configs like DP1/TP4 and DP2/TP2.

I left a few comments about details. Please consider addressing them. Also, please fix the commit to pass DCO check.

tohtana · 2026-08-04T01:00:23Z

Can you also share your plan for the the next step as this PR has 1/2 in the title.
I think it would be good to start with the combination with ZeRO. Probably we could do further optimizations based on profiling.

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

Results of comparing AutoTP module injection against the DeepCompile autotp pass (deepspeedai/DeepSpeed#8204) on dense Qwen3.5 (618M, hybrid GatedDeltaNet/attention, 4xA100): 1.31x speedup at DP2/TP2 and 1.45x at DP1/TP4 with loss agreement at the bf16 noise floor. Includes the writeup, plots, per-run metrics, and pinned environment. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

autotp works now

tohtana

Hi @therealnaveenkamal,
Thank you for the update! I think the earlier issues have been addressed.
I also found some correctness issues in the current code. Can you check them?

tohtana · 2026-08-10T23:18:35Z

+    Any tensor-parallel layer the pass cannot rewrite is rejected rather than left on the
+    module-level path.
+    """
+    for name, module in model.named_modules():


Can we set the flag defer_collectives_to_compiler only all modules passed the check?
If this raises an error in the loop, only some modules have defer_collectives_to_compiler=True. But the outer code might catch the error and fallback to eager. In that case, some communication collectives will be skipped.

got it. now we go through all the modules and if any module is incompatible, we raise an error

tohtana · 2026-08-10T23:30:48Z

+
+        assert specs is not None
+        by_type = {spec.partition_type for spec in specs}
+        assert PartitionType.ROW in by_type, "the supported entry should still be applied"


I think this expectation (and the code) is wrong. To make TP work, COLUMN and ROW should be paired. Partially skipping the conversion breaks it.

The issue with Llama4 MoE router should be addressed in another PR. How about making it fail when any unsupported style is found.

fixed it now.

tohtana · 2026-08-10T23:50:25Z

+# already records the partitioning decision the pass needs. 
+
+COLUMN_PARALLEL_LAYERS = (LinearLayer, SubParamLinearLayer)
+ROW_PARALLEL_LAYERS = (LinearAllreduce, SubParamLinearAllreduce)


SubParamLinearAllreduce.forward() always executes its module-level row all-reduce, while the compiler pass classifies that layer as row parallel and inserts another graph all-reduce.

thanks for this, @tohtana. I've made a check for defer_collectives_to_compiler in layers.py

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

… transformers not 5.x fixed for partial compiler flags, now raises error

therealnaveenkamal · 2026-08-11T23:55:50Z

Hi @tohtana - thanks for pointing out the bugs. I've made the changes. Let me know if this is okay. Thanks.

tohtana · 2026-08-12T06:33:47Z

Hi @therealnaveenkamal, thank you for the update!
Overall the change looks okay to me, but this test failed in my env. Can you check it?

tests/unit/compile/test_tp_compile.py::TestAutoTPCompileMoE::test_mixtral_matches_module_injection

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

added guards for test_tp_compile

tohtana · 2026-08-12T21:55:06Z

Hi @therealnaveenkamal,
I investigated the test failure. I think it was a bug on HF side that existed from v5.8.0 to v5.10.0.

It started with huggingface/transformers#45621 and fixed by huggingface/transformers#45634. I found you already pushed the version guard, but can we be more specific about versions?
I don't see any other issues. After we confirm the new version guard works, let's merge this PR.

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

added warning for transformers bug

therealnaveenkamal · 2026-08-12T22:51:05Z

Thanks @tohtana, Guard is now exact: the range [5.8.0, 5.10.1) lives as a constant in init_tp.py, and the test skips it.

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

tohtana · 2026-08-12T23:38:10Z

Hi @therealnaveenkamal,
Great, thank you for the update!

I just found the code should work even with the affected versions as long as we choose other than batched_mm. I opened a small PR to your branch: https://github.com/therealnaveenkamal/DeepSpeed/pull/5/changes
Can you review it? Feel free to modifty it if you want.

therealnaveenkamal · 2026-08-12T23:43:47Z

@tohtana sorry about that. I've merged the PR.

tohtana

@therealnaveenkamal Thank you for your great work! This is definitely a significant step for DeepCompile.
I really appreciate your contribution to DeepSpeed.

…eepspeedai#8294) ## Problem huggingface/transformers#47579 (on `main` since `861f4c41`, 2026-08-21) makes `PretrainedConfig.__init__` inject `"embed_tokens": "embedding_rowwise"` into `base_model_tp_plan` whenever `tie_word_embeddings` is true. `SUPPORTED_STYLES` is a strict allowlist and `convert()` raises on any style outside it, rejecting the whole plan, so AutoTP plan conversion now fails for every tied-embedding model (Qwen2/Qwen3, Llama, Gemma) built against transformers `main`. ## Fix Recognize `embedding_rowwise` and convert it to a `SKIP` spec, which is option 2 in deepspeedai#8290: the entry is understood and the embedding is deliberately left replicated. What follows is the behaviour DeepSpeed already implements rather than a new policy. `lm_head` still converts to a gathered column spec, and `_configure_gathered_column_tie_fallbacks` then sees that `lm_head.weight is embed_tokens.weight` and leaves both modules replicated, logging that coupled vocabulary-parallel embedding is not supported yet. A tied model is therefore left in the shape it has on a transformers release without the injection, with both modules replicated and the tie intact. The entry maps to `SKIP` with `grad_allreduce` left false, unlike `replicated_with_grad_allreduce`. The parameter is never split, so `register_replicated_grad_hooks` must not register an all-reduce for it. Styles that are still unknown continue to reject the whole plan. `test_unsupported_style_rejects_whole_plan` is unchanged and still passes. ## Verification Run on CPU in a container at `edaa7221`, against transformers `main` (5.16.0.dev0) and torch 2.13.0+cpu. - The two added tests fail on master with the reported `ValueError` and pass with this change. - `tests/unit/module_inject/` and `tests/unit/runtime/test_tp_plan_extraction.py`: 47 passed, on Python 3.11 and on 3.12. - `pre-commit run --files` on the three changed files passes yapf, check-torchdist, check-license and codespell; flake8 5.0.4 exits 0 on them under Python 3.11. - Not verified here: `test_qwen2_tied_lm_head_falls_back_to_replicated`, which needs 2 GPUs. That is the test deepspeedai#8290 reports as failing and the one this change is meant to restore. ## Two things worth deciding separately Scoping this to `embedding_rowwise` leaves the next transformers-side style to fail the same way, since the injection is unconditional and the allowlist is deny-by-default against a vocabulary DeepSpeed does not own. A general rule for unknown styles looks like a maintainer call rather than something to settle here. Related to that, the `convert()` docstring says entries with an unsupported style become SKIP specs instead of invalidating the plan, but no code path does that, and none does after this change either: an unsupported style still raises before the loop is reached. The docstring and the raise arrived together in deepspeedai#8204, so I have left both alone. Happy to follow up once you have picked the policy. Refs deepspeedai#8290. This covers the conversion failure only, and does not implement vocabulary-parallel tied embeddings (option 1 or 3 in that issue), so I have not used a closing keyword. --------- Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>

therealnaveenkamal requested review from hwchen2017, loadams, tjruwase and tohtana as code owners August 2, 2026 05:36

chatgpt-codex-connector Bot reviewed Aug 2, 2026

View reviewed changes

Comment thread deepspeed/compile/passes/tp_compile.py Outdated

Comment thread deepspeed/compile/passes/tp_compile.py Outdated

tohtana requested changes Aug 4, 2026

View reviewed changes

Comment thread deepspeed/compile/custom_ops/tp_collectives.py Outdated

Comment thread tests/unit/compile/test_tp_compile.py Outdated

therealnaveenkamal added 3 commits August 5, 2026 13:51

added support for autotp

8abe7ed

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

fix for 2x residuals

d442001

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

restore AutoTP torch version check

478c1e9

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

therealnaveenkamal force-pushed the feature/autotp branch from a3a2bb4 to 478c1e9 Compare August 5, 2026 23:53

therealnaveenkamal added 2 commits August 7, 2026 06:28

autotp works now

e7762c3

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

cleaning up

4738778

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

Merge pull request #1 from therealnaveenkamal/feature/autotp-dev

3f6e179

autotp works now

tohtana requested changes Aug 10, 2026

View reviewed changes

therealnaveenkamal added 4 commits August 11, 2026 18:43

fixed for partial compiler flags, now raises error

40d92d6

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

skip pytest if transformers not 5.x

a816df0

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

precommit

84ec6ca

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

fixed for partial compiler flags, now raises error and skip pytest if…

445632e

… transformers not 5.x fixed for partial compiler flags, now raises error

therealnaveenkamal added 2 commits August 12, 2026 07:01

added guards for test_tp_compile

218cf7b

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

added guards for test_tp_compile

196e62b

added guards for test_tp_compile

therealnaveenkamal added 2 commits August 12, 2026 18:38

added warning for transformers bug

0f71bad

Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>

added warning for transformers bug

66a5b2b

added warning for transformers bug

Fix AutoTP MoE compatibility guard

6cb92ef

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Narrow AutoTP Transformers MoE guard to batched_mm

611ed98

tohtana approved these changes Aug 12, 2026

View reviewed changes

tohtana enabled auto-merge August 12, 2026 23:51

Merge branch 'master' into feature/autotp

8377dda

tohtana added this pull request to the merge queue Aug 13, 2026

Merged via the queue into deepspeedai:master with commit 7904603 Aug 13, 2026
14 of 15 checks passed

ebarkhordar mentioned this pull request Aug 23, 2026

[AutoTP] Convert HF embedding_rowwise tp_plan entries to SKIP specs #8294

Merged

This was referenced Aug 29, 2026

Fix AutoTP + deep compile collectives silently drop when AC is on #8355

Merged

[BUG]AutoTP compile pass silently dropping tensor-parallel collectives #8357

Closed

Conversation

therealnaveenkamal commented Aug 2, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

Uh oh!

tohtana left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

tohtana commented Aug 4, 2026

Uh oh!

tohtana left a comment

Choose a reason for hiding this comment

Uh oh!

tohtana Aug 10, 2026

Choose a reason for hiding this comment

Uh oh!

therealnaveenkamal Aug 11, 2026

Choose a reason for hiding this comment

Uh oh!

tohtana Aug 10, 2026

Choose a reason for hiding this comment

Uh oh!

tohtana Aug 10, 2026

Choose a reason for hiding this comment

Uh oh!

therealnaveenkamal Aug 11, 2026

Choose a reason for hiding this comment

Uh oh!

tohtana Aug 10, 2026

Choose a reason for hiding this comment

Uh oh!

therealnaveenkamal Aug 11, 2026

Choose a reason for hiding this comment

Uh oh!

therealnaveenkamal commented Aug 11, 2026

Uh oh!

tohtana commented Aug 12, 2026

Uh oh!

tohtana commented Aug 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

therealnaveenkamal commented Aug 12, 2026

Uh oh!

tohtana commented Aug 12, 2026

Uh oh!

therealnaveenkamal commented Aug 12, 2026

Uh oh!

tohtana left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

tohtana left a comment •

edited

Loading

tohtana commented Aug 12, 2026 •

edited

Loading