10000 Pass `torch.load(weights_only=)` internally to avoid FutureWarning by awaelchli · Pull Request #130663 · pytorch/pytorch · GitHub
[go: up one dir, main page]
Skip to content

Pass torch.load(weights_only=) internally to avoid FutureWarning #130663

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 4 commits into from
Closed

Pass torch.load(weights_only=) internally to avoid FutureWarning #130663

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
345cd72
Avoid weights_only warning
awaelchli Jul 13, 2024
da424bb
lintrunner
awaelchli Jul 15, 2024
b43aca5
Merge branch 'main' into bugfix-torch-load-warning
awaelchli Jul 15, 2024
e0275f8
switch to weights_only=True in filesystem.py
awaelchli Jul 15, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions torch/distributed/checkpoint/default_planner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -210,10 +210,12 @@ def load_bytes(self, read_item: ReadItem, value: io.BytesIO) -> None:
set_element(
self.original_state_dict,
self.mappings[read_item.dest_index.fqn],
torch.load(value),
torch.load(value, weights_only=False),
)
else:
self.state_dict[read_item.dest_index.fqn] = torch.load(value)
self.state_dict[read_item.dest_index.fqn] = torch.load(
value, weights_only=False
)

def resolve_tensor(self, read_item: ReadItem):
tensor = self.lookup_tensor(read_item.dest_index)
Expand Down
6 changes: 5 additions & 1 deletion torch/distributed/checkpoint/filesystem.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -654,7 +654,11 @@ def read_data(self, plan: LoadPlan, planner: LoadPlanner) -> Future[None]:
else:
tensor = cast(
Tensor,
torch.load(cast(IO[bytes], file_slice), map_location="cpu"),
torch.load(
cast(IO[bytes], file_slice),
map_location="cpu",
weights_only=True,
),
)
tensor = narrow_tensor_by_index(
tensor, req.storage_offsets, req.lengths
Expand Down
6 changes: 4 additions & 2 deletions torch/distributed/checkpoint/format_utils.py
8000
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,9 @@ def read_data(self, plan: LoadPlan, planner: LoadPlanner) -> Future[None]:
# TODO: read on each host, instead of only the coordinator
if self.is_coordinator:
assert self.checkpoint_id is not None
torch_state_dict = torch.load(self.checkpoint_id, map_location="cpu")
torch_state_dict = torch.load(
self.checkpoint_id, map_location="cpu", weights_only=False
)
if planner.flatten_state_dict:
torch_state_dict, _ = flatten_state_dict(torch_state_dict)
else:
Expand Down Expand Up @@ -231,7 +233,7 @@ def torch_save_to_dcp(
To avoid OOM, it's recommended to only run this function on a single rank.
"""

state_dict = torch.load(torch_save_path)
state_dict = torch.load(torch_save_path, weights_only=False)
# we don't need stateful behavior here because the expectation is anything loaded by
# torch.load would not contain stateful objects.
_save_state_dict(
Expand Down
2 changes: 1 addition & 1 deletion torch/distributed/checkpoint/planner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -330,7 +330,7 @@ class LoadPlanner:
>>>
>>> def load_bytes(self, read_item, value):
>>> # Remove the "foo_" prefix
>>> self.original_state_dict[read_item.dest_index.fqn[4:]] = torch.load(value)
>>> self.original_state_dict[read_item.dest_index.fqn[4:]] = torch.load(value, weights_only=False)


Modifying resolve_tensor and commit_tensor to handle load time transformation.
Expand Down
2 changes: 1 addition & 1 deletion torch/distributed/optim/zero_redundancy_optimizer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

return obj


Expand Down
Loading
0