Vectorize local shard retrieval #5826
Merged
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
jonb377
force-pushed
the
jonbolin/copy-pool
branch
from
b9d8a7c to
8992668
Compare
jonb377
force-pushed
the
jonbolin/vectorize-local-shards
branch
from
8d0b5bf to
bad9356
Compare
yeounoh
reviewed
Nov 23, 2023
| XLATensorPtr xtensor = bridge::GetXlaTensor(input); | ||
| m.def( | ||
| "_get_local_shard_replica_and_indices", | ||
| [](const std::vector<at::Tensor>& input) |
There was a problem hiding this comment.
nit. rename input to inputs or input_tensors
jonb377
force-pushed
the
jonbolin/copy-pool
branch
from
8992668 to
fa65980
Compare
jonb377
force-pushed
the
jonbolin/vectorize-local-shards
branch
from
bad9356 to
7ecc38d
Compare
jonb377
force-pushed
the
jonbolin/copy-pool
branch
from
fa65980 to
c8f7315
Compare
jonb377
force-pushed
the
jonbolin/vectorize-local-shards
branch
from
7ecc38d to
48e9943
Compare
jonb377
force-pushed
the
jonbolin/vectorize-local-shards
branch
from
48e9943 to
b673c07
Compare
ManfeiBai
pushed a commit
to ManfeiBai/PyTorchXLA
that referenced
this pull request
Dec 1, 2023
chunnienc
pushed a commit
to chunnienc/xla
that referenced
this pull request
Dec 14, 2023
golechwierowicz
pushed a commit
that referenced
this pull request
Jan 12, 2024
bhavya01
pushed a commit
that referenced
this pull request
Apr 22, 2024
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.
To capitalize on the improvements in #5824 and #5825, moving tensor shards to CPU should be batched. This change does the following:
_get_local_shardsand_get_local_shard_replica_and_indicesoperate on lists of tensors instead of individual tensors._sharded_cpu_state_dictto use the batched method across all sharded tensors in the state_dict.With all three changes applied, the amount of time spent transferring the state_dict to CPU for a 2B parameter model decreases from >10s to 3.4s, which unblocks training much more quickly.