Safetensors offload #20321
Merged
Safetensors offload #20321
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
The documentation is not available anymore as the PR was closed or merged. |
sgugger
commented
Nov 21, 2022
| @@ -28,7 +28,7 @@ | |||
|
|
|||
| import torch | |||
| from packaging import version | |||
| from torch import Tensor, device, nn | |||
There was a problem hiding this comment.
This import was dangerous with so many local variables named device. It was only used in a type annotation.
sgugger
commented
Nov 21, 2022
Comment on lines
+2539
to
+2543
| offload_index = { | ||
| p: {"safetensors_file": os.path.join(folder, f), "weight_name": p, "dtype": str_dtype} | ||
| for p, f in sharded_metadata["weight_map"].items() | ||
| if param_device_map[p] == "disk" | ||
| } |
There was a problem hiding this comment.
We build the offload_index here when the checkpoin is safetensors, this way we can skip loading all the shards and gain time.
LysandreJik
approved these changes
Nov 25, 2022
mpierrau
pushed a commit
to mpierrau/transformers
that referenced
this pull request
Dec 15, 2022
* INtegrate safetensos in weight offloading * Use safetensors checkpoint for offload when available * Make naming consistent * Make load faster * Quality * Add default
This was referenced Mar 7, 2023
4 tasks
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
What does this PR do?
This PRs make offload to disk more efficient for models that have a checkpoint in safetensors format: instead of re-saving everything as Numpy memory-mapped array, it uses directly the fact we can access the tensor in the checkpoint without loading the rest.
Goes with huggingface/accelerate#873