Skip to content

Conversation

d4l3k
Copy link
Member

@d4l3k d4l3k commented Feb 5, 2025

Added in a new streaming_save/load implementations that don't use an extra layer of zipfile serialization

This is a rework of #54 that supports DTensor

Test plan:

pytest torchft/serialization_test.py torchft/checkpointing_test.py -x

Added in a new streaming_save/load implementations that don't use an
extra layer of zipfile serialization

Co-authored-by: Tristan Rice <[email protected]>
Co-authored-by: Krishn Parasar <[email protected]>
@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Meta Open Source bot. label Feb 5, 2025
@d4l3k d4l3k closed this Feb 7, 2025
pytorchmergebot pushed a commit to pytorch/pytorch that referenced this pull request Feb 7, 2025
… methods (#146555)

Summary:

This is intended for use with torchft when we need to do a streaming state dict transfer. This is strictly superior to the prior streaming method in torchft as this supports all tensor subclasses such as DTensor.

This supports 100% of the inputs to torch.save/load but is not wire compatible nor intended to have any backwards compatibility.

Security wise this fully supports weights_only and defaults to True. It does use pickle for some metadata but uses weights_only for the metadata.

Adapted from:

meta-pytorch/torchft#101

meta-pytorch/torchft#54

Test Plan:

pytest test/distributed/test_serialization.py

Pull Request resolved: #146555
Approved by: https://github.com/fegin, https://github.com/mikaylagawarecki

Co-authored-by: Krishn Parasar <[email protected]>
Raymo111 pushed a commit to pytorch/pytorch that referenced this pull request Feb 20, 2025
… methods (#146555)

Summary:

This is intended for use with torchft when we need to do a streaming state dict transfer. This is strictly superior to the prior streaming method in torchft as this supports all tensor subclasses such as DTensor.

This supports 100% of the inputs to torch.save/load but is not wire compatible nor intended to have any backwards compatibility.

Security wise this fully supports weights_only and defaults to True. It does use pickle for some metadata but uses weights_only for the metadata.

Adapted from:

meta-pytorch/torchft#101

meta-pytorch/torchft#54

Test Plan:

pytest test/distributed/test_serialization.py

Pull Request resolved: #146555
Approved by: https://github.com/fegin, https://github.com/mikaylagawarecki

Co-authored-by: Krishn Parasar <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants