add dsv3 from paddlenlp-sft

Author: Difers

examples/run_finetune.py

@@ -140,6 +140,10 @@ def main():
     model_config.max_sequence_length = training_args.max_seq_len
     model_config.num_nextn_predict_layers = model_args.num_nextn_predict_layers
     model_config._attn_implementation = model_args.attn_impl
+    model_config.gradient_accumulation_steps = training_args.gradient_accumulation_steps
+    model_config.using_flex_token = model_args.using_flex_token
+    model_config.using_fake_gate = model_args.using_fake_gate
+    model_config.aux_loss_alpha = model_args.aux_loss_alpha
     logger.info(f"Final model config: {model_config}")
     logger.info("Creating model")
 
@@ -278,13 +282,16 @@ def neft_post_hook(module, input, output):
         training_args.logging_steps = int(training_args.max_steps / training_args.num_train_epochs)
 
     callbacks = []
+
     if getattr(model_config, "topk_method", None) == "noaux_tc":
-        callbacks += [MoECorrectionBiasAdjustCallback(lr=0)]
+        # deepseek_v3 finetune do not update the bias, so set lr to 0.0
+        callbacks += [MoECorrectionBiasAdjustCallback(lr=0.0)]
 
     if training_args.use_expert_parallel:
         callbacks += [MoeExpertsGradScaleCallback(training_args)]
 
-    print("callbacks:", callbacks, flush=True)
+    logger.info(f"callbacks: {callbacks}")
+
     trainer = SFTTrainer(
         model=model,
         args=training_args,
@@ -295,6 +302,7 @@ def neft_post_hook(module, input, output):
         data_collator=data_collator,
         do_generation=data_args.eval_with_do_generation,
         data_args=data_args,
+        callbacks=callbacks,
     )
     trainable_parameters = [
         p for p in model.parameters() if not p.stop_gradient or ("quantization_linear" in p.name and "w_1" in p.name)

paddleformers/nn/linear.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import contextlib
+
 import paddle.nn as nn
 from paddle.incubate.nn import FusedLinear
 
@@ -22,6 +24,7 @@
     RowParallelLinear,
     RowSequenceParallelLinear,
 )
+from ..transformers.model_utils import dtype_guard
 from .general import GeneralInterface
 
 __all__ = ["Linear"]
@@ -51,6 +54,12 @@ def create(
         input_is_parallel: bool = True,
         fuse_matmul_bias: bool = False,
     ):
+        def linear_type_gaurd():
+            if config.use_fp8:
+                return dtype_guard("float8_e4m3fn")
+            else:
+                return contextlib.nullcontext()
+
         if linear_type is None and config is None:
             raise ValueError("linear_type or config must be specified")
 
@@ -59,7 +68,9 @@ def create(
 
         linear_cls = self._global_mapping[linear_type]
         kwargs = self.get_linear_kwargs(linear_type, has_bias, gather_output, input_is_parallel, fuse_matmul_bias)
-        return linear_cls(in_features=in_features, out_features=out_features, weight_attr=weight_attr, **kwargs)
+
+        with linear_type_gaurd():
+            return linear_cls(in_features=in_features, out_features=out_features, weight_attr=weight_attr, **kwargs)
 
     @classmethod
     def get_linear_type(self, config: PretrainedConfig, tp_plan=None):

paddleformers/nn/mlp.py

@@ -45,6 +45,10 @@ 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(
@@ -57,6 +61,7 @@ def __init__(
                     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)
@@ -72,6 +77,7 @@ def __init__(
                     config=config,
                     fuse_matmul_bias=config.fuse_linear,
                     tp_plan="colwise",
+                    linear_type=linear_type,
                 ),
             )
             self.gate_proj = getattr(self, gate_proj_name)
@@ -87,6 +93,7 @@ def __init__(
                     config=config,
                     fuse_matmul_bias=config.fuse_linear,
                     tp_plan="colwise",
+                    linear_type=linear_type,
                 ),
             )
             self.up_proj = getattr(self, up_proj_name)
@@ -102,6 +109,7 @@ def __init__(
                 config=config,
                 fuse_matmul_bias=config.fuse_linear,
                 tp_plan="rowwise",
+                linear_type=linear_type,
             ),
         )
         self.down_proj = getattr(self, down_proj_name)

paddleformers/nn/pp_model.py

@@ -507,12 +507,28 @@ 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
+    _rms_norm_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
+        MTPLayerPipeCls = self._mtp_layer_pipe_cls if self._mtp_layer_pipe_cls is not None else None
+        RMSNormPipeCls = self._rms_norm_pipe_cls if self._rms_norm_pipe_cls is not None else RMSNormPipe
 
         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 +575,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 +589,12 @@ def __init__(self, config: PretrainedConfig, **kwargs):
                 ),
                 f"model.layers.{i}",
             )
+        for i in range(config.num_nextn_predict_layers):
+            if MTPLayerPipeCls is not None:
+                self.add_sequential_layer(
+                    LayerDesc(MTPLayerPipeCls, 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(
@@ -582,22 +604,22 @@ def __init__(self, config: PretrainedConfig, **kwargs):
             )
 
         self.add_sequential_layer(
-            LayerDesc(RMSNormPipe if config.use_rmsnorm else LayerNormPipe, config=config),
+            LayerDesc(RMSNormPipeCls if config.use_rmsnorm else LayerNormPipe, config=config),
             "model.norm",
         )
 
         if config.tie_word_embeddings:
             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 +652,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