improve

Author: li126com

internlm/model/modules/linear.py

@@ -349,9 +349,12 @@ def __init__(
         else:
             super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype)
 
-    def forward(self, input: torch.Tensor, no_communication=False) -> torch.Tensor:  # pylint: disable=W0622
+        self.last_block_layer = False
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:  # pylint: disable=W0622
         _class_name = self.__class__.__name__
         assert self._communicator is not None, f"{_class_name} should register with a communicator first."
+        no_communication = bool(gpc.recompute_forward_no_comm and self.last_block_layer)
         return fused_dense_func(
             input,
             self.weight,
@@ -414,6 +417,7 @@ def __init__(
         multiple_of: int = 1,
         device: torch.device = None,
         dtype: torch.dtype = None,
+        layer_name: str = "default",
     ) -> None:
         if in_features % multiple_of:
             raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}")
@@ -430,6 +434,9 @@ def __init__(
             split_mode="row",
         )
 
+        if layer_name == "w2":
+            self.last_block_layer = True
+
 
 class ScaleColumnParallelLinear(ParallelLinearWithCommExt):
     """
@@ -471,7 +478,7 @@ def __init__(
         self.first_eval_flag = True
         self.tmp_weight = None
 
-    def forward(self, input, no_communication=False):  # pylint: disable=W0622
+    def forward(self, input):  # pylint: disable=W0622
         _class_name = self.__class__.__name__
         assert self._communicator is not None, f"{_class_name} should register with a communicator first."
 
@@ -502,7 +509,6 @@ def forward(self, input, no_communication=False):  # pylint: disable=W0622
             communicator=self._communicator,
             module=self,
             bias=self.bias,
-            no_communication=no_communication,
         )
 
 
@@ -603,6 +609,7 @@ def new_linear(
             multiple_of,
             device,
             dtype,
+            layer_name=name,
         )
     else:
         err_msg = (

internlm/model/modules/mlp.py

@@ -6,7 +6,6 @@
 import torch
 from torch import nn
 
-from internlm.core.context.parallel_context import global_context as gpc
 from internlm.model.modules.linear import new_linear
 from internlm.model.modules.utils import Silu
 from internlm.utils.logger import get_logger
@@ -99,7 +98,7 @@ def forward(self, x):
         else:
             fussed_out = self.fused_w1_w3(x)
             w1_o, w3_o = torch.split(fussed_out, fussed_out.shape[-1] // 2, dim=-1)
-        out = self.w2(Silu(w1_o, w3_o), no_communication=gpc.recompute_forward_no_comm)
+        out = self.w2(Silu(w1_o, w3_o))
         return out
 
 

internlm/model/utils.py

@@ -4,7 +4,6 @@
 
 from internlm.core.context import ParallelMode
 from internlm.core.context.parallel_context import global_context as gpc
-from internlm.core.parallel.comm.tensor import _GATHER_DIM
 from internlm.model.modules.mha import MHA
 
 
@@ -60,21 +59,16 @@ def convert_attn_args_to_kwargs(args, kwargs) -> Dict[str, Any]:
 
 def padding_residual(residual):
     requires_grad = residual.requires_grad
-    pad_before = gpc.get_local_rank(ParallelMode.TENSOR) * residual.shape[_GATHER_DIM]
-    pad_after = (
-        gpc.get_world_size(ParallelMode.TENSOR) - gpc.get_local_rank(ParallelMode.TENSOR) - 1
-    ) * residual.shape[_GATHER_DIM]
-
-    pad_before_tensor = torch.zeros(
-        (*residual.shape[:_GATHER_DIM], pad_before, *residual.shape[_GATHER_DIM + 1 :]),
-        dtype=residual.dtype,
-        device=residual.device,
-    )
-    pad_after_tensor = torch.zeros(
-        (*residual.shape[:_GATHER_DIM], pad_after, *residual.shape[_GATHER_DIM + 1 :]),
+    _GATHER_DIM = 1
+    total_size = gpc.get_world_size(ParallelMode.TENSOR) * residual.shape[_GATHER_DIM]
+    zero_padding_tensor = torch.zeros(
+        (*residual.shape[:_GATHER_DIM], total_size, *residual.shape[_GATHER_DIM + 1 :]),
         dtype=residual.dtype,
         device=residual.device,
     )
-    residual = torch.cat([pad_before_tensor, residual, pad_after_tensor], dim=1).requires_grad_(requires_grad)
+    start_idx = gpc.get_local_rank(ParallelMode.TENSOR) * residual.shape[_GATHER_DIM]
+    end_idx = start_idx + residual.shape[_GATHER_DIM]
+    zero_padding_tensor[:, start_idx:end_idx, :] = residual
+    residual = zero_padding_tensor.requires_grad_(requires_grad)
 
     return residual

internlm/solver/activation_checkpoint.py

@@ -151,7 +151,7 @@ def backward(ctx, *args):
         for i, idx in enumerate(tensor_indices):
             inputs[idx] = tensors[i]
 
-        # no_communication
+        # when checkpoint_tp_no_comm==True, we use TP recomputation communication optimization
         no_communication = getattr(gpc.config.model, "checkpoint_tp_no_comm", False)
 
         detached_inputs = detach_variable(tuple(inputs))