Support DualPipeV for DeepSeek

llm/config/deepseek-v3/pretrain_argument.json

@@ -4,11 +4,11 @@
   "input_dir": "./data",
   "output_dir": "./checkpoints/pretrain_ckpts",
   "per_device_train_batch_size": 1,
-  "gradient_accumulation_steps": 4,
+  "gradient_accumulation_steps": 10,
   "per_device_eval_batch_size": 1,
   "tensor_parallel_degree": 1,
-  "pipeline_parallel_degree": 1,
-  "sharding_parallel_degree": 8,
+  "pipeline_parallel_degree": 2,
+  "sharding_parallel_degree": 4,
   "expert_parallel_degree": 4,
   "sharding": "stage1",
   "virtual_pp_degree": 1,
@@ -33,9 +33,10 @@
   "do_eval": true,
   "do_predict": false,
   "disable_tqdm": true,
-  "recompute": true,
+  "recompute": false,
   "distributed_dataloader": 1,
   "recompute_granularity": "full",
   "unified_checkpoint": true,
-  "save_total_limit": 2
-}
+  "save_total_limit": 2,
+  "pipeline_parallel_config": "use_dualpipev"
+}

paddlenlp/transformers/deepseek_v2/configuration.py

@@ -179,6 +179,7 @@
         attention_dropout=0.0,
         speculate_model_type=False,
         using_flex_token=False,
+        use_dualpipev=False,
         **kwargs,
     ):
         self.vocab_size = vocab_size
@@ -227,6 +228,7 @@
         self.speculate_model_type = speculate_model_type
         self.use_fp8 = False
         self.using_flex_token = using_flex_token
+        self.use_dualpipev = use_dualpipev
 
         super().__init__(
             pad_token_id=pad_token_id,

paddlenlp/transformers/deepseek_v2/modeling.py

@@ -130,7 +130,7 @@
     return assignment_list
 
 
-def parallel_matmul(x: Tensor, y: Tensor, tensor_parallel_output=True):
+def parallel_matmul(x: Tensor, y: Tensor, transpose_y=False, tensor_parallel_output=True):
     is_fleet_init = True
     tensor_parallel_degree = 1
     try:
@@ -148,15 +148,15 @@
     if is_fleet_init and tensor_parallel_degree > 1 and y_is_distributed:
         # if not running under distributed.launch, it will raise AttributeError: 'Fleet' object has no attribute '_hcg'
         input_parallel = paddle.distributed.collective._c_identity(x, group=model_parallel_group)
-        logits = paddle.matmul(input_parallel, y, transpose_y=False)
+        logits = paddle.matmul(input_parallel, y, transpose_y=transpose_y)
 
         if tensor_parallel_output:
             return logits
 
         return paddle.distributed.collective._c_concat(logits, group=model_parallel_group)
 
     else:
-        logits = paddle.matmul(x, y, transpose_y=False)
+        logits = paddle.matmul(x, y, transpose_y=transpose_y)
         return logits
 
 
@@ -826,6 +826,10 @@
 
     def forward(self, hidden_states):
         final_hidden_states, l_aux, l_zloss = super().forward(hidden_states)
+        final_hidden_states = self.auxilibaryloss_and_shared_expert_compute(hidden_states, final_hidden_states, l_aux)
+        return final_hidden_states
+
+    def auxilibaryloss_and_shared_expert_compute(self, hidden_states, final_hidden_states, l_aux):
         if self.training and self.alpha > 0.0:
             l_aux = l_aux * self.alpha
             final_hidden_states = AddAuxiliaryLoss.apply(final_hidden_states, l_aux)
@@ -1145,6 +1149,48 @@
         self.input_layernorm = DeepseekV2RMSNorm(config)
         self.post_attention_layernorm = DeepseekV2RMSNorm(config)
 
+    def self_attn_and_gate_compute(
+        self,
+        hidden_states: paddle.Tensor,
+        position_ids: Optional[paddle.Tensor] = None,
+        attention_mask: Optional[paddle.Tensor] = None,
+        output_attentions: Optional[bool] = False,
+        past_key_value: Optional[Tuple[paddle.Tensor]] = None,
+        use_cache: Optional[bool] = False,
+        attn_mask_startend_row_indices: Optional[paddle.Tensor] = None,
+        **kwargs,
+    ):
+        hidden_states, residual = self.self_attn_compute(
+            hidden_states,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            past_key_value=past_key_value,
+            use_cache=use_cache,
+            attn_mask_startend_row_indices=attn_mask_startend_row_indices,
+            **kwargs,
+        )
+        probs, routing_map, l_aux, l_zloss = self.mlp.gate_compute(hidden_states)
+        return probs, routing_map, l_aux, l_zloss
+
+    def auxilibaryloss_and_shared_expert_compute(self, residual, hidden_states, expert_output, l_aux):
+        hidden_states = self.mlp.auxilibaryloss_and_shared_expert_compute(hidden_states, expert_output, l_aux)
+        hidden_states = residual + hidden_states
+
+    def post_process_output(self, hidden_states, output_attentions, use_cache, self_attn_weights, present_key_value):
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        if use_cache:
+            outputs += (present_key_value,)
+
+        if type(outputs) is tuple and len(outputs) == 1:
+            outputs = outputs[0]
+
+        return outputs
+
     def forward(
         self,
         hidden_states: paddle.Tensor,
@@ -1170,10 +1216,6 @@
                 (see `past_key_values`).
             past_key_value (`Tuple(paddle.Tensor)`, *optional*): cached past key and value projection states
         """
-        if "padding_mask" in kwargs:
-            warnings.warn(
-                "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
-            )
         residual = hidden_states
 
         hidden_states = self.input_layernorm(hidden_states)
@@ -1216,18 +1258,60 @@
         hidden_states = self.mlp(hidden_states)
         hidden_states = residual + hidden_states
 
-        outputs = (hidden_states,)
+        return self.post_process_output(
+            hidden_states, output_attentions, use_cache, self_attn_weights, present_key_value
+        )
 
-        if output_attentions:
-            outputs += (self_attn_weights,)
+    def self_attn_compute(
+        self,
+        hidden_states: paddle.Tensor,
+        position_ids: Optional[paddle.Tensor] = None,
+        attention_mask: Optional[paddle.Tensor] = None,
+        output_attentions: Optional[bool] = False,
+        past_key_value: Optional[Tuple[paddle.Tensor]] = None,
+        use_cache: Optional[bool] = False,
+        attn_mask_startend_row_indices: Optional[paddle.Tensor] = None,
+        **kwargs
+    ):
+        residual = hidden_states
 
-        if use_cache:
-            outputs += (present_key_value,)
+        hidden_states = self.input_layernorm(hidden_states)
 
-        if type(outputs) is tuple and len(outputs) == 1:
-            outputs = outputs[0]
+        # Self Attention
+        has_gradient = not hidden_states.stop_gradient
+        if (
+            self.enable_recompute
+            and self.layerwise_recompute
+            and has_gradient
+            and self.recompute_granularity == "full_attn"
+        ):
+            hidden_states, self_attn_weights, present_key_value = recompute(
+                self.self_attn,
+                hidden_states=hidden_states,
+                position_ids=position_ids,
+                attention_mask=attention_mask,
+                output_attentions=output_attentions,
+                past_key_value=past_key_value,
+                use_cache=use_cache,
+                attn_mask_startend_row_indices=attn_mask_startend_row_indices,
+                **kwargs,
+            )
+        else:
+            hidden_states, self_attn_weights, present_key_value = self.self_attn(
+                hidden_states=hidden_states,
+                position_ids=position_ids,
+                attention_mask=attention_mask,
+                output_attentions=output_attentions,
+                past_key_value=past_key_value,
+                use_cache=use_cache,
+                attn_mask_startend_row_indices=attn_mask_startend_row_indices,
+                **kwargs,
+            )
+        hidden_states = residual + hidden_states
 
-        return outputs
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        return hidden_states, residual
 
 
 class DeepseekV2MTPLayer(DeepseekV2DecoderLayer):
@@ -1892,7 +1976,7 @@
 
 
 class DeepseekV2LMHead(nn.Layer):
-    def __init__(self, config: DeepseekV2Config):
+    def __init__(self, config: DeepseekV2Config, embedding_weight=None):
         super(DeepseekV2LMHead, self).__init__()
         self.config = config
 
@@ -1906,11 +1990,16 @@
         else:
             vocab_size = config.vocab_size
 
-        self.weight = self.create_parameter(
-            shape=[config.hidden_size, vocab_size],
-            dtype=paddle.get_default_dtype(),
-            default_initializer=nn.initializer.XavierNormal(1.0),
-        )
+        if embedding_weight is not None:
+            self.transpose_y = True
+            self.weight = embedding_weight
+        else:
+            self.transpose_y = False
+            self.weight = self.create_parameter(
+                shape=[config.hidden_size, vocab_size],
+                dtype=paddle.get_default_dtype(),
+                default_initializer=nn.initializer.XavierNormal(1.0),
+            )
         # Must set distributed attr for Tensor Parallel !
         self.weight.is_distributed = True if (vocab_size != config.vocab_size) else False
 
@@ -1922,7 +2011,9 @@
         if tensor_parallel_output is None:
             tensor_parallel_output = self.config.tensor_parallel_output
 
-        logits = parallel_matmul(hidden_states, self.weight, tensor_parallel_output=tensor_parallel_output)
+        logits = parallel_matmul(
+            hidden_states, self.weight, transpose_y=self.transpose_y, tensor_parallel_output=tensor_parallel_output
+        )
         return logits
 
     def extra_repr(self):

paddlenlp/transformers/deepseek_v2/modeling_pp.py

@@ -39,13 +39,18 @@
     DeepseekV2RMSNorm,
 )
 
+try:
+    from paddle.distributed.fleet.meta_parallel import LocalSharedLayerDesc
+except:
+    LocalSharedLayerDesc = None
+
 __all__ = [
     "DeepseekV2ForCausalLMPipe",
 ]
 
 
 def parse_args(args):
-    if isinstance(args, tuple):
+    if isinstance(args, (tuple, list)):
         if len(args) == 4:
             hidden_states, attention_mask, attn_mask_startend_row_indices, position_ids = args
 
@@ -55,6 +60,9 @@
         elif len(args) == 2:
             hidden_states, attention_mask = args
             attn_mask_startend_row_indices, position_ids = None, None
+        else:  # len(args) == 1:
+            hidden_states = args[0]
+            attention_mask, attn_mask_startend_row_indices, position_ids = None, None, None
     else:
         hidden_states = args
         attention_mask, attn_mask_startend_row_indices, position_ids = None, None, None
@@ -321,8 +329,8 @@
 
 
 class DeepseekV2LMHeadPipe(DeepseekV2LMHead):
-    def __init__(self, config):
-        super(DeepseekV2LMHeadPipe, self).__init__(config)
+    def __init__(self, config, embedding_weight=None):
+        super(DeepseekV2LMHeadPipe, self).__init__(config, embedding_weight=embedding_weight)
 
     @property
     def embedding_weight(self):
@@ -406,6 +414,10 @@
             assert len(self.no_recompute_layers) == 0, "for pp with full recompute, no_recompute_layers is not support"
 
         virtual_pp_degree = getattr(self.config, "virtual_pp_degree", 1)
+        use_dualpipev = getattr(self.config, "use_dualpipev", False)
+        if use_dualpipev:
+            assert LocalSharedLayerDesc is not None, "LocalSharedLayerDesc is None, please update your paddle."
+        shared_class = LocalSharedLayerDesc if use_dualpipev else SharedLayerDesc
 
         def get_hcg():
             return fleet.get_hybrid_communicate_group()
@@ -420,7 +432,7 @@
 
         if config.tie_word_embeddings:
             self.add_sequential_layer(
-                SharedLayerDesc(
+                shared_class(
                     "DeepseekV2_shared_weight",
                     DeepseekV2EmbeddingPipe,
                     shared_weight_attr="embedding_weight",
@@ -453,12 +465,11 @@
 
         if config.tie_word_embeddings:
             self.add_sequential_layer(
-                SharedLayerDesc(
+                shared_class(
                     "DeepseekV2_shared_weight",
                     DeepseekV2LMHeadPipe,
                     shared_weight_attr="embedding_weight",
                     config=config,
-                    **{"transpose_y": True},
                 ),
                 "lm_head",
             )
@@ -489,6 +500,7 @@
                 "partition": False,
             },
             num_virtual_pipeline_stages=virtual_pp_degree,
+            use_dualpipev=use_dualpipev,
         )
         # You should call init here, since there is a  diamond inheritance problem
         self.apply(self._init_weights)
@@ -497,3 +509,39 @@
 
     def get_loss_fn(self, config):
         return DeepseekV2PretrainingCriterionPipe(config)
+
+    def overlapped_forward_backward(
+        self,
+        module0,  # the module of the forward chunk
+        inputs0,
+        criterion0,
+        labels0,
+        module1,  # the module of the backward chunk, maybe not used
+        loss1,
+        outputs1,
+        output_grads1,
+        scaler,
+    ):
+        outputs0 = inputs0
+        for layer in module0:
+            outputs0 = layer(outputs0)
+
+        outputs0 = [outputs0] if isinstance(outputs0, paddle.Tensor) else outputs0
+
+        if labels0 is not None:
+            loss0 = criterion0(outputs0, labels0)
+        else:
+            loss0 = None
+
+        if loss1 is not None:
+            if scaler:
+                paddle.autograd.backward(scaler.scale(loss1))
+            else:
+                paddle.autograd.backward(loss1)
+        else:
+            paddle.autograd.backward(
+                tensors=outputs1,
+                grad_tensors=output_grads1,
+            )
+
+        return outputs0, loss0

paddlenlp/transformers/gpt/modeling.py

@@ -1350,6 +1350,8 @@
             Tensor: The pretraining loss. Its data type should be float32 and its shape is [1].
 
         """
+        if isinstance(prediction_scores, list):
+            prediction_scores = prediction_scores[0]
         with paddle.amp.auto_cast(False):
             masked_lm_loss = self.loss_func(prediction_scores.astype("float32"), masked_lm_labels.unsqueeze(2))
             # skip ignore_index which loss == 0