diff --git a/tests/test_moe/test_moe_checkpoint.py b/tests/test_moe/test_moe_checkpoint.py index 5b6fe441112d..d2cff44ad9bd 100644 --- a/tests/test_moe/test_moe_checkpoint.py +++ b/tests/test_moe/test_moe_checkpoint.py @@ -14,7 +14,7 @@ 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 23ad1a3dc6a4..4826d87ac044 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.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 5987e31f71c9..18b50eb5c482 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.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 d38f66fef658..49c452938e25 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 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 7e140bf862f2..b43e52bb4c6a 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.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 641136718161..03a1a609b672 100644 --- a/tests/test_zero/test_found_inf.py +++ b/tests/test_zero/test_legacy/test_found_inf.py @@ -4,6 +4,7 @@ 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_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 6085de3c8919..b76648321451 100644 --- a/tests/test_zero/test_shard_param.py +++ b/tests/test_zero/test_legacy/test_shard_param.py @@ -4,6 +4,7 @@ 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.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 c435d9bb1ef7..40d2820d800a 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