Pipeline parallelism for BMTrain TransformerBlockList

bmtrain/__init__.py

@@ -10,7 +10,7 @@
 from .block_layer import CheckpointBlock, TransformerBlockList
 from .backward import optim_step
 from .wrapper import BMTrainModelWrapper
-
+from .pipe_layer import PipelineTransformerBlockList
 from . import debug
 from .store import save, load
 

bmtrain/benchmark/__init__.py

@@ -1,2 +1,3 @@
 from .all_gather import all_gather
-from .reduce_scatter import reduce_scatter
+from .reduce_scatter import reduce_scatter
+from .send_recv import send_recv

bmtrain/benchmark/send_recv.py

@@ -0,0 +1,31 @@
+from .. import nccl
+from .shape import SHAPES
+from ..global_var import config
+from ..utils import print_rank
+from .utils import format_size
+import torch
+def send_recv():
+    current_stream = torch.cuda.current_stream()
+    for shape in SHAPES:
+        send_size = shape
+
+        send_buffer = torch.empty( send_size // 2, dtype=torch.half, device="cuda" )
+        recv_buffer = torch.empty( send_size // 2, dtype=torch.half, device="cuda" )
+
+        start_evt = torch.cuda.Event(enable_timing=True)
+        end_evt = torch.cuda.Event(enable_timing=True)
+
+        current_stream.record_event(start_evt)
+        nccl.groupStart()
+        if config['rank'] in [0,2,4,6]:
+            nccl.send(send_buffer.storage(), config['rank']+1, config['comm'])
+        else:
+            nccl.recv(recv_buffer.storage(), config['rank']-1, config['comm'])
+        nccl.groupEnd()
+        current_stream.record_event(end_evt)
+        current_stream.synchronize()
+        time_usage = start_evt.elapsed_time(end_evt)
+
+        bw = shape / 1024 / 1024 / 1024 * 1000 / time_usage
+        print_rank("Send Recv:\tsize {}\ttime: {:4.3f}\tbw: {:2.6f} GB/s".format(format_size(send_size), time_usage, bw))
+

bmtrain/block_layer.py

@@ -1,6 +1,6 @@
-from typing import Dict, Iterable, Iterator, Union
-
+from typing import Dict, Iterable, Iterator, Union, List
 
+from .utils import round_up
 from .global_var import config
 import torch
 from . import nccl
@@ -10,8 +10,6 @@
 from . import debug
 import copy
 
-def round_up(x, d):
-    return (x + d - 1) // d * d
 
 # the flag is used to control the zero level , 0 means normal zero3 , 1 means forward without release parameter ,2 means backward without gather parameter
 class OpCheckpointBlock(torch.autograd.Function):
@@ -120,16 +118,19 @@ def backward(ctx, *grad_outputs):
         return (None, None, None, None) + tuple(grads)
 
 class CheckpointBlockContext:
-    def __init__(self ,block : 'CheckpointBlock',ctx_dict : dict = None, flag : int = 0) -> None:
+    def __init__(self, block : 'CheckpointBlock', ctx_dict : dict = None, flag : int = 0, pipe = False) -> None:
         self.block = block
-        self.ctx_dict=ctx_dict
+        self.ctx_dict = ctx_dict
         self._param_buffer = {}
         self._grad_buffer = {}
         self._param_tensor = {}
         self._grad_tensor = {}
         self.flag = flag
         self._need_release = False
-
+        if pipe:
+            self.comm = config["zero_comm"] 
+        else:
+            self.comm = config["comm"]
     def enter(self):
         """
         gather parameters
@@ -146,23 +147,23 @@ def enter(self):
                 assert self.block._storage_params[kw].is_cuda
                 assert kw not in self._grad_buffer
                 assert kw not in self._param_buffer
-                local_param = self.block._storage_params[kw]
-
+                local_param = self.block._storage_params[kw]    
+           
                 storage_type = local_param.storage_type()
                 if self.flag != 2:
-                    self._param_buffer[kw] = storage_type(val["partition_size"] * config["world_size"])
+                    self._param_buffer[kw] = storage_type(val["partition_size"] * val["world_size"])
                     self._param_tensor[kw] = torch.tensor([], dtype=self._param_buffer[kw].dtype, device=self._param_buffer[kw].device).set_(self._param_buffer[kw])
 
                 if requires_grad and local_param.requires_grad:
-                    self._grad_buffer[kw] = storage_type(val["partition_size"] * config["world_size"])
+                    self._grad_buffer[kw] = storage_type(val["partition_size"] * val["world_size"])
                     self._grad_tensor[kw] = torch.tensor([], dtype=self._grad_buffer[kw].dtype, device=self._grad_buffer[kw].device).set_(self._grad_buffer[kw]).zero_()
             if self.flag != 2:
                 nccl.groupStart()
                 for kw, val in self.block._storage_info.items():
                     nccl.allGather(
                         self.block._storage_params[kw].storage(),
                         self._param_buffer[kw],
-                        config["comm"]
+                        self.comm
                     )
                 nccl.groupEnd()
 
@@ -177,25 +178,25 @@ def enter(self):
                 self._grad_tensor[kw].record_stream(current_stream)
 
         # update parameters in block
-            for param in self.block._param_info:
-                kw_name = param["kw_name"]
-                offset = param["offset"]
-                shape = param["shape"]
+        for param in self.block._param_info:
+            kw_name = param["kw_name"]
+            offset = param["offset"]
+            shape = param["shape"]
 
-                if self.flag != 2:
-                    dtype = self._param_buffer[kw_name].dtype
-                    device = self._param_buffer[kw_name].device
-                    param["parameter"].data = torch.tensor([], dtype=dtype, device=device).set_(self._param_buffer[kw_name], offset, shape)                
-                else:
-                    dtype = param["parameter"].data.dtype
-                    device = param["parameter"].data.device
-                    param["parameter"].data = torch.tensor([], dtype=dtype, device=device).set_(self.ctx_dict[kw_name], offset, shape)
+            if self.flag != 2:
+                dtype = self._param_buffer[kw_name].dtype
+                device = self._param_buffer[kw_name].device
+                param["parameter"].data = torch.tensor([], dtype=dtype, device=device).set_(self._param_buffer[kw_name], offset, shape)                
+            else:
+                dtype = param["parameter"].data.dtype
+                device = param["parameter"].data.device
+                param["parameter"].data = torch.tensor([], dtype=dtype, device=device).set_(self.ctx_dict[kw_name], offset, shape)
 
-                if requires_grad and kw_name in self._grad_buffer:
-                    param["parameter"].requires_grad_(True)
-                    param["parameter"].grad = torch.tensor([], dtype=dtype, device=device).set_(self._grad_buffer[kw_name], offset, shape)
-                else:
-                    param["parameter"].requires_grad_(False)
+            if requires_grad and kw_name in self._grad_buffer:
+                param["parameter"].requires_grad_(True)
+                param["parameter"].grad = torch.tensor([], dtype=dtype, device=device).set_(self._grad_buffer[kw_name], offset, shape)
+            else:
+                param["parameter"].requires_grad_(False)
 
     def __enter__(self):
         self.enter()
@@ -209,7 +210,6 @@ def exit(self):
             return
         self._need_release = False
         self.block._ready = False
-
         requires_grad = torch.is_grad_enabled()
         if requires_grad:
             for kw, val in self.block._storage_info.items():
@@ -237,7 +237,7 @@ def exit(self):
                             self._grad_buffer[kw],
                             local_param.grad.storage(),
                             "sum",
-                            config["comm"]
+                            self.comm
                         )
                 nccl.groupEnd()
 
@@ -355,7 +355,8 @@ def __init__(self, inner_module : torch.nn.Module):
         offsets = {}
         # intialize storage buffers
         for kw, val in self._storage_info.items():
-            partition_size = round_up(val["total"], config['world_size']) // config['world_size']
+            val["world_size"] = config["world_size"]
+            partition_size = round_up(val["total"], val["world_size"]) // val["world_size"]
             val["partition_size"] = partition_size
             val["begin"] = config['rank'] * partition_size
             val["end"] = (config['rank'] + 1) * partition_size
@@ -538,7 +539,6 @@ def init_parameters(self):
                 # initialzie here
                 tmp_tensor = torch.empty(it["shape"], device=param.device, dtype=param.dtype)
                 param._init_method(tmp_tensor)
-
                 param_st = it["offset"]
                 param_end = it["offset"] + it["size"]
                 kw_name = it["kw_name"]
@@ -604,7 +604,6 @@ def named_modules(self, memo = None, prefix: str = '', remove_duplicate: bool =
             1 -> ('0', Linear(in_features=2, out_features=2, bias=True))
 
         """
-        # print("here in named_modules hahaha")
 
         if memo is None:
             memo = set()
@@ -654,7 +653,7 @@ def forward(ctx, placeholder, self : 'TransformerBlockList', save_list, hidden_s
         with torch.no_grad():
             for i in range(len(self)):
                 if save_list[i][0] == i:
-                    layer_inputs.append(hidden_state)
+                    layer_inputs.append(hidden_state.detach())
                 cuda_rng_state.append( torch.cuda.get_rng_state() )
                 if config['zero_level']==2:
                     flag = 1
@@ -674,13 +673,20 @@ def forward(ctx, placeholder, self : 'TransformerBlockList', save_list, hidden_s
         ctx.layer_inspector = layer_inspector
         ctx.cuda_rng_state = cuda_rng_state
 
-
         ctx.save_for_backward(*layer_inputs, *tensors)
-        return hidden_state
+
+        if self.return_hidden_states:
+            middle_hiddens = layer_inputs 
+            for mid in middle_hiddens:
+                mid.requires_grad_()
+            middle_hiddens = torch.stack(middle_hiddens, dim=0)
+            return hidden_state, middle_hiddens
+        else:
+            return hidden_state, None
 
 
     @staticmethod
-    def backward(ctx, grad_hidden_state : torch.Tensor):
+    def backward(ctx, grad_hidden_state : torch.Tensor, grad_middle: List[torch.Tensor]):
         def exit_prev(prev_ctx, prev_grad):
             if prev_ctx is not None:
                 if prev_grad:
@@ -738,7 +744,7 @@ def exit_prev(prev_ctx, prev_grad):
                                 ctx.save_list[j+1][0] = j+1
 
                     torch.cuda.set_rng_state(ctx.cuda_rng_state[i])
-                    ipt = layer_inputs[ctx.save_list[i][1]].detach().requires_grad_()
+                    ipt = layer_inputs[ctx.save_list[i][1]].requires_grad_()
                     if config['zero_level'] == 2:
                         flag = 2
                     else:
@@ -759,12 +765,15 @@ def exit_prev(prev_ctx, prev_grad):
                         assert it["group"] == ctx.layer_inspector[i][j]["group"], "Backward step changed"
 
                         # change the tensor in placeholder
+                        ctx.layer_inspector[i][j]["requires_grad"] = it["requires_grad"]
                         ctx.layer_inspector[i][j]["tensor"] = it["tensor"]
                     torch.autograd.backward(
                         [output],
                         [grad_hidden_state]
                     )
                     grad_hidden_state = ipt.grad
+                    if grad_middle is not None:
+                        grad_hidden_state = grad_hidden_state + grad_middle[i]
 
                 exit_prev(prev_ctx, prev_grad)
 
@@ -836,6 +845,11 @@ def __iter__(self) -> Iterator[CheckpointBlock]:
     def __getitem__(self, index: Union[int, str]) -> CheckpointBlock:
         return self._modules[str(index)]
 
-    def forward(self, hidden_state, *args):
+    def forward(self, hidden_state, *args, return_hidden_states = False):
+        self.return_hidden_states = return_hidden_states
         placeholder = torch.tensor([], requires_grad=torch.is_grad_enabled())
-        return OpTransformerBlockList.apply(placeholder, self, self.save_list, hidden_state, *args)
+        last_hidden, middle_hiddens = OpTransformerBlockList.apply(placeholder, self, self.save_list, hidden_state, *args)
+        if return_hidden_states:
+            return last_hidden, middle_hiddens
+        else:
+            return last_hidden

bmtrain/distributed/__init__.py

@@ -1 +1 @@
-from .ops import all_gather, all_reduce
+from .ops import all_gather, all_reduce, broadcast, recv_activations, send_activations