Offload activation async support

.gitignore

@@ -2,7 +2,7 @@
 __pycache__/
 *.py[cod]
 *$py.class
-
+*nsys-rep
 # C extensions
 *.so
 
@@ -150,4 +150,4 @@ log
 .vscode
 
 !bmtrain/dist
-tests/test_log.txt
+tests/test_log.txt

bmtrain/block_layer.py

@@ -66,7 +66,7 @@ class CheckpointBlock(torch.nn.Module):
         >>> y2, ... = transformer_block(x)
         >>> assert torch.allclose(y1, y2)
     """
-    def __init__(self, inner_module : torch.nn.Module, use_checkpoint=True, zero_level=3):
+    def __init__(self, inner_module : torch.nn.Module, use_checkpoint=True, use_offload=False, offload_level=0, zero_level=3):
         super().__init__()
         self._module = inner_module
         self._inputs = None
@@ -78,9 +78,9 @@ def __init__(self, inner_module : torch.nn.Module, use_checkpoint=True, zero_lev
         self._storage_params : Dict[str, torch.nn.Parameter] = {}
         self._storage_info = {}
         self._ready = False
-        # sort parameters by name
+        # sort parameters by nam_next_modulee
         ordered_parameters = list(self._module.named_parameters())
-
+        assert not (use_checkpoint and use_offload), "It does not make sense to use offload and checkpointing at the same time" 
         # calc total number of parameters
         for name, param in ordered_parameters:
             if not isinstance(param, DistributedParameter):
@@ -202,6 +202,11 @@ def __init__(self, inner_module : torch.nn.Module, use_checkpoint=True, zero_lev
         self._pre_module = [] #save the pre module of self
         self._ref_count = 0 #incremental in forward and  decreasing in backward
         self._mode = "BLOCK" #BLOCK or ZERO or PIPE
+        if use_offload and offload_level != 0:
+            self._mode = "OFFLOAD"
+            self._on_device = False
+            self.offload_level = offload_level
+
         self.all_input_no_grad = False
         self.all_param_no_grad = False
         self._zero_level = zero_level
@@ -536,19 +541,23 @@ def __init__(self, modules: Iterable[CheckpointBlock], num_hidden=1, sqrt=False)
 
         self._modules = {}
         pre_module = None
+        offload = 0
         for i, module in enumerate(modules):
             if not isinstance(module, CheckpointBlock):
                 module = CheckpointBlock(module)
-
-            module._mode = "ZERO"
+            module._mode = "ZERO" if module._mode == "BLOCK" else module._mode
             module.set_pre_module(pre_module)
             pre_module = module
             module._is_first_layer = False
             module._is_last_layer = False
-
+            if module._mode == "OFFLOAD":
+                offload+=1
+                module.calc_event = torch.cuda.Event()
+                module.offload_event = torch.cuda.Event()
             self._modules[str(i)] = module
+            module._idx = i
             self.add_module(str(i), module)
-
+        print(f"offload layer: {offload}")
         self._modules[str(0)]._is_first_layer = True
         self._modules[str(len(modules)-1)]._is_last_layer = True
 
@@ -575,7 +584,7 @@ def __init__(self, modules: Iterable[CheckpointBlock], num_hidden=1, sqrt=False)
             self.save_list = save_list
         else:
             self.save_list = [(i, i) for i in range(len(self))]
-
+             
     def __len__(self) -> int:
         return len(self._modules)
 

bmtrain/hook_func.py

@@ -1,10 +1,157 @@
 import torch
 from .global_var import config
 from .checkpointing import CheckpointBlockContext
+from .distributed import all_gather, broadcast, all_reduce, send_activations, recv_activations 
+from collections import deque,OrderedDict
+from contextlib import contextmanager
+from .utils import round_up
+
+class Offload_Dict:
+
+    def __init__(self):
+        self._offload_dict = OrderedDict()
+
+    def add(self, tensor):
+        tensor = tensor.contiguous()
+        tensor_id = id(tensor)
+        data_ptr = tensor.storage().data_ptr()
+        if data_ptr not in self._offload_dict:
+            self._offload_dict[data_ptr] = {}
+            self._offload_dict[data_ptr]["stor"] = tensor.storage()
+            self._offload_dict[data_ptr]["size"] = tensor.storage().size()
+            self._offload_dict[data_ptr]["dtype"] = tensor.storage().dtype
+            self._offload_dict[data_ptr]["tensors"] = {}
+
+        self._offload_dict[data_ptr]["tensors"][id(tensor)] = {}
+        self._offload_dict[data_ptr]["tensors"][id(tensor)]["numel"] = tensor.numel()
+        self._offload_dict[data_ptr]["tensors"][id(tensor)]['dtype'] = tensor.dtype
+        self._offload_dict[data_ptr]["tensors"][id(tensor)]['offset'] = tensor.storage_offset()
+        self._offload_dict[data_ptr]["tensors"][id(tensor)]['tensor'] = tensor
+        self._offload_dict[data_ptr]["tensors"][id(tensor)]["shape"] = tensor.shape
+        self._device = "cuda"
+        return (data_ptr,tensor_id)
+
+    def get_total(self):
+        fp16_total = sum([v['size'] for v in self._offload_dict.values() if v['dtype'] == torch.float16])
+        fp32_total = sum([v['size'] for v in self._offload_dict.values() if v['dtype'] == torch.float32])        
+        return fp16_total,fp32_total
+
+    def make_cpu_storage(self):
+        fp16_total = sum([v['size'] for v in self._offload_dict.values() if v['dtype'] == torch.float16])
+        fp32_total = sum([v['size'] for v in self._offload_dict.values() if v['dtype'] == torch.float32])
+        fp16_storage = torch.HalfStorage(fp16_total).pin_memory()
+        fp32_storage = torch.FloatStorage(fp32_total).pin_memory()
+        self.fp16_storage = fp16_storage
+        self.fp32_storage = fp32_storage 
+        self.fp16_total = fp16_total
+        self.fp32_total = fp32_total
+
+    def get(self, key):
+        data_ptr, tensor_id = key
+        return self._offload_dict[data_ptr]['tensors'][tensor_id]["tensor"]
+
+    def pop_all(self):
+        self._offload_dict.clear()
+
+    def h2d_memcpy(self):
+        fp16_storage_cuda = self.fp16_storage.cuda(non_blocking=True)
+        fp32_storage_cuda = self.fp32_storage.cuda(non_blocking=True)
+        for key,val in self._offload_dict.items():
+            for id_val in val['tensors'].values():
+                id_val['tensor'] = torch.tensor([], dtype=id_val['dtype'],device=fp16_storage_cuda.device)
+                if id_val['dtype'] == torch.float16:
+                    id_val['tensor'].set_(fp16_storage_cuda, id_val['abs_offset'], id_val['shape'])
+                elif id_val['dtype'] == torch.float32:
+                    id_val['tensor'].set_(fp32_storage_cuda, id_val['abs_offset'], id_val['shape'])
+
+    def record_stream(self, stream):
+        for key, val in self._offload_dict.items():
+            for id_val in val['tensors'].values():
+                id_val['tensor'].record_stream(stream)
+
+    def d2h_memcpy(self):   
+        fp16_offset = 0
+        fp32_offset = 0
+        fp16_total = sum([v['size'] for v in self._offload_dict.values() if v['dtype'] == torch.float16])
+        fp32_total = sum([v['size'] for v in self._offload_dict.values() if v['dtype'] == torch.float32])
+        assert fp16_total <= self.fp16_total
+        assert fp32_total <= self.fp32_total
+        fp16_storage = self.fp16_storage
+        fp32_storage = self.fp32_storage
+        for key,val in self._offload_dict.items():
+            assert val['dtype'] in [torch.float16, torch.float32]
+            storage = fp16_storage if val['dtype'] == torch.float16 else fp32_storage
+            offset = fp16_offset if val['dtype'] == torch.float16 else fp32_offset
+            for id_val in val['tensors'].values():
+                cpu_tensor = torch.tensor([], dtype=id_val['dtype'], device="cpu") \
+                    .set_(storage, offset+id_val['offset'], id_val['shape'])
+                id_val["abs_offset"] = offset+id_val['offset']
+                id_val['tensor'] = cpu_tensor.copy_(id_val['tensor'], non_blocking=True)
+            if val['dtype'] == torch.float16:
+                fp16_offset += val['size']
+            else:
+                fp32_offset += val['size']
+            val['stor'] = None
+
+def find_pre_module_helper(m):
+    if len(m) == 0:
+        return None
+    if m._mode == "OFFLOAD":
+        return m
+    else:
+        return find_pre_module_helper(m.pre_module())
+
+def offload_wrapper(offload_dict):
+    def pack_hook(tensor):
+        if isinstance(tensor, torch.nn.Parameter):
+            return (tensor,) 
+        elif tensor.dtype not in [torch.float16]:
+            return (tensor,)
+        else:
+            key = offload_dict.add(tensor)
+            return (tensor.device, key)
+    def unpack_hook(packed):
+        if len(packed) == 2:
+            device, key = packed
+            tensor = offload_dict.get(key)
+            assert tensor.device == device
+            return tensor
+        else:
+            tensor, = packed
+            return tensor
+    return pack_hook, unpack_hook
+
+def offload_pre_hook(module, input):
+   if hasattr(module, "_offload_hook"):
+        pack_hook, unpack_hook = module._offload_hook
+        torch._C._autograd._push_saved_tensors_default_hooks(
+            pack_hook, unpack_hook
+        ) 
+
+def offload_post_hook(module, input, output):
+    if hasattr(module, "_offload_hook"):
+        torch._C._autograd._pop_saved_tensors_default_hooks()
 
 def zero_pre_forward(module, inputs):
     enter = True
     pipe = False
+    if module._mode == "OFFLOAD":
+        if not hasattr(module, "_offload_dict"):
+            module._offload_dict = Offload_Dict()
+        pack_hook, unpack_hook = offload_wrapper(module._offload_dict)
+        if module.offload_level == 1:
+            for n, m in module.named_modules():
+                if m.__class__.__name__ == "Linear" and not hasattr(m, "_offload_hook"):
+                    m._offload_hook = (pack_hook, unpack_hook)
+                    m.register_forward_pre_hook(offload_pre_hook)
+                    m.register_forward_hook(offload_post_hook)
+        elif module.offload_level == 2:
+            if not hasattr(module, "_offload_hook"):
+                module._offload_hook = (pack_hook, unpack_hook)
+            torch._C._autograd._push_saved_tensors_default_hooks(
+                pack_hook, unpack_hook
+            )
+
     if module._mode == "PIPE":
         enter = module._micro_idx == 0
         pipe = True
@@ -25,14 +172,42 @@ def zero_post_forward(module, inputs, outputs):
     exit = True
     if module._mode == "PIPE":
         exit = module._micro_idx == config['micros'] - 1
-
+    elif module._mode == "OFFLOAD":
+        torch.cuda.current_stream().record_event(module.calc_event)
+        pre_offload_module = find_pre_module_helper(module.pre_module())
+        if pre_offload_module is not None:
+            torch.cuda.current_stream().wait_event(pre_offload_module.offload_event)
+        with torch.cuda.stream(config["offload_stream"]):
+            config["offload_stream"].wait_event(module.calc_event)
+            if not hasattr(module._offload_dict, "fp16_storage"):
+                module._offload_dict.make_cpu_storage()
+            module._offload_dict.record_stream(config["offload_stream"])
+            module._offload_dict.d2h_memcpy()
+            if len(module._next_module) > 0:
+                config["offload_stream"].record_event(module.offload_event)
+        if module.offload_level == 2:
+            torch._C._autograd._pop_saved_tensors_default_hooks()
     if exit:
         module._forward_block_ctx.exit(forward_flag)
         module._ref_count += 1
 
 def zero_pre_backward(module, grad_outputs):
     backward_flag = 2 if module._zero_level == 2 else 0
     if module._mode != "PIPE":
+        if module._mode == "OFFLOAD" or (len(module._next_module) == 0):
+            if len(module._next_module) != 0:
+                current_stream = torch.cuda.current_stream()
+                current_stream.wait_event(module.offload_event)
+            pre_module = find_pre_module_helper(module.pre_module())
+            if pre_module is not None:
+                pre_module._on_device = True
+                with torch.cuda.stream(config["offload_stream"]):
+                    if (len(module._next_module) != 0):
+                        torch.cuda.current_stream().wait_event(module.calc_event)
+                    pre_module._offload_dict.h2d_memcpy()
+                    torch.cuda.current_stream().record_event(pre_module.offload_event)
+            if (len(module._next_module) != 0):
+                module._offload_dict.record_stream(current_stream)
         module._backward_block_ctx = CheckpointBlockContext(module, module._layer_dict)
         module._backward_block_ctx.enter(backward_flag, True)
         if not module._is_last_layer: 
@@ -45,6 +220,10 @@ def zero_pre_backward(module, grad_outputs):
 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 == "OFFLOAD":
+            module._on_device = False
+            module._offload_dict.pop_all()
+            torch.cuda.current_stream().record_event(module.calc_event)
         if module._is_first_layer: 
             module.backward_release(backward_flag)
     else:

bmtrain/init.py

@@ -74,6 +74,7 @@ def init_distributed(
     config["calc_stream"] = torch.cuda.current_stream()
     config["load_stream"] = torch.cuda.Stream(priority=-1)
     config["tp_comm_stream"] = torch.cuda.Stream(priority=-1)
+    config["offload_stream"] = torch.cuda.Stream()
     config["pp_comm_stream"] = torch.cuda.Stream(priority=-1)
     config['barrier_stream'] = torch.cuda.Stream()
     config["load_event"] = torch.cuda.Event()

bmtrain/nn/__init__.py

@@ -3,4 +3,4 @@
 from .row_parallel_linear import RowParallelLinear
 from .parallel_embedding import ParallelEmbedding 
 from .parallel_cross_entropy_func import parallel_cross_entropy_func
-from .parallel_linear_func import OpParallelLinear
+from .parallel_linear_func import OpParallelLinear

example/layers/attention.py

@@ -32,6 +32,7 @@ def __init__(self,
 
         self.softmax = torch.nn.Softmax(dim=-1)
         self.num_heads = num_heads
+        self.num_kv_heads = num_heads
         self.dim_head = dim_head
         self.dim_model = dim_model
 
@@ -44,18 +45,20 @@ def forward(self,
         batch_size, seq_q, dim_model = hidden_q.size()
         seq_kv = hidden_kv.size(1)
 
-        assert hidden_q.data_ptr() == hidden_kv.data_ptr()
-
-        hidden_q = bmt.nn.OpParallelLinear.apply(
-			hidden_q,
-			torch.cat([self.project_q.weight, self.project_k.weight, self.project_v.weight], dim=0),
-			torch.cat([self.project_q.bias, self.project_k.bias, self.project_v.bias], dim=0),
-			True, False,
-			False, None
-		)
-
-        h_q, h_k, h_v = hidden_q.chunk(3, dim=-1)
-
+        if isinstance(self.project_q, ColumnParallelLinear):
+            assert hidden_q.data_ptr() == hidden_kv.data_ptr()
+            hidden_q = bmt.nn.OpParallelLinear.apply(
+                hidden_q,
+                torch.cat([self.project_q.weight, self.project_k.weight, self.project_v.weight], dim=0),
+                torch.cat([self.project_q.bias, self.project_k.bias, self.project_v.bias], dim=0) if self.project_q.bias is not None else None,
+                True, False,
+                False, None
+            )
+            h_q, h_k, h_v = hidden_q.chunk(3, dim=-1)
+        else:
+            h_q : torch.Tensor = self.project_q(hidden_q)
+            h_k : torch.Tensor = self.project_k(hidden_q)
+            h_v : torch.Tensor = self.project_v(hidden_q)
         if config['tp_size'] > 1:
             #batch_size  will changed in TensorParallel
             batch_size = h_v.shape[0]