Implement pipeline-aware shard estimator and turn it on by default for prefetch pipeline #1924
Conversation
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
facebook-github-bot
added
the
CLA Signed
|
This pull request was exported from Phabricator. Differential Revision: D56444328 |
…r prefetch pipeline (pytorch#1924) Summary: Lots of rank load imbalance comes from underestimation of sparse, which is not only a function of parameter, but the HBM usage during input and output. This is especially bad if we have multi-stage pipeline, which keeps multiple copies of input for each table. Prefetched embedding has the worst performance. By heuristically analyzing the current memory snapshot, we noticed 4~7x of extra input. This diff use a new formulas to calculate HBM usage: - Multiple input are considered depending on the pipeline - Output and input are added using max, since output tensor is not occupying extra HBM before a2a communication, at that time input is no longer used For backward compatibility, we haven't roll out the change to ally user as it may incur extra failure of jobs. Instead, we set up different pipeline-awareness mode: - None: Use old formulas regardless - Prefetch-Only (default): Use new formulas if prefetch pipeline is on - All: Use new formulas. Differential Revision: D56444328
levythu
force-pushed
the
export-D56444328
branch
from
59db2a4 to
a133b90
Compare
|
This pull request was exported from Phabricator. Differential Revision: D56444328 |
levythu
added a commit
to levythu/torchrec
that referenced
this pull request
Apr 24, 2024
…r prefetch pipeline (pytorch#1924) Summary: Lots of rank load imbalance comes from underestimation of sparse, which is not only a function of parameter, but the HBM usage during input and output. This is especially bad if we have multi-stage pipeline, which keeps multiple copies of input for each table. Prefetched embedding has the worst performance. By heuristically analyzing the current memory snapshot, we noticed 4~7x of extra input. This diff use a new formulas to calculate HBM usage: - Multiple input are considered depending on the pipeline - Output and input are added using max, since output tensor is not occupying extra HBM before a2a communication, at that time input is no longer used For backward compatibility, we haven't roll out the change to ally user as it may incur extra failure of jobs. Instead, we set up different pipeline-awareness mode: - None: Use old formulas regardless - Prefetch-Only (default): Use new formulas if prefetch pipeline is on - All: Use new formulas. Differential Revision: D56444328
levythu
added a commit
to levythu/torchrec
that referenced
this pull request
Apr 24, 2024
… than HBM-only Summary: Checked pytorch#1924 for context. Basically, under some specific case, a table could use even more hbm when offloaded and prefetched. in this case, we'd rather not offload them. Differential Revision: D56505315
facebook-github-bot
pushed a commit
that referenced
this pull request
Apr 25, 2024
… than HBM-only Summary: Check #1924 for context. Basically, under some specific case, a table could use even more hbm when offloaded and prefetched. in this case, we'd rather not offload them. Reviewed By: henrylhtsang Differential Revision: D56505315 fbshipit-source-id: f8cdf03568f79ad7e2ef842b8939280de9ec9c84
levythu
added a commit
to levythu/torchrec
that referenced
this pull request
Sep 12, 2024
Summary: Prior context: pytorch#1924 or D56444328. From the in-depth analysis, the output of embedding didn't really contribute to the peak memory. Instead, they're quite ephemeral and deallocated right after output-dist. (See the analysis of original PR/diff for more details). Therefore, this diff remove output out of the memory accounting for an embedding shard. This behavior is not always safe though: in rare case, e.g. a rank contain only sparse , or for some reason sparse forward/backward happens when HBM usage hits peak, output size will contribute to peak memory instead -- the change leave the open interface to still count in output size and caller will set it as needed. For backward compatibility, this change only applies to any estimation using "new formula" (the pipeline aware formula). The old formula will keep untouched and always include output. {F1862246115} Differential Revision: D62540207
facebook-github-bot
pushed a commit
that referenced
this pull request
Sep 12, 2024
Summary: Pull Request resolved: #2382 Prior context: #1924 or D56444328. From the in-depth analysis, the output of embedding didn't really contribute to the peak memory. Instead, they're quite ephemeral and deallocated right after output-dist. (See the analysis of original PR/diff for more details). Therefore, this diff remove output out of the memory accounting for an embedding shard. This behavior is not always safe though: in rare case, e.g. a rank contain only sparse , or for some reason sparse forward/backward happens when HBM usage hits peak, output size will contribute to peak memory instead -- the change leave the open interface to still count in output size and caller will set it as needed. For backward compatibility, this change only applies to any estimation using "new formula" (the pipeline aware formula). The old formula will keep untouched and always include output. {F1862246115} Reviewed By: sarckk Differential Revision: D62540207 fbshipit-source-id: 0af679817ab61b508b5ce06fec4e0136e7a63dea
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.
Summary:
Lots of rank load imbalance comes from underestimation of sparse, which is not only a function of parameter, but the HBM usage during input and output. This is especially bad if we have multi-stage pipeline, which keeps multiple copies of input for each table.
Prefetched embedding has the worst performance. By heuristically analyzing the current memory snapshot, we noticed 4~7x of extra input.
This diff use a new formulas to calculate HBM usage:
For backward compatibility, we haven't roll out the change to ally user as it may incur extra failure of jobs. Instead, we set up different pipeline-awareness mode:
Differential Revision: D56444328