ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: 统一state_dict 处理方式

Browse Source
This commit is contained in:
ViperEkura
2026-03-13 22:41:56 +08:00
parent e35cb0d84a
commit c4feab96fe
1 changed files with 4 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
khaosz/data/checkpoint.py
+4 -6
View File
@@ -1,5 +1,5 @@
import json import json
import torch import safetensors.torch as st
import torch.distributed as dist import torch.distributed as dist
from pathlib import Path from pathlib import Path
@@ -34,9 +34,8 @@ class Checkpoint:
} }
with open(save_path / "meta.json", "w") as f: with open(save_path / "meta.json", "w") as f:
json.dump(meta, f, indent=2) json.dump(meta, f, indent=2)
with open(save_path / f"state_dict.pt", "wb") as f: st.save_file(self.state_dict, save_path / f"state_dict.safetensors")
torch.save(self.state_dict, f)
@classmethod @classmethod
def load( def load(
@@ -57,8 +56,7 @@ class Checkpoint:
dist.broadcast_object_list(meta_list, src=0) dist.broadcast_object_list(meta_list, src=0)
meta = meta_list[0] meta = meta_list[0]
with open(save_path / f"state_dict.pt", "rb") as f: state_dict = st.load_file(save_path / f"state_dict.safetensors")
state_dict = torch.load(f)
return cls( return cls(
state_dict=state_dict, state_dict=state_dict,