BMTrain now supports 1F1B Pipeline schedule!

.github/workflows/build.yml

@@ -6,14 +6,15 @@ on:
     branches:
       - 'dev'
       - 'main'
+  push:
+    branches:
+      - 'dev'
 
 jobs:
   build-archive-wheel:
 
     uses: OpenBMB/BMTrain/.github/workflows/build_whl.yml@main
-    secrets:
-      DOCKERHUB_TOKEN: ${{ secrets.DOCKERHUB_TOKEN }}
-      DOCKERHUB_USERNAME: ${{ secrets.DOCKERHUB_USERNAME }}
+    secrets: inherit
 
   publish:
     needs: build-archive-wheel

bmtrain/__init__.py

@@ -1,4 +1,4 @@
-from .utils import print_block, print_dict, print_rank, see_memory, load_nccl_pypi
+from .utils import print_block, print_dict, print_rank, print_rank_pp, see_memory, load_nccl_pypi
 try:
     from . import nccl
 except:
@@ -10,11 +10,11 @@
 from .layer import DistributedModule
 from .param_init import init_parameters, grouped_parameters
 from .synchronize import synchronize, sum_loss, wait_loader, gather_result
-from .block_layer import Block, TransformerBlockList
+from .block_layer import Block, TransformerBlockList, PipeDreamBlockList
 from .wrapper import BMTrainModelWrapper
 from .pipe_layer import PipelineTransformerBlockList
 from . import debug
-from .store import save, load
+from .store import save, load, clean
 
 from . import loss
 from . import distributed

bmtrain/benchmark/all_reduce.py

@@ -0,0 +1,27 @@
+from .. import nccl
+from .shape import SHAPES
+from ..global_var import config
+from ..utils import round_up, print_rank
+from .utils import format_size
+import torch
+
+def all_reduce():
+    current_stream = torch.cuda.current_stream()
+    for shape in SHAPES:
+        global_size = round_up(shape, config['world_size'] * 2)
+
+        partition_tensor = torch.empty( global_size // 2, dtype=torch.half, device="cuda" )
+        global_tensor = torch.empty( global_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.allReduce(partition_tensor.storage(), global_tensor.storage(),"sum", config['comm'])
+        current_stream.record_event(end_evt)
+        current_stream.synchronize()
+        time_usage = start_evt.elapsed_time(end_evt)
+
+        bw = global_size / 1024 / 1024 / 1024 * 1000 / time_usage * 2
+        print_rank("All reduce:\tsize {}\ttime: {:4.3f}\tbw: {:2.6f} GB/s".format(format_size(global_size), time_usage, bw))
+

bmtrain/distributed/__init__.py

@@ -1 +1 @@
-from .ops import all_gather, all_reduce, broadcast, recv_activations, send_activations, reduce_scatter
+from .ops import all_gather, all_reduce, broadcast, recv_tensor, send_tensor, groupcall, send_object, recv_object, reduce_scatter

bmtrain/distributed/dtype.py

@@ -0,0 +1,12 @@
+import torch
+DTYPE_LIST = [
+    torch.float64,
+    torch.float32,
+    torch.float16,
+    torch.int64,
+    torch.int32,
+    torch.int16,
+    torch.int8,
+    torch.bfloat16,
+    torch.bool
+]

bmtrain/distributed/ops.py

@@ -1,48 +1,13 @@
 import torch
-from ..global_var import config
-from ..nccl import allGather as ncclAllGather, recv
+import bmtrain as bmt
+from ..global_var import config, rank
+from ..nccl import allGather as ncclAllGather
 from ..nccl import allReduce as ncclAllReduce
 from ..nccl import broadcast as ncclBroadcast
 from ..nccl import reduceScatter as ncclReduceScatter
-from ..nccl import send as ncclSend
-from ..nccl import recv as ncclRecv
-from ..nccl import commCount,commRank,NCCLCommunicator
-DTYPE_LIST = [
-    torch.float64,
-    torch.float32,
-    torch.float16,
-    torch.int64,
-    torch.int32,
-    torch.int16,
-    torch.int8,
-    torch.bfloat16,
-    torch.bool
-]
-def send_activations(hidden_state, next_rank, comm):
-    send_meta(hidden_state, next_rank, comm)
-    ncclSend(hidden_state.storage(), next_rank, comm)
-
-def recv_activations(prev_rank, comm):
-    dtype, shape = recv_meta(prev_rank, comm)
-    hidden_state = torch.empty(shape, dtype=dtype, device="cuda")
-    ncclRecv(hidden_state.storage(), prev_rank, comm)
-    return hidden_state
-
-def send_meta(x, next_rank, comm):
-    meta_data = torch.tensor(data=[0]*50, device="cuda", dtype=torch.int)
-    meta_data[0] = len(x.size())
-    meta_data[1] = DTYPE_LIST.index(x.dtype)
-    meta_data[2:len(x.size())+2] = torch.tensor(x.size(), device="cuda", dtype=torch.int)
-    meta_data = meta_data.contiguous()
-    ncclSend(meta_data.storage(), next_rank, comm)
-
-def recv_meta(prev_rank, comm):
-    meta_data = torch.tensor(data=[0]*50, device="cuda", dtype=torch.int)
-    ncclRecv(meta_data.storage(), prev_rank, comm)
-    n_dims = meta_data[0].item()
-    dtype = DTYPE_LIST[meta_data[1].item()]
-    shape = meta_data[2:n_dims+2].tolist()
-    return dtype,shape
+from ..nccl import commCount, commRank, NCCLCommunicator, groupStart, groupEnd
+from .p2p_ops import *
+
 
 class OpBroadcast(torch.autograd.Function):
 

bmtrain/distributed/p2p_ops.py

@@ -0,0 +1,159 @@
+import torch
+from bmtrain import config
+from ..nccl import reduceScatter as ncclReduceScatter
+from ..nccl import send as ncclSend
+from ..nccl import recv as ncclRecv
+from ..nccl import groupStart,groupEnd
+from .dtype import DTYPE_LIST
+import pickle
+import contextlib
+
+_p2p_stream = {}
+_p2p_events = {}
+
+@contextlib.contextmanager
+def groupcall():
+    groupStart()
+    yield
+    groupEnd()
+class handler:
+    def __init__(self, event):
+        self.event= event
+
+    def wait(self):
+        torch.cuda.current_stream().wait_event(self.event)
+
+def send_object(obj, peer_rank, comm):
+    data_bytes: bytes = pickle.dumps(obj)
+    data_length: int = len(data_bytes)
+
+    gpu_data_length = torch.tensor([data_length], device="cuda", dtype=torch.long)
+    ncclSend(gpu_data_length.storage(), peer_rank, comm)
+    byte_storage = torch.ByteStorage.from_buffer(data_bytes).cuda()
+    ncclSend(byte_storage, peer_rank, comm)
+
+def recv_object(peer_rank, comm):
+    data_length = torch.tensor([0], device="cuda", dtype=torch.long)
+    ncclRecv(data_length.storage(), peer_rank, comm)
+    data_bytes_stor = torch.cuda.ByteStorage(data_length.item())
+    ncclRecv(data_bytes_stor, peer_rank, comm)
+    tensor = torch.ByteTensor(data_bytes_stor.cpu())
+    data = pickle.loads(tensor.numpy().tobytes())
+    return data
+
+def record_stream_helper(tensor_list, stream):
+    for t in tensor_list:
+        t.record_stream(stream)
+
+def send_tensors(tensor_list, peer_rank, comm):
+    handler = _send_tensors(tensor_list, peer_rank, comm)
+    handler.wait()
+
+def isend_tensor(tensor_list, peer_rank, comm):
+    return _send_tensors(tensor_list, peer_rank, comm)
+
+def _send_tensors(tensor_list, peer_rank, comm):
+    p2p_key = f"send {peer_rank}"
+    if p2p_key not in _p2p_stream:
+        _p2p_stream[p2p_key] = torch.cuda.Stream()
+    if p2p_key not in _p2p_events: 
+        _p2p_events[p2p_key] = torch.cuda.Event()
+    stream = _p2p_stream[p2p_key]
+    event = _p2p_events[p2p_key]
+    event.record(torch.cuda.current_stream()) 
+    stream.wait_event(event)
+    with torch.cuda.stream(stream):
+        length = torch.tensor(data=[len([h for h in tensor_list ])], device="cuda", dtype=torch.int)
+        flags = torch.tensor(data=[0 for _ in range(len(tensor_list))], device="cuda",dtype=torch.int)
+        for i in range(len(tensor_list)):
+            if tensor_list[i] is None:
+                flag = -1
+            elif torch.is_tensor(tensor_list[i]):
+                flag = 0
+            else:
+                flag = 1
+            flags[i] = flag
+        ncclSend(length.storage(), peer_rank, comm)
+        ncclSend(flags.contiguous().storage(), peer_rank, comm)
+        for i in range(len(tensor_list)):
+            if flags[i] == 0:
+                tensor_list[i].record_stream(stream)
+                send_tensor(tensor_list[i], peer_rank, comm)
+            elif flags[i] == 1:
+                send_object(tensor_list[i], peer_rank, comm)
+        event.record(stream)
+    return handler(event)
+
+def recv_tensors(peer_rank, comm):
+    tensors, handle = _recv_tensors(peer_rank, comm)
+    handle.wait()
+    return tensors
+
+def irecv_tensors(peer_rank, comm):
+    tensors, handle = _recv_tensors(peer_rank, comm)
+    return tensors, handle
+
+def _recv_tensors(peer_rank, comm):
+    p2p_key = f"recv {peer_rank}"
+    if p2p_key not in _p2p_stream:
+        _p2p_stream[p2p_key] = torch.cuda.Stream()
+    if p2p_key not in _p2p_events:
+        _p2p_events[p2p_key] = torch.cuda.Event()
+    stream = _p2p_stream[p2p_key]
+    event = _p2p_events[p2p_key]
+    with torch.cuda.stream(stream):
+        length = torch.tensor(data=[0], device="cuda", dtype=torch.int)
+        tensor_list = []
+        ncclRecv(length.storage(), peer_rank, comm)
+        flags = torch.tensor(data=[0 for _ in range(length)], device="cuda",dtype=torch.int)
+        ncclRecv(flags.storage(), peer_rank, comm)
+        for i in range(length[0].item()):
+            flag = flags[i].item()
+            if flag == -1:
+                tensor_list.append(None)
+            elif flag == 0:
+                recv = recv_tensor(peer_rank, comm)
+                tensor_list.append(recv)
+            elif flag == 1:
+                recv = recv_object(peer_rank, comm)
+                tensor_list.append(recv)
+    event.record(stream)
+    record_stream_helper([tensor_list[i] for i in range(length[0].item()) if flags[i].item() != -1], torch.cuda.current_stream())
+    return tensor_list, handler(event)
+
+def send_tensor(hidden_state, peer_rank, comm):
+    hidden_state = hidden_state.contiguous()
+    send_meta(hidden_state, peer_rank, comm)
+    ncclSend(hidden_state.storage(), peer_rank, comm)
+
+def send_tensor_inplace(hidden_state, peer_rank, comm):
+    hidden_state = hidden_state.contiguous()
+    ncclSend(hidden_state.storage(), peer_rank, comm)
+
+def recv_tensor_inplace(hidden_state, peer_rank, comm):
+    hidden_state = hidden_state.contiguous()
+    ncclRecv(hidden_state.storage(), peer_rank, comm)
+    return hidden_state
+
+def recv_tensor(peer_rank, comm):
+    dtype, shape = recv_meta(peer_rank, comm)
+    hidden_state = torch.empty(shape, dtype=dtype, device="cuda")
+    ncclRecv(hidden_state.storage(), peer_rank, comm)
+    return hidden_state
+
+def send_meta(x, peer_rank, comm):
+    meta_data = torch.tensor(data=[0]*50, device="cuda", dtype=torch.int)
+    meta_data[0] = len(x.size())
+    meta_data[1] = DTYPE_LIST.index(x.dtype)
+    meta_data[2:len(x.size())+2] = torch.tensor(x.size(), device="cuda", dtype=torch.int)
+    meta_data = meta_data.contiguous()
+    ncclSend(meta_data.storage(), peer_rank, comm)
+
+def recv_meta(peer_rank, comm):
+    meta_data = torch.tensor(data=[0]*50, device="cuda", dtype=torch.int)
+    ncclRecv(meta_data.storage(), peer_rank, comm)
+    n_dims = meta_data[0].item()
+    dtype = DTYPE_LIST[meta_data[1].item()]
+    shape = meta_data[2:n_dims+2].tolist()
+
+    return dtype,shape

bmtrain/hook_func.py

@@ -4,51 +4,75 @@
 
 def zero_pre_forward(module, inputs):
     enter = True
-    pipe = False
-    if module._mode == "PIPE":
-        enter = module._micro_idx == 0
-        pipe = True
+    if module._mode == "PIPE" or module._mode == "1F1B":
+        if not hasattr(module, "_micro_forward_idx") or module._micro_forward_idx == -1:
+            module._micro_forward_idx = 0
+            enter = True
+        else:
+            enter = False
+            module._micro_forward_idx += 1
     if enter:
         zero_level = module._zero_level 
         forward_flag = 1 if zero_level == 2 else 0
         if zero_level == 2 and not module._need_release:
             forward_flag = 2 # repeating forward in same layer
         if module.all_param_no_grad: #only forward
             forward_flag = 0
-        module._forward_block_ctx = ZeroContext(module, module._layer_dict, pipe=pipe)
-        module._forward_block_ctx.enter(forward_flag)
+        if module._mode == "1F1B":
+            module._block_ctx = ZeroContext(module, module._layer_dict)
+            module._block_ctx.enter(0, requires_grad=True)
+        else:
+            module._forward_block_ctx = ZeroContext(module, module._layer_dict)
+            module._forward_block_ctx.enter(forward_flag)
 
 def zero_post_forward(module, inputs, outputs):
     forward_flag = 1 if module._zero_level == 2 else 0
     if module.all_param_no_grad:
         forward_flag = 0
     exit = True
-    if module._mode == "PIPE":
-        exit = module._micro_idx == config['micros'] - 1
+    if module._mode == "PIPE" or module._mode == "1F1B":
+        if module._micro_forward_idx == config["micros"] - 1:
+            module._micro_forward_idx = -1
+            if module._mode == "1F1B":
+                exit = False
+            else:
+                exit = True
+        else:
+            exit = False
 
     if exit:
         module._forward_block_ctx.exit(forward_flag)
 
 def zero_pre_backward(module, grad_outputs):
     backward_flag = 2 if module._zero_level == 2 else 0
-    if module._mode != "PIPE":
+    if module._mode != "PIPE" and module._mode != "1F1B":
         module._backward_block_ctx = ZeroContext(module, module._layer_dict)
         module._backward_block_ctx.enter(backward_flag, True)
         module.release_next_module(backward_flag)
     else:
-        if module._micro_idx == config['micros'] - 1:
-            module._backward_block_ctx = ZeroContext(module, module._layer_dict, pipe=True)
-            module._backward_block_ctx.enter(backward_flag, True)
+        if not hasattr(module, "_micro_backward_idx") or module._micro_backward_idx == -1:
+            if module._mode == "1F1B":
+                module._micro_backward_idx = 0
+            else:
+                module._micro_backward_idx = 0
+                module._backward_block_ctx = ZeroContext(module, module._layer_dict)
+                module._backward_block_ctx.enter(backward_flag,requires_grad=True)
+        else:
+            module._micro_backward_idx += 1
 
 def zero_post_backward(module, grad_inputs, grad_outputs):
     backward_flag = 2 if module._zero_level == 2 else 0
-    if module._mode != "PIPE":
+    if module._mode != "PIPE" and module._mode != "1F1B":
         if module._is_first_layer: 
             module.release(backward_flag)
     else:
-        if module._micro_idx == 0:
-            module.release(backward_flag)
-        module._micro_idx -= 1
+        if module._micro_backward_idx == config["micros"] - 1:
+            if module._mode == "1F1B":
+                module._block_ctx.exit(0, backward=True)
+                config['load_stream'].record_event(config['load_event'])
+            else:
+                module.release(backward_flag)
+            module._micro_backward_idx = -1
 
 class OneStepNoGradFunc(torch.autograd.Function):
     """