Skip to content

Conversation

thiagocrepaldi
Copy link
Collaborator

@thiagocrepaldi thiagocrepaldi commented Jan 11, 2024

Stack from ghstack (oldest at bottom):

During ONNXProgram.save, the implicit/explicit state_dict passed in must
be loaded in memory in order to read each initializer and create an
external tensor proto with them

This PR ensures torch.load uses memory-map to support large models that
cannot fit in memory

During ONNXProgram.save, the implicit/explicit state_dict passed in must
be loaded in memory in order to read each initializer and create an
external tensor proto with them

This PR ensures torch.load uses memory-map to support large models that
cannot fit in memory

[ghstack-poisoned]
Copy link

pytorch-bot bot commented Jan 11, 2024

🔗 Helpful Links

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

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

✅ No Failures

As of commit d5c3f17 with merge base b4a3563 (image):
💚 Looks good so far! There are no failures yet. 💚

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

@pytorch-bot pytorch-bot bot added the release notes: onnx torch.onnx related changes that should show up in the release notes label Jan 11, 2024
@thiagocrepaldi thiagocrepaldi added module: onnx Related to torch.onnx onnx-triaged triaged by ONNX team labels Jan 11, 2024
… large models"

During ONNXProgram.save, the implicit/explicit state_dict passed in must
be loaded in memory in order to read each initializer and create an
external tensor proto with them

This PR ensures torch.load uses memory-map to support large models that
cannot fit in memory

[ghstack-poisoned]
… large models"

During ONNXProgram.save, the implicit/explicit state_dict passed in must
be loaded in memory in order to read each initializer and create an
external tensor proto with them

This PR ensures torch.load uses memory-map to support large models that
cannot fit in memory

[ghstack-poisoned]
thiagocrepaldi pushed a commit that referenced this pull request Jan 11, 2024
During ONNXProgram.save, the implicit/explicit state_dict passed in must
be loaded in memory in order to read each initializer and create an
external tensor proto with them

This PR ensures torch.load uses memory-map to support large models that
cannot fit in memory

ghstack-source-id: 5b1b31e
Pull Request resolved: #117295
Copy link
Collaborator

@BowenBao BowenBao left a comment

Choose a reason for hiding this comment

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

Very nice. Do you have a case study show casing the efficiency/effectiveness before and after?

extra_state_dict = torch.load(path)
# Loads checkpoint using memory-map on CPU to succeed with large models
extra_state_dict = torch.load(
path, map_location="cpu", mmap=True, weights_only=True
Copy link
Collaborator

Choose a reason for hiding this comment

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

noticed weights_only=True is added. is this intended?

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

It is a security feature. when weights_only=False, malicious pickled checkpoints can execute code on the machine.
If the checkpoint can be loaded with weights_only=True, we should - but I am experimenting with it

@thiagocrepaldi
Copy link
Collaborator Author

Very nice. Do you have a case study show casing the efficiency/effectiveness before and after?

Not yet.

@thiagocrepaldi
Copy link
Collaborator Author

@pytorchbot merge

@pytorch-bot pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Jan 12, 2024
@pytorchmergebot
Copy link
Collaborator

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

@facebook-github-bot facebook-github-bot deleted the gh/thiagocrepaldi/19/head branch January 15, 2024 15:23
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/trunk Trigger trunk jobs on your pull request Merged module: onnx Related to torch.onnx onnx-triaged triaged by ONNX team open source release notes: onnx torch.onnx related changes that should show up in the release notes

Projects

Status: Done

Development

Successfully merging this pull request may close these issues.

4 participants