From 52736205d9b7141e1d812e7d1ba618d1528cbe95 Mon Sep 17 00:00:00 2001 From: Jiarui Fang Date: Wed, 6 Jul 2022 17:39:46 +0800 Subject: [PATCH] [checkpoint] make unitest faster (#1217) --- .../utils/checkpoint/module_checkpoint.py | 35 ++++++----------- tests/test_utils/test_colo_checkpoint.py | 39 +++++++++++-------- 2 files changed, 34 insertions(+), 40 deletions(-) diff --git a/colossalai/utils/checkpoint/module_checkpoint.py b/colossalai/utils/checkpoint/module_checkpoint.py index c4de1c5ea..0cdb17d6c 100644 --- a/colossalai/utils/checkpoint/module_checkpoint.py +++ b/colossalai/utils/checkpoint/module_checkpoint.py @@ -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) diff --git a/tests/test_utils/test_colo_checkpoint.py b/tests/test_utils/test_colo_checkpoint.py index 6e7d4441d..48742fc18 100644 --- a/tests/test_utils/test_colo_checkpoint.py +++ b/tests/test_utils/test_colo_checkpoint.py @@ -1,21 +1,20 @@ from abc import ABC, abstractmethod -import os, sys, shutil +import os, shutil import torch import torch.nn as nn import pytest import copy -import operator -import colossalai -from colossalai.context.parallel_mode import ParallelMode +from functools import partial + import torch.multiprocessing as mp import torch.distributed as dist + +import colossalai from colossalai.testing import rerun_if_address_is_in_use from colossalai.utils.cuda import get_current_device from colossalai.utils import free_port from colossalai.utils.model.colo_init_context import ColoInitContext -from colossalai.tensor import ColoTensorSpec, ComputePattern, ComputeSpec, DistSpecManager, distspec, ProcessGroup, ColoTensor -from colossalai.core import global_context as gpc -from functools import partial +from colossalai.tensor import ComputePattern, ComputeSpec, DistSpecManager, distspec, ProcessGroup from colossalai.nn.parallel.data_parallel import ColoDDP from colossalai.utils.checkpoint import save_checkpoint, load_checkpoint from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR @@ -46,15 +45,17 @@ class DummyDataGenerator(ABC): class DummyDataLoader(DummyDataGenerator): - batch_size = 128 - category = 16 - feature_size = 256 + + def __init__(self, batch_size, category, feature_size, length=10): + super().__init__(length) + self.batch_size = batch_size + self.category = category + self.feature_size = feature_size def generate(self): image_dict = {} - image_dict['pixel_values'] = torch.rand( - DummyDataLoader.batch_size, DummyDataLoader.feature_size, device=get_current_device()) * 2 - 1 - image_dict['label'] = torch.randint(DummyDataLoader.category, (DummyDataLoader.batch_size,), + image_dict['pixel_values'] = torch.rand(self.batch_size, self.feature_size, device=get_current_device()) * 2 - 1 + image_dict['label'] = torch.randint(self.category, (self.batch_size,), dtype=torch.int64, device=get_current_device()) return image_dict @@ -102,11 +103,15 @@ def remove(path): def run_checkpoint(init_spec_func, use_ddp, test_epoch, pg): - train_dataloader = DummyDataLoader(length=16) + batch = 3 + feature = 32 + category = 16 + train_dataloader = DummyDataLoader(batch, category, feature, length=16) with ColoInitContext(device=get_current_device()): - model = MLP(256, 16, 64) - model_reload = MLP(256, 16, 64) - model_ref = MLP(256, 16, 64) + model = MLP(feature, category) + model_reload = MLP(feature, category) + model_ref = MLP(feature, category) + model = model.cuda() model_reload = model_reload.cuda() model_ref = model_ref.cuda()