[Checkpoint][2D][3/N] Add nested_tensors for distributed checkpoint to core distributed #89501
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/89501
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 9f6be82: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
wz337
requested review from
H-Huang,
awgu,
kwen2501,
mrshenli,
pritamdamania87,
rohan-varma and
zhaojuanmao
as code owners
wz337
changed the title
wanchaol
left a comment
wanchaol
left a comment
There was a problem hiding this comment.
Some questions about why the placement is always on cuda:0? stamp to unblock the migration
There was a problem hiding this comment.
this seems not used at all?
There was a problem hiding this comment.
It will be used in the optimizer.py, which will be upstreamed in the following PR.
There was a problem hiding this comment.
why this is cuda:0 only?
There was a problem hiding this comment.
I believe this is a placeholder, because calling _init_from_local_shards_and_global_metadata() in line 110 requires it as an input arg.
There was a problem hiding this comment.
is there a test with regard to this function?
There was a problem hiding this comment.
There is an integration test once the planner is updated with 2D functionality. There is no individual test for this yet. Will add a test as test improvement once everything is moved over.
|
@pytorchmergebot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Merge failedReason: This PR is too stale; the last push date was more than 3 days ago. Please rebase and try again. You can rebase by leaving the following comment on this PR: Details for Dev Infra teamRaised by workflow job |
|
@pytorchmergebot rebase |
|
@pytorchbot successfully started a rebase job. Check the current status here |
|
Successfully rebased |
|
@pytorchmergebot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
| from .utils import _element_wise_add | ||
|
|
||
|
|
||
| # TODO: update docstring for nested_tensor.py |
There was a problem hiding this comment.
Why are we creating another file with a name that's matching something existing? Is there a plan to use the actual NestedTensor here eventually? cc @drisspg
There was a problem hiding this comment.
Agree, that this naming clash is confusing. For reference this is documentation on NestedTensor: https://pytorch.org/docs/stable/nested.html
…o core distributed (pytorch#89501) This PR moves nested_tensors to torch.distributed.checkpoint. This is a pre-req for enabling 2D checkpoint. This flattens sharded tensors in state_dict. It is used when saving and loading FSDP SHARDED_STATE_DICT. Docstring, individual and integration test will be added in the following PRs. Pull Request resolved: pytorch#89501 Approved by: https://github.com/wanchaol
…92705) Fixes #90350. Pull Request resolved: #92705 Approved by: https://github.com/kumpera
5 participants
This PR moves nested_tensors to torch.distributed.checkpoint. This is a pre-req for enabling 2D checkpoint.
This flattens sharded tensors in state_dict. It is used when saving and loading FSDP SHARDED_STATE_DICT.
Docstring, individual and integration test will be added in the following PRs.