Make block can accept kwargs

Author: MayDomine

bmtrain/block_layer.py

@@ -311,9 +311,13 @@ def post_hook(self, out):
         post_out = tuple(post_out)
         return post_out
 
-    def forward(self, *args):
+    def forward(self, *args, **kwargs):
+        signature = inspect.signature(self._module.forward)
+        bound_args = signature.bind(*args, **kwargs)
+        args = bound_args.args
         arg_list = self.pre_hook(*args)
 
+
         if self.all_input_no_grad and not self.all_param_no_grad:
             placeholder = torch.tensor([], requires_grad=torch.is_grad_enabled())
             return hook_func.OneStepNoGradFunc.apply(self, placeholder, *arg_list)

bmtrain/wrapper.py

@@ -15,11 +15,15 @@ def make_distributed(model: torch.nn.Module):
     for kw in list(model._buffers.keys()):
         if model._buffers[kw] is not None:
             model._buffers[kw] = model._buffers[kw].cuda()
-
+    is_module_list = isinstance(model, torch.nn.ModuleList)
+    pre_module = None
     for kw in list(model._modules.keys()):
-        if isinstance(model, torch.nn.ModuleList):
+        if is_module_list:
             if not isinstance(model._modules[kw], Block):
                 model._modules[kw] = Block(model_wrapper_dispatch(model._modules[kw]))
+                if pre_module is not None:
+                    model._modules[kw].set_pre_module(pre_module)
+                pre_module = model._modules[kw]
         else:
             model._modules[kw] = model_wrapper_dispatch(model._modules[kw])