Pass torch.load(weights_only=) internally to avoid FutureWarning
#130663
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/130663
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit e0275f8 with merge base 9df4bc6 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
There is also: |
There was a problem hiding this comment.
For weights_only=True, if there is a known list of GLOBALs that might be encountered (e.g. for DTensor) that we expect to be safe we can use the
with torch.serialization.safe_globals([bla]): context manager
e.g. for DTensor, [DTensor, DeviceMesh, Shard, DTensorSpec, TensorMeta] might be needed
| @@ -108,7 +108,7 @@ def _broadcast_object( | |||
| ) | |||
| dist.broadcast(data_recv_tensor, src=src_rank, group=group, async_op=False) | |||
| buffer = io.BytesIO(data_recv_tensor.cpu().numpy()) | |||
| obj = torch.load(buffer, map_location=device) | |||
| obj = torch.load(buffer, map_location=device, weights_only=False) | |||
There was a problem hiding this comment.
From the code above, we can see this is broadcasting an object pickled to bytes, so loading here means unpickling. The function doesn't make any assumptions like obj being a tensor, so we need to set weights_only=False here.
|
@mikaylagawarecki @malfet The reason I set them to False is simply to ensure I don't break anything, since False was implicitly the default before. For the distributed checkpoint utilities, I don't know whether we can make the assumptions that only tensors are in the checkpoint. |
Switching to |
|
Could someone approve the CI workflow again, thx. I didn't know my push would disable it again. |
|
@pytorchbot merge |
Merge failedReason: This PR needs a If not, please add the To add a label, you can comment to pytorchbot, for example For more information, see Details for Dev Infra teamRaised by workflow job |
|
@pytorchbot label "release notes: python_frontend" |
|
@pytorchbot 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 |
…130663) Fixes #130658 Pull Request resolved: #130663 Approved by: https://github.com/malfet, https://github.com/LucasLLC
…ytorch#130663) Fixes pytorch#130658 Pull Request resolved: pytorch#130663 Approved by: https://github.com/malfet, https://github.com/LucasLLC
|
@pytorchbot cherry-pick --onto release/2.4 -c critical --fixes #130658 |
…130663) Fixes #130658 Pull Request resolved: #130663 Approved by: https://github.com/malfet, https://github.com/LucasLLC (cherry picked from commit ad314a2)
Cherry picking #130663The cherry pick PR is at #133594 and it is linked with issue #130658. The following tracker issues are updated: Details for Dev Infra teamRaised by workflow job |
…133594) Pass `torch.load(weights_only=)` internally to avoid FutureWarning (#130663) Fixes #130658 Pull Request resolved: #130663 Approved by: https://github.com/malfet, https://github.com/LucasLLC (cherry picked from commit ad314a2) Co-authored-by: Adrian Wälchli <aedu.waelchli@gmail.com>
Fixes #130658
cc @XilunWu @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @LucasLLC @MeetVadakkanchery @mhorowitz