(1)recover bos download (2) move dsv3_fast_pretrain from env to arg (3) move load_hf_ckpt

Author: hushenwei2000

paddleformers/trainer/utils/load_hf_ckpt.py renamed to paddleformers/examples/deepseek_v3/load_hf_ckpt.py

@@ -21,23 +21,13 @@
 
 from __future__ import annotations
 
-import contextlib
-import math
-import os
-import warnings
 from functools import partial
-from typing import List, Optional, Tuple, Union
 
 import paddle
-import paddle.distributed as dist
-import paddle.distributed.fleet.meta_parallel as mpu
 import paddle.nn.functional as F
 from paddle import Tensor, nn
 from paddle.distributed import fleet
-from paddle.distributed.fleet.meta_parallel import get_rng_state_tracker
-from paddle.distributed.fleet.recompute.recompute import recompute
 from paddle.jit import to_static
-from paddle.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from paddle.utils import try_import
 
 try:
@@ -48,7 +38,6 @@
 try:
     from paddle.distributed.fleet.utils.sequence_parallel_utils import (
         GatherOp,
-        ScatterOp,
         mark_as_sequence_parallel_parameter,
     )
 except:
@@ -62,43 +51,20 @@
     flash_attention = None
 
 from config.configuration import DeepseekV2FastConfig
-from moe_gate import PretrainedMoEGate
-from moe_layer import MoEFlexTokenLayer, MoELayer
 from paddle.distributed.fleet.meta_parallel.zero_bubble_utils import WeightGradStore
 
-from paddleformers.transformers.activations import ACT2FN
-from paddleformers.transformers.conversion_utils import (
-    StateDictNameMapping,
-    init_name_mappings,
-)
-from paddleformers.transformers.deepseek_v2 import DeepseekV2RotaryEmbedding, Linear
-from paddleformers.transformers.deepseek_v2 import fp8_linear as linear_utils
 from paddleformers.transformers.deepseek_v2 import (
+    DeepseekV2RotaryEmbedding,
     yarn_find_correction_range,
     yarn_get_mscale,
     yarn_linear_ramp_mask,
 )
 from paddleformers.transformers.fp8_utils import (
-    FP8Linear,
     FP8LinearFunctionBase,
     cache_fp8_weight,
     set_parameter_color,
 )
-from paddleformers.transformers.llama import fusion_ops
-from paddleformers.transformers.llama.modeling import get_use_casual_mask
-from paddleformers.transformers.model_outputs import (
-    BaseModelOutputWithPastAndMTP,
-    CausalLMOutputWithPast,
-    SequenceClassifierOutputWithPast,
-)
-from paddleformers.transformers.model_utils import (
-    PretrainedModel,
-    dtype_guard,
-    register_base_model,
-)
-from paddleformers.transformers.utils import cast_if_needed, device_guard
-from paddleformers.utils.initializer import kaiming_uniform_
-from paddleformers.utils.log import logger
+from paddleformers.transformers.utils import device_guard
 from paddleformers.utils.tools import get_env_device
 
 try:
@@ -117,13 +83,7 @@ def swiglu(x, y=None):
 except ImportError:
     fused_partial_rope = None
 
-from paddleformers.transformers.deepseek_v2 import (
-    DeepseekV2ForCausalLM,
-    DeepseekV2ForSequenceClassification,
-    DeepseekV2Model,
-    DeepseekV2PretrainedModel,
-    DeepseekV2PretrainingCriterion,
-)
+from paddleformers.transformers.deepseek_v2 import rotate_half
 
 __all__ = [
     "DeepseekV2LMHead",
@@ -153,6 +113,13 @@ def rms_norm_fused(x_in, w, eps, use_fast_ln=False):
         return fused_ln.fused_rms_norm(x_in, w, eps)[0]
 
 
+def cast_if_needed(x, dtype):
+    """
+    cast_if_needed
+    """
+    return x.cast(dtype) if x.dtype != dtype else x
+
+
 def fusion_rms_norm(hidden_states, weight, variance_epsilon, use_fast_ln=False):
     if get_env_device() == "npu":
         return paddle.base.core.eager._run_custom_op("rms_norm_npu", hidden_states, weight, variance_epsilon)[0]

paddleformers/examples/deepseek_v3/modeling.py

@@ -23,7 +23,6 @@
 
 import contextlib
 import math
-import os
 import warnings
 from functools import partial
 from typing import List, Optional, Tuple, Union
@@ -65,7 +64,7 @@
     init_name_mappings,
 )
 from paddleformers.transformers.deepseek_v2 import fp8_linear as linear_utils
-from paddleformers.transformers.deepseek_v2.fp8_linear import Linear
+from paddleformers.transformers.deepseek_v2.fp8_linear import Linear as Linear_
 from paddleformers.transformers.fp8_utils import (
     FP8KeepXLinear,
     FP8Linear,
@@ -100,7 +99,6 @@ def swiglu(x, y=None):
     fused_partial_rope = None
 from modeling import (
     AddAuxiliaryLoss,
-    DeepseekV2PretrainingCriterion,
     DeepseekV2RMSNorm,
     DeepseekV2RotaryEmbedding,
     DeepseekV2YarnRotaryEmbedding,
@@ -117,6 +115,7 @@ def swiglu(x, y=None):
 from paddleformers.transformers.deepseek_v2 import (
     DeepseekV2DynamicNTKScalingRotaryEmbedding,
     DeepseekV2LinearScalingRotaryEmbedding,
+    DeepseekV2PretrainingCriterion,
     _expand_2d_mask,
     _make_causal_mask,
     is_casual_mask,
@@ -168,7 +167,7 @@ def __init__(self, config: DeepseekV2FastConfig, hidden_size=None, intermediate_
         self.hidden_size = config.hidden_size if hidden_size is None else hidden_size
         self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size
         self.fuse_attention_ffn = config.fuse_attention_ffn
-        Linear = FP8Linear if self.config.dsv3_use_fp8_gemm else Linear
+        Linear = FP8Linear if self.config.dsv3_use_fp8_gemm else Linear_
 
         def linear_dtype_gaurd():
             if config.use_fp8:
@@ -534,7 +533,7 @@ def linear_dtype_gaurd():
 
         self._init_rope()
         self.softmax_scale = self.q_head_dim ** (-0.5)
-        Linear = FP8Linear if self.config.dsv3_use_fp8_gemm else Linear
+        Linear = FP8Linear if self.config.dsv3_use_fp8_gemm else Linear_
 
         # fmt: off
         if self.config.tensor_parallel_degree > 1:
@@ -1260,7 +1259,7 @@ def get_tensor_parallel_split_mappings(num_layers):
     def _init_weights(self, layer):
         if self.config.tensor_parallel_degree > 1:
             rng_tracker = get_rng_state_tracker().rng_state
-        Linear = FP8Linear if self.config.dsv3_use_fp8_gemm else Linear
+        Linear = FP8Linear if self.config.dsv3_use_fp8_gemm else Linear_
 
         if isinstance(
             layer,

paddleformers/examples/deepseek_v3/modeling_fast.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 import math
-import os
 from typing import OrderedDict, Tuple, Union
 
 import paddle
@@ -37,7 +36,6 @@
 from config.configuration import DeepseekV2FastConfig
 from modeling import (
     DeepseekV2LMHead,
-    DeepseekV2PretrainingCriterion,
     DeepseekV2RMSNorm,
     TemporaryVarContext,
     set_global_step,
@@ -49,7 +47,7 @@
 from paddle.distributed.fleet.recompute.recompute import recompute
 from paddle.distributed.fleet.utils.sequence_parallel_utils import ScatterOp
 
-from paddleformers.transformers.deepseek_v2 import DeepseekV2RotaryEmbedding
+from paddleformers.transformers.deepseek_v2 import DeepseekV2PretrainingCriterion
 from paddleformers.transformers.model_utils import PipelinePretrainedModel
 from paddleformers.utils.log import logger
 from paddleformers.utils.tools import get_env_device

paddleformers/examples/deepseek_v3/modeling_pp.py

@@ -19,5 +19,5 @@
 
 # mpirun sh script/kill_process.sh 
 # mpirun rm -rf output
-nohup bash script/train_gpu.sh ./config/pretrain_argument.json  > run.log 2>&1 &
+nohup bash script/train_gpu.sh ./config/pretrain_argument.json --dsv3_fast_pretrain=True > run.log 2>&1 &
 

paddleformers/examples/deepseek_v3/run.sh

@@ -21,7 +21,7 @@
 
 import paddle
 from config.configuration import DeepseekV2FastConfig
-from modeling import DeepseekV2ForCausalLM
+from load_hf_ckpt import load_huggingface_ckpt
 from modeling_pp import DeepseekV2ForCausalLMPipe
 
 from paddleformers.data.causal_dataset import (
@@ -40,14 +40,12 @@
     speed_metrics,
 )
 from paddleformers.transformers import (
-    AutoConfig,
-    AutoModelForCausalLM,
-    AutoModelForCausalLMPipe,
     AutoTokenizer,
     CosineAnnealingWithWarmupDecay,
     LinearAnnealingWithWarmupDecay,
 )
 from paddleformers.transformers.configuration_utils import LlmMetaConfig, llmmetaclass
+from paddleformers.transformers.deepseek_v2 import DeepseekV2ForCausalLM
 from paddleformers.utils.batch_sampler import DistributedBatchSampler
 from paddleformers.utils.log import logger
 from paddleformers.utils.tools import get_env_device
@@ -413,8 +411,7 @@ def main():
                 "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
             )
 
-    tokenizer = AutoTokenizer.from_pretrained(model_args.tokenizer_name_or_path, **{"download_hub": "bos"})
-    # config = AutoConfig.from_pretrained("./")
+    tokenizer = AutoTokenizer.from_pretrained(model_args.tokenizer_name_or_path, download_hub="huggingface")
     config = DeepseekV2FastConfig.from_pretrained("./config/config.json")
 
     # set all llm config
@@ -583,6 +580,12 @@ def main():
 
     callbacks = [StepFlexToken(), FP8QuantWeightCallback()]
 
+    def resume_from_custom_func(model):
+        if training_args.resume_from_huggingface_ckpt:
+            load_huggingface_ckpt(model, training_args.resume_from_huggingface_ckpt)
+        else:
+            logger.info("No resume from checkpoint since training args 'resume_from_huggingface_ckpt' is None.")
+
     trainer = PretrainingTrainer(
         model=model,
         args=training_args,
@@ -592,6 +595,7 @@ def main():
         optimizers=(None, lr_scheduler),
         tokenizer=tokenizer,
         callbacks=callbacks,
+        resume_from_custom_func=resume_from_custom_func,
     )
 
     checkpoint = None

paddleformers/examples/deepseek_v3/run_pretrain.py

@@ -86,15 +86,12 @@
 )
 from ..peft import LoKrModel, LoRAModel, PrefixModelForCausalLM, ReFTModel, VeRAModel
 from ..peft.lora import QuantizationLoRABaseLinear
+from ..quantization.quantization_linear import (
+    ColumnParallelQuantizationLinear,
+    QuantizationLinear,
+    RowParallelQuantizationLinear,
+)
 
-try:
-    from ..quantization.quantization_linear import (
-        ColumnParallelQuantizationLinear,
-        QuantizationLinear,
-        RowParallelQuantizationLinear,
-    )
-except:
-    QuantizationLinear = None
 try:
     from paddle.distributed.fleet.utils.sequence_parallel_utils import (
         register_sequence_parallel_allreduce_hooks,
@@ -199,7 +196,6 @@
     nested_numpify,
     nested_truncate,
 )
-from .utils.load_hf_ckpt import load_huggingface_ckpt
 from .utils.sharding_io import ShardingIO
 
 DEFAULT_CALLBACKS = [DefaultFlowCallback]
@@ -307,6 +303,7 @@ def __init__(
         optimizers: Tuple[paddle.optimizer.Optimizer, paddle.optimizer.lr.LRScheduler] = (None, None),
         preprocess_logits_for_metrics: Callable[[paddle.Tensor, paddle.Tensor], paddle.Tensor] = None,
         processing_class: Optional[ImageProcessingMixin] = None,
+        resume_from_custom_func: Optional[Callable] = None,
     ):
 
         if args is None:
@@ -361,6 +358,7 @@ def __init__(
         self.train_dataset = train_dataset
         self.eval_dataset = eval_dataset
         self.tokenizer = tokenizer
+        self.resume_from_custom_func = resume_from_custom_func
         if not args.skip_profile_timer:
             set_timers()
         self.timers = get_timers()
@@ -1138,8 +1136,8 @@ def _inner_training_loop(
         if self.args.ignore_data_skip:
             self.timers and self.timers("read-data").start()
 
-        if self.args.resume_from_huggingface_ckpt is not None:
-            load_huggingface_ckpt(model, self.args.resume_from_huggingface_ckpt)
+        if self.resume_from_custom_func is not None:
+            self.resume_from_custom_func(self.model)
 
         for epoch in range(epochs_trained, num_train_epochs):
             if isinstance(train_dataloader, paddle.io.DataLoader) and isinstance(

paddleformers/trainer/trainer.py

@@ -1092,6 +1092,10 @@ class TrainingArguments:
         default=False,
         metadata={"help": "Save model to HuggingFace safetensors."},
     )
+    dsv3_fast_pretrain: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Use fast pretrain version of DeepSeekV3."},
+    )
 
     def __post_init__(self):
         world_size = paddle.distributed.get_world_size()
@@ -1409,7 +1413,7 @@ def is_segment_parallel_supported():
                     else:
                         order = ["dp", "sharding", "pp", "mp"]
                 if self.use_expert_parallel:
-                    if not os.getenv("DSV3_FAST_PRETRAIN", "False"):
+                    if not self.dsv3_fast_pretrain:
                         if self.moe_sharding_parallel_degree >= 1 and self.expert_parallel_degree > 1:
                             order.insert(-1, "ep")
                             sd_idx = order.index("sharding")
@@ -1571,7 +1575,8 @@ def is_segment_parallel_supported():
                 fleet.init(is_collective=True, strategy=strategy)
                 logger.info(strategy)
 
-                if os.getenv("DSV3_FAST_PRETRAIN", "False"):
+                # if os.getenv("DSV3_FAST_PRETRAIN", "False"):
+                if self.dsv3_fast_pretrain:
                     if self.expert_parallel_degree > 1:
                         self.add_moe_comm_group()
 

paddleformers/trainer/training_args.py

@@ -35,7 +35,6 @@
     [
         ("bert", "BertConfig"),
         ("deepseek_v2", "DeepseekV2Config"),
-        ("deepseek_v2_fast", "DeepseekV2FastConfig"),
         ("deepseek_v3", "DeepseekV3Config"),
         ("ernie4_5", "Ernie4_5Config"),
         ("llama", "LlamaConfig"),

paddleformers/transformers/auto/configuration.py

@@ -56,8 +56,6 @@
         "yarn_find_correction_range",
         "get_triangle_upper_mask",
         "DeepseekV2LinearScalingRotaryEmbedding",
-        "set_global_step",
-        "get_global_step",
     ],
     "modeling_auto": [
         "DeepseekV2LMHeadAuto",