refine many code

Author: Difers

examples/config/deepseek_v3/sft_4k_argument_dsv3.json

@@ -8,6 +8,7 @@
   "eval_dataset_path": "/root/paddlejob/tmpspace/chenzhichao/PaddleNLP-SFT/llm/en_data/dev.json",
   "eval_dataset_prob": "1.0",
   "eval_dataset_type": "erniekit",
+  "packing": true,
   "per_device_train_batch_size": 1,
   "gradient_accumulation_steps": 16,
   "per_device_eval_batch_size": 1,
@@ -55,6 +56,7 @@
   "sharding_parallel_config": "split_param",
   "tensor_parallel_output": true,
   "num_nextn_predict_layers": 1,
-  "convert_from_hf": true
+  "convert_from_hf": true,
+  "use_attn_mask_startend_row_indices": true
 }
 

examples/run_finetune.py

@@ -312,7 +312,7 @@ def neft_post_hook(module, input, output):
     if training_args.use_expert_parallel:
         callbacks += [MoeExpertsGradScaleCallback(training_args)]
 
-    logger.info("callbacks:", callbacks, flush=True)
+    logger.info(f"callbacks: {callbacks}")
 
     trainer = SFTTrainer(
         model=model,

paddleformers/nn/lm_head.py

@@ -68,7 +68,7 @@ def _set_distributed_attr(self, param):
         if param.is_distributed:
             param.split_axis = 0
 
-    def forward(self, hidden_states, tensor_parallel_output=None):
+    def forward(self, hidden_states, tensor_parallel_output=None, gather_hidden_states=True):
         """Project hidden states to vocabulary logits.
 
         Args:
@@ -103,7 +103,7 @@ def forward(self, hidden_states, tensor_parallel_output=None):
             self.bias,
             tensor_parallel_output,
             training=self.training,
-            gather_hidden_states=True,
+            gather_hidden_states=gather_hidden_states,
         )
 
     def extra_repr(self):

paddleformers/nn/mlp.py

@@ -12,10 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import paddle
+
+import contextlib
+
 import paddle.nn as nn
 from paddle.incubate.nn.functional import swiglu as fused_swiglu
 
 from ..generation.configuration_utils import PretrainedConfig
+from ..transformers.model_utils import dtype_guard
 from .activation import ACT2FN
 from .linear import Linear
 
@@ -45,66 +49,81 @@ def __init__(
         self.act_type = config.get("hidden_act", "silu")
         self.act_fn = ACT2FN[self.act_type]
         self.fuse_up_gate = fuse_up_gate
+        self.is_moe = kwargs.get("is_moe", False)
+        linear_type = None
+        if self.is_moe:
+            linear_type = "default"
 
-        if self.fuse_up_gate:
-            setattr(
-                self,
-                gate_up_proj_name,
-                Linear.create(
-                    self.hidden_size,
-                    self.intermediate_size * 2,
-                    has_bias=self.has_bias,
-                    config=config,
-                    fuse_matmul_bias=config.fuse_linear,
-                    tp_plan="colwise",
-                ),
-            )
-            self.up_gate_proj = getattr(self, gate_up_proj_name)
-        else:
-            # set attr for gate_proj
-            setattr(
-                self,
-                gate_proj_name,
-                Linear.create(
-                    self.hidden_size,
-                    self.intermediate_size,
-                    has_bias=self.has_bias,
-                    config=config,
-                    fuse_matmul_bias=config.fuse_linear,
-                    tp_plan="colwise",
-                ),
-            )
-            self.gate_proj = getattr(self, gate_proj_name)
+        def linear_type_gaurd():
+            if config.use_fp8:
+                return dtype_guard("float8_e4m3fn")
+            else:
+                return contextlib.nullcontext()
+
+        with linear_type_gaurd():
+            if self.fuse_up_gate:
+                setattr(
+                    self,
+                    gate_up_proj_name,
+                    Linear.create(
+                        self.hidden_size,
+                        self.intermediate_size * 2,
+                        has_bias=self.has_bias,
+                        config=config,
+                        fuse_matmul_bias=config.fuse_linear,
+                        tp_plan="colwise",
+                        linear_type=linear_type,
+                    ),
+                )
+                self.up_gate_proj = getattr(self, gate_up_proj_name)
+            else:
+                # set attr for gate_proj
+                setattr(
+                    self,
+                    gate_proj_name,
+                    Linear.create(
+                        self.hidden_size,
+                        self.intermediate_size,
+                        has_bias=self.has_bias,
+                        config=config,
+                        fuse_matmul_bias=config.fuse_linear,
+                        tp_plan="colwise",
+                        linear_type=linear_type,
+                    ),
+                )
+                self.gate_proj = getattr(self, gate_proj_name)
+
+                # set attr for up_proj
+                setattr(
+                    self,
+                    up_proj_name,
+                    Linear.create(
+                        self.hidden_size,
+                        self.intermediate_size,
+                        has_bias=self.has_bias,
+                        config=config,
+                        fuse_matmul_bias=config.fuse_linear,
+                        tp_plan="colwise",
+                        linear_type=linear_type,
+                    ),
+                )
+                self.up_proj = getattr(self, up_proj_name)
 
-            # set attr for up_proj
+            # set attr for down_proj
             setattr(
                 self,
-                up_proj_name,
+                down_proj_name,
                 Linear.create(
-                    self.hidden_size,
                     self.intermediate_size,
+                    self.hidden_size,
                     has_bias=self.has_bias,
                     config=config,
                     fuse_matmul_bias=config.fuse_linear,
-                    tp_plan="colwise",
+                    tp_plan="rowwise",
+                    linear_type=linear_type,
                 ),
             )
-            self.up_proj = getattr(self, up_proj_name)
-
-        # set attr for down_proj
-        setattr(
-            self,
-            down_proj_name,
-            Linear.create(
-                self.intermediate_size,
-                self.hidden_size,
-                has_bias=self.has_bias,
-                config=config,
-                fuse_matmul_bias=config.fuse_linear,
-                tp_plan="rowwise",
-            ),
-        )
-        self.down_proj = getattr(self, down_proj_name)
+            self.down_proj = getattr(self, down_proj_name)
 
     def forward(self, x):
         if self.fuse_up_gate:

paddleformers/nn/norm.py

@@ -17,7 +17,9 @@
 from paddle.incubate.nn.functional import fused_rms_norm_ext
 
 from ..generation.configuration_utils import PretrainedConfig
+from ..transformers.llama import fusion_ops
 from ..utils.log import logger
+from ..utils.tools import get_env_device
 from .general import GeneralInterface
 
 try:

paddleformers/nn/pp_model.py

@@ -343,8 +343,19 @@ def __init__(self, *args, **kwargs):
 
     def forward(self, args):
         hidden_states, _, _, _, _ = parse_args(args)
-        hidden_states = super().forward(hidden_states)
-        return hidden_states
+
+        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 :]
+
+            output_list = [super().forward(hidden_states)]
+            for hidden_states in hidden_states_mtp:
+                output_list.append(super().forward(hidden_states))
+            return output_list
+        else:
+            hidden_states = super().forward(hidden_states)
+            return hidden_states
 
 
 class LayerNormPipe(LayerNorm):
@@ -389,6 +400,12 @@ def forward(self, args):
                 [batch_size, sequence_length, vocab_size]
                 representing unnormalized log probabilities for each token
         """
+        if self.config.num_nextn_predict_layers > 0:
+            logits = []
+            for _hidden_states in args:
+                logits.append(super().forward(_hidden_states))
+            return logits
+
         hidden_states, _, _, _, _ = parse_args(args)
         logits = super().forward(hidden_states)
         return logits
@@ -507,12 +524,25 @@ class GeneralModelForCausalLMPipe(PipelinePretrainedModel, PipelineLayer):
     transpose_weight_keys = None
     _embed_cls = None
     _rotary_emb_cls = None
+    _mtp_layer_pipe_cls = None
+    _embedding_pipe_cls = None
+    _decoder_layer_pipe_cls = None
+    _criterion_pipe_cls = None
+    _lmhead_pipe_cls = None
 
     def __init__(self, config: PretrainedConfig, **kwargs):
         # dynamic inherit DecoderLayer
         if self._decoder_layer_cls is None:
             raise ValueError("_decoder_layer_cls must be set before init.")
-        DecoderLayerPipe = make_decoder_layer_pipe(self._decoder_layer_cls)
+
+        EmbeddingPipeCls = self._embedding_pipe_cls if self._embedding_pipe_cls is not None else Embedding
+
+        if self._decoder_layer_pipe_cls is None:
+            DecoderLayerPipe = make_decoder_layer_pipe(self._decoder_layer_cls)
+        else:
+            DecoderLayerPipe = self._decoder_layer_pipe_cls
+
+        LMHeadPipeCls = self._lmhead_pipe_cls if self._lmhead_pipe_cls is not None else LMHeadPipe
 
         new_initializer_range = math.sqrt(0.3333 / config.hidden_size)
         logger.info(f"change initializer-range from {config.initializer_range} to {new_initializer_range}")
@@ -559,7 +589,7 @@ def __init__(self, config: PretrainedConfig, **kwargs):
         else:
             self.add_sequential_layer(
                 LayerDesc(
-                    EmbeddingPipe, config=config, embed_cls=self._embed_cls, rotary_emb_cls=self._rotary_emb_cls
+                    EmbeddingPipeCls, config=config, embed_cls=self._embed_cls, rotary_emb_cls=self._rotary_emb_cls
                 ),
                 "model",
             )
@@ -573,6 +603,11 @@ def __init__(self, config: PretrainedConfig, **kwargs):
                 ),
                 f"model.layers.{i}",
             )
+        for i in range(config.num_nextn_predict_layers):
+            self.add_sequential_layer(
+                LayerDesc(self._mtp_layer_pipe_cls, config=config, layer_idx=config.num_hidden_layers + i),
+                f"model.layers.{config.num_hidden_layers + i}",
+            )
         for i in range(config.add_tail_layers):
             self.add_sequential_layer(
                 LayerDesc(
@@ -590,14 +625,14 @@ def __init__(self, config: PretrainedConfig, **kwargs):
             self.add_sequential_layer(
                 SharedLayerDesc(
                     "model_shared_weight",
-                    LMHeadPipe,
+                    LMHeadPipeCls,
                     shared_weight_attr="embedding_weight",
                     config=config,
                 ),
                 "lm_head",
             )
         else:
-            self.add_sequential_layer(LayerDesc(LMHeadPipe, config=config), "lm_head")
+            self.add_sequential_layer(LayerDesc(LMHeadPipeCls, config=config), "lm_head")
         recompute_interval = 0
 
         seg_method = config.pp_seg_method if hasattr(config, "pp_seg_method") else "layer:DecoderLayer|EmptyLayer"
@@ -630,10 +665,12 @@ def __init__(self, config: PretrainedConfig, **kwargs):
         )
 
     def get_loss_fn(self, config):
+        CriterionPipeCls = self._criterion_pipe_cls if self._criterion_pipe_cls is not None else CriterionLayerPipe
+
         if config.get("dpo_config", None) is not None:
-            loss_fn = CriterionLayerPipe(config, use_infohub=True)
+            loss_fn = CriterionPipeCls(config, use_infohub=True)
         else:
-            loss_fn = CriterionLayerPipe(config)
+            loss_fn = CriterionPipeCls(config)
 
         return loss_fn
 

paddleformers/transformers/__init__.py

@@ -123,20 +123,18 @@
         "DeepseekV2DynamicNTKScalingRotaryEmbedding",
         "DeepseekV2MLP",
         "yarn_get_mscale",
-        "DeepseekV2LMHead",
         "DeepseekV2DecoderLayer",
-        "DeepseekV2PretrainingCriterion",
         "yarn_find_correction_range",
         "get_triangle_upper_mask",
         "DeepseekV2LinearScalingRotaryEmbedding",
+        "DeepseekV2ForCausalLMPipe",
     ],
     "deepseek_v2.modeling_auto": [
         "DeepseekV2LMHeadAuto",
         "DeepseekV2ForCausalLMAuto",
         "DeepseekV2ModelAuto",
         "DeepseekV2PretrainedModelAuto",
     ],
-    "deepseek_v2.modeling_pp": ["DeepseekV2ForCausalLMPipe"],
     "deepseek_v2.mfu_utils": ["DeepSeekProjection"],
     "deepseek_v2.kernel": [
         "act_quant",
@@ -160,14 +158,14 @@
         "DeepseekV3ForSequenceClassification",
         "DeepseekV3Model",
         "DeepseekV3PretrainedModel",
+        "DeepseekV3ForCausalLMPipe",
     ],
     "deepseek_v3.modeling_auto": [
         "DeepseekV3LMHeadAuto",
         "DeepseekV3ForCausalLMAuto",
         "DeepseekV3ModelAuto",
         "DeepseekV3PretrainedModelAuto",
     ],
-    "deepseek_v3.modeling_pp": ["DeepseekV3ForCausalLMPipe"],
     "ernie4_5.configuration": ["Ernie4_5Config"],
     "ernie4_5.modeling": ["Ernie4_5Model", "Ernie4_5ForCausalLM", "Ernie4_5ForCausalLMPipe"],
     "ernie4_5.tokenizer": ["Ernie4_5Tokenizer"],

paddleformers/transformers/deepseek_v2/__init__.py

@@ -50,20 +50,18 @@
         "DeepseekV2DynamicNTKScalingRotaryEmbedding",
         "DeepseekV2MLP",
         "yarn_get_mscale",
-        "DeepseekV2LMHead",
         "DeepseekV2DecoderLayer",
-        "DeepseekV2PretrainingCriterion",
         "yarn_find_correction_range",
         "get_triangle_upper_mask",
         "DeepseekV2LinearScalingRotaryEmbedding",
+        "DeepseekV2ForCausalLMPipe",
     ],
     "modeling_auto": [
         "DeepseekV2LMHeadAuto",
         "DeepseekV2ForCausalLMAuto",
         "DeepseekV2ModelAuto",
         "DeepseekV2PretrainedModelAuto",
     ],
-    "modeling_pp": ["DeepseekV2ForCausalLMPipe"],
     "mfu_utils": ["DeepSeekProjection"],
     "kernel": [
         "act_quant",

paddleformers/transformers/deepseek_v2/configuration.py

@@ -160,12 +160,13 @@ def __init__(
         first_k_dense_replace=0,
         norm_topk_prob=False,
         scoring_func="softmax",
-        aux_loss_alpha=0.001,
+        aux_loss_alpha=0.0001,
         seq_aux=True,
         hidden_act="silu",
         max_position_embeddings=2048,
         seq_length=32768,
         initializer_range=0.02,
+        use_rmsnorm=True,
         rms_norm_eps=1e-6,
         use_cache=True,
         pad_token_id=None,
@@ -217,6 +218,7 @@ def __init__(
         self.num_key_value_heads = num_key_value_heads
         self.hidden_act = hidden_act
         self.initializer_range = initializer_range
+        self.use_rmsnorm = use_rmsnorm
         self.rms_norm_eps = rms_norm_eps
         self.pretraining_tp = pretraining_tp
         self.use_cache = use_cache