added allgather and allreduce

pretrain_ict.py

@@ -57,6 +57,73 @@ def model_provider():
     return general_model_provider(False, False)
 
 
+
+def get_group_world_size_rank():
+
+    group = mpu.get_data_parallel_group()
+    rank = torch.distributed.get_rank(group=group)
+    world_size = torch.distributed.get_world_size(group=group)
+
+    return group, rank, world_size
+
+
+def get_rank_chunk_along_first_dim(tensor):
+
+    group, rank, world_size = get_group_world_size_rank()
+
+    assert tensor.shape[0] % world_size == 0
+    dim_size = tensor.shape[0] // world_size
+    output_list = torch.split(tensor, dim_size, dim=0)
+    
+    output = output_list[rank].contiguous()
+    return output
+
+
+class AllgatherFromDataParallelRegion(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, input_):
+
+        assert input_.dim() == 2
+        group, rank, world_size = get_group_world_size_rank()
+
+        tensor_list = [torch.empty_like(input_) for _ in range(world_size)]
+        tensor_list[rank] = input_
+        torch.distributed.all_gather(tensor_list, input_, group=group)
+
+        output = torch.cat(tensor_list, dim=0).contiguous()
+
+        return output
+
+
+    @staticmethod
+    def backward(ctx, grad_output):
+
+        return get_rank_chunk_along_first_dim(grad_output)
+
+
+class AllReduceFromDataParallelRegion(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, input_):
+
+        assert input_.dim() == 2
+        group, rank, world_size = get_group_world_size_rank()
+
+        tensor_list = [torch.zero_like(input_) for _ in range(world_size)]
+        tensor_list[rank] = input_
+        output = torch.cat(tensor_list, dim=0).contiguous() 
+        torch.distributed.all_reduce(output, group=group)
+
+        return output
+
+
+    @staticmethod
+    def backward(ctx, grad_output):
+
+        return get_rank_chunk_along_first_dim(grad_output)
+
+
 def get_batch(data_iterator):
     # Items and their type.
     keys = ['query_tokens', 'query_pad_mask',
@@ -95,21 +162,35 @@ def forward_step(data_iterator, model):
     # Forward model.
     query_logits, block_logits = model(query_tokens, query_pad_mask, block_tokens, block_pad_mask)
 
-    data_parallel_size = dist.get_world_size() / args.model_parallel_size
-    batch_size = query_logits.shape[0]
-    global_batch_size = int(batch_size * data_parallel_size)
+    IMPLEMENTATION = 'original'
 
-    all_logits_shape = (int(global_batch_size), int(query_logits.shape[1]))
-    all_query_logits = torch.cuda.FloatTensor(*all_logits_shape).type(query_logits.dtype).fill_(0.0)
-    all_block_logits = all_query_logits.clone()
+    if IMPLEMENTATION == 'original':
+        data_parallel_size = dist.get_world_size() / args.model_parallel_size
+        batch_size = query_logits.shape[0]
+        global_batch_size = int(batch_size * data_parallel_size)
+        
+        all_logits_shape = (int(global_batch_size), int(query_logits.shape[1]))
+        all_query_logits = torch.cuda.FloatTensor(*all_logits_shape).type(query_logits.dtype).fill_(0.0)
+        all_block_logits = all_query_logits.clone()
 
-    # record this processes' data
-    all_query_logits[args.rank * batch_size:(args.rank + 1) * batch_size] = query_logits
-    all_block_logits[args.rank * batch_size:(args.rank + 1) * batch_size] = block_logits
+        # record this processes' data
+        all_query_logits[args.rank * batch_size:(args.rank + 1) * batch_size] = query_logits
+        all_block_logits[args.rank * batch_size:(args.rank + 1) * batch_size] = block_logits
 
-    # merge data from all processes
-    dist.all_reduce(all_query_logits)
-    dist.all_reduce(all_block_logits)
+        # merge data from all processes
+        dist.all_reduce(all_query_logits)
+        dist.all_reduce(all_block_logits)
+
+    elif IMPLEMENTATION == 'allgather':
+        all_query_logits = AllgatherFromDataParallelRegion.apply(query_logits)
+        all_block_logits = AllgatherFromDataParallelRegion.apply(block_logits)
+
+    elif IMPLEMENTATION == 'allreduce':
+        all_query_logits = AllReduceFromDataParallelRegion.apply(query_logits)
+        all_block_logits = AllReduceFromDataParallelRegion.apply(block_logits)
+
+    else:
+        raise Exception('should not be here.')
 
     # scores are inner products between query and block embeddings
     retrieval_scores = all_query_logits.float().matmul(torch.transpose(all_block_logits, 0, 1).float())