[LLM] Add expert parallel (#9368)

* add expert parallel utils

* update gates

* update

* update base methods

* update moe_layer

* update moebase

* add moe_gate and moe_layer for qwen2moe

* add config

* update

* update gate dtype

* update moe gate and layer

* update moe_layer.py

* update

* update

* update token_priority method

* update data type

* remove old moe

* fix moe capacity reduce.Max

* update comment

* lint

llm/config/qwen2moe/lora_argument.json

@@ -0,0 +1,34 @@
+{
+    "model_name_or_path": "Qwen/Qwen2-57B-A14B",
+    "dataset_name_or_path": "./data",
+    "output_dir": "./checkpoints/lora_ckpts",
+    "per_device_train_batch_size": 4,
+    "gradient_accumulation_steps": 4,
+    "per_device_eval_batch_size": 8,
+    "eval_accumulation_steps":16,
+    "num_train_epochs": 3,
+    "learning_rate": 3e-04,
+    "warmup_steps": 30,
+    "logging_steps": 1,
+    "evaluation_strategy": "epoch",
+    "save_strategy": "epoch",
+    "src_length": 1024,
+    "max_length": 2048,
+    "bf16": true,
+    "fp16_opt_level": "O2",
+    "do_train": true,
+    "do_eval": true,
+    "disable_tqdm": true,
+    "load_best_model_at_end": true,
+    "eval_with_do_generation": false,
+    "metric_for_best_model": "accuracy",
+    "recompute": true,
+    "save_total_limit": 1,
+    "tensor_parallel_degree": 1,
+    "pipeline_parallel_degree": 1,
+    "lora": true,
+    "unified_checkpoint": true,
+    "zero_padding": false,
+    "use_flash_attention": true,
+    "pissa": false
+  }

llm/config/qwen2moe/pretrain_argument.json

@@ -0,0 +1,40 @@
+{
+    "model_name_or_path": "Qwen/Qwen2-57B-A14B",
+    "tokenizer_name_or_path": "Qwen/Qwen2-57B-A14B",
+    "input_dir": "./data",
+    "output_dir": "./checkpoints/pretrain_ckpts",
+    "per_device_train_batch_size": 2,
+    "gradient_accumulation_steps": 1,
+    "per_device_eval_batch_size": 2,
+    "tensor_parallel_degree": 1,
+    "pipeline_parallel_degree": 1,
+    "sharding": "stage2",
+    "virtual_pp_degree": 1,
+    "sequence_parallel": 0,   
+    "use_flash_attention": true,
+    "use_fused_rms_norm": true,
+    "max_seq_length": 4096,
+    "learning_rate": 3e-05,
+    "min_learning_rate": 3e-06,
+    "warmup_steps": 30,
+    "logging_steps": 1,
+    "max_steps": 10000,
+    "save_steps": 5000,
+    "eval_steps": 1000,
+    "weight_decay": 0.01,
+    "bf16": true,
+    "fp16_opt_level": "O2",
+    "warmup_ratio": 0.01,
+    "max_grad_norm": 1.0,
+    "dataloader_num_workers": 1,
+    "continue_training": 1,
+    "do_train": true,
+    "do_eval": true,
+    "do_predict": true,
+    "disable_tqdm": true,
+    "recompute": true,
+    "distributed_dataloader": 1,
+    "recompute_granularity": "full",
+    "unified_checkpoint": true,
+    "save_total_limit": 2
+  }

llm/config/qwen2moe/sft_argument.json

@@ -0,0 +1,33 @@
+{
+    "model_name_or_path": "Qwen/Qwen2-57B-A14B",
+    "dataset_name_or_path": "./data",
+    "output_dir": "./checkpoints/sft_ckpts",
+    "per_device_train_batch_size": 1,
+    "gradient_accumulation_steps": 4,
+    "per_device_eval_batch_size": 8,
+    "eval_accumulation_steps":16,
+    "num_train_epochs": 3,
+    "learning_rate": 3e-05,
+    "warmup_steps": 30,
+    "logging_steps": 1,
+    "evaluation_strategy": "epoch",
+    "save_strategy": "epoch",
+    "src_length": 1024,
+    "max_length": 2048,
+    "bf16": true,
+    "fp16_opt_level": "O2",
+    "do_train": true,
+    "do_eval": true,
+    "disable_tqdm": true,
+    "load_best_model_at_end": true,
+    "eval_with_do_generation": false,
+    "metric_for_best_model": "accuracy",
+    "recompute": true,
+    "save_total_limit": 1,
+    "tensor_parallel_degree": 1,
+    "pipeline_parallel_degree": 1,
+    "sharding": "stage2",
+    "zero_padding": false,
+    "unified_checkpoint": true,
+    "use_flash_attention": true
+  }

paddlenlp/trainer/training_args.py

@@ -883,6 +883,14 @@ class TrainingArguments:
         default=False,
         metadata={"help": "Enable MoE (Mixture of Experts) expert parallel training"},
     )
+    expert_max_capacity: Optional[int] = field(
+        default=pow(2, 32),
+        metadata={"help": "Enable MoE (Mixture of Experts) expert max token capacity"},
+    )
+    expert_min_capacity: Optional[int] = field(
+        default=1,
+        metadata={"help": "Enable MoE (Mixture of Experts) expert min token capacity"},
+    )
     release_grads: Optional[bool] = field(
         default=False, metadata={"help": "Whether to release gradients during training. Default is `False`."}
     )

paddlenlp/transformers/__init__.py

@@ -32,6 +32,8 @@
 from .attention_utils import create_bigbird_rand_mask_idx_list
 from .sequence_parallel_utils import AllGatherVarlenOp, sequence_parallel_sparse_mask_labels
 from .tensor_parallel_utils import parallel_matmul, parallel_linear, fused_head_and_loss_fn
+from .moe_gate import *
+from .moe_layer import *
 
 try:
     from paddle.distributed.fleet.utils.sequence_parallel_utils import (