diff --git a/colossalai/nn/__init__.py b/colossalai/nn/__init__.py index 7c37650ed..13ae5187f 100644 --- a/colossalai/nn/__init__.py +++ b/colossalai/nn/__init__.py @@ -5,6 +5,3 @@ from .metric import * from .model import * from .optimizer import * from ._ops import * - -from .modules import ColoLinear, ColoEmbedding -from .module_utils import register_colo_module, is_colo_module, get_colo_module, init_colo_module, check_colo_module diff --git a/colossalai/nn/modules/__init__.py b/colossalai/nn/modules/__init__.py deleted file mode 100644 index 3d6a0e69b..000000000 --- a/colossalai/nn/modules/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .colo_module import ColoModule -from .linear import ColoLinear -from .embedding import ColoEmbedding \ No newline at end of file diff --git a/colossalai/nn/parallel/__init__.py b/colossalai/nn/parallel/__init__.py new file mode 100644 index 000000000..c22b027ca --- /dev/null +++ b/colossalai/nn/parallel/__init__.py @@ -0,0 +1,3 @@ +from .data_parallel import ColoDDP, ColoDDPV2 + +__all__ = ['ColoDDP', 'ColoDDPV2'] diff --git a/colossalai/nn/parallel.py b/colossalai/nn/parallel/data_parallel.py similarity index 99% rename from colossalai/nn/parallel.py rename to colossalai/nn/parallel/data_parallel.py index a257a2377..668964592 100644 --- a/colossalai/nn/parallel.py +++ b/colossalai/nn/parallel/data_parallel.py @@ -7,8 +7,6 @@ from colossalai.zero.utils.zero_hook_v2 import ZeROHookV2 from colossalai.tensor.chunk import ChunkManager, TensorState from colossalai.tensor.param_op_hook import use_param_op_hooks -__all__ = ['ColoDDP', 'ColoDDPV2'] - def free_storage(data: torch.Tensor) -> None: """Free underlying storage of a Tensor.""" diff --git a/colossalai/nn/parallel/layers/__init__.py b/colossalai/nn/parallel/layers/__init__.py new file mode 100644 index 000000000..6d205828d --- /dev/null +++ b/colossalai/nn/parallel/layers/__init__.py @@ -0,0 +1,15 @@ +from .colo_module import ColoModule +from .linear import ColoLinear +from .embedding import ColoEmbedding +from .module_utils import register_colo_module, is_colo_module, get_colo_module, init_colo_module, check_colo_module + +__all__ = [ + 'ColoModule', + 'register_colo_module', + 'is_colo_module', + 'get_colo_module', + 'init_colo_module', + 'check_colo_module', + 'ColoLinear', + 'ColoEmbedding', +] diff --git a/colossalai/nn/modules/colo_module.py b/colossalai/nn/parallel/layers/colo_module.py similarity index 100% rename from colossalai/nn/modules/colo_module.py rename to colossalai/nn/parallel/layers/colo_module.py diff --git a/colossalai/nn/modules/embedding.py b/colossalai/nn/parallel/layers/embedding.py similarity index 100% rename from colossalai/nn/modules/embedding.py rename to colossalai/nn/parallel/layers/embedding.py diff --git a/colossalai/nn/modules/linear.py b/colossalai/nn/parallel/layers/linear.py similarity index 100% rename from colossalai/nn/modules/linear.py rename to colossalai/nn/parallel/layers/linear.py diff --git a/colossalai/nn/module_utils.py b/colossalai/nn/parallel/layers/module_utils.py similarity index 99% rename from colossalai/nn/module_utils.py rename to colossalai/nn/parallel/layers/module_utils.py index 158bd1062..800c9226d 100644 --- a/colossalai/nn/module_utils.py +++ b/colossalai/nn/parallel/layers/module_utils.py @@ -1,6 +1,6 @@ from typing import Dict from colossalai.tensor import ColoParameter, ParallelAction, TensorSpec -from .modules import ColoModule +from . import ColoModule import torch _COLOSSAL_MODULES: Dict[type, ColoModule] = {} diff --git a/colossalai/utils/__init__.py b/colossalai/utils/__init__.py index 15ef96de3..875b5a93b 100644 --- a/colossalai/utils/__init__.py +++ b/colossalai/utils/__init__.py @@ -11,8 +11,6 @@ from .memory import (report_memory_usage, colo_device_memory_used, colo_set_proc colo_device_memory_capacity, colo_set_cpu_memory_capacity, colo_get_cpu_memory_capacity) from .timer import MultiTimer, Timer from .tensor_detector import TensorDetector -from .model.utils import InsertPostInitMethodToModuleSubClasses -from .model.colo_init_context import ColoInitContext __all__ = [ 'checkpoint', @@ -52,6 +50,4 @@ __all__ = [ 'disposable', 'colo_set_cpu_memory_capacity', 'colo_get_cpu_memory_capacity', - 'InsertPostInitMethodToModuleSubClasses', - 'ColoInitContext', ] diff --git a/colossalai/utils/model/colo_init_context.py b/colossalai/utils/model/colo_init_context.py index 87489f0ce..5eaeff19f 100644 --- a/colossalai/utils/model/colo_init_context.py +++ b/colossalai/utils/model/colo_init_context.py @@ -2,7 +2,7 @@ from .utils import InsertPostInitMethodToModuleSubClasses import torch from colossalai.tensor import ColoTensor, ColoParameter -from colossalai.nn import register_colo_module, init_colo_module, \ +from colossalai.nn.parallel.layers import register_colo_module, \ ColoLinear, ColoEmbedding from torch import nn diff --git a/colossalai/utils/model/pipelinable.py b/colossalai/utils/model/pipelinable.py index aab066b34..14d19a7ae 100644 --- a/colossalai/utils/model/pipelinable.py +++ b/colossalai/utils/model/pipelinable.py @@ -1,10 +1,7 @@ import torch -import functools import inspect -from colossalai.amp.naive_amp import NaiveAMPModel -from colossalai.utils.model.utils import _substitute_init_recursively, InsertPostInitMethodToModuleSubClasses, call_to_str +from colossalai.utils.model.utils import InsertPostInitMethodToModuleSubClasses, call_to_str from colossalai.builder.pipeline import partition_uniform, partition_balanced -from colossalai.core import global_context as gpc from colossalai.nn.layer.utils import CheckpointModule from colossalai.tensor import ColoTensor diff --git a/colossalai/zero/init_ctx/init_context.py b/colossalai/zero/init_ctx/init_context.py index 8e44c0632..93d5a455f 100644 --- a/colossalai/zero/init_ctx/init_context.py +++ b/colossalai/zero/init_ctx/init_context.py @@ -1,9 +1,12 @@ import contextlib import functools from typing import Optional +from contextlib import AbstractContextManager + import torch import torch.nn as nn import torch.distributed as dist + from colossalai.context.parallel_mode import ParallelMode from colossalai.core import global_context as gpc from colossalai.context.singleton_meta import SingletonMeta @@ -12,8 +15,7 @@ from colossalai.zero.shard_utils import BaseShardStrategy from colossalai.zero.sharded_model._utils import cast_tensor_to_fp16 from colossalai.zero.sharded_model.sharded_model_v2 import ShardedModelV2 from colossalai.zero.sharded_param import ShardedParamV2 -from contextlib import AbstractContextManager -from colossalai.utils import InsertPostInitMethodToModuleSubClasses +from colossalai.utils.model.utils import InsertPostInitMethodToModuleSubClasses class ZeroContextConfig(object): diff --git a/colossalai/zero/zero_optimizer.py b/colossalai/zero/zero_optimizer.py index 401d7257e..fa5b84456 100644 --- a/colossalai/zero/zero_optimizer.py +++ b/colossalai/zero/zero_optimizer.py @@ -2,7 +2,7 @@ import torch import torch.distributed as dist from enum import Enum from torch.optim import Optimizer -from colossalai.nn.parallel import ColoDDPV2 +from colossalai.nn.parallel.data_parallel import ColoDDPV2 from typing import Dict from colossalai.amp.naive_amp.grad_scaler import DynamicGradScaler from colossalai.core import global_context as gpc diff --git a/tests/test_tensor/test_context.py b/tests/test_tensor/test_context.py index 59d0e8498..8171ebfab 100644 --- a/tests/test_tensor/test_context.py +++ b/tests/test_tensor/test_context.py @@ -1,10 +1,7 @@ import pytest -from colossalai.utils import ColoInitContext +from colossalai.utils.model.colo_init_context import ColoInitContext -from numpy import allclose, require import torch -from colossalai.tensor import ColoTensor -from copy import deepcopy from colossalai.utils.cuda import get_current_device diff --git a/tests/test_tensor/test_gpt.py b/tests/test_tensor/test_gpt.py index 85f8de367..4c438a22d 100644 --- a/tests/test_tensor/test_gpt.py +++ b/tests/test_tensor/test_gpt.py @@ -5,14 +5,14 @@ import torch.multiprocessing as mp 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 import ColoInitContext +from colossalai.utils.model.colo_init_context import ColoInitContext from colossalai.tensor import TensorSpec, ComputePattern, ParallelAction, DistSpecManager, distspec from colossalai.core import global_context as gpc from functools import partial from _utils import tensor_equal, tensor_shard_equal, set_seed from tests.components_to_test.registry import non_distributed_component_funcs from torch.nn.parallel import DistributedDataParallel as DDP -from colossalai.nn.parallel import ColoDDP +from colossalai.nn.parallel.data_parallel import ColoDDP def init_1d_row_spec(model): diff --git a/tests/test_tensor/test_hybrid_device.py b/tests/test_tensor/test_hybrid_device.py index f1dcd8f20..a6d5cf14f 100644 --- a/tests/test_tensor/test_hybrid_device.py +++ b/tests/test_tensor/test_hybrid_device.py @@ -1,13 +1,14 @@ -from colossalai.utils import free_port, ColoInitContext, get_current_device +from colossalai.utils import free_port, get_current_device +from colossalai.utils.model.colo_init_context import ColoInitContext from colossalai.testing import rerun_if_address_is_in_use -from colossalai.tensor import TensorSpec, ComputePattern, ParallelAction +from colossalai.tensor import ComputePattern, ParallelAction from functools import partial from colossalai.core import global_context as gpc from colossalai.context import ParallelMode -from colossalai.nn import init_colo_module -from colossalai.nn.parallel import ColoDDP +from colossalai.nn.parallel.layers import init_colo_module +from colossalai.nn.parallel.data_parallel import ColoDDP import colossalai import torch diff --git a/tests/test_tensor/test_model.py b/tests/test_tensor/test_model.py index 019152aa6..65f633a53 100644 --- a/tests/test_tensor/test_model.py +++ b/tests/test_tensor/test_model.py @@ -5,11 +5,11 @@ import colossalai import pytest import torch import torch.multiprocessing as mp -from colossalai.testing import parameterize, rerun_if_address_is_in_use +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 import ColoInitContext -from colossalai.tensor import distspec, named_params_with_colotensor, TensorSpec, ComputePattern, \ +from colossalai.utils.model.colo_init_context import ColoInitContext +from colossalai.tensor import distspec, TensorSpec, ComputePattern, \ ParallelAction, ColoTensor, DistSpecManager from colossalai.context import ParallelMode from colossalai.core import global_context as gpc diff --git a/tests/test_tensor/test_module_spec.py b/tests/test_tensor/test_module_spec.py index 0a4e7c9d2..e8ddcd884 100644 --- a/tests/test_tensor/test_module_spec.py +++ b/tests/test_tensor/test_module_spec.py @@ -6,7 +6,7 @@ import torch import torch.multiprocessing as mp from colossalai.tensor import TensorSpec, ComputePattern, ParallelAction -from colossalai.nn import init_colo_module, check_colo_module +from colossalai.nn.parallel.layers import init_colo_module, check_colo_module from _utils import tensor_equal, tensor_shard_equal, set_seed import colossalai diff --git a/tests/test_tensor/test_zero.py b/tests/test_tensor/test_zero.py index a87c73359..7a862c5b4 100644 --- a/tests/test_tensor/test_zero.py +++ b/tests/test_tensor/test_zero.py @@ -5,14 +5,14 @@ import torch.multiprocessing as mp 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 import ColoInitContext -from colossalai.tensor import TensorSpec, ComputePattern, ParallelAction, DistSpecManager, distspec, ColoParameter, ChunkManager +from colossalai.utils.model.colo_init_context import ColoInitContext +from colossalai.tensor import ChunkManager from colossalai.core import global_context as gpc from functools import partial -from _utils import tensor_equal, tensor_shard_equal, set_seed +from _utils import tensor_equal, set_seed from tests.components_to_test.registry import non_distributed_component_funcs from torch.nn.parallel import DistributedDataParallel as DDP -from colossalai.nn.parallel import ColoDDP, ColoDDPV2 +from colossalai.nn.parallel import ColoDDPV2 from colossalai.testing import parameterize diff --git a/tests/test_tensor/test_zero_optim.py b/tests/test_tensor/test_zero_optim.py index 6c908870c..66072dca7 100644 --- a/tests/test_tensor/test_zero_optim.py +++ b/tests/test_tensor/test_zero_optim.py @@ -6,11 +6,11 @@ from colossalai.context.parallel_mode import ParallelMode 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 import ColoInitContext +from colossalai.utils.model.colo_init_context import ColoInitContext from colossalai.tensor import ChunkManager from colossalai.core import global_context as gpc from functools import partial -from _utils import tensor_equal, tensor_shard_equal, set_seed +from _utils import tensor_equal, set_seed from tests.components_to_test.registry import non_distributed_component_funcs from torch.nn.parallel import DistributedDataParallel as DDP from colossalai.nn.parallel import ColoDDPV2 diff --git a/tests/test_utils/test_pipelinable.py b/tests/test_utils/test_pipelinable.py index 2be3b264c..09d815632 100644 --- a/tests/test_utils/test_pipelinable.py +++ b/tests/test_utils/test_pipelinable.py @@ -1,13 +1,8 @@ -import os.path as osp - -import pytest import torch import torch.multiprocessing as mp from colossalai.utils.model.pipelinable import PipelinableContext -from functools import partial -from colossalai.utils import free_port from colossalai.testing import rerun_on_exception NUM_CHUNKS = 1