From 86be7440076de238bacdd4bf401513d44025c388 Mon Sep 17 00:00:00 2001 From: jiaruifang Date: Wed, 6 Jul 2022 17:34:24 +0800 Subject: [PATCH 1/4] make it faster --- tests/test_utils/test_colo_checkpoint.py | 39 +++++++++++++----------- 1 file changed, 22 insertions(+), 17 deletions(-) diff --git a/tests/test_utils/test_colo_checkpoint.py b/tests/test_utils/test_colo_checkpoint.py index 6e7d4441d760..48742fc18a58 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 @@ def __len__(self): 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() From 9c504afabd4db00b254c96b44014a4ef3bcd0276 Mon Sep 17 00:00:00 2001 From: jiaruifang Date: Tue, 12 Jul 2022 23:01:25 +0800 Subject: [PATCH 2/4] [hotfix] torchvison fx tests --- tests/test_fx/test_pipeline/test_torchvision/test_torchvision.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_fx/test_pipeline/test_torchvision/test_torchvision.py b/tests/test_fx/test_pipeline/test_torchvision/test_torchvision.py index dab485063ab2..e52889e3be6c 100644 --- a/tests/test_fx/test_pipeline/test_torchvision/test_torchvision.py +++ b/tests/test_fx/test_pipeline/test_torchvision/test_torchvision.py @@ -10,6 +10,7 @@ import random import numpy as np import inspect +import pytest MANUAL_SEED = 0 random.seed(MANUAL_SEED) From 0488f3a6cf2968acf67bddf0a201dfa1db53ce34 Mon Sep 17 00:00:00 2001 From: jiaruifang Date: Tue, 12 Jul 2022 23:27:13 +0800 Subject: [PATCH 3/4] [hotfix] rename duplicated named test_gpt.py --- tests/test_tensor/{test_gpt.py => test_gpt2.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/test_tensor/{test_gpt.py => test_gpt2.py} (100%) diff --git a/tests/test_tensor/test_gpt.py b/tests/test_tensor/test_gpt2.py similarity index 100% rename from tests/test_tensor/test_gpt.py rename to tests/test_tensor/test_gpt2.py From 6960701ff8c59ce3279280d8250c28e23b9c1c87 Mon Sep 17 00:00:00 2001 From: jiaruifang Date: Wed, 13 Jul 2022 00:16:31 +0800 Subject: [PATCH 4/4] [hotfix] dist mgr torch version --- colossalai/tensor/dist_spec_mgr.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/colossalai/tensor/dist_spec_mgr.py b/colossalai/tensor/dist_spec_mgr.py index 51b7bfb918ec..f1dc241a8968 100644 --- a/colossalai/tensor/dist_spec_mgr.py +++ b/colossalai/tensor/dist_spec_mgr.py @@ -88,11 +88,13 @@ def _gather(tensor: torch.Tensor, old_dist_spec: _DistSpec, pg: ProcessGroup) -> torch.Tensor: a replicated tensor. """ assert old_dist_spec.placement.value == 's', f"The old_dist_spec of DistSpecManager._gather must be SHARD!" - if version.parse(torch.__version__) < version.parse("1.11.0"): + is_cpu_tensor = False + if tensor.device.type == 'cpu': # pytorch lower than 1.11 dose not support gather a cpu tensor. # Therefore, we transfer tensor to GPU before gather. saved_dev = tensor.device tensor.data = tensor.data.cuda() + is_cpu_tensor = True buffer = [torch.empty_like(tensor) for _ in range(pg.tp_world_size())] assert tensor.device.type == 'cuda' @@ -106,7 +108,7 @@ def _gather(tensor: torch.Tensor, old_dist_spec: _DistSpec, pg: ProcessGroup) -> buffer = new_buffer assert len(buffer) == 1 - if version.parse(torch.__version__) < version.parse("1.11.0"): + if is_cpu_tensor: buffer[0].data = buffer[0].data.to(saved_dev) return buffer[0]