Support magic send for mtp (#10916)

* support fp8 weight quant cache

* fix bug

* fix confilct

* support magic send

Author: zhangbo9674

llm/run_pretrain.py

@@ -27,6 +27,7 @@
     print_rank_0,
 )
 from paddlenlp.trainer import (
+    FP8QuantWeightCallback,
     PdArgumentParser,
     StepFlexToken,
     Trainer,
@@ -568,7 +569,7 @@ def main():
         * data_args.max_seq_length
     )
 
-    callbacks = [StepFlexToken()]
+    callbacks = [StepFlexToken(), FP8QuantWeightCallback()]
 
     trainer = PretrainingTrainer(
         model=model,

paddlenlp/trainer/trainer_callback.py

@@ -20,6 +20,7 @@
 """
 import dataclasses
 import json
+import os
 from dataclasses import dataclass
 from typing import Dict, List, Optional, Union
 
@@ -41,6 +42,7 @@
     "PrinterCallback",
     "EarlyStoppingCallback",
     "StepFlexToken",
+    "FP8QuantWeightCallback",
 ]
 
 
@@ -615,3 +617,34 @@ def on_step_begin(
         model = kwargs.pop("model")
         if hasattr(model, "step_flex_token"):
             model.step_flex_token(state.global_step)
+
+
+g_shard_bypass_dygraph_optimizer = int(os.environ.get("FLAGS_shard_bypass_dygraph_optimizer", 0))
+
+
+def enable_in_dict_config(config, key):
+    """enable_in_dict_config"""
+    return key in config and config[key]
+
+
+skip_count = 0
+
+
+class FP8QuantWeightCallback(TrainerCallback):
+    """
+    FP8QuantWeightCallback
+    """
+
+    def on_step_begin(self, args, state, control, **kwargs):
+        """
+        每个step开始前把专家参数quant成fp8q
+        """
+        model = kwargs["model"]
+        optimizer = kwargs["optimizer"]
+        global skip_count
+
+        if not g_shard_bypass_dygraph_optimizer or skip_count == 0:
+            model.fp8_quant_weight(True)
+            optimizer.clear_param_storage("moe_expert")
+
+        skip_count += 1

paddlenlp/transformers/deepseek_v2/modeling.py

@@ -981,6 +981,49 @@ def __init__(self, config: DeepseekV2Config, norm_weight=None, norm_eps=None):
                     config=config, intermediate_size=intermediate_size, is_moe=False
                 )
 
+    def fp8_quant_weight(self, batch_mode=False):
+        """Quantize weights in FP8 format.
+
+        Args:
+            batch_mode: If True, quantize all weights in batch mode using the first expert's weights.
+                    If False, quantize each expert's weights individually.
+        """
+
+        def quantize_weights(weight_list, weight_obj=None):
+            """Helper function to quantize a list of weights."""
+            if weight_obj is None:
+                weight_obj = weight_list[0]
+
+            # Quantize without transpose
+            fp8_weight, fp8_scale = paddle.incubate.nn.functional.fused_stack_transpose_quant(
+                weight_list, transpose=False
+            )
+            setattr(weight_obj, "fp8_weight_stacked", fp8_weight)
+            setattr(weight_obj, "fp8_scale_stacked", fp8_scale)
+
+            # Quantize with transpose
+            fp8_weight_t, fp8_scale_t = paddle.incubate.nn.functional.fused_stack_transpose_quant(
+                weight_list, transpose=True
+            )
+            setattr(weight_obj, "fp8_weight_stacked_transpose", fp8_weight_t)
+            setattr(weight_obj, "fp8_scale_stacked_transpose", fp8_scale_t)
+
+        if batch_mode:
+            # Batch mode: process all experts' weights together
+            expert_w1_list = [expert.w1 for expert in self.experts if expert is not None]
+            expert_w2_list = [expert.w2 for expert in self.experts if expert is not None]
+
+            if expert_w1_list:
+                quantize_weights(expert_w1_list, expert_w1_list[0])
+            if expert_w2_list:
+                quantize_weights(expert_w2_list, expert_w2_list[0])
+        else:
+            # Individual mode: process each expert's weights separately
+            for expert in self.experts:
+                if expert is not None:
+                    quantize_weights([expert.w1])
+                    quantize_weights([expert.w1])
+
     def forward(self, hidden_states):
         if self.using_post_norm_recompute:
             super().update_flex_token()
@@ -1928,6 +1971,12 @@ def __init__(self, config: DeepseekV2Config, layer_idx: int, layerwise_recompute
         else:
             self.mlp = DeepseekV2MLPClass(config)
 
+    def fp8_quant_weight(self, batch_mode=False):
+        """fp8_quant_weight"""
+        if isinstance(self.mlp, DeepseekV2MoE):
+            logger.info(f"fp8 quant weight for mlp {type(self.mlp)}")
+            self.mlp.fp8_quant_weight(batch_mode)
+
     def forward(
         self,
         hidden_states: paddle.Tensor,

paddlenlp/transformers/deepseek_v2/modeling_pp.py

@@ -65,6 +65,10 @@
     "DeepseekV2ForCausalLMPipe",
 ]
 
+import queue
+
+global_inputs_embeds_mtp_queue = queue.Queue()
+
 
 DSV3_USE_FP8_GEMM = os.getenv("DSV3_USE_FP8_GEMM", "False").lower() == "true"
 DSV3_USE_FP8_DISPATCH = os.getenv("DSV3_USE_FP8_DISPATCH", "False").lower() == "true"
@@ -1019,7 +1023,7 @@ def forward(self, args):
         inputs_embeds = self.embed_tokens(input_ids)
 
         batch_size, seq_length = input_ids.shape
-        if self.config.send_mtp_embed:
+        if self.config.num_nextn_predict_layers > 0:
             seq_length -= self.config.num_nextn_predict_layers
 
             if attention_mask is not None:
@@ -1042,7 +1046,7 @@ def forward(self, args):
             attention_mask = paddle.tril(paddle.ones((seq_length, seq_length), dtype="bool"))
             attention_mask.stop_gradient = True
 
-        if self.config.send_mtp_embed:
+        if self.config.num_nextn_predict_layers > 0:
             inputs_embeds_extra = inputs_embeds[:, -self.config.num_nextn_predict_layers :, :]  # [B, S, D]
             inputs_embeds = inputs_embeds[:, : -self.config.num_nextn_predict_layers, :]
             inputs_embeds_ori = inputs_embeds
@@ -1054,6 +1058,7 @@ def forward(self, args):
                 # [seq_len * bs / n, num_head * head_dim] (n is mp parallelism)
                 inputs_embeds = ScatterOp.apply(inputs_embeds)
             embeds_res = [inputs_embeds]
+            mtp_embeds = []
             for depth in range(self.config.num_nextn_predict_layers):
                 inputs_embeds_mtp = paddle.concat(
                     [
@@ -1065,12 +1070,19 @@ def forward(self, args):
                 if self.sequence_parallel:
                     inputs_embeds_mtp = inputs_embeds_mtp.reshape([-1, inputs_embeds_mtp.shape[-1]])
                     inputs_embeds_mtp = ScatterOp.apply(inputs_embeds_mtp)
-                embeds_res.append(inputs_embeds_mtp)
-            # if not self.sequence_parallel
-            # mtp_embeds: [B*num_nextn_predict_layers, seq_len, hidden_size]
-            # else:
-            # mtp_embeds: [B*seq_len*num_nextn_predict_layers, hidden_size]
-            inputs_embeds = paddle.concat(embeds_res, axis=-1)
+                mtp_embeds.append(inputs_embeds_mtp)
+
+            if self.config.send_mtp_embed:
+                embeds_res.extend(mtp_embeds)
+                # if not self.sequence_parallel
+                # mtp_embeds: [B*num_nextn_predict_layers, seq_len, hidden_size]
+                # else:
+                # mtp_embeds: [B*seq_len*num_nextn_predict_layers, hidden_size]
+                inputs_embeds = paddle.concat(embeds_res, axis=-1)
+            else:
+                global global_inputs_embeds_mtp_queue
+                cloned_mtp_embeds = [t.detach() for t in mtp_embeds]
+                global_inputs_embeds_mtp_queue.put(cloned_mtp_embeds)
             return return_args(inputs_embeds, attention_mask, attn_mask_startend_row_indices, position_ids)
         else:
             if self.sequence_parallel:
@@ -1359,9 +1371,15 @@ class DeepseekV2MTPLayerPipe(DeepseekV2MTPLayer):
     def forward(self, args):
         hidden_states, attention_mask, attn_mask_startend_row_indices, position_ids = parse_args(args)
 
-        hidden_states_list = paddle.split(hidden_states, self.config.num_nextn_predict_layers + 1, axis=-1)
-        hidden_states_main_model = hidden_states_list[0]
-        inputs_embeds_cur_depth_list = hidden_states_list[1:]
+        if self.config.send_mtp_embed:
+            hidden_states_list = paddle.split(hidden_states, self.config.num_nextn_predict_layers + 1, axis=-1)
+            hidden_states_main_model = hidden_states_list[0]
+            inputs_embeds_cur_depth_list = hidden_states_list[1:]
+        else:
+            hidden_states_main_model = hidden_states
+            global global_inputs_embeds_mtp_queue
+            inputs_embeds_cur_depth_list = global_inputs_embeds_mtp_queue.get()
+
         has_gradient = not hidden_states_main_model.stop_gradient
 
         if attention_mask is not None and attention_mask.dtype == paddle.int32:
@@ -1426,7 +1444,7 @@ def __init__(self, config):
     def forward(self, args):
         hidden_states, attention_mask, attn_mask_startend_row_indices, position_ids = parse_args(args)
 
-        if self.config.send_mtp_embed:
+        if self.config.num_nextn_predict_layers > 0:
             hidden_states_list = paddle.split(hidden_states, self.config.num_nextn_predict_layers + 1, axis=-1)
             hidden_states = hidden_states_list[0]
             hidden_states_mtp = hidden_states_list[-self.config.num_nextn_predict_layers :]
@@ -1451,7 +1469,7 @@ def embedding_weight(self):
         return get_attr(self, "weight")
 
     def forward(self, args: Union[Tuple, paddle.Tensor]):
-        if self.config.send_mtp_embed:
+        if self.config.num_nextn_predict_layers > 0:
             logits = []
             for _hidden_states in args:
                 logits.append(super().forward(_hidden_states))
@@ -1466,7 +1484,7 @@ def build_schedule_node(self):
 
 class DeepseekV2PretrainingCriterionPipe(DeepseekV2PretrainingCriterion):
     def forward(self, logits, labels):
-        if self.config.send_mtp_embed:
+        if self.config.num_nextn_predict_layers > 0:
             mtp_logits = logits[1:]
             logits = logits[0]
             loss = super().forward(logits, labels, mtp_logits=mtp_logits)
@@ -1669,6 +1687,19 @@ def compute_recompute_fwd_gate_up_list(pp_nums, all_dl_nums, dense_dl_nums, reco
         # DON'T init PipelinePretrainedModel
         # PipelinePretrainedModel.__init__(self.super(), config=config)
 
+    def fp8_quant_weight(self, batch_mode=False):
+        """fp8_quant_weight"""
+        with paddle.no_grad():
+            for i, layer in self._sub_layers.items():
+                if isinstance(
+                    layer, paddle.distributed.fleet.meta_parallel.parallel_layers.pp_layers.PipelineLayerChunk
+                ):
+                    for i, sub_layer in layer.named_sublayers():
+                        if isinstance(sub_layer, DeepseekV2DecoderLayer) and hasattr(sub_layer, "fp8_quant_weight"):
+                            sub_layer.fp8_quant_weight(batch_mode)
+                if isinstance(layer, DeepseekV2DecoderLayer) and hasattr(layer, "fp8_quant_weight"):
+                    layer.fp8_quant_weight(batch_mode)
+
     def get_loss_fn(self, config):
         return DeepseekV2PretrainingCriterionPipe(config)
 

paddlenlp/transformers/fp8_utils.py

@@ -46,6 +46,29 @@ def swiglu(x, y=None):
 ]
 
 
+def _get_fp8_weight_and_scale(weight, stacked=False, transpose=False):
+    """_get_fp8_weight_and_scale"""
+    if stacked:
+        if transpose:
+            fp8_weight, fp8_scale = weight.fp8_weight_stacked_transpose, weight.fp8_scale_stacked_transpose
+        else:
+            fp8_weight, fp8_scale = weight.fp8_weight_stacked, weight.fp8_scale_stacked
+    else:
+        if transpose:
+            fp8_weight, fp8_scale = weight.fp8_weight_transpose, weight.fp8_scale_transpose
+        else:
+            fp8_weight, fp8_scale = weight.fp8_weight, weight.fp8_scale
+    return fp8_weight, fp8_scale
+
+
+def fused_stack_quant(expert_weight_list, transpose=False):
+    if hasattr(expert_weight_list[0], "fp8_weight_stacked"):
+        w, scale = _get_fp8_weight_and_scale(expert_weight_list[0], stacked=True, transpose=transpose)
+    else:
+        w, scale = paddle.incubate.nn.functional.fused_stack_transpose_quant(expert_weight_list, transpose=transpose)
+    return w, scale
+
+
 class FP8LinearFunctionBase:
     @staticmethod
     def dequantize_fp8_to_fp32(fp8_tensor, scale):
@@ -524,7 +547,9 @@ def backward(ctx, do3):
         )
 
         # ===== call func common_fp8_mlp_bwd =====
-        d_norm_output, dw1, dw2 = FP8LinearFunctionBase.common_fp8_mlp_bwd(do3, x_fp8, x_scale, x_t_fp8, x_t_scale, w1, w2)
+        d_norm_output, dw1, dw2 = FP8LinearFunctionBase.common_fp8_mlp_bwd(
+            do3, x_fp8, x_scale, x_t_fp8, x_t_scale, w1, w2
+        )
 
         # ===== reshape to origin shape =====
         if len(x_orig_shape) > 2:
@@ -723,7 +748,7 @@ def fwd_gate_up(self, x, expert_w1, num_expert, tokens_per_expert):
         if not self.is_split_group_gemm:
             self.m_indices = self.gen_m_indices(tokens_per_expert)
         # concat w1, shape is [num_groups, n, k]
-        w1_t_quant, w1_t_scale = paddle.incubate.nn.functional.fused_stack_transpose_quant(expert_w1, transpose=True)
+        w1_t_quant, w1_t_scale = fused_stack_quant(expert_w1, transpose=True)
         w1_t_quant = w1_t_quant.reshape([num_expert, -1, w1_t_quant.shape[-1]])
         w1_t_scale = w1_t_scale.reshape([num_expert, -1, w1_t_scale.shape[-1]])
 
@@ -765,7 +790,7 @@ def fwd_down(self, o1, unzipped_probs, expert_w2, num_expert, o3=None, clear_o1=
         [m_sum, k] = [m_sum, n] * [num_groups, n, k]
         """
         # concat and transpose w2
-        w2_quant, w2_scale = paddle.incubate.nn.functional.fused_stack_transpose_quant(expert_w2, transpose=True)
+        w2_quant, w2_scale = fused_stack_quant(expert_w2, transpose=True)
         w2_quant = w2_quant.reshape([num_expert, -1, w2_quant.shape[-1]])
         w2_scale = w2_scale.reshape([num_expert, -1, w2_scale.shape[-1]])
 
@@ -801,9 +826,7 @@ def bwd_dowm_input(self, expert_w2, unzipped_grad, o1, inplace_swiglu_prob=False
         [m_sum, n] = [m_sum, k] * [num_groups, k, n]
         """
         # recompute concated_w2_2d
-        bw_w2_quant, bw_w2_scale = paddle.incubate.nn.functional.fused_stack_transpose_quant(
-            expert_w2, transpose=False
-        )
+        bw_w2_quant, bw_w2_scale = fused_stack_quant(expert_w2, transpose=False)
         bw_w2_quant = bw_w2_quant.reshape([len(expert_w2), -1, bw_w2_quant.shape[-1]])
         bw_w2_scale = bw_w2_scale.reshape([len(expert_w2), -1, bw_w2_scale.shape[-1]])
 
@@ -849,9 +872,7 @@ def bwd_gate_up_input(self, do1, expert_w1, dx=None):
         [m_sum, k] = [m_sum, n] * [num_groups, n, k]
         """
         # recompute concated_w1_t
-        bw_w1_quant, bw_w1_scale = paddle.incubate.nn.functional.fused_stack_transpose_quant(
-            expert_w1, transpose=False
-        )
+        bw_w1_quant, bw_w1_scale = fused_stack_quant(expert_w1, transpose=False)
         bw_w1_quant = bw_w1_quant.reshape([len(expert_w1), -1, bw_w1_quant.shape[-1]])
         bw_w1_scale = bw_w1_scale.reshape([len(expert_w1), -1, bw_w1_scale.shape[-1]])