github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-11 04:41:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

[hotfix] fix no optimizer in save/load (#1363)

Browse Source
This commit is contained in:
HELSON 2022-07-26 10:53:53 +08:00 committed by GitHub
parent cd063ac37f
commit 943a96323e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 38 additions and 36 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/tensor/__init__.py
View File

@ -13,6 +13,6 @@ from . import distspec
__all__ = [ __all__ = [
'ColoTensor', 'convert_parameter', 'ComputePattern', 'ComputeSpec', 'named_params_with_colotensor', 'ColoParameter', 'ColoTensor', 'convert_parameter', 'ComputePattern', 'ComputeSpec', 'named_params_with_colotensor', 'ColoParameter',
'distspec', 'DistSpecManager', 'ParamOpHook', 'ParamOpHookManager', 'ChunkManager', 'TensorState', 'ProcessGroup', 'distspec', 'DistSpecManager', 'ParamOpHook', 'ParamOpHookManager', 'ProcessGroup', 'ColoTensorSpec',
'ColoTensorSpec', 'TensorSpec', 'ShardSpec', 'ReplicaSpec' 'ShardSpec', 'ReplicaSpec'
] ]

2
colossalai/utils/checkpoint/module_checkpoint.py
View File

@ -46,6 +46,7 @@ def save_checkpoint(dire: str,
# synchronize all the processes # synchronize all the processes
dist.barrier() dist.barrier()
if optimizer is not None:
mapping = dict() mapping = dict()
optim_state = optimizer.state_dict() optim_state = optimizer.state_dict()
for k, v in optim_state['state'].items(): for k, v in optim_state['state'].items():
@ -108,6 +109,7 @@ def load_checkpoint(dire,
delattr(p, 'save_ready') delattr(p, 'save_ready')
del mapping del mapping
if optimizer is not None:
mapping = dict() mapping = dict()
for k, v in optimizer.state_dict()['state'].items(): for k, v in optimizer.state_dict()['state'].items():
for n, t in v.items(): for n, t in v.items():