[Tensor] remove ParallelAction, use ComputeSpec instread (#1166)

tests/test_tensor/test_addmm_tp.py

@@ -5,7 +5,7 @@ import torch.nn as nn
 import torch.multiprocessing as mp
 from colossalai.tensor import ColoTensor
 from colossalai.tensor import distspec
-from colossalai.tensor import TensorSpec, ComputePattern, ParallelAction, DistSpecManager
+from colossalai.tensor import TensorSpec, ComputePattern, ComputeSpec, DistSpecManager
 from colossalai.context import ParallelMode
 from colossalai.testing import rerun_if_address_is_in_use
 from colossalai.utils import free_port
@@ -41,7 +41,7 @@ class Conv1D(nn.Module):
 def init_1d_row(weight, bias):
     spec = TensorSpec(
         distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [0], [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
-        ParallelAction(ComputePattern.TP1D))
+        ComputeSpec(ComputePattern.TP1D))
     with DistSpecManager.no_grad():
         weight.set_spec(spec)
 
@@ -49,7 +49,7 @@ def init_1d_row(weight, bias):
 def init_1d_col(weight, bias):
     spec = TensorSpec(
         distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [-1], [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
-        ParallelAction(ComputePattern.TP1D))
+        ComputeSpec(ComputePattern.TP1D))
     with DistSpecManager.no_grad():
         weight.set_spec(spec)
         bias.set_spec(spec)

tests/test_tensor/test_embedding_bag_tp.py

@@ -11,14 +11,14 @@ import torch.multiprocessing as mp
 from colossalai.testing import rerun_if_address_is_in_use
 from colossalai.utils import free_port
 from colossalai.core import global_context as gpc
-from colossalai.tensor import TensorSpec, ComputePattern, ParallelAction, DistSpecManager
+from colossalai.tensor import TensorSpec, ComputePattern, ComputeSpec, DistSpecManager
 from _utils import tensor_equal, tensor_shard_equal
 
 
 def init_1d_col(weight):
     spec = TensorSpec(
         distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [-1], [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
-        ParallelAction(ComputePattern.TP1D))
+        ComputeSpec(ComputePattern.TP1D))
     with DistSpecManager.no_grad():
         weight.set_spec(spec)
 

tests/test_tensor/test_embedding_tp.py

@@ -11,14 +11,14 @@ import torch.multiprocessing as mp
 from colossalai.testing import rerun_if_address_is_in_use
 from colossalai.utils import free_port
 from colossalai.core import global_context as gpc
-from colossalai.tensor import TensorSpec, ComputePattern, ParallelAction, DistSpecManager
+from colossalai.tensor import TensorSpec, ComputePattern, ComputeSpec, DistSpecManager
 from _utils import tensor_equal, tensor_shard_equal
 
 
 def init_1d_row(weight):
     spec = TensorSpec(
         distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [0], [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
-        ParallelAction(ComputePattern.TP1D))
+        ComputeSpec(ComputePattern.TP1D))
     with DistSpecManager.no_grad():
         weight.set_spec(spec)
 
@@ -26,7 +26,7 @@ def init_1d_row(weight):
 def init_1d_col(weight):
     spec = TensorSpec(
         distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [-1], [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
-        ParallelAction(ComputePattern.TP1D))
+        ComputeSpec(ComputePattern.TP1D))
     with DistSpecManager.no_grad():
         weight.set_spec(spec)
 

tests/test_tensor/test_gpt.py

@@ -6,7 +6,7 @@ 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 TensorSpec, ComputePattern, ParallelAction, DistSpecManager, distspec
+from colossalai.tensor import TensorSpec, ComputePattern, ComputeSpec, DistSpecManager, distspec
 from colossalai.core import global_context as gpc
 from functools import partial
 from _utils import tensor_equal, tensor_shard_equal, set_seed
@@ -18,7 +18,7 @@ from colossalai.nn.parallel.data_parallel import ColoDDP
 def init_1d_row_spec(model):
     spec = TensorSpec(
         distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [0], [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
-        ParallelAction(ComputePattern.TP1D))
+        ComputeSpec(ComputePattern.TP1D))
     with DistSpecManager.no_grad():
         for n, p in model.named_parameters():
             if 'weight' in n and 'ln' not in n:
@@ -28,7 +28,7 @@ def init_1d_row_spec(model):
 def init_1d_col_spec(model):
     spec = TensorSpec(
         distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [-1], [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
-        ParallelAction(ComputePattern.TP1D))
+        ComputeSpec(ComputePattern.TP1D))
     with DistSpecManager.no_grad():
         for n, p in model.named_parameters():
             if 'ln' not in n and ('weight' in n or 'bias' in n):

tests/test_tensor/test_hybrid_device.py

@@ -1,7 +1,7 @@
 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 ComputePattern, ParallelAction
+from colossalai.tensor import ComputePattern, ComputeSpec
 
 from functools import partial
 from colossalai.core import global_context as gpc
@@ -46,7 +46,7 @@ def run_hybrid_device(use_ddp, mode):
 
     print(f'embedding weight size: {real_model.embed.weight.size()} | device: {real_model.embed.weight.device}')
     #print(f'linear weight size: {real_model.proj.weight.size()} | device: {real_model.proj.weight.device}')
-    parallel_action = ParallelAction(ComputePattern.TP1D)
+    parallel_action = ComputeSpec(ComputePattern.TP1D)
     init_colo_module(model, parallel_action, recursive=True, mode=mode)
 
     # use cpu gloo to handle embedding
@@ -63,6 +63,7 @@ def run_hybrid_device(use_ddp, mode):
     out.sum().backward()
     optimizer.step()
 
+
 def run_dist(rank, world_size, port, use_ddp, mode):
     if use_ddp and world_size == 1:
         return
@@ -71,6 +72,7 @@ def run_dist(rank, world_size, port, use_ddp, mode):
     colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
     run_hybrid_device(use_ddp, mode)
 
+
 @pytest.mark.dist
 @pytest.mark.parametrize('world_size', [1, 4])
 @pytest.mark.parametrize('use_ddp', [False, True])
@@ -78,7 +80,7 @@ def run_dist(rank, world_size, port, use_ddp, mode):
 @rerun_if_address_is_in_use()
 # Working for simulate the embedding(CPU DP+TP) -> nn(GPU DP+TP)
 def _test_hybrid_device(world_size, use_ddp, mode):
-    run_func = partial(run_dist, world_size=world_size, port=free_port(), use_ddp=use_ddp ,mode=mode)
+    run_func = partial(run_dist, world_size=world_size, port=free_port(), use_ddp=use_ddp, mode=mode)
     mp.spawn(run_func, nprocs=world_size)
 
 

tests/test_tensor/test_linear_tp.py

@@ -12,14 +12,14 @@ import torch.nn.functional as F
 from colossalai.testing import rerun_if_address_is_in_use
 from colossalai.utils import free_port
 from colossalai.core import global_context as gpc
-from colossalai.tensor import TensorSpec, ComputePattern, ParallelAction, DistSpecManager
+from colossalai.tensor import TensorSpec, ComputePattern, ComputeSpec, DistSpecManager
 from _utils import tensor_equal, tensor_shard_equal
 
 
 def init_1d_row(weight, bias):
     spec = TensorSpec(
         distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [-1], [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
-        ParallelAction(ComputePattern.TP1D))
+        ComputeSpec(ComputePattern.TP1D))
     with DistSpecManager.no_grad():
         weight.set_spec(spec)
 
@@ -27,7 +27,7 @@ def init_1d_row(weight, bias):
 def init_1d_col(weight, bias):
     spec = TensorSpec(
         distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [0], [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
-        ParallelAction(ComputePattern.TP1D))
+        ComputeSpec(ComputePattern.TP1D))
     with DistSpecManager.no_grad():
         weight.set_spec(spec)
         bias.set_spec(spec)

tests/test_tensor/test_model.py

@@ -10,7 +10,7 @@ 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 distspec, TensorSpec, ComputePattern, \
-    ParallelAction, ColoTensor, DistSpecManager
+    ComputeSpec, ColoTensor, DistSpecManager
 from colossalai.context import ParallelMode
 from colossalai.core import global_context as gpc
 from colossalai.nn.optimizer import ColoOptimizer
@@ -21,7 +21,7 @@ from _utils import tensor_equal, tensor_shard_equal, set_seed
 def init_1d_row_linear(weight):
     spec = TensorSpec(
         distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [-1], [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
-        ParallelAction(ComputePattern.TP1D))
+        ComputeSpec(ComputePattern.TP1D))
     with DistSpecManager.no_grad():
         weight.set_spec(spec)
 
@@ -29,7 +29,7 @@ def init_1d_row_linear(weight):
 def init_1d_col_linear(weight):
     spec = TensorSpec(
         distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [0], [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
-        ParallelAction(ComputePattern.TP1D))
+        ComputeSpec(ComputePattern.TP1D))
     with DistSpecManager.no_grad():
         weight.set_spec(spec)
 
@@ -37,7 +37,7 @@ def init_1d_col_linear(weight):
 def init_1d_row_embedding(weight):
     spec = TensorSpec(
         distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [0], [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
-        ParallelAction(ComputePattern.TP1D))
+        ComputeSpec(ComputePattern.TP1D))
     with DistSpecManager.no_grad():
         weight.set_spec(spec)
 
@@ -45,7 +45,7 @@ def init_1d_row_embedding(weight):
 def init_1d_col_embedding(weight):
     spec = TensorSpec(
         distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [-1], [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
-        ParallelAction(ComputePattern.TP1D))
+        ComputeSpec(ComputePattern.TP1D))
     with DistSpecManager.no_grad():
         weight.set_spec(spec)
 

tests/test_tensor/test_module_spec.py

@@ -5,7 +5,7 @@ from functools import partial
 import torch
 import torch.multiprocessing as mp
 
-from colossalai.tensor import TensorSpec, ComputePattern, ParallelAction
+from colossalai.tensor import TensorSpec, ComputePattern, ComputeSpec
 from colossalai.nn.parallel.layers import init_colo_module, check_colo_module
 from _utils import tensor_equal, tensor_shard_equal, set_seed
 
@@ -40,7 +40,7 @@ def run_model_with_spec(mode, model_name):
         for p1, p2 in zip(model.parameters(), model_seq.parameters()):
             p2.data.copy_(p1.data)
 
-    parallel_action = ParallelAction(ComputePattern.TP1D)
+    parallel_action = ComputeSpec(ComputePattern.TP1D)
     # Not all layers in Bert can be mod by 4.
     # e.g. row shard for all layers is invalid because the first dim of some layer is the classification type size 2.
     if 'bert' == model_name:
@@ -114,7 +114,7 @@ def run_linear_with_spec(mode):
 
     model_handy = copy(model)
 
-    parallel_action = ParallelAction(ComputePattern.TP1D)
+    parallel_action = ComputeSpec(ComputePattern.TP1D)
     init_colo_module(model, parallel_action, recursive=True, mode=mode)
 
     x = torch.rand(2, 4).cuda()
@@ -148,7 +148,7 @@ def run_check_shared_param():
         model = BertForMaskedLM(config)
 
     model = model.cuda()
-    parallel_action = ParallelAction(ComputePattern.TP1D)
+    parallel_action = ComputeSpec(ComputePattern.TP1D)
     # model.cls.predictions.decoder and model.cls.predictions share the bias, so they should have the same spec
     assert len(model.cls.predictions.decoder.bias.shared_param_modules) == 2
     # They are all Linear, so both row is allowed. This should pass check.
@@ -156,7 +156,7 @@ def run_check_shared_param():
     # This should be detected by check because you can not set weight as row while set bias as col.
     col_spec = TensorSpec(
         distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [0], [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
-        ParallelAction(ComputePattern.TP1D))
+        ComputeSpec(ComputePattern.TP1D))
     model.cls.predictions.bias.set_spec(col_spec)
     try:
         check_colo_module(model.cls.predictions.decoder, recursive=False)

tests/test_tensor/test_zero_optim.py

@@ -19,7 +19,7 @@ from colossalai.zero import ZeroOptimizer
 from colossalai.testing import parameterize
 from colossalai.amp import convert_to_apex_amp
 from colossalai.gemini.gemini_mgr import GeminiManager
-from colossalai.tensor import TensorSpec, ComputePattern, ParallelAction, DistSpecManager, distspec
+from colossalai.tensor import TensorSpec, ComputePattern, ComputeSpec, DistSpecManager, distspec
 
 
 def check_param_equal(model, torch_model):
@@ -47,7 +47,7 @@ def run_fwd_bwd(model, criterion, optimizer, input_ids, attn_mask):
 def init_1d_row_spec(model):
     spec = TensorSpec(
         distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [0], [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
-        ParallelAction(ComputePattern.TP1D))
+        ComputeSpec(ComputePattern.TP1D))
     with DistSpecManager.no_grad():
         for n, p in model.named_parameters():
             if 'weight' in n and 'ln' not in n:
@@ -57,7 +57,7 @@ def init_1d_row_spec(model):
 def init_1d_col_spec(model):
     spec = TensorSpec(
         distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [-1], [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
-        ParallelAction(ComputePattern.TP1D))
+        ComputeSpec(ComputePattern.TP1D))
     with DistSpecManager.no_grad():
         for n, p in model.named_parameters():
             if 'ln' not in n and ('weight' in n or 'bias' in n):