[Tensor] add embedding tp1d row (#904)

tests/components_to_test/simple_net.py

@@ -5,7 +5,6 @@ from .utils.dummy_data_generator import DummyDataGenerator
 from .registry import non_distributed_component_funcs
 from colossalai.utils.cuda import get_current_device
 
-
 class SimpleNet(CheckpointModule):
     """
     In this no-leaf module, it has subordinate nn.modules and a nn.Parameter.
@@ -13,12 +12,14 @@ class SimpleNet(CheckpointModule):
 
     def __init__(self, checkpoint=False) -> None:
         super().__init__(checkpoint=checkpoint)
+        self.embed = nn.Embedding(20, 4)
         self.proj1 = nn.Linear(4, 8)
         self.ln1 = nn.LayerNorm(8)
         self.proj2 = nn.Linear(8, 4)
         self.ln2 = nn.LayerNorm(4)
 
     def forward(self, x):
+        x = self.embed(x)
         x = self.proj1(x)
         x = self.ln1(x)
         x = self.proj2(x)
@@ -26,11 +27,12 @@ class SimpleNet(CheckpointModule):
         return x
 
 
+
 class DummyDataLoader(DummyDataGenerator):
 
     def generate(self):
-        data = torch.rand(16, 4, device=get_current_device())
-        label = torch.randint(low=0, high=2, size=(16,), device=get_current_device())
+        data = torch.randint(low=0, high=20, size=(16,20), device=get_current_device())
+        label = torch.randint(low=0, high=2, size=(16,4), device=get_current_device())
         return data, label
 
 

tests/test_tensor/test_embedding_tp.py

@@ -65,10 +65,60 @@ def run_embedding_tp1d_col_test():
     W_grad = torch.chunk(W_grad, DEPTH, dim=-1)[local_rank]
     check_equal(W_grad, layer.weight.grad)
 
+def run_embedding_tp1d_row_test():
+    device = get_current_device()
+    dtype = torch.float32
+    DEPTH = gpc.get_world_size(ParallelMode.PARALLEL_1D)
+    num_embeddings = 12
+    embedding_dim = 32
+
+    local_rank = gpc.get_local_rank(ParallelMode.PARALLEL_1D)
+
+    layer_master = torch.nn.Embedding(num_embeddings, embedding_dim)
+    layer = torch.nn.Embedding(num_embeddings, embedding_dim)
+
+    A_master = torch.tensor((0,3,6,9), device=device)
+    A = broadcast_tensor_chunk(A_master, chunk_size=1)
+
+    W_shape = (num_embeddings, embedding_dim)
+    W_master = torch.randn(W_shape, dtype=dtype, device=device)
+    W = broadcast_tensor_chunk(W_master, chunk_size=1)
+    W.requires_grad = True
+
+    # replace the torch nn.Parameters with ColoTensor
+    sharded_weight = ColoTensor.init_from_torch_tensor(W)
+    parallel_action_list = [
+        ParallelAction(priority=1, compute_pattern=ComputePattern.TP1DRow_Embedding, 
+        parallel_mode=ParallelMode.PARALLEL_1D)
+    ]
+    spec = TensorSpec(parallel_action_list)
+    sharded_weight.set_spec(spec) # reshard
+    replace_parameter_add_grad(layer, sharded_weight)
+    out = layer(A)
+
+    replace_parameter_add_grad(layer_master, W_master)
+    C_master = layer_master(A_master)
+    C = C_master.clone()
+
+    check_equal(out, C)
+
+    grad_shape = C_master.shape
+    grad_master = torch.randn(grad_shape, dtype=dtype, device=get_current_device())
+    grad = broadcast_tensor_chunk(grad_master, chunk_size=1)
+    out.backward(grad)
+
+    grad_master = grad_master.clone()
+    C_master.backward(grad_master)
+
+    W_grad = W_master.grad
+    W_grad = torch.chunk(W_grad, DEPTH, dim=0)[local_rank]
+    check_equal(W_grad, layer.weight.grad)
+
 def run_dist(rank, world_size, port):
     config = dict(parallel=dict(tensor=dict(mode="1d", size=world_size),))
     colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
     run_embedding_tp1d_col_test()
+    run_embedding_tp1d_row_test()
 
 @pytest.mark.dist
 @parameterize('world_size', [1, 4])

tests/test_tensor/test_model.py

@@ -47,6 +47,11 @@ def run_1d_col_tp():
     ]
     spec_col = TensorSpec(parallel_action_list_col)
 
+    parallel_action_list_embedding_col = [
+        ParallelAction(priority=1, compute_pattern=ComputePattern.TP1DCol_Embedding, parallel_mode=ParallelMode.PARALLEL_1D)
+    ]
+    spec_embedding_col = TensorSpec(parallel_action_list_embedding_col)
+
     set_seed(1)
     if rank == 0:
         model_torch = model_builder(checkpoint=True)
@@ -60,6 +65,8 @@ def run_1d_col_tp():
             p.set_spec(spec_col)
         if 'proj2' in name and 'weight' in name:
             p.set_spec(spec_row)
+        if 'embed' in name and 'weight' in name:
+            p.set_spec(spec_embedding_col)
 
     model = model.cuda()
 
@@ -172,6 +179,11 @@ def run_1d_row_tp():
     ]
     spec = TensorSpec(parallel_action_list)
 
+    parallel_action_list_embedding_row = [
+        ParallelAction(priority=1, compute_pattern=ComputePattern.TP1DRow_Embedding, parallel_mode=ParallelMode.PARALLEL_1D)
+    ]
+    spec_embedding_row = TensorSpec(parallel_action_list_embedding_row)
+
     set_seed(1)
     if rank == 0:
         model_torch = model_builder(checkpoint=True)
@@ -183,6 +195,8 @@ def run_1d_row_tp():
             continue
         if 'weight' in name and 'LayerNorm' not in name and 'ln' not in name and 'embed' not in name:
             p.set_spec(spec)
+        if 'embed' in name and 'weight' in name:
+            p.set_spec(spec_embedding_row)
 
     model = model.cuda()
 
@@ -227,7 +241,7 @@ def run_dist(rank, world_size, port):
     config = dict(parallel=dict(tensor=dict(mode="1d", size=world_size),))
     colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
     run_1d_row_tp()
-
+    run_1d_col_tp()
 
 @pytest.mark.dist
 @parameterize('world_size', [1, 4])
@@ -238,6 +252,6 @@ def test_simple_net(world_size):
 
 
 if __name__ == '__main__':
-    # test_simple_net()
+    test_simple_net()
     # test_model_parameters()
-    test_colo_optimizer()
+    # test_colo_optimizer()