[PyTorch] Adding TP overlap support for te.Linear with parallel_mode="column"
#1343
Conversation
This file contains 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
… in sequence-parallel Linear backward Signed-off-by: Alp Dener <[email protected]>
Signed-off-by: Alp Dener <[email protected]>
…dated unit tests Signed-off-by: Alp Dener <[email protected]>
for more information, see https://pre-commit.ci
denera
force-pushed
the
linear-tp-overlap-ag-fprop-rs-dgrad
branch
from
90458d4 to
4e3e61a
Compare
|
/te-ci pytorch L1 |
timmoon10
approved these changes
Nov 21, 2024
Comment on lines
866
to
870
| ub_overlap_ag: bool = False, | ||
| ub_overlap_rs: bool = False, | ||
| ub_bulk_dgrad: bool = False, | ||
| ub_bulk_wgrad: bool = False, | ||
| ub_name: Optional[str] = None, |
There was a problem hiding this comment.
We should seriously consider deprecating these UB options and just passing in a dict. The UB interface is unstable and will likely be so for some while. A dict would be better for backward compatibility (reinterpret old options) and forward compatibility (ignore unknown options). This would be especially helpful for Mcore integration.
For example, the operation-based API passes in UB options with a dict:
Comment on lines
+931
to
+935
| assert not (self.ub_overlap_rs_fprop and self.ub_overlap_ag_fprop), "Internal TE error!" | ||
| assert not (self.ub_overlap_ag_dgrad and self.ub_overlap_rs_dgrad), "Internal TE error!" | ||
| assert not ( | ||
| self.ub_overlap_rs_dgrad and (self.ub_bulk_dgrad or self.ub_bulk_wgrad) | ||
| ), "Internal TE error!" |
There was a problem hiding this comment.
More descriptive error messages would be helpful.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
te.Linearcurrently only supports TP overlap withparallel_mode="row"where it overlaps reduce-scatter in the forward pass, and all-gather with dgrad in the backward pass.This PR adds new options to enable all-gather overlap in the forward pass, and reduce-scatter overlap with dgrad in the backward pass, when
parallel_mode="column".Fixes #1312
Type of change
Checklist: