[colotensor] add Tensor.view op and its unit test (#1343)

[colotensor] add megatron initialization for gpt2

tests/test_tensor/common_utils/_utils.py

@@ -73,3 +73,9 @@ def split_param_row_tp1d(param, pg):
 
 def split_param_col_tp1d(param, pg):
     split_param_single_dim_tp1d(-1, param, pg)
+
+
+def debug_print(ranks, *args):
+    if dist.get_rank() in ranks:
+        print(*args)
+    dist.barrier()

tests/test_tensor/core/test_tensor.py

@@ -75,7 +75,7 @@ def _run_view(world_size):
     assert t.size_global(1) == 5
     assert t.size_global() == torch.Size([4 * world_size, 5])
 
-    t = t.view_global(4 * 5 * world_size)
+    t = t.view(4 * 5 * world_size)
     assert t.shape == torch.Size([4 * 5 * world_size])
 
 
@@ -129,9 +129,9 @@ def _run_set_tensor_spec(world_size):
     spec1 = ColoTensorSpec(pg)
     t1 = ColoTensor.from_torch_tensor(torch.randn(2, 3, 4), spec1)
 
-    dist_spec2 = (ShardSpec([-1], [pg.tp_world_size()]), None)
+    dist_spec2 = ShardSpec([-1], [pg.tp_world_size()])
     assert t1.is_replicate()
-    t1.set_dist_spec(*dist_spec2)
+    t1.set_dist_spec(dist_spec2)
     assert t1.is_shard_1dcol()
 
 

tests/test_tensor/model/test_gpt2.py

@@ -15,6 +15,7 @@ from colossalai.utils.model.colo_init_context import ColoInitContext
 from colossalai.tensor import ShardSpec, ComputePattern, ComputeSpec, ProcessGroup, ColoTensor, ColoTensorSpec
 from colossalai.nn.parallel.data_parallel import ColoDDP
 from tests.components_to_test.registry import non_distributed_component_funcs
+from tests.test_tensor.common_utils import split_param_col_tp1d, split_param_row_tp1d, debug_print
 
 
 def init_1d_row_spec(model, pg: ProcessGroup):
@@ -34,6 +35,32 @@ def init_1d_col_spec(model, pg: ProcessGroup):
             p.set_tensor_spec(*spec)
 
 
+def init_megatron_spec(model, pg: ProcessGroup):
+    for mn, module in model.named_modules():
+        # debug_print([0], mn)
+        for pn, param in module.named_parameters(recurse=False):
+            # debug_print([0], '\t', pn, param.compute_spec, param.shape)
+            param.set_process_group(pg)
+
+            if 'mlp.c_fc' in mn:
+                if 'weight' in pn or 'bias' in pn:
+                    split_param_col_tp1d(param, pg)
+                    param.compute_spec.set_output_replicate(False)
+                else:
+                    raise RuntimeError
+            elif 'mlp.c_proj' in mn:
+                if 'weight' in pn:
+                    split_param_row_tp1d(param, pg)
+                else:
+                    assert 'bias' in pn
+            elif 'wte' in mn or 'wpe' in mn:
+                assert 'weight' in pn
+                split_param_col_tp1d(param, pg)
+            elif 'c_fc' in mn or 'c_proj' in mn:
+                split_param_col_tp1d(param, pg)
+            # debug_print([0], '\t', param.compute_spec, param.shape)
+
+
 def check_param_equal(model, torch_model, pg: ProcessGroup):
     for p, torch_p in zip(model.parameters(), torch_model.parameters()):
         assert pg.tp_local_rank() is not None, f"{pg.rank()} {pg.tp_world_size()} {pg._tp_degree} {pg.tp_local_rank()}1"
@@ -102,8 +129,10 @@ def run_dist(rank, world_size, port, use_ddp):
     if use_ddp and world_size == 1:
         return
     colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
-    run_gpt(init_1d_row_spec, use_ddp)
-    run_gpt(init_1d_col_spec, use_ddp)
+    # Comments below tests for speed concern
+    # run_gpt(init_1d_row_spec, use_ddp)
+    # run_gpt(init_1d_col_spec, use_ddp)
+    run_gpt(init_megatron_spec, use_ddp)
 
 
 @pytest.mark.dist
@@ -116,4 +145,4 @@ def test_gpt(world_size, use_ddp):
 
 
 if __name__ == '__main__':
-    test_gpt(4, use_ddp=True)
+    test_gpt(4, use_ddp=False)

tests/test_tensor/ops/test_view.py

@@ -0,0 +1,100 @@
+from functools import partial
+
+import colossalai
+import pytest
+import torch
+import torch.multiprocessing as mp
+import torch.distributed as dist
+from colossalai.testing import rerun_if_address_is_in_use
+from colossalai.utils import free_port, get_current_device
+from colossalai.tensor import ColoTensorSpec, ProcessGroup, ColoTensor, ShardSpec
+from colossalai.tensor.distspec import DistPlacementPattern
+from tests.test_tensor.common_utils import split_param_row_tp1d, split_param_col_tp1d, debug_print
+
+
+def exam_view_core(pg):
+    # the case of replicated ColoTensors
+    x = torch.randn(4, 4).cuda()
+    x_colo = ColoTensor(x, ColoTensorSpec(pg))
+
+    y = x.view(2, -1, 2)
+    y_colo = x_colo.view(2, -1, 2)
+
+    assert torch.all(y == y_colo)
+    assert y_colo.dist_spec.placement == DistPlacementPattern.REPLICATE
+    # the perfect case of col-sliced ColoTensors
+    split_param_col_tp1d(x_colo, pg)
+
+    z = x.view(torch.Size((2, 1, 2, -1)))
+    z_colo = x_colo.view(torch.Size((2, 1, 2, -1)))
+    if dist.get_rank() == 0:
+        z = z[:, :, :, 0:2]
+    else:
+        z = z[:, :, :, 2:]
+    assert torch.all(z == z_colo)
+    assert z_colo.dist_spec == x_colo.dist_spec
+    # the perfect case of row-sliced ColoTensors
+    split_param_row_tp1d(x_colo, pg)
+
+    z = x.view(torch.Size((-1, 2, 2)))
+    z_colo = x_colo.view(torch.Size((-1, 2, 2)))
+    if dist.get_rank() == 0:
+        z = z[0:2, :, :]
+    else:
+        z = z[2:, :, :]
+    assert torch.all(z == z_colo)
+    assert z_colo.dist_spec == x_colo.dist_spec
+    # the normal case of row-sliced ColoTensors
+    z = x.view(-1, 2, 2, 2)
+    z_colo = x_colo.view(-1, 2, 2, 2)
+    assert torch.all(z == z_colo)
+    assert y_colo.dist_spec.placement == DistPlacementPattern.REPLICATE
+
+
+def exam_view_autograd(pg):
+    x = torch.randn(8, 2, device=get_current_device(), requires_grad=True)
+    y = torch.randn(8, 2, device=get_current_device(), requires_grad=True)
+    with torch.no_grad():
+        y.copy_(x)
+    y = ColoTensor(y, ColoTensorSpec(pg))
+    y_slice = y.redistribute(ShardSpec([-1], [pg.tp_world_size()]))
+
+    xx = x.view(2, 2, -1)
+    yy_slice = y_slice.view(2, 2, -1)
+    yy = yy_slice.to_replicate()
+    grad = torch.randn(2, 2, 4, device=get_current_device())
+
+    xx.backward(grad)
+    yy.backward(grad)
+    assert torch.all(x.grad == y.grad)
+
+
+def exam_view_errors(pg):
+    x = torch.randn(8, 2, device=get_current_device())
+    x = ColoTensor(x, ColoTensorSpec(pg))
+    split_param_row_tp1d(x, pg)
+
+    x.view('a', 'b', 'c')
+    x.view(8, -1)
+    x.view([-2, -2, -2])
+    x.view((-1, -1, -1))
+
+
+def run_dist(rank, world_size, port):
+    colossalai.launch(config=dict(), rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
+    pg = ProcessGroup(tp_degree=torch.distributed.get_world_size())
+    exam_view_core(pg)
+    exam_view_autograd(pg)
+    # exam_view_errors(pg)
+
+
+@pytest.mark.dist
+@pytest.mark.parametrize('world_size', [2])
+@rerun_if_address_is_in_use()
+def test_view(world_size):
+    run_func = partial(run_dist, world_size=world_size, port=free_port())
+    mp.spawn(run_func, nprocs=world_size)
+
+
+if __name__ == '__main__':
+    test_view(2)

tests/test_tensor/test_colo_checkpoint_tools.py

@@ -11,7 +11,7 @@ from colossalai.utils.cuda import get_current_device
 from colossalai.utils import free_port
 from colossalai.tensor import ComputePattern, ComputeSpec, ColoTensor, ShardSpec, ProcessGroup, ColoTensorSpec
 from colossalai.utils.checkpoint.utils import gather_tensor, scatter_tensor
-from tests.test_tensor._utils import tensor_shard_equal
+from tests.test_tensor.common_utils import tensor_shard_equal
 
 
 def run_dist(rank, world_size, port, dp_degree, tp_degree):
@@ -24,7 +24,7 @@ def run_dist(rank, world_size, port, dp_degree, tp_degree):
 
     gather_tensor(param)
     if dist.get_rank() == 0:
-        assert torch.allclose(x, param.data, rtol=0, atol=0)
+        assert torch.all(x == param)
     else:
         assert tensor_shard_equal(x, param.data, pg.tp_local_rank(), pg.tp_world_size())
     dist.barrier()

tests/test_zero/test_sharded_optim_state_dict.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 functools import partial
-from tests.test_tensor._utils import set_seed
+from tests.test_tensor.common_utils import set_seed
 from tests.components_to_test.registry import non_distributed_component_funcs
 from colossalai.testing import parameterize
 from colossalai.nn.optimizer import HybridAdam

tests/test_zero/test_zero_optim_state_dict.py

@@ -9,7 +9,7 @@ from colossalai.utils import free_port
 from colossalai.utils.model.colo_init_context import ColoInitContext
 from colossalai.core import global_context as gpc
 from functools import partial
-from tests.test_tensor._utils import set_seed
+from tests.test_tensor.common_utils import set_seed
 from tests.components_to_test.registry import non_distributed_component_funcs
 from colossalai.nn.parallel.data_parallel import ZeroDDP
 from colossalai.gemini import ChunkManager, GeminiManager