github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[checkpoint] use args, kwargs in save_checkpoint, load_checkpoint (#1368)

pull/1370/head
HELSON 2022-07-26 14:41:53 +08:00 committed by GitHub
parent c491c2a948
commit 8463290642
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -39,7 +39,7 @@ def save_checkpoint(dire: str,
delattr(v, 'save_ready') delattr(v, 'save_ready')
# model saving # model saving
save_state = {'epoch': epoch, 'model': model_state} save_state = {'epoch': epoch, 'model': model_state}
torch.save(save_state, dire + '/epoch_{}_model.pth'.format(epoch)) torch.save(save_state, dire + '/epoch_{}_model.pth'.format(epoch), *args, **kwargs)
# delete old dicts # delete old dicts
del model_state del model_state
@ -57,7 +57,7 @@ def save_checkpoint(dire: str,
if rank == 0: if rank == 0:
save_state = {'epoch': epoch, 'optim': optim_state} save_state = {'epoch': epoch, 'optim': optim_state}
torch.save(save_state, dire + '/epoch_{}_optim.pth'.format(epoch)) torch.save(save_state, dire + '/epoch_{}_optim.pth'.format(epoch), *args, **kwargs)
# recover colo tensors in rank0 # recover colo tensors in rank0
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():
@ -96,7 +96,7 @@ def load_checkpoint(dire,
gather_tensor(p) gather_tensor(p)
if rank == 0: if rank == 0:
load_state = torch.load(dire + '/epoch_{}_model.pth'.format(epoch)) load_state = torch.load(dire + '/epoch_{}_model.pth'.format(epoch), *args, **kwargs)
model.load_state_dict(load_state['model']) model.load_state_dict(load_state['model'])
dist.barrier() dist.barrier()
@ -118,7 +118,7 @@ def load_checkpoint(dire,
gather_tensor(t) gather_tensor(t)
if rank == 0: if rank == 0:
colo_checkpoint = torch.load(dire + '/epoch_{}_optim.pth'.format(epoch)) colo_checkpoint = torch.load(dire + '/epoch_{}_optim.pth'.format(epoch), *args, **kwargs)
optimizer.load_state_dict(colo_checkpoint['optim']) optimizer.load_state_dict(colo_checkpoint['optim'])
dist.barrier() dist.barrier()