(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".
| # finding one here means this activation has been handled. | ||
| if activation.op == "call_function" and activation.target is COLUMN_PARALLEL_OP: | ||
| continue | ||
| _insert_after(gm, activation, COLUMN_PARALLEL_OP) |
There was a problem hiding this comment.
When a LinearLayer has gather_output=True (for example the default/HF lm_head path), defer_collectives_to_compiler() intentionally leaves its module-level ColumnParallel/gather collectives enabled, but the graph pass still reaches this insertion because it only checks the layer type. That adds an extra copy_to_tp_region before the matmul, so the layer's input gradient is all-reduced once by the compiled op and once by the existing ColumnParallel.apply, silently scaling upstream gradients. Skip modules whose defer_collectives_to_compiler flag is false, or at least gather_output layers.
Useful? React with 👍 / 👎.
| collective_node.meta["val"] = node.meta.get("val") | ||
| # Steal every consumer first, then hand the original back as this node's own input; doing it in | ||
| # the other order would leave the new node feeding itself. | ||
| node.replace_all_uses_with(collective_node) |
There was a problem hiding this comment.
When the same activation also feeds any non-column consumer, this helper's replace_all_uses_with rewrites those consumers to read from copy_to_tp_region as well. Its forward is an identity, but its backward all-reduces the combined gradient from every rewritten consumer, whereas module-level AutoTP only all-reduces the gradient from column-parallel branches. This silently corrupts upstream gradients; only the matching column matmul users should be retargeted.
Useful? React with 👍 / 👎.
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.
| @@ -0,0 +1,72 @@ | |||
| # Copyright (c) Microsoft Corporation. | |||
There was a problem hiding this comment.
For new files, please show DeepSpeed's copyright. You can find an example in other files.
|
|
||
| from unit.common import DistributedTest | ||
|
|
||
| pytestmark = pytest.mark.skipif(not required_torch_version(min_version=2.9), |
There was a problem hiding this comment.
Can you clarify what limits this feature to v2.9+? Also, can we have the same check in the main code?
|
Can you also share your plan for the the next step as this PR has 1/2 in the title. |
2 participants
Working on #8104
cc @tohtana