From bf5066fba7ad845294a4c5cba13e65fc50fd12a3 Mon Sep 17 00:00:00 2001 From: HELSON Date: Tue, 19 Jul 2022 15:46:24 +0800 Subject: [PATCH] [refactor] refactor ColoTensor's unit tests (#1340) --- tests/test_tensor/_utils/__init__.py | 1 - tests/test_tensor/common_utils/__init__.py | 1 + .../test_tensor/{_utils/_util.py => common_utils/_utils.py} | 0 tests/test_tensor/{ => core}/test_dist_spec_mgr.py | 0 tests/test_tensor/{ => core}/test_tensor.py | 0 tests/test_tensor/{ => model}/test_gpt2.py | 2 +- tests/test_tensor/{ => model}/test_model.py | 6 ++---- tests/test_tensor/{ => model}/test_module_spec.py | 2 +- tests/test_tensor/{ => ops}/test_addmm_tp.py | 2 +- tests/test_tensor/{ => ops}/test_embedding_bag_tp.py | 2 +- tests/test_tensor/{ => ops}/test_embedding_tp.py | 2 +- tests/test_tensor/{ => ops}/test_linear_tp.py | 2 +- tests/test_tensor/{ => ops}/test_loss_func.py | 0 tests/test_tensor/{ => ops}/test_op.py | 0 tests/test_tensor/test_parameter.py | 2 +- tests/test_tensor/test_zero_optim.py | 2 +- 16 files changed, 11 insertions(+), 13 deletions(-) delete mode 100644 tests/test_tensor/_utils/__init__.py create mode 100644 tests/test_tensor/common_utils/__init__.py rename tests/test_tensor/{_utils/_util.py => common_utils/_utils.py} (100%) rename tests/test_tensor/{ => core}/test_dist_spec_mgr.py (100%) rename tests/test_tensor/{ => core}/test_tensor.py (100%) rename tests/test_tensor/{ => model}/test_gpt2.py (98%) rename tests/test_tensor/{ => model}/test_model.py (98%) rename tests/test_tensor/{ => model}/test_module_spec.py (99%) rename tests/test_tensor/{ => ops}/test_addmm_tp.py (95%) rename tests/test_tensor/{ => ops}/test_embedding_bag_tp.py (94%) rename tests/test_tensor/{ => ops}/test_embedding_tp.py (92%) rename tests/test_tensor/{ => ops}/test_linear_tp.py (93%) rename tests/test_tensor/{ => ops}/test_loss_func.py (100%) rename tests/test_tensor/{ => ops}/test_op.py (100%) diff --git a/tests/test_tensor/_utils/__init__.py b/tests/test_tensor/_utils/__init__.py deleted file mode 100644 index 8b2ce749a..000000000 --- a/tests/test_tensor/_utils/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from ._util import * \ No newline at end of file diff --git a/tests/test_tensor/common_utils/__init__.py b/tests/test_tensor/common_utils/__init__.py new file mode 100644 index 000000000..5387db704 --- /dev/null +++ b/tests/test_tensor/common_utils/__init__.py @@ -0,0 +1 @@ +from ._utils import * diff --git a/tests/test_tensor/_utils/_util.py b/tests/test_tensor/common_utils/_utils.py similarity index 100% rename from tests/test_tensor/_utils/_util.py rename to tests/test_tensor/common_utils/_utils.py diff --git a/tests/test_tensor/test_dist_spec_mgr.py b/tests/test_tensor/core/test_dist_spec_mgr.py similarity index 100% rename from tests/test_tensor/test_dist_spec_mgr.py rename to tests/test_tensor/core/test_dist_spec_mgr.py diff --git a/tests/test_tensor/test_tensor.py b/tests/test_tensor/core/test_tensor.py similarity index 100% rename from tests/test_tensor/test_tensor.py rename to tests/test_tensor/core/test_tensor.py diff --git a/tests/test_tensor/test_gpt2.py b/tests/test_tensor/model/test_gpt2.py similarity index 98% rename from tests/test_tensor/test_gpt2.py rename to tests/test_tensor/model/test_gpt2.py index 5c1d33cdd..19c6716b0 100644 --- a/tests/test_tensor/test_gpt2.py +++ b/tests/test_tensor/model/test_gpt2.py @@ -1,7 +1,7 @@ import pytest from functools import partial -from _utils import tensor_equal, tensor_shard_equal, set_seed +from tests.test_tensor.common_utils import tensor_equal, tensor_shard_equal, set_seed import torch from torch.nn.parallel import DistributedDataParallel as DDP diff --git a/tests/test_tensor/test_model.py b/tests/test_tensor/model/test_model.py similarity index 98% rename from tests/test_tensor/test_model.py rename to tests/test_tensor/model/test_model.py index 4f03c0f7e..c50393467 100644 --- a/tests/test_tensor/test_model.py +++ b/tests/test_tensor/model/test_model.py @@ -1,7 +1,5 @@ import pytest from functools import partial -from _utils import tensor_shard_equal, set_seed - import torch import torch.multiprocessing as mp @@ -15,7 +13,8 @@ from colossalai.tensor import ColoTensor, ProcessGroup from colossalai.nn.optimizer import ColossalaiOptimizer from tests.components_to_test.registry import non_distributed_component_funcs -from _utils import split_param_row_tp1d, split_param_col_tp1d +from tests.test_tensor.common_utils import tensor_shard_equal, check_equal, set_seed, \ + split_param_row_tp1d, split_param_col_tp1d def run_1d_hybrid_tp(model_name): @@ -264,7 +263,6 @@ def run_1d_row_tp(model_name: str): def _run_pretrain_load(): - from _utils import check_equal from transformers import BertForMaskedLM set_seed(1) model_pretrained = BertForMaskedLM.from_pretrained('bert-base-uncased') diff --git a/tests/test_tensor/test_module_spec.py b/tests/test_tensor/model/test_module_spec.py similarity index 99% rename from tests/test_tensor/test_module_spec.py rename to tests/test_tensor/model/test_module_spec.py index b51d9df42..a3eda1d8a 100644 --- a/tests/test_tensor/test_module_spec.py +++ b/tests/test_tensor/model/test_module_spec.py @@ -7,7 +7,7 @@ import torch.multiprocessing as mp from colossalai.tensor import ColoTensor, ComputePattern, ComputeSpec, ShardSpec, ColoTensorSpec from colossalai.nn.parallel.layers import init_colo_module, check_colo_module -from _utils import tensor_equal, tensor_shard_equal, set_seed +from tests.test_tensor.common_utils import tensor_equal, tensor_shard_equal, set_seed import colossalai from colossalai.utils.cuda import get_current_device diff --git a/tests/test_tensor/test_addmm_tp.py b/tests/test_tensor/ops/test_addmm_tp.py similarity index 95% rename from tests/test_tensor/test_addmm_tp.py rename to tests/test_tensor/ops/test_addmm_tp.py index 83f996517..5182868b5 100644 --- a/tests/test_tensor/test_addmm_tp.py +++ b/tests/test_tensor/ops/test_addmm_tp.py @@ -8,7 +8,7 @@ from colossalai.tensor import ColoTensorSpec from colossalai.testing import rerun_if_address_is_in_use from colossalai.utils import free_port from functools import partial -from _utils import tensor_shard_equal, tensor_equal, split_param_row_tp1d, split_param_col_tp1d +from tests.test_tensor.common_utils import tensor_shard_equal, tensor_equal, split_param_row_tp1d, split_param_col_tp1d class Conv1D(nn.Module): diff --git a/tests/test_tensor/test_embedding_bag_tp.py b/tests/test_tensor/ops/test_embedding_bag_tp.py similarity index 94% rename from tests/test_tensor/test_embedding_bag_tp.py rename to tests/test_tensor/ops/test_embedding_bag_tp.py index f241bd9c8..c7a1604e5 100644 --- a/tests/test_tensor/test_embedding_bag_tp.py +++ b/tests/test_tensor/ops/test_embedding_bag_tp.py @@ -8,7 +8,7 @@ import torch.multiprocessing as mp from colossalai.testing import rerun_if_address_is_in_use from colossalai.utils import free_port from colossalai.tensor import ColoParameter, ColoTensorSpec, ProcessGroup -from _utils import tensor_equal, tensor_shard_equal, split_param_col_tp1d +from tests.test_tensor.common_utils import tensor_equal, tensor_shard_equal, split_param_col_tp1d def run_with_spec(spec_init_func): diff --git a/tests/test_tensor/test_embedding_tp.py b/tests/test_tensor/ops/test_embedding_tp.py similarity index 92% rename from tests/test_tensor/test_embedding_tp.py rename to tests/test_tensor/ops/test_embedding_tp.py index 8e0a27cb5..541dc5c09 100644 --- a/tests/test_tensor/test_embedding_tp.py +++ b/tests/test_tensor/ops/test_embedding_tp.py @@ -8,7 +8,7 @@ import torch.multiprocessing as mp from colossalai.testing import rerun_if_address_is_in_use from colossalai.utils import free_port from colossalai.tensor import ColoTensorSpec, ProcessGroup, ColoTensor -from _utils import tensor_equal, tensor_shard_equal, split_param_col_tp1d, split_param_row_tp1d +from tests.test_tensor.common_utils import tensor_equal, tensor_shard_equal, split_param_col_tp1d, split_param_row_tp1d def run_with_spec(spec_init_func, pg: ProcessGroup): diff --git a/tests/test_tensor/test_linear_tp.py b/tests/test_tensor/ops/test_linear_tp.py similarity index 93% rename from tests/test_tensor/test_linear_tp.py rename to tests/test_tensor/ops/test_linear_tp.py index 734f57160..603e98564 100644 --- a/tests/test_tensor/test_linear_tp.py +++ b/tests/test_tensor/ops/test_linear_tp.py @@ -8,7 +8,7 @@ import torch.nn.functional as F from colossalai.testing import rerun_if_address_is_in_use from colossalai.utils import free_port from colossalai.tensor import ColoTensorSpec, ProcessGroup, ColoTensor -from _utils import tensor_equal, tensor_shard_equal, split_param_col_tp1d, split_param_row_tp1d +from tests.test_tensor.common_utils import tensor_equal, tensor_shard_equal, split_param_col_tp1d, split_param_row_tp1d def run_with_spec(spec_init_func, split_bias): diff --git a/tests/test_tensor/test_loss_func.py b/tests/test_tensor/ops/test_loss_func.py similarity index 100% rename from tests/test_tensor/test_loss_func.py rename to tests/test_tensor/ops/test_loss_func.py diff --git a/tests/test_tensor/test_op.py b/tests/test_tensor/ops/test_op.py similarity index 100% rename from tests/test_tensor/test_op.py rename to tests/test_tensor/ops/test_op.py diff --git a/tests/test_tensor/test_parameter.py b/tests/test_tensor/test_parameter.py index 37b549aa9..7c3c4b213 100644 --- a/tests/test_tensor/test_parameter.py +++ b/tests/test_tensor/test_parameter.py @@ -1,7 +1,7 @@ from colossalai.tensor import ColoParameter, ColoTensor, ColoTensorSpec, ProcessGroup import torch import pytest -from _utils import tensor_equal +from common_utils import tensor_equal import colossalai from colossalai.utils import free_port diff --git a/tests/test_tensor/test_zero_optim.py b/tests/test_tensor/test_zero_optim.py index ca5bf94e1..3b374fa1e 100644 --- a/tests/test_tensor/test_zero_optim.py +++ b/tests/test_tensor/test_zero_optim.py @@ -8,7 +8,7 @@ from colossalai.utils import free_port from colossalai.utils.model.colo_init_context import ColoInitContext from colossalai.gemini import ChunkManager from functools import partial -from _utils import tensor_equal, set_seed, tensor_shard_equal +from tests.test_tensor.common_utils import tensor_equal, set_seed, tensor_shard_equal from tests.components_to_test.registry import non_distributed_component_funcs from torch.nn.parallel import DistributedDataParallel as DDP from colossalai.nn.parallel import ZeroDDP