github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-16 06:30:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

[checkpoint] checkpoint for ColoTensor Model (#1196)

Browse Source
This commit is contained in:
Jiarui Fang
2022-07-06 17:22:03 +08:00
committed by GitHub
parent 291e22aac6
commit f38006ea83
4 changed files with 292 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
colossalai/utils/model/colo_init_context.py
View File

@@ -38,15 +38,18 @@ def colo_state_dict(self, destination=None, prefix='', keep_vars=False, state_di
# build param to spec mapping
mapping1 = dict()
mapping2 = dict()
mapping3 = dict()
# gather all params
has_dist_parameter = False
with torch.no_grad():
for param in self.parameters():
if isinstance(param, ColoParameter) and param.has_compute_spec():
if isinstance(param, ColoParameter):
has_dist_parameter = True
mapping1[id(param)] = copy(param.dist_spec)
mapping2[id(param)] = copy(param.compute_spec)
mapping3[id(param)] = param.get_process_group()
param.set_dist_spec(distspec.replicate())
param.process_group = None
# TODO: fix when keep_vars = True
# when keep_vars = False, the state_dict_func will call detach to create
@@ -64,6 +67,7 @@ def colo_state_dict(self, destination=None, prefix='', keep_vars=False, state_di
if param_id in mapping1:
dist_spec = mapping1[id(param)]
compute_spec = mapping2[id(param)]
param.process_group = mapping3[id(param)]
param.set_tensor_spec(dist_spec, compute_spec)
return ret