Use mmap option to load_state_dict #28331
Conversation
|
|
||
| return torch.load(checkpoint_file, map_location=map_location, weights_only=True) | ||
| extra_args = {} | ||
| # mmap can only be used with files serialized with zipfile-based format. |
There was a problem hiding this comment.
how often does this come up? (zipfile-based format)
There was a problem hiding this comment.
zip file is used often as it's the default for torch.save (https://pytorch.org/docs/stable/generated/torch.save.html).
Sure. Please let me know if my test case is correct. |
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
Do you mean add a test case for device_map == "meta' ? |
|
That would be a good way of making sure this will be fine with DeepSpeedZero (the if branch) 🤗 |
Currently none of existing tests covers the branch: Anyway, I added a guard for "meta" device. |
Use mmap option to load_state_dict (huggingface#28331)
Use mmap option to load_state_dict (huggingface#28331)
Use mmap option to load_state_dict (huggingface#28331)
Use mmap option to load_state_dict (huggingface#28331)
Use torch.load(mmap=True) to accelerate checkpoint loading
#28332
cc @SunMarc @sgugger