(1/2) Implementing Compiler Pass for AutoTP - #8204
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 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".
There was a problem hiding this comment.
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.
|
Can you also share your plan for the the next step as this PR has 1/2 in the title. |
Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
a3a2bb4 to
478c1e9
Compare
Signed-off-by: Naveenraj Kamalakannan <therealnaveenkamal@gmail.com>
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>
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
left a comment
tohtana
left a comment
There was a problem hiding this comment.
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?
| 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(): |
There was a problem hiding this comment.
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.
|
|
||
| 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" |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
The issue with Llama4 MoE router should be addressed in another PR. How about making it fail when any unsupported style is found.
| # already records the partitioning decision the pass needs. | ||
|
|
||
| COLUMN_PARALLEL_LAYERS = (LinearLayer, SubParamLinearLayer) | ||
| ROW_PARALLEL_LAYERS = (LinearAllreduce, SubParamLinearAllreduce) |
There was a problem hiding this comment.
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.
2 participants
Working on #8104
cc @tohtana