From 933048ad3e68f243d0cfeda1e7a8cf1b55a4cf88 Mon Sep 17 00:00:00 2001 From: ver217 Date: Thu, 6 Apr 2023 09:38:25 +0800 Subject: [PATCH] [test] reorganize zero/gemini tests (#3445) --- tests/test_moe/test_moe_checkpoint.py | 2 +- tests/test_moe/test_moe_colo_init.py | 2 +- tests/test_moe/test_moe_zero_init.py | 2 +- tests/test_moe/test_moe_zero_model.py | 2 +- tests/test_moe/test_moe_zero_optim.py | 2 +- .../update => test_zero/test_gemini}/test_chunk_mgrv2.py | 0 .../update => test_zero/test_gemini}/test_chunkv2.py | 0 .../update => test_zero/test_gemini}/test_fwd_bwd.py | 0 .../update => test_zero/test_gemini}/test_gemini_use_rmt.py | 0 .../update => test_zero/test_gemini}/test_get_torch_model.py | 0 .../update => test_zero/test_gemini}/test_grad_clip.py | 0 .../update => test_zero/test_gemini}/test_inference.py | 0 .../{test_gemini/update => test_zero/test_gemini}/test_optim.py | 0 tests/{ => test_zero}/test_gemini/test_runtime_mem_tracer.py | 0 .../update => test_zero/test_gemini}/test_search.py | 0 .../update => test_zero/test_gemini}/test_zeroddp_state_dict.py | 0 .../test_gemini}/test_zerooptim_state_dict.py | 0 tests/test_zero/{ => test_legacy}/common.py | 0 tests/test_zero/{ => test_legacy}/test_found_inf.py | 2 +- .../test_legacy}/test_gemini_manager.py | 0 tests/test_zero/{ => test_legacy}/test_init_context.py | 0 tests/{test_gemini => test_zero/test_legacy}/test_param_op.py | 0 tests/test_zero/{ => test_legacy}/test_shard_model_v2.py | 0 tests/test_zero/{ => test_legacy}/test_shard_param.py | 2 +- .../{ => test_legacy}/test_sharded_optim_state_dict.py | 0 tests/test_zero/{ => test_legacy}/test_sharded_optim_v2.py | 0 .../{ => test_legacy}/test_sharded_optim_with_sync_bn.py | 0 tests/test_zero/{ => test_legacy}/test_state_dict.py | 1 - tests/test_zero/{ => test_legacy}/test_tensor_utils.py | 0 tests/test_zero/{ => test_legacy}/test_zero_engine.py | 0 .../{low_level_zero => test_low_level}/test_grad_acc.py | 0 .../{low_level_zero => test_low_level}/test_zero1_2.py | 0 .../{low_level_zero => test_low_level}/test_zero_init.py | 0 .../{low_level_zero => test_low_level}/test_zero_tp.py | 0 34 files changed, 7 insertions(+), 8 deletions(-) rename tests/{test_gemini/update => test_zero/test_gemini}/test_chunk_mgrv2.py (100%) rename tests/{test_gemini/update => test_zero/test_gemini}/test_chunkv2.py (100%) rename tests/{test_gemini/update => test_zero/test_gemini}/test_fwd_bwd.py (100%) rename tests/{test_gemini/update => test_zero/test_gemini}/test_gemini_use_rmt.py (100%) rename tests/{test_gemini/update => test_zero/test_gemini}/test_get_torch_model.py (100%) rename tests/{test_gemini/update => test_zero/test_gemini}/test_grad_clip.py (100%) rename tests/{test_gemini/update => test_zero/test_gemini}/test_inference.py (100%) rename tests/{test_gemini/update => test_zero/test_gemini}/test_optim.py (100%) rename tests/{ => test_zero}/test_gemini/test_runtime_mem_tracer.py (100%) rename tests/{test_gemini/update => test_zero/test_gemini}/test_search.py (100%) rename tests/{test_gemini/update => test_zero/test_gemini}/test_zeroddp_state_dict.py (100%) rename tests/{test_gemini/update => test_zero/test_gemini}/test_zerooptim_state_dict.py (100%) rename tests/test_zero/{ => test_legacy}/common.py (100%) rename tests/test_zero/{ => test_legacy}/test_found_inf.py (97%) rename tests/{test_gemini => test_zero/test_legacy}/test_gemini_manager.py (100%) rename tests/test_zero/{ => test_legacy}/test_init_context.py (100%) rename tests/{test_gemini => test_zero/test_legacy}/test_param_op.py (100%) rename tests/test_zero/{ => test_legacy}/test_shard_model_v2.py (100%) rename tests/test_zero/{ => test_legacy}/test_shard_param.py (98%) rename tests/test_zero/{ => test_legacy}/test_sharded_optim_state_dict.py (100%) rename tests/test_zero/{ => test_legacy}/test_sharded_optim_v2.py (100%) rename tests/test_zero/{ => test_legacy}/test_sharded_optim_with_sync_bn.py (100%) rename tests/test_zero/{ => test_legacy}/test_state_dict.py (98%) rename tests/test_zero/{ => test_legacy}/test_tensor_utils.py (100%) rename tests/test_zero/{ => test_legacy}/test_zero_engine.py (100%) rename tests/test_zero/{low_level_zero => test_low_level}/test_grad_acc.py (100%) rename tests/test_zero/{low_level_zero => test_low_level}/test_zero1_2.py (100%) rename tests/test_zero/{low_level_zero => test_low_level}/test_zero_init.py (100%) rename tests/test_zero/{low_level_zero => test_low_level}/test_zero_tp.py (100%) diff --git a/tests/test_moe/test_moe_checkpoint.py b/tests/test_moe/test_moe_checkpoint.py index 5b6fe4411..d2cff44ad 100644 --- a/tests/test_moe/test_moe_checkpoint.py +++ b/tests/test_moe/test_moe_checkpoint.py @@ -14,7 +14,7 @@ from colossalai.utils import free_port, get_current_device from colossalai.zero import ColoInitContext from tests.test_moe.test_moe_zero_init import MoeModel from tests.test_tensor.common_utils import debug_print -from tests.test_zero.common import CONFIG +from tests.test_zero.test_legacy.common import CONFIG def exam_moe_checkpoint(): diff --git a/tests/test_moe/test_moe_colo_init.py b/tests/test_moe/test_moe_colo_init.py index 23ad1a3dc..4826d87ac 100644 --- a/tests/test_moe/test_moe_colo_init.py +++ b/tests/test_moe/test_moe_colo_init.py @@ -13,7 +13,7 @@ from colossalai.utils import free_port, get_current_device from colossalai.zero import ColoInitContext from tests.test_moe.test_moe_zero_init import MoeModel from tests.test_tensor.common_utils import debug_print -from tests.test_zero.common import CONFIG +from tests.test_zero.test_legacy.common import CONFIG @parameterize("init_device_type", ['cpu', 'cuda']) diff --git a/tests/test_moe/test_moe_zero_init.py b/tests/test_moe/test_moe_zero_init.py index 5987e31f7..18b50eb5c 100644 --- a/tests/test_moe/test_moe_zero_init.py +++ b/tests/test_moe/test_moe_zero_init.py @@ -14,7 +14,7 @@ from colossalai.testing import parameterize, rerun_if_address_is_in_use from colossalai.utils import free_port, get_current_device from colossalai.zero.legacy.init_ctx import ZeroInitContext from colossalai.zero.legacy.shard_utils import BucketTensorShardStrategy, TensorShardStrategy -from tests.test_zero.common import CONFIG +from tests.test_zero.test_legacy.common import CONFIG class MoeModel(nn.Module): diff --git a/tests/test_moe/test_moe_zero_model.py b/tests/test_moe/test_moe_zero_model.py index d38f66fef..49c452938 100644 --- a/tests/test_moe/test_moe_zero_model.py +++ b/tests/test_moe/test_moe_zero_model.py @@ -17,7 +17,7 @@ from colossalai.zero.legacy.sharded_model._utils import cast_tensor_to_fp16 from colossalai.zero.legacy.sharded_model.utils import col_model_deepcopy from tests.components_to_test.registry import non_distributed_component_funcs from tests.test_moe.test_moe_zero_init import MoeModel -from tests.test_zero.common import CONFIG, check_grads_padding, run_fwd_bwd +from tests.test_zero.test_legacy.common import CONFIG, check_grads_padding, run_fwd_bwd @parameterize("enable_autocast", [False]) diff --git a/tests/test_moe/test_moe_zero_optim.py b/tests/test_moe/test_moe_zero_optim.py index 7e140bf86..b43e52bb4 100644 --- a/tests/test_moe/test_moe_zero_optim.py +++ b/tests/test_moe/test_moe_zero_optim.py @@ -20,7 +20,7 @@ from colossalai.zero.legacy.sharded_optim import ShardedOptimizerV2 from colossalai.zero.low_level._utils import has_inf_or_nan from tests.components_to_test.registry import non_distributed_component_funcs from tests.test_moe.test_moe_zero_init import MoeModel -from tests.test_zero.common import CONFIG, check_sharded_model_params +from tests.test_zero.test_legacy.common import CONFIG, check_sharded_model_params def _run_step(model, optimizer, data, label, criterion, grad_handler): diff --git a/tests/test_gemini/update/test_chunk_mgrv2.py b/tests/test_zero/test_gemini/test_chunk_mgrv2.py similarity index 100% rename from tests/test_gemini/update/test_chunk_mgrv2.py rename to tests/test_zero/test_gemini/test_chunk_mgrv2.py diff --git a/tests/test_gemini/update/test_chunkv2.py b/tests/test_zero/test_gemini/test_chunkv2.py similarity index 100% rename from tests/test_gemini/update/test_chunkv2.py rename to tests/test_zero/test_gemini/test_chunkv2.py diff --git a/tests/test_gemini/update/test_fwd_bwd.py b/tests/test_zero/test_gemini/test_fwd_bwd.py similarity index 100% rename from tests/test_gemini/update/test_fwd_bwd.py rename to tests/test_zero/test_gemini/test_fwd_bwd.py diff --git a/tests/test_gemini/update/test_gemini_use_rmt.py b/tests/test_zero/test_gemini/test_gemini_use_rmt.py similarity index 100% rename from tests/test_gemini/update/test_gemini_use_rmt.py rename to tests/test_zero/test_gemini/test_gemini_use_rmt.py diff --git a/tests/test_gemini/update/test_get_torch_model.py b/tests/test_zero/test_gemini/test_get_torch_model.py similarity index 100% rename from tests/test_gemini/update/test_get_torch_model.py rename to tests/test_zero/test_gemini/test_get_torch_model.py diff --git a/tests/test_gemini/update/test_grad_clip.py b/tests/test_zero/test_gemini/test_grad_clip.py similarity index 100% rename from tests/test_gemini/update/test_grad_clip.py rename to tests/test_zero/test_gemini/test_grad_clip.py diff --git a/tests/test_gemini/update/test_inference.py b/tests/test_zero/test_gemini/test_inference.py similarity index 100% rename from tests/test_gemini/update/test_inference.py rename to tests/test_zero/test_gemini/test_inference.py diff --git a/tests/test_gemini/update/test_optim.py b/tests/test_zero/test_gemini/test_optim.py similarity index 100% rename from tests/test_gemini/update/test_optim.py rename to tests/test_zero/test_gemini/test_optim.py diff --git a/tests/test_gemini/test_runtime_mem_tracer.py b/tests/test_zero/test_gemini/test_runtime_mem_tracer.py similarity index 100% rename from tests/test_gemini/test_runtime_mem_tracer.py rename to tests/test_zero/test_gemini/test_runtime_mem_tracer.py diff --git a/tests/test_gemini/update/test_search.py b/tests/test_zero/test_gemini/test_search.py similarity index 100% rename from tests/test_gemini/update/test_search.py rename to tests/test_zero/test_gemini/test_search.py diff --git a/tests/test_gemini/update/test_zeroddp_state_dict.py b/tests/test_zero/test_gemini/test_zeroddp_state_dict.py similarity index 100% rename from tests/test_gemini/update/test_zeroddp_state_dict.py rename to tests/test_zero/test_gemini/test_zeroddp_state_dict.py diff --git a/tests/test_gemini/update/test_zerooptim_state_dict.py b/tests/test_zero/test_gemini/test_zerooptim_state_dict.py similarity index 100% rename from tests/test_gemini/update/test_zerooptim_state_dict.py rename to tests/test_zero/test_gemini/test_zerooptim_state_dict.py diff --git a/tests/test_zero/common.py b/tests/test_zero/test_legacy/common.py similarity index 100% rename from tests/test_zero/common.py rename to tests/test_zero/test_legacy/common.py diff --git a/tests/test_zero/test_found_inf.py b/tests/test_zero/test_legacy/test_found_inf.py similarity index 97% rename from tests/test_zero/test_found_inf.py rename to tests/test_zero/test_legacy/test_found_inf.py index 641136718..03a1a609b 100644 --- a/tests/test_zero/test_found_inf.py +++ b/tests/test_zero/test_legacy/test_found_inf.py @@ -4,6 +4,7 @@ import pytest import torch import torch.multiprocessing as mp from common import CONFIG +from test_sharded_optim_v2 import _run_step import colossalai from colossalai.nn.optimizer import HybridAdam @@ -16,7 +17,6 @@ from colossalai.zero.legacy.sharded_model import ShardedModelV2 from colossalai.zero.legacy.sharded_optim import ShardedOptimizerV2 from colossalai.zero.low_level._utils import has_inf_or_nan from tests.components_to_test.registry import non_distributed_component_funcs -from tests.test_zero.test_sharded_optim_v2 import _run_step @parameterize("cpu_offload", [True, False]) diff --git a/tests/test_gemini/test_gemini_manager.py b/tests/test_zero/test_legacy/test_gemini_manager.py similarity index 100% rename from tests/test_gemini/test_gemini_manager.py rename to tests/test_zero/test_legacy/test_gemini_manager.py diff --git a/tests/test_zero/test_init_context.py b/tests/test_zero/test_legacy/test_init_context.py similarity index 100% rename from tests/test_zero/test_init_context.py rename to tests/test_zero/test_legacy/test_init_context.py diff --git a/tests/test_gemini/test_param_op.py b/tests/test_zero/test_legacy/test_param_op.py similarity index 100% rename from tests/test_gemini/test_param_op.py rename to tests/test_zero/test_legacy/test_param_op.py diff --git a/tests/test_zero/test_shard_model_v2.py b/tests/test_zero/test_legacy/test_shard_model_v2.py similarity index 100% rename from tests/test_zero/test_shard_model_v2.py rename to tests/test_zero/test_legacy/test_shard_model_v2.py diff --git a/tests/test_zero/test_shard_param.py b/tests/test_zero/test_legacy/test_shard_param.py similarity index 98% rename from tests/test_zero/test_shard_param.py rename to tests/test_zero/test_legacy/test_shard_param.py index 6085de3c8..b76648321 100644 --- a/tests/test_zero/test_shard_param.py +++ b/tests/test_zero/test_legacy/test_shard_param.py @@ -4,6 +4,7 @@ from functools import partial import pytest import torch import torch.multiprocessing as mp +from common import CONFIG, allclose import colossalai from colossalai.testing import parameterize, rerun_if_address_is_in_use @@ -12,7 +13,6 @@ from colossalai.zero.legacy.gemini.stateful_tensor import StatefulTensor from colossalai.zero.legacy.shard_utils import BucketTensorShardStrategy, TensorShardStrategy from colossalai.zero.legacy.sharded_param import ShardedTensor from colossalai.zero.legacy.sharded_param.sharded_param import ShardedParamV2 -from tests.test_zero.common import CONFIG, allclose @parameterize("shard_strategy_class", [TensorShardStrategy, BucketTensorShardStrategy]) diff --git a/tests/test_zero/test_sharded_optim_state_dict.py b/tests/test_zero/test_legacy/test_sharded_optim_state_dict.py similarity index 100% rename from tests/test_zero/test_sharded_optim_state_dict.py rename to tests/test_zero/test_legacy/test_sharded_optim_state_dict.py diff --git a/tests/test_zero/test_sharded_optim_v2.py b/tests/test_zero/test_legacy/test_sharded_optim_v2.py similarity index 100% rename from tests/test_zero/test_sharded_optim_v2.py rename to tests/test_zero/test_legacy/test_sharded_optim_v2.py diff --git a/tests/test_zero/test_sharded_optim_with_sync_bn.py b/tests/test_zero/test_legacy/test_sharded_optim_with_sync_bn.py similarity index 100% rename from tests/test_zero/test_sharded_optim_with_sync_bn.py rename to tests/test_zero/test_legacy/test_sharded_optim_with_sync_bn.py diff --git a/tests/test_zero/test_state_dict.py b/tests/test_zero/test_legacy/test_state_dict.py similarity index 98% rename from tests/test_zero/test_state_dict.py rename to tests/test_zero/test_legacy/test_state_dict.py index c435d9bb1..40d2820d8 100644 --- a/tests/test_zero/test_state_dict.py +++ b/tests/test_zero/test_legacy/test_state_dict.py @@ -1,7 +1,6 @@ #!/usr/bin/env python # -*- encoding: utf-8 -*- -from copy import deepcopy from functools import partial import pytest diff --git a/tests/test_zero/test_tensor_utils.py b/tests/test_zero/test_legacy/test_tensor_utils.py similarity index 100% rename from tests/test_zero/test_tensor_utils.py rename to tests/test_zero/test_legacy/test_tensor_utils.py diff --git a/tests/test_zero/test_zero_engine.py b/tests/test_zero/test_legacy/test_zero_engine.py similarity index 100% rename from tests/test_zero/test_zero_engine.py rename to tests/test_zero/test_legacy/test_zero_engine.py diff --git a/tests/test_zero/low_level_zero/test_grad_acc.py b/tests/test_zero/test_low_level/test_grad_acc.py similarity index 100% rename from tests/test_zero/low_level_zero/test_grad_acc.py rename to tests/test_zero/test_low_level/test_grad_acc.py diff --git a/tests/test_zero/low_level_zero/test_zero1_2.py b/tests/test_zero/test_low_level/test_zero1_2.py similarity index 100% rename from tests/test_zero/low_level_zero/test_zero1_2.py rename to tests/test_zero/test_low_level/test_zero1_2.py diff --git a/tests/test_zero/low_level_zero/test_zero_init.py b/tests/test_zero/test_low_level/test_zero_init.py similarity index 100% rename from tests/test_zero/low_level_zero/test_zero_init.py rename to tests/test_zero/test_low_level/test_zero_init.py diff --git a/tests/test_zero/low_level_zero/test_zero_tp.py b/tests/test_zero/test_low_level/test_zero_tp.py similarity index 100% rename from tests/test_zero/low_level_zero/test_zero_tp.py rename to tests/test_zero/test_low_level/test_zero_tp.py