[checkpoint] checkpoint for ColoTensor Model (#1196)

colossalai/utils/checkpoint/__init__.py

@@ -0,0 +1,3 @@
+from .module_checkpoint import save_checkpoint, load_checkpoint
+
+__all__ = ['save_checkpoint', 'load_checkpoint']

colossalai/utils/checkpoint/module_checkpoint.py

@@ -0,0 +1,73 @@
+import torch
+import torch.nn as nn
+import torch.distributed as dist
+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,
+                    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 (_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.
+    """
+    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
+    }
+    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']
+    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.load_state_dict(lr_scheduler_dict)

colossalai/utils/model/colo_init_context.py

@@ -38,15 +38,18 @@ def colo_state_dict(self, destination=None, prefix='', keep_vars=False, state_di
     # build param to spec mapping
     mapping1 = dict()
     mapping2 = dict()
+    mapping3 = dict()
     # gather all params
     has_dist_parameter = False
     with torch.no_grad():
         for param in self.parameters():
-            if isinstance(param, ColoParameter) and param.has_compute_spec():
+            if isinstance(param, ColoParameter):
                 has_dist_parameter = True
                 mapping1[id(param)] = copy(param.dist_spec)
                 mapping2[id(param)] = copy(param.compute_spec)
+                mapping3[id(param)] = param.get_process_group()
                 param.set_dist_spec(distspec.replicate())
+                param.process_group = None
 
     # TODO: fix when keep_vars = True
     # when keep_vars = False, the state_dict_func will call detach to create
@@ -64,6 +67,7 @@ def colo_state_dict(self, destination=None, prefix='', keep_vars=False, state_di
             if param_id in mapping1:
                 dist_spec = mapping1[id(param)]
                 compute_spec = mapping2[id(param)]
+                param.process_group = mapping3[id(param)]
                 param.set_tensor_spec(dist_spec, compute_spec)
     return ret