OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

[DCP] Introduce modules metadata in the storage_meta (#146654)

Browse Source 
Summary: Introduce the list of modules in the storage_meta which is shared between the planner and the storage writer. We will use it to let the storage writer know about the modules in the state dict and create module directories in the checkpoint.

Test Plan: UTs

Pull Request resolved: https://github.com/pytorch/pytorch/pull/146654
Approved by: https://github.com/MeetVadakkanchery
This commit is contained in:
Saurabh Mishra 2025-02-13 17:44:30 +00:00 committed by PyTorch MergeBot
parent 938209fb6f
commit 7077d0ac8c
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
torch/distributed/checkpoint/metadata.py
View File

@ -131,6 +131,7 @@ class StorageMeta:
checkpoint_id: Union[str, os.PathLike, None] = None
save_id: Optional[str] = None
load_id: Optional[str] = None
modules: list[str] = field(default_factory=list)
@dataclass