github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-12 12:47:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

[checkpoint] support huggingface style sharded checkpoint (#3461)

Browse Source 
* [checkpoint] support huggingface style sharded checkpoint

* [checkpoint] support huggingface style sharded checkpoint

* [checkpoint] support huggingface style sharded checkpoint

* [checkpoint] support huggingface style sharded checkpoint

* [checkpoint] support huggingface style sharded checkpoint

---------

Co-authored-by: luchen <luchen@luchendeMBP.lan>
This commit is contained in:
jiangmingyan
2023-04-06 16:23:39 +08:00
committed by GitHub
parent 6afeb1202a
commit 52a933e175
4 changed files with 291 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
colossalai/checkpoint_io/index_file.py
View File

@@ -148,3 +148,9 @@ class CheckpointIndexFile:
"""
ckpt_path = self.weight_map[param_name]
return ckpt_path
def get_all_param_names(self):
"""
Get all the weight keys.
"""
return list(self.weight_map.keys())