[Cherry-pick] Fix safetensors shape

docs/trainer.md

@@ -705,4 +705,8 @@ Trainer 是一个简单，但功能完整的 Paddle训练和评估模块，并
                        Whether use flatten_param_grads method in optimizer,
                        only used on NPU devices.(default:False)
 
+  --use_expert_parallel
+                       Whether to enable MoE (Mixture of Experts) expert parallel training.
+                       (default: False)
+
 ```

llm/finetune_generation.py

@@ -16,6 +16,7 @@
 import sys
 from dataclasses import dataclass, field
 from functools import partial
+from typing import Optional
 
 import paddle
 from argument import (
@@ -45,6 +46,7 @@
     AutoConfig,
     AutoModelForCausalLM,
     AutoTokenizer,
+    Llama3Tokenizer,
     LlamaTokenizer,
 )
 from paddlenlp.utils.log import logger
@@ -65,6 +67,10 @@ class FinetuneArguments(TrainingArguments):
         default=0,
         metadata={"help": "The steps use to control the learing rate."},
     )
+    tensor_parallel_output: Optional[bool] = field(
+        default=False,
+        metadata={"help": "whether to output logits in distributed status"},
+    )
 
 
 def read_local_dataset(path):
@@ -139,7 +145,7 @@ def main():
         if not training_args.autotuner_benchmark:
             model = AutoModelForCausalLMPipe.from_pretrained(
                 model_args.model_name_or_path,
-                tensor_parallel_output=False,
+                tensor_parallel_output=training_args.tensor_parallel_output,
                 tensor_parallel_degree=training_args.tensor_parallel_degree,
                 tensor_parallel_rank=training_args.tensor_parallel_rank,
                 use_flash_attention=model_args.use_flash_attention,
@@ -151,7 +157,7 @@ def main():
             # NOTE(gongenlei): new add autotuner_benchmark
             model_config = AutoConfig.from_pretrained(
                 model_args.model_name_or_path,
-                tensor_parallel_output=False,
+                tensor_parallel_output=training_args.tensor_parallel_output,
                 tensor_parallel_degree=training_args.tensor_parallel_degree,
                 tensor_parallel_rank=training_args.tensor_parallel_rank,
                 dtype=dtype,
@@ -162,7 +168,7 @@ def main():
     else:
         model_config = AutoConfig.from_pretrained(
             model_args.model_name_or_path,
-            tensor_parallel_output=False,
+            tensor_parallel_output=training_args.tensor_parallel_output,
             tensor_parallel_degree=training_args.tensor_parallel_degree,
             tensor_parallel_rank=training_args.tensor_parallel_rank,
             dtype=dtype,
@@ -232,7 +238,7 @@ def neft_post_hook(module, input, output):
     if tokenizer.chat_template is not None:
         data_args.eval_with_do_generation = False
 
-    if isinstance(tokenizer, LlamaTokenizer):
+    if isinstance(tokenizer, LlamaTokenizer) or isinstance(tokenizer, Llama3Tokenizer):
         tokenizer.pad_token_id = tokenizer.eos_token_id
 
     if data_args.dataset_name_or_path is None:

llm/run_pretrain.py

@@ -46,6 +46,7 @@
 )
 from paddlenlp.utils.batch_sampler import DistributedBatchSampler
 from paddlenlp.utils.log import logger
+from paddlenlp.utils.tools import get_env_device
 
 
 def add_start_docstrings(*docstr):
@@ -483,6 +484,16 @@ def main():
         config.num_attention_heads % config.sep_parallel_degree == 0
     ), f"num_attention_heads:{config.num_attention_heads} must be divisible by sep_parallel_degree {config.sep_parallel_degree}"
 
+    if get_env_device() == "xpu" and training_args.gradient_accumulation_steps > 1:
+        try:
+            from paddle_xpu.layers.nn.linear import LinearConfig  # noqa: F401
+
+            LinearConfig.enable_accumulate_steps_opt()
+            LinearConfig.set_accumulate_steps(training_args.gradient_accumulation_steps)
+        except ImportError:
+            # It's OK, not use accumulate_steps optimization
+            pass
+
     print("Final pre-training config:", config)
 
     # Set the dtype for loading model

llm/utils.py

@@ -125,9 +125,11 @@ def get_lora_target_modules(model):
             ".*v_proj.*",
             ".*k_proj.*",
             ".*o_proj.*",
+            ".*qkv_proj.*",
             ".*gate_proj.*",
             ".*down_proj.*",
             ".*up_proj.*",
+            ".*gate_up_fused_proj.*",
         ]
     elif model.base_model_prefix == "opt":
         target_modules = [
@@ -209,6 +211,13 @@ def prediction_step(
             # keepdim in order to maintain the same shape as logits
             if isinstance(logits, (list, tuple)):
                 logits = logits[0]
+            # all gather logits when enabling tensor_parallel_output
+            if self.args.tensor_parallel_degree > 1 and getattr(self.args, "tensor_parallel_output", False):
+                hcg = fleet.get_hybrid_communicate_group()
+                model_parallel_group = hcg.get_model_parallel_group()
+                gathered_logits = []
+                dist.all_gather(gathered_logits, logits, group=model_parallel_group)
+                logits = paddle.concat(gathered_logits, axis=-1)
             return (loss, logits.argmax(axis=-1, keepdim=True), labels)
 
         loss = None

model_zoo/gpt-3/ppfleetx/models/language_model/gpt/dygraph/hybrid_model.py

@@ -48,13 +48,16 @@
     MinLengthLogitsProcessor,
     RepetitionPenaltyLogitsProcessor,
 )
-from paddle.distributed.fleet.utils.sequence_parallel_utils import (
-    ColumnSequenceParallelLinear,
-    GatherOp,
-    RowSequenceParallelLinear,
-    ScatterOp,
-    mark_as_sequence_parallel_parameter,
-)
+try:
+    from paddle.distributed.fleet.utils.sequence_parallel_utils import (
+        ColumnSequenceParallelLinear,
+        GatherOp,
+        RowSequenceParallelLinear,
+        ScatterOp,
+        mark_as_sequence_parallel_parameter,
+    )
+except:
+    pass
 
 from paddlenlp.transformers.segment_parallel_utils  import ReshardLayer
 

model_zoo/gpt-3/ppfleetx/models/language_model/language_module.py

@@ -24,9 +24,12 @@
 from ppfleetx.core.module.basic_module import BasicModule
 from ppfleetx.data.tokenizers import GPTTokenizer
 from ppfleetx.distributed.apis import env
-from paddle.distributed.fleet.utils.sequence_parallel_utils import (
-    register_sequence_parallel_allreduce_hooks,
-)
+try:
+    from paddle.distributed.fleet.utils.sequence_parallel_utils import (
+        register_sequence_parallel_allreduce_hooks,
+    )
+except:
+    pass
 from ppfleetx.utils.log import logger
 
 # TODO(haohongxiang): to solve the problem of cross-reference

paddlenlp/__init__.py

@@ -18,7 +18,7 @@
 PADDLENLP_STABLE_VERSION = "PADDLENLP_STABLE_VERSION"
 
 
-__version__ = "2.7.1.post"
+__version__ = "2.8.1.post"
 if os.getenv(PADDLENLP_STABLE_VERSION):
     __version__ = __version__.replace(".post", "")