frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[DCP] Avoid multiple storage writer resets in async save (#159448)

Browse Source 
Summary: Avoid multiple storage writer resets in async save. Currently the reset gets called by the async_save method and then again in the save method. In the async path, async_save should only do the staging and the reset should only happen in the synchronous save path.

Test Plan:
```
buck test 'fbcode//mode/opt' //aiplatform/modelstore/experimental/DCP/tests:checkpoint_dist_client_test
```
https://www.internalfb.com/intern/testinfra/testrun/15199648841705052

Rollback Plan:

Differential Revision: D79230339

Pull Request resolved: https://github.com/pytorch/pytorch/pull/159448
Approved by: https://github.com/meetv18
This commit is contained in:
Saurabh Mishra
2025-09-10 00:43:03 +00:00
committed by PyTorch MergeBot
parent 5539916fe1
commit 33589374b6
1 changed files with 0 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
torch/distributed/checkpoint/state_dict_saver.py
View File

@ -312,10 +312,6 @@ def async_save(
)
)
storage_writer = cast(
StorageWriter, _storage_setup(storage_writer, checkpoint_id, reader=False)
)
state_dict = _stateful_to_state_dict(state_dict)
@_dcp_method_logger(log_exceptions=True)