github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-17 23:18:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

[checkpoint] make unitest faster (#1217)

Browse Source
This commit is contained in:
Jiarui Fang
2022-07-06 17:39:46 +08:00
committed by GitHub
parent f38006ea83
commit 52736205d9
2 changed files with 34 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -5,7 +5,8 @@ import collections
from torch.optim.lr_scheduler import CosineAnnealingLR as _CosineAnnealingLR
from colossalai.utils.model.colo_init_context import colo_state_dict
def save_checkpoint(dire,
def save_checkpoint(dire: str,
epoch: int,
model: torch.nn.Module,
optimizer: torch.optim.Optimizer = None,
@@ -15,30 +16,21 @@ def save_checkpoint(dire,
"""save_checkpoint
save a model, whose parameters are `ColoTensor`s.
Args:
dire (_type_): _description_
epoch (int): _description_
model (torch.nn.Module): _description_
optimizer (torch.optim.Optimizer, optional): _description_. Defaults to None.
lr_scheduler (torch.optim.lr_scheduler._LRScheduler, optional): _description_. Defaults to None.
dire (str): directory to save the checkpoint files.
epoch (int): the number of epoch
model (torch.nn.Module): a torch module initialized by ColoInitContext
optimizer (torch.optim.Optimizer, optional): optimizers. Defaults to None.
lr_scheduler (torch.optim.lr_scheduler._LRScheduler, optional): lr schedule. Defaults to None.
"""
model_state = {
'epoch': epoch,
'model': colo_state_dict(model, state_dict_func=nn.Module.state_dict)
}
model_state = {'epoch': epoch, 'model': colo_state_dict(model, state_dict_func=nn.Module.state_dict)}
if dist.get_rank() == 0:
torch.save(model_state, dire + '/epoch_{}_model.pth'.format(epoch))
lr_scheduler_dict = lr_scheduler.state_dict()
lr_scheduler_dict['after_scheduler'] = lr_scheduler_dict['after_scheduler'].state_dict()
optim_state = {
'epoch': epoch,
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler_dict
}
optim_state = {'epoch': epoch, 'optimizer': optimizer.state_dict(), 'lr_scheduler': lr_scheduler_dict}
torch.save(optim_state, dire + '/epoch_{}_optim_rank_{}.pth'.format(epoch, dist.get_rank()))
def load_checkpoint(dire,
epoch: int,
rank: int,
@@ -64,10 +56,7 @@ def load_checkpoint(dire,
optimizer.load_state_dict(optim_state['optimizer'])
lr_scheduler_dict = optim_state['lr_scheduler']
after_scheduler_dict = lr_scheduler_dict['after_scheduler']
lr_scheduler_dict['after_scheduler'] = _CosineAnnealingLR(
optimizer,
after_scheduler_dict['T_max'],
after_scheduler_dict['eta_min'],
after_scheduler_dict['last_epoch']
)
lr_scheduler_dict['after_scheduler'] = _CosineAnnealingLR(optimizer, after_scheduler_dict['T_max'],
after_scheduler_dict['eta_min'],
after_scheduler_dict['last_epoch'])
lr_scheduler.load_state_dict(lr_scheduler_dict)