mirror of https://github.com/hpcaitech/ColossalAI
76 lines
3.4 KiB
Python
76 lines
3.4 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.distributed as dist
|
|
import collections
|
|
import inspect
|
|
from colossalai.utils.model.colo_init_context import colo_state_dict
|
|
|
|
|
|
def filter_dict(dict_to_filter, thing_with_kwargs):
|
|
sig = inspect.signature(thing_with_kwargs)
|
|
filter_keys = [param.name for param in sig.parameters.values() if param.kind == param.POSITIONAL_OR_KEYWORD]
|
|
filter_dict = {}
|
|
for filter_key in filter_keys:
|
|
if filter_key in dict_to_filter:
|
|
filter_dict[filter_key] = dict_to_filter[filter_key]
|
|
return filter_dict
|
|
|
|
|
|
def save_checkpoint(dire: str,
|
|
epoch: int,
|
|
model: torch.nn.Module,
|
|
optimizer: torch.optim.Optimizer = None,
|
|
lr_scheduler: torch.optim.lr_scheduler._LRScheduler = None,
|
|
*args,
|
|
**kwargs):
|
|
"""save_checkpoint
|
|
save a model, whose parameters are `ColoTensor`s.
|
|
Args:
|
|
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)}
|
|
if dist.get_rank() == 0:
|
|
torch.save(model_state, dire + '/epoch_{}_model.pth'.format(epoch))
|
|
optim_state = {'epoch': epoch, 'optimizer': optimizer.state_dict(), 'lr_scheduler': lr_scheduler.state_dict()}
|
|
torch.save(optim_state, dire + '/epoch_{}_optim_rank_{}.pth'.format(epoch, dist.get_rank()))
|
|
|
|
|
|
def load_checkpoint(dire,
|
|
epoch: int,
|
|
rank: int,
|
|
model: torch.nn.Module,
|
|
optimizer: torch.optim.Optimizer = None,
|
|
lr_scheduler: torch.optim.lr_scheduler._LRScheduler = None,
|
|
*args,
|
|
**kwargs):
|
|
"""load_checkpoint
|
|
load a model, whose parameters are `ColoTensor`s.
|
|
Args:
|
|
dire (_type_): _description_
|
|
epoch (int): _description_
|
|
rank (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.
|
|
"""
|
|
model_state = torch.load(dire + '/epoch_{}_model.pth'.format(epoch))
|
|
model_state['model'] = collections.OrderedDict([(k.split('.', 1)[1], v) for k, v in model_state['model'].items()])
|
|
model.load_state_dict(model_state['model'])
|
|
optim_state = torch.load(dire + '/epoch_{}_optim_rank_{}.pth'.format(epoch, rank))
|
|
optimizer.load_state_dict(optim_state['optimizer'])
|
|
lr_scheduler_dict = optim_state['lr_scheduler']
|
|
if 'after_scheduler_type' in lr_scheduler_dict:
|
|
after_scheduler_type = lr_scheduler_dict.pop('after_scheduler_type')
|
|
after_scheduler_dict = lr_scheduler_dict.pop('after_scheduler_dict')
|
|
reload_scheduler = getattr(torch.optim.lr_scheduler, after_scheduler_type)
|
|
filtered_dict = filter_dict(after_scheduler_dict, reload_scheduler)
|
|
lr_scheduler_dict['after_scheduler'] = reload_scheduler(
|
|
optimizer,
|
|
**filtered_dict,
|
|
)
|
|
lr_scheduler.load_state_dict(lr_scheduler_dict)
|