Add torch.serialization.skip_data context manager #134504

mikaylagawarecki · 2024-08-26T20:55:27Z

Semantic

The semantic is
(1) By default torch.serialization.skip_data(materialize_fake_tensors=False) will make torch.save skip writing storages (but reserve space for them in the checkpoint).

import torch
import torch.nn as nn

sd = nn.Linear(3, 5).state_dict()
with torch.serialization.skip_data():
    torch.save(sd, 'foo.pt')
print(torch.load('foo.pt', weights_only=True))

(2) With torch.serialization.skip_data(materialize_fake_tensors=True)If FakeTensor is passed to torch.save the pickler will treat these FakeTensors as being "materialized" space will be reserved in the checkpoint for the associated storage bytes, and when loading the type will be Tensor instead of FakeTensor)

import torch
import torch.nn as nn
from torch._subclasses.fake_tensor import FakeTensorMode 

with FakeTensorMode():
    m = nn.Linear(3, 5, dtype=torch.float16, device='cuda')

sd = m.state_dict()
with torch.serialization.skip_data(materialize_fake_tensors=True):
    torch.save(sd, 'bla.pt')
print(torch.load('bla.pt', weights_only=True))
# OrderedDict([('weight', tensor([[0., 0., 0.],
#        [0., 0., 0.],
#        [0., 0., 0.],
#        [0., 0., 0.],
#        [0., 0., 0.]], device='cuda:0', dtype=torch.float16)), ('bias', tensor([0., 0., 0., 0., 0.], device='cuda:0', dtype=torch.float16))])

Follow Ups

torch.load semantic for skip_data context manager
Mechanism for getting offsets of storages saved via this method (for writing in a separate pass)

Stack from ghstack (oldest at bottom):

-> Add torch.serialization.skip_data context manager #134504

Differential Revision: D62238610

[ghstack-poisoned]

pytorch-bot · 2024-08-26T20:55:29Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/134504

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 0cc1afb with merge base 5a0e7a4 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

ghstack-source-id: c7ec2503006744fc967a26d7623f107f1952e1b1 Pull Request resolved: #134504

[ghstack-poisoned]

ghstack-source-id: d877e269960a0984896931ad6b61ae05688eaa9a Pull Request resolved: #134504

The semantic is (1) If real tensors are passed, storages bytes will not be written but zipfile metadata + sufficient space will be saved in the checkpoint for ```python import torch import torch.nn as nn sd = nn.Linear(3, 5).state_dict() torch.save(sd, 'foo.pt', metadata_only=True) print(torch.load('foo.pt', weights_only=True)) # OrderedDict([('weight', tensor([[0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.]])), ('bias', tensor([0., 0., 0., 0., 0.]))]) ``` (2) If FakeTensor is passed, space will be saved in the checkpoint for the associated storage ```python from torch._subclasses.fake_tensor import FakeTensorMode with FakeTensorMode(): m = nn.Linear(3, 5, dtype=torch.float16, device='cuda') sd = m.state_dict() torch.save(sd, 'bla.pt', metadata_only=True) print(torch.load('bla.pt', weights_only=True)) # OrderedDict([('weight', tensor([[0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.]], device='cuda:0', dtype=torch.float16)), ('bias', tensor([0., 0., 0., 0., 0.], device='cuda:0', dtype=torch.float16))]) ``` [ghstack-poisoned]

ghstack-source-id: f462eb65e4620006092e4c8446ed2b99b3325f2f Pull Request resolved: #134504

torch/_tensor.py

mikaylagawarecki · 2024-08-26T21:41:04Z

torch/serialization.py

@@ -978,7 +1004,7 @@ def persistent_id(obj):
            # If storage is allocated, ensure that any other saved storages
            # pointing to the same data all have the same dtype. If storage is
            # not allocated, don't perform this check
-            if storage.data_ptr() != 0:
+            if str(storage.device) != "meta" and storage.data_ptr() != 0:


Using str(storage.device) != "meta" to avoid the deprecation message mentioned above when accessing FakeTensor data_ptr

...are there other cases when device is not meta but storage data_ptr is 0?

ok, think this fix is good then to avoid accessing data_ptr specifically for FakeTensor storage

The semantic is (1) If real tensors are passed, storages bytes will not be written but zipfile metadata + sufficient space will be saved in the checkpoint for ```python import torch import torch.nn as nn sd = nn.Linear(3, 5).state_dict() torch.save(sd, 'foo.pt', metadata_only=True) print(torch.load('foo.pt', weights_only=True)) # OrderedDict([('weight', tensor([[0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.]])), ('bias', tensor([0., 0., 0., 0., 0.]))]) ``` (2) If FakeTensor is passed to `torch.save(metadata_only=True)` space will be reserved in the checkpoint for the associated storage bytes, and when loading the type will be Tensor instead of FakeTensor) ```python from torch._subclasses.fake_tensor import FakeTensorMode with FakeTensorMode(): m = nn.Linear(3, 5, dtype=torch.float16, device='cuda') sd = m.state_dict() torch.save(sd, 'bla.pt', metadata_only=True) print(torch.load('bla.pt', weights_only=True)) # OrderedDict([('weight', tensor([[0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.]], device='cuda:0', dtype=torch.float16)), ('bias', tensor([0., 0., 0., 0., 0.], device='cuda:0', dtype=torch.float16))]) ``` [ghstack-poisoned]

ghstack-source-id: 2ded66d41b110f0c9bf44ad0b36688201a9d85b4 Pull Request resolved: #134504

The semantic is (1) If real tensors are passed, storages bytes will not be written but zipfile metadata + sufficient space will be saved in the checkpoint for ```python import torch import torch.nn as nn sd = nn.Linear(3, 5).state_dict() torch.save(sd, 'foo.pt', metadata_only=True) print(torch.load('foo.pt', weights_only=True)) # OrderedDict([('weight', tensor([[0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.]])), ('bias', tensor([0., 0., 0., 0., 0.]))]) ``` (2) If FakeTensor is passed to `torch.save(metadata_only=True)` space will be reserved in the checkpoint for the associated storage bytes, and when loading the type will be Tensor instead of FakeTensor) ```python from torch._subclasses.fake_tensor import FakeTensorMode with FakeTensorMode(): m = nn.Linear(3, 5, dtype=torch.float16, device='cuda') sd = m.state_dict() torch.save(sd, 'bla.pt', metadata_only=True) print(torch.load('bla.pt', weights_only=True)) # OrderedDict([('weight', tensor([[0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.]], device='cuda:0', dtype=torch.float16)), ('bias', tensor([0., 0., 0., 0., 0.], device='cuda:0', dtype=torch.float16))]) ``` [ghstack-poisoned]

ghstack-source-id: 4590955bb49646f3b461245fbc88c2051297bbd3 Pull Request resolved: #134504

The semantic is (1) By default `torch.serialization.skip_data(materialize_fake_tensors=False)` will make torch.save skip writing storages (but reserve space for them in the checkpoint). ```python import torch import torch.nn as nn sd = nn.Linear(3, 5).state_dict() with torch.serialization.skip_data(): torch.save(sd, 'foo.pt') print(torch.load('foo.pt', weights_only=True)) ``` (2) With `torch.serialization.skip_data(materialize_fake_tensors=True)`If FakeTensor is passed to `torch.save()` the pickler will treat these FakeTensors as being "materialized" space will be reserved in the checkpoint for the associated storage bytes, and when loading the type will be Tensor instead of FakeTensor) ```python import torch import torch.nn as nn from torch._subclasses.fake_tensor import FakeTensorMode with FakeTensorMode(): m = nn.Linear(3, 5, dtype=torch.float16, device='cuda') sd = m.state_dict() with torch.serialization.skip_data(materialize_fake_tensors=True): torch.save(sd, 'bla.pt') print(torch.load('bla.pt', weights_only=True)) # OrderedDict([('weight', tensor([[0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.]], device='cuda:0', dtype=torch.float16)), ('bias', tensor([0., 0., 0., 0., 0.], device='cuda:0', dtype=torch.float16))]) ``` [ghstack-poisoned]

ghstack-source-id: 04f6268af51de244f7c7f1a267ae85902d9c490a Pull Request resolved: #134504

The semantic is (1) By default `torch.serialization.skip_data(materialize_fake_tensors=False)` will make torch.save skip writing storages (but reserve space for them in the checkpoint). ```python import torch import torch.nn as nn sd = nn.Linear(3, 5).state_dict() with torch.serialization.skip_data(): torch.save(sd, 'foo.pt') print(torch.load('foo.pt', weights_only=True)) ``` (2) With `torch.serialization.skip_data(materialize_fake_tensors=True)`If FakeTensor is passed to `torch.save()` the pickler will treat these FakeTensors as being "materialized" space will be reserved in the checkpoint for the associated storage bytes, and when loading the type will be Tensor instead of FakeTensor) ```python import torch import torch.nn as nn from torch._subclasses.fake_tensor import FakeTensorMode with FakeTensorMode(): m = nn.Linear(3, 5, dtype=torch.float16, device='cuda') sd = m.state_dict() with torch.serialization.skip_data(materialize_fake_tensors=True): torch.save(sd, 'bla.pt') print(torch.load('bla.pt', weights_only=True)) # OrderedDict([('weight', tensor([[0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.]], device='cuda:0', dtype=torch.float16)), ('bias', tensor([0., 0., 0., 0., 0.], device='cuda:0', dtype=torch.float16))]) ``` [ghstack-poisoned]

ghstack-source-id: 353853b426919cac893a28a8a83968a92cf36aa7 Pull Request resolved: #134504

The semantic is (1) By default `torch.serialization.skip_data(materialize_fake_tensors=False)` will make torch.save skip writing storages (but reserve space for them in the checkpoint). ```python import torch import torch.nn as nn sd = nn.Linear(3, 5).state_dict() with torch.serialization.skip_data(): torch.save(sd, 'foo.pt') print(torch.load('foo.pt', weights_only=True)) ``` (2) With `torch.serialization.skip_data(materialize_fake_tensors=True)`If FakeTensor is passed to `torch.save()` the pickler will treat these FakeTensors as being "materialized" space will be reserved in the checkpoint for the associated storage bytes, and when loading the type will be Tensor instead of FakeTensor) ```python import torch import torch.nn as nn from torch._subclasses.fake_tensor import FakeTensorMode with FakeTensorMode(): m = nn.Linear(3, 5, dtype=torch.float16, device='cuda') sd = m.state_dict() with torch.serialization.skip_data(materialize_fake_tensors=True): torch.save(sd, 'bla.pt') print(torch.load('bla.pt', weights_only=True)) # OrderedDict([('weight', tensor([[0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.]], device='cuda:0', dtype=torch.float16)), ('bias', tensor([0., 0., 0., 0., 0.], device='cuda:0', dtype=torch.float16))]) ``` [ghstack-poisoned]

ghstack-source-id: 508e473d845f34db9c7c0fd7c3f1fa1f8a30b46a Pull Request resolved: #134504

The semantic is (1) By default `torch.serialization.skip_data(materialize_fake_tensors=False)` will make `torch.save` skip writing storages (but reserve space for them in the checkpoint). ```python import torch import torch.nn as nn sd = nn.Linear(3, 5).state_dict() with torch.serialization.skip_data(): torch.save(sd, 'foo.pt') print(torch.load('foo.pt', weights_only=True)) ``` (2) With `torch.serialization.skip_data(materialize_fake_tensors=True)`If FakeTensor is passed to `torch.save` the pickler will treat these FakeTensors as being "materialized" space will be reserved in the checkpoint for the associated storage bytes, and when loading the type will be Tensor instead of FakeTensor) ```python import torch import torch.nn as nn from torch._subclasses.fake_tensor import FakeTensorMode with FakeTensorMode(): m = nn.Linear(3, 5, dtype=torch.float16, device='cuda') sd = m.state_dict() with torch.serialization.skip_data(materialize_fake_tensors=True): torch.save(sd, 'bla.pt') print(torch.load('bla.pt', weights_only=True)) # OrderedDict([('weight', tensor([[0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.]], device='cuda:0', dtype=torch.float16)), ('bias', tensor([0., 0., 0., 0., 0.], device='cuda:0', dtype=torch.float16))]) ``` [ghstack-poisoned]

ghstack-source-id: 5257374670569afd9696e23dfa7c100e00e63399 Pull Request resolved: #134504

albanD · 2024-08-27T20:31:37Z

torch/_tensor.py

        state = torch._utils._get_obj_state(self)
-        if type(self) is Tensor and not state:
+        # Ignore all state when using FakeTensor with skip_data(materialize_fake_tensors) because FakeTensor has


I think this is sub-optimal in the long term. We should have a clearer contract on how our FakeTensor -> real Tensor with no data.
Dropping every single field from it might not be acceptable for everyone and might be interesting if there was a method to extract such Tensor from it easily.

This sounds ok as a temporary solution here I think

I agree, technically what we want is to drop every attribute that is in the dict of a "normal" FakeTensor, but keep anything else, is that right?

But agreed that for now it seems tricky to handle this because I'm not sure what is in a "normal" FakeTensor's dict is stable

torch/_tensor.py

albanD · 2024-08-27T20:44:03Z

torch/_tensor.py

-                    ),
+                (
+                    isinstance(
+                        self, torch._subclasses.functional_tensor.FunctionalTensor


Should we just split functional + fake Tensor to another condition just above? I have to admit I can't read this condition anymore :D

albanD · 2024-08-27T20:45:46Z

torch/serialization.py

+# (1) map_location (needed for wrapper subclasses/third party devices to torch._utils)
+# (2) skip_data (needed for torch.Tensor.__reduce_ex__ for skip_data ctx)
+# (3) materialize_fake_tensors (needed for torch.Tensor.__reduce_ex__ for skip_data ctx)
+_serialization_tls = threading.local()


Ho I would have expected you just did torch._C._stash_obj_in_tls("_serialization_tls", _serialization_tls) here and nothing else below?

Or if we don't know how to do it. It's ok to remote it all and open an issue as a TODO for later.

Removed and opened issue here #134680

torch/serialization.py

The semantic is (1) By default `torch.serialization.skip_data(materialize_fake_tensors=False)` will make `torch.save` skip writing storages (but reserve space for them in the checkpoint). ```python import torch import torch.nn as nn sd = nn.Linear(3, 5).state_dict() with torch.serialization.skip_data(): torch.save(sd, 'foo.pt') print(torch.load('foo.pt', weights_only=True)) ``` (2) With `torch.serialization.skip_data(materialize_fake_tensors=True)`If FakeTensor is passed to `torch.save` the pickler will treat these FakeTensors as being "materialized" space will be reserved in the checkpoint for the associated storage bytes, and when loading the type will be Tensor instead of FakeTensor) ```python import torch import torch.nn as nn from torch._subclasses.fake_tensor import FakeTensorMode with FakeTensorMode(): m = nn.Linear(3, 5, dtype=torch.float16, device='cuda') sd = m.state_dict() with torch.serialization.skip_data(materialize_fake_tensors=True): torch.save(sd, 'bla.pt') print(torch.load('bla.pt', weights_only=True)) # OrderedDict([('weight', tensor([[0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.]], device='cuda:0', dtype=torch.float16)), ('bias', tensor([0., 0., 0., 0., 0.], device='cuda:0', dtype=torch.float16))]) ``` [ghstack-poisoned]

ghstack-source-id: ddcd957cc6258a1d5eb7bf54910deee2ee02527d Pull Request resolved: #134504

The semantic is (1) By default `torch.serialization.skip_data(materialize_fake_tensors=False)` will make `torch.save` skip writing storages (but reserve space for them in the checkpoint). ```python import torch import torch.nn as nn sd = nn.Linear(3, 5).state_dict() with torch.serialization.skip_data(): torch.save(sd, 'foo.pt') print(torch.load('foo.pt', weights_only=True)) ``` (2) With `torch.serialization.skip_data(materialize_fake_tensors=True)`If FakeTensor is passed to `torch.save` the pickler will treat these FakeTensors as being "materialized" space will be reserved in the checkpoint for the associated storage bytes, and when loading the type will be Tensor instead of FakeTensor) ```python import torch import torch.nn as nn from torch._subclasses.fake_tensor import FakeTensorMode with FakeTensorMode(): m = nn.Linear(3, 5, dtype=torch.float16, device='cuda') sd = m.state_dict() with torch.serialization.skip_data(materialize_fake_tensors=True): torch.save(sd, 'bla.pt') print(torch.load('bla.pt', weights_only=True)) # OrderedDict([('weight', tensor([[0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.]], device='cuda:0', dtype=torch.float16)), ('bias', tensor([0., 0., 0., 0., 0.], device='cuda:0', dtype=torch.float16))]) ``` [ghstack-poisoned]

ghstack-source-id: 507667c8e1f8956f09c236d486d66f744fc89233 Pull Request resolved: #134504

## Semantic The semantic is (1) By default `torch.serialization.skip_data(materialize_fake_tensors=False)` will make `torch.save` skip writing storages (but reserve space for them in the checkpoint). ```python import torch import torch.nn as nn sd = nn.Linear(3, 5).state_dict() with torch.serialization.skip_data(): torch.save(sd, 'foo.pt') print(torch.load('foo.pt', weights_only=True)) ``` (2) With `torch.serialization.skip_data(materialize_fake_tensors=True)`If FakeTensor is passed to `torch.save` the pickler will treat these FakeTensors as being "materialized" space will be reserved in the checkpoint for the associated storage bytes, and when loading the type will be Tensor instead of FakeTensor) ```python import torch import torch.nn as nn from torch._subclasses.fake_tensor import FakeTensorMode with FakeTensorMode(): m = nn.Linear(3, 5, dtype=torch.float16, device='cuda') sd = m.state_dict() with torch.serialization.skip_data(materialize_fake_tensors=True): torch.save(sd, 'bla.pt') print(torch.load('bla.pt', weights_only=True)) # OrderedDict([('weight', tensor([[0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.]], device='cuda:0', dtype=torch.float16)), ('bias', tensor([0., 0., 0., 0., 0.], device='cuda:0', dtype=torch.float16))]) ``` ## Follow Ups - [ ] `torch.load` semantic for skip_data context manager - [ ] Mechanism for getting offsets of storages saved via this method (for writing in a separate pass) [ghstack-poisoned]

ghstack-source-id: f0ebde9fd35937451885a562b05b7c02d4f91dac Pull Request resolved: #134504

mikaylagawarecki · 2024-09-05T13:23:08Z

@mikaylagawarecki has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

## Semantic The semantic is (1) By default `torch.serialization.skip_data(materialize_fake_tensors=False)` will make `torch.save` skip writing storages (but reserve space for them in the checkpoint). ```python import torch import torch.nn as nn sd = nn.Linear(3, 5).state_dict() with torch.serialization.skip_data(): torch.save(sd, 'foo.pt') print(torch.load('foo.pt', weights_only=True)) ``` (2) With `torch.serialization.skip_data(materialize_fake_tensors=True)`If FakeTensor is passed to `torch.save` the pickler will treat these FakeTensors as being "materialized" space will be reserved in the checkpoint for the associated storage bytes, and when loading the type will be Tensor instead of FakeTensor) ```python import torch import torch.nn as nn from torch._subclasses.fake_tensor import FakeTensorMode with FakeTensorMode(): m = nn.Linear(3, 5, dtype=torch.float16, device='cuda') sd = m.state_dict() with torch.serialization.skip_data(materialize_fake_tensors=True): torch.save(sd, 'bla.pt') print(torch.load('bla.pt', weights_only=True)) # OrderedDict([('weight', tensor([[0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.]], device='cuda:0', dtype=torch.float16)), ('bias', tensor([0., 0., 0., 0., 0.], device='cuda:0', dtype=torch.float16))]) ``` ## Follow Ups - [ ] `torch.load` semantic for skip_data context manager - [ ] Mechanism for getting offsets of storages saved via this method (for writing in a separate pass) Differential Revision: [D62238610](https://our.internmc.facebook.com/intern/diff/D62238610) [ghstack-poisoned]

ghstack-source-id: edd5529c1a59d84dcc63562874a73c7927f524d2 Pull Request resolved: #134504

## Semantic The semantic is (1) By default `torch.serialization.skip_data(materialize_fake_tensors=False)` will make `torch.save` skip writing storages (but reserve space for them in the checkpoint). ```python import torch import torch.nn as nn sd = nn.Linear(3, 5).state_dict() with torch.serialization.skip_data(): torch.save(sd, 'foo.pt') print(torch.load('foo.pt', weights_only=True)) ``` (2) With `torch.serialization.skip_data(materialize_fake_tensors=True)`If FakeTensor is passed to `torch.save` the pickler will treat these FakeTensors as being "materialized" space will be reserved in the checkpoint for the associated storage bytes, and when loading the type will be Tensor instead of FakeTensor) ```python import torch import torch.nn as nn from torch._subclasses.fake_tensor import FakeTensorMode with FakeTensorMode(): m = nn.Linear(3, 5, dtype=torch.float16, device='cuda') sd = m.state_dict() with torch.serialization.skip_data(materialize_fake_tensors=True): torch.save(sd, 'bla.pt') print(torch.load('bla.pt', weights_only=True)) # OrderedDict([('weight', tensor([[0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.]], device='cuda:0', dtype=torch.float16)), ('bias', tensor([0., 0., 0., 0., 0.], device='cuda:0', dtype=torch.float16))]) ``` ## Follow Ups - [ ] `torch.load` semantic for skip_data context manager - [ ] Mechanism for getting offsets of storages saved via this method (for writing in a separate pass) Differential Revision: [D62238610](https://our.internmc.facebook.com/intern/diff/D62238610) [ghstack-poisoned]

ghstack-source-id: 9258c783fbb455fa0caa507c00e1a74e834d7cf0 Pull Request resolved: #134504

mikaylagawarecki · 2024-09-05T13:52:22Z

@mikaylagawarecki has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

mikaylagawarecki · 2024-09-05T14:05:52Z

Repro given no longer fails, see D62238610

albanD

SGTM!

mikaylagawarecki · 2024-09-05T14:42:54Z

@pytorchbot merge

pytorchmergebot · 2024-09-05T14:44:45Z

Merge started

Your 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

Advanced Debugging

Check the merge workflow status
here

mikaylagawarecki · 2024-09-05T14:57:52Z

@pytorchbot merge

pytorchmergebot · 2024-09-05T14:58:12Z

The merge job was canceled or timed out. This most often happen if two merge requests were issued for the same PR, or if merge job was waiting for more than 6 hours for tests to finish. In later case, please do not hesitate to reissue the merge command
For more information see pytorch-bot wiki.

pytorchmergebot · 2024-09-05T14:59:59Z

Merge started

Your 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

Advanced Debugging

Check the merge workflow status
here

## Semantic The semantic is (1) By default `torch.serialization.skip_data(materialize_fake_tensors=False)` will make `torch.save` skip writing storages (but reserve space for them in the checkpoint). ```python import torch import torch.nn as nn sd = nn.Linear(3, 5).state_dict() with torch.serialization.skip_data(): torch.save(sd, 'foo.pt') print(torch.load('foo.pt', weights_only=True)) ``` (2) With `torch.serialization.skip_data(materialize_fake_tensors=True)`If FakeTensor is passed to `torch.save` the pickler will treat these FakeTensors as being "materialized" space will be reserved in the checkpoint for the associated storage bytes, and when loading the type will be Tensor instead of FakeTensor) ```python import torch import torch.nn as nn from torch._subclasses.fake_tensor import FakeTensorMode with FakeTensorMode(): m = nn.Linear(3, 5, dtype=torch.float16, device='cuda') sd = m.state_dict() with torch.serialization.skip_data(materialize_fake_tensors=True): torch.save(sd, 'bla.pt') print(torch.load('bla.pt', weights_only=True)) # OrderedDict([('weight', tensor([[0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.]], device='cuda:0', dtype=torch.float16)), ('bias', tensor([0., 0., 0., 0., 0.], device='cuda:0', dtype=torch.float16))]) ``` ## Follow Ups - [ ] `torch.load` semantic for skip_data context manager - [ ] Mechanism for getting offsets of storages saved via this method (for writing in a separate pass) Pull Request resolved: pytorch#134504 Approved by: https://github.com/albanD

…4504)" This reverts commit 202600b. Reverted pytorch#134504 on behalf of https://github.com/mikaylagawarecki due to This is breaking Windows docs tests due to NamedTemporaryFile on Windows not working well ([comment](pytorch#134504 (comment)))

## Semantic The semantic is (1) By default `torch.serialization.skip_data(materialize_fake_tensors=False)` will make `torch.save` skip writing storages (but reserve space for them in the checkpoint). ```python import torch import torch.nn as nn sd = nn.Linear(3, 5).state_dict() with torch.serialization.skip_data(): torch.save(sd, 'foo.pt') print(torch.load('foo.pt', weights_only=True)) ``` (2) With `torch.serialization.skip_data(materialize_fake_tensors=True)`If FakeTensor is passed to `torch.save` the pickler will treat these FakeTensors as being "materialized" space will be reserved in the checkpoint for the associated storage bytes, and when loading the type will be Tensor instead of FakeTensor) ```python import torch import torch.nn as nn from torch._subclasses.fake_tensor import FakeTensorMode with FakeTensorMode(): m = nn.Linear(3, 5, dtype=torch.float16, device='cuda') sd = m.state_dict() with torch.serialization.skip_data(materialize_fake_tensors=True): torch.save(sd, 'bla.pt') print(torch.load('bla.pt', weights_only=True)) # OrderedDict([('weight', tensor([[0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.]], device='cuda:0', dtype=torch.float16)), ('bias', tensor([0., 0., 0., 0., 0.], device='cuda:0', dtype=torch.float16))]) ``` ## Follow Ups - [ ] `torch.load` semantic for skip_data context manager - [ ] Mechanism for getting offsets of storages saved via this method (for writing in a separate pass) Pull Request resolved: pytorch#134504 Approved by: https://github.com/albanD

…4504)" This reverts commit 94db935. Reverted pytorch#134504 on behalf of https://github.com/kit1980 due to See D62082697 ([comment](pytorch#134504 (comment)))

## Semantic The semantic is (1) By default `torch.serialization.skip_data(materialize_fake_tensors=False)` will make `torch.save` skip writing storages (but reserve space for them in the checkpoint). ```python import torch import torch.nn as nn sd = nn.Linear(3, 5).state_dict() with torch.serialization.skip_data(): torch.save(sd, 'foo.pt') print(torch.load('foo.pt', weights_only=True)) ``` (2) With `torch.serialization.skip_data(materialize_fake_tensors=True)`If FakeTensor is passed to `torch.save` the pickler will treat these FakeTensors as being "materialized" space will be reserved in the checkpoint for the associated storage bytes, and when loading the type will be Tensor instead of FakeTensor) ```python import torch import torch.nn as nn from torch._subclasses.fake_tensor import FakeTensorMode with FakeTensorMode(): m = nn.Linear(3, 5, dtype=torch.float16, device='cuda') sd = m.state_dict() with torch.serialization.skip_data(materialize_fake_tensors=True): torch.save(sd, 'bla.pt') print(torch.load('bla.pt', weights_only=True)) # OrderedDict([('weight', tensor([[0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.]], device='cuda:0', dtype=torch.float16)), ('bias', tensor([0., 0., 0., 0., 0.], device='cuda:0', dtype=torch.float16))]) ``` ## Follow Ups - [ ] `torch.load` semantic for skip_data context manager - [ ] Mechanism for getting offsets of storages saved via this method (for writing in a separate pass) Differential Revision: [D62238610](https://our.internmc.facebook.com/intern/diff/D62238610) Pull Request resolved: pytorch#134504 Approved by: https://github.com/albanD

## Semantic The semantic is (1) By default `torch.serialization.skip_data(materialize_fake_tensors=False)` will make `torch.save` skip writing storages (but reserve space for them in the checkpoint). ```python import torch import torch.nn as nn sd = nn.Linear(3, 5).state_dict() with torch.serialization.skip_data(): torch.save(sd, 'foo.pt') print(torch.load('foo.pt', weights_only=True)) ``` (2) With `torch.serialization.skip_data(materialize_fake_tensors=True)`If FakeTensor is passed to `torch.save` the pickler will treat these FakeTensors as being "materialized" space will be reserved in the checkpoint for the associated storage bytes, and when loading the type will be Tensor instead of FakeTensor) ```python import torch import torch.nn as nn from torch._subclasses.fake_tensor import FakeTensorMode with FakeTensorMode(): m = nn.Linear(3, 5, dtype=torch.float16, device='cuda') sd = m.state_dict() with torch.serialization.skip_data(materialize_fake_tensors=True): torch.save(sd, 'bla.pt') print(torch.load('bla.pt', weights_only=True)) # OrderedDict([('weight', tensor([[0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.]], device='cuda:0', dtype=torch.float16)), ('bias', tensor([0., 0., 0., 0., 0.], device='cuda:0', dtype=torch.float16))]) ``` ## Follow Ups - [ ] `torch.load` semantic for skip_data context manager - [ ] Mechanism for getting offsets of storages saved via this method (for writing in a separate pass) Pull Request resolved: pytorch#134504 Approved by: https://github.com/albanD

…4504)" This reverts commit 202600b. Reverted pytorch#134504 on behalf of https://github.com/mikaylagawarecki due to This is breaking Windows docs tests due to NamedTemporaryFile on Windows not working well ([comment](pytorch#134504 (comment)))

## Semantic The semantic is (1) By default `torch.serialization.skip_data(materialize_fake_tensors=False)` will make `torch.save` skip writing storages (but reserve space for them in the checkpoint). ```python import torch import torch.nn as nn sd = nn.Linear(3, 5).state_dict() with torch.serialization.skip_data(): torch.save(sd, 'foo.pt') print(torch.load('foo.pt', weights_only=True)) ``` (2) With `torch.serialization.skip_data(materialize_fake_tensors=True)`If FakeTensor is passed to `torch.save` the pickler will treat these FakeTensors as being "materialized" space will be reserved in the checkpoint for the associated storage bytes, and when loading the type will be Tensor instead of FakeTensor) ```python import torch import torch.nn as nn from torch._subclasses.fake_tensor import FakeTensorMode with FakeTensorMode(): m = nn.Linear(3, 5, dtype=torch.float16, device='cuda') sd = m.state_dict() with torch.serialization.skip_data(materialize_fake_tensors=True): torch.save(sd, 'bla.pt') print(torch.load('bla.pt', weights_only=True)) # OrderedDict([('weight', tensor([[0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.]], device='cuda:0', dtype=torch.float16)), ('bias', tensor([0., 0., 0., 0., 0.], device='cuda:0', dtype=torch.float16))]) ``` ## Follow Ups - [ ] `torch.load` semantic for skip_data context manager - [ ] Mechanism for getting offsets of storages saved via this method (for writing in a separate pass) Pull Request resolved: pytorch#134504 Approved by: https://github.com/albanD

…4504)" This reverts commit 94db935. Reverted pytorch#134504 on behalf of https://github.com/kit1980 due to See D62082697 ([comment](pytorch#134504 (comment)))

## Semantic The semantic is (1) By default `torch.serialization.skip_data(materialize_fake_tensors=False)` will make `torch.save` skip writing storages (but reserve space for them in the checkpoint). ```python import torch import torch.nn as nn sd = nn.Linear(3, 5).state_dict() with torch.serialization.skip_data(): torch.save(sd, 'foo.pt') print(torch.load('foo.pt', weights_only=True)) ``` (2) With `torch.serialization.skip_data(materialize_fake_tensors=True)`If FakeTensor is passed to `torch.save` the pickler will treat these FakeTensors as being "materialized" space will be reserved in the checkpoint for the associated storage bytes, and when loading the type will be Tensor instead of FakeTensor) ```python import torch import torch.nn as nn from torch._subclasses.fake_tensor import FakeTensorMode with FakeTensorMode(): m = nn.Linear(3, 5, dtype=torch.float16, device='cuda') sd = m.state_dict() with torch.serialization.skip_data(materialize_fake_tensors=True): torch.save(sd, 'bla.pt') print(torch.load('bla.pt', weights_only=True)) # OrderedDict([('weight', tensor([[0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.], # [0., 0., 0.]], device='cuda:0', dtype=torch.float16)), ('bias', tensor([0., 0., 0., 0., 0.], device='cuda:0', dtype=torch.float16))]) ``` ## Follow Ups - [ ] `torch.load` semantic for skip_data context manager - [ ] Mechanism for getting offsets of storages saved via this method (for writing in a separate pass) Differential Revision: [D62238610](https://our.internmc.facebook.com/intern/diff/D62238610) Pull Request resolved: pytorch#134504 Approved by: https://github.com/albanD

Add metadata_only flag to torch.save

920d15f

[ghstack-poisoned]

mikaylagawarecki mentioned this pull request Aug 26, 2024

Prototype changes to create fake checkpoints with empty storages #134503

Closed

mikaylagawarecki added a commit that referenced this pull request Aug 26, 2024

Add metadata_only flag to torch.save

9ba4180

ghstack-source-id: c7ec2503006744fc967a26d7623f107f1952e1b1 Pull Request resolved: #134504

Update on "Add metadata_only flag to torch.save"

fe1f39e

[ghstack-poisoned]

mikaylagawarecki added a commit that referenced this pull request Aug 26, 2024

Add metadata_only flag to torch.save

f2082a8

ghstack-source-id: d877e269960a0984896931ad6b61ae05688eaa9a Pull Request resolved: #134504

mikaylagawarecki added a commit that referenced this pull request Aug 26, 2024

Add metadata_only flag to torch.save

abcc351

ghstack-source-id: f462eb65e4620006092e4c8446ed2b99b3325f2f Pull Request resolved: #134504

mikaylagawarecki commented Aug 26, 2024

View reviewed changes

torch/_tensor.py Outdated Show resolved Hide resolved

mikaylagawarecki commented Aug 26, 2024

View reviewed changes

mikaylagawarecki added a commit that referenced this pull request Aug 26, 2024

Add metadata_only flag to torch.save

600b7a5

ghstack-source-id: 2ded66d41b110f0c9bf44ad0b36688201a9d85b4 Pull Request resolved: #134504

mikaylagawarecki added a commit that referenced this pull request Aug 27, 2024

Add metadata_only flag to torch.save

6ef50bd

ghstack-source-id: 4590955bb49646f3b461245fbc88c2051297bbd3 Pull Request resolved: #134504

mikaylagawarecki changed the title ~~Add metadata_only flag to torch.save~~ Add torch.serialization.skip_data context manager Aug 27, 2024

mikaylagawarecki added a commit that referenced this pull request Aug 27, 2024

Add metadata_only flag to torch.save

8f1925f

ghstack-source-id: 04f6268af51de244f7c7f1a267ae85902d9c490a Pull Request resolved: #134504

mikaylagawarecki requested a review from albanD August 27, 2024 19:11

mikaylagawarecki added a commit that referenced this pull request Aug 27, 2024

Add metadata_only flag to torch.save

a449983

ghstack-source-id: 353853b426919cac893a28a8a83968a92cf36aa7 Pull Request resolved: #134504

mikaylagawarecki added a commit that referenced this pull request Aug 27, 2024

Add metadata_only flag to torch.save

31ce170

ghstack-source-id: 508e473d845f34db9c7c0fd7c3f1fa1f8a30b46a Pull Request resolved: #134504

mikaylagawarecki added a commit that referenced this pull request Aug 27, 2024

Add metadata_only flag to torch.save

f5b5901

ghstack-source-id: 5257374670569afd9696e23dfa7c100e00e63399 Pull Request resolved: #134504

albanD reviewed Aug 27, 2024

View reviewed changes

mikaylagawarecki added a commit that referenced this pull request Aug 28, 2024

Add metadata_only flag to torch.save

a132d66

ghstack-source-id: ddcd957cc6258a1d5eb7bf54910deee2ee02527d Pull Request resolved: #134504

mikaylagawarecki mentioned this pull request Aug 28, 2024

Use torch._C._stash_obj_in_tls for global state in serialization #134680

Open

mikaylagawarecki added a commit that referenced this pull request Aug 28, 2024

Add metadata_only flag to torch.save

6734652

ghstack-source-id: 507667c8e1f8956f09c236d486d66f744fc89233 Pull Request resolved: #134504

mikaylagawarecki added a commit that referenced this pull request Sep 5, 2024

Add metadata_only flag to torch.save

a4e3076

ghstack-source-id: f0ebde9fd35937451885a562b05b7c02d4f91dac Pull Request resolved: #134504

mikaylagawarecki added a commit that referenced this pull request Sep 5, 2024

Add metadata_only flag to torch.save

c940867

ghstack-source-id: edd5529c1a59d84dcc63562874a73c7927f524d2 Pull Request resolved: #134504

mikaylagawarecki added a commit that referenced this pull request Sep 5, 2024

Add metadata_only flag to torch.save

112b44c

ghstack-source-id: 9258c783fbb455fa0caa507c00e1a74e834d7cf0 Pull Request resolved: #134504

albanD approved these changes Sep 5, 2024

View reviewed changes

pytorchmergebot added the merging label Sep 5, 2024

pytorchmergebot closed this in a096f28 Sep 5, 2024

pytorchmergebot removed the merging label Sep 5, 2024

github-actions bot deleted the gh/mikaylagawarecki/261/head branch October 6, 2024 02:09

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add torch.serialization.skip_data context manager #134504

Add torch.serialization.skip_data context manager #134504

mikaylagawarecki commented Aug 26, 2024 •

edited

Loading

pytorch-bot bot commented Aug 26, 2024 •

edited

Loading

mikaylagawarecki Aug 26, 2024 •

edited

Loading

albanD Aug 26, 2024

mikaylagawarecki Aug 27, 2024

albanD Aug 27, 2024

mikaylagawarecki Aug 28, 2024 •

edited

Loading

albanD Aug 27, 2024

albanD Aug 27, 2024

mikaylagawarecki Aug 28, 2024

mikaylagawarecki commented Sep 5, 2024

mikaylagawarecki commented Sep 5, 2024

mikaylagawarecki commented Sep 5, 2024

albanD left a comment

mikaylagawarecki commented Sep 5, 2024

pytorchmergebot commented Sep 5, 2024

mikaylagawarecki commented Sep 5, 2024

pytorchmergebot commented Sep 5, 2024

pytorchmergebot commented Sep 5, 2024

Add torch.serialization.skip_data context manager #134504

Add torch.serialization.skip_data context manager #134504

Conversation

mikaylagawarecki commented Aug 26, 2024 • edited Loading

Semantic

Follow Ups

pytorch-bot bot commented Aug 26, 2024 • edited Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/134504

✅ No Failures

mikaylagawarecki Aug 26, 2024 • edited Loading

Choose a reason for hiding this comment

albanD Aug 26, 2024

Choose a reason for hiding this comment

mikaylagawarecki Aug 27, 2024

Choose a reason for hiding this comment

albanD Aug 27, 2024

Choose a reason for hiding this comment

mikaylagawarecki Aug 28, 2024 • edited Loading

Choose a reason for hiding this comment

albanD Aug 27, 2024

Choose a reason for hiding this comment

albanD Aug 27, 2024

Choose a reason for hiding this comment

mikaylagawarecki Aug 28, 2024

Choose a reason for hiding this comment

mikaylagawarecki commented Sep 5, 2024

mikaylagawarecki commented Sep 5, 2024

mikaylagawarecki commented Sep 5, 2024

albanD left a comment

Choose a reason for hiding this comment

mikaylagawarecki commented Sep 5, 2024

pytorchmergebot commented Sep 5, 2024

Merge started

mikaylagawarecki commented Sep 5, 2024

pytorchmergebot commented Sep 5, 2024

pytorchmergebot commented Sep 5, 2024

Merge started

mikaylagawarecki commented Aug 26, 2024 •

edited

Loading

pytorch-bot bot commented Aug 26, 2024 •

edited

Loading

mikaylagawarecki Aug 26, 2024 •

edited

Loading

mikaylagawarecki Aug 28, 2024 •

edited

Loading