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]support generalized scheduler (#1222)

Browse Source
This commit is contained in:
Yi Zhao
2022-07-07 18:16:38 +08:00
committed by GitHub
parent a98319f023
commit 04537bf83e
4 changed files with 85 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -2,10 +2,20 @@ import torch
import torch.nn as nn
import torch.distributed as dist
import collections
from torch.optim.lr_scheduler import CosineAnnealingLR as _CosineAnnealingLR
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,
@@ -25,9 +35,7 @@ def save_checkpoint(dire: str,
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.state_dict()}
torch.save(optim_state, dire + '/epoch_{}_optim_rank_{}.pth'.format(epoch, dist.get_rank()))
@@ -55,8 +63,13 @@ def load_checkpoint(dire,
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']
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'])
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)