[LLM] support trainer for gpt and llama pre-training. (#6053)

* support trainer for gpt and llama pretrain.

* update

* fix

* fix use flash.

* support llama pre-train and post-train

* fix

* support use virtual pp degree.

Author: ZHUI

examples/language_model/llama/README.md

@@ -13,6 +13,65 @@ Llama 模型的权重的使用则需要遵循[License](../../../paddlenlp/transf
 
 <a name="1"></a>
 
+## 预训练
+
+预训练数据制作参考[此处](../../../model_zoo/ernie-1.0/preprocess/docs/OpenWebText2.md)
+
+为了方便用户运行测试本模型，本项目提供了处理好的100k条doc的训练样本：
+```shell
+wget https://bj.bcebos.com/paddlenlp/models/transformers/llama/data/llama_openwebtext_100k_ids.npy
+wget https://bj.bcebos.com/paddlenlp/models/transformers/llama/data/llama_openwebtext_100k_idx.npz
+```
+
+将所有预处理得到的文件统一放入一个文件夹中，以备训练使用：
+
+```
+mkdir data
+mv llama_openwebtext_100k_ids.npy ./data
+mv llama_openwebtext_100k_idx.npz ./data
+```
+
+使用下面脚本,即可在llama-7b的基础上,继续训练.
+```shell
+python -u  -m paddle.distributed.launch \
+    --gpus "0,1,2,3,4,5,6,7" \
+    --log_dir "output/$task_name""_log" \
+    run_pretrain.py \
+    --model_type "llama" \
+    --model_name_or_path "facebook/llama-7b" \
+    --tokenizer_name_or_path "facebook/llama-7b" \
+    --input_dir "./data" \
+    --output_dir "output/$task_name" \
+    --split 949,50,1 \
+    --max_seq_length 2048 \
+    --per_device_train_batch_size 1 \
+    --per_device_eval_batch_size 1 \
+    --use_flash_attention 1 \
+    --use_fused_rms_norm 0 \
+    --fp16  \
+    --fp16_opt_level "O2"  \
+    --scale_loss 1024 \
+    --learning_rate 0.00001 \
+    --min_learning_rate 0.000005 \
+    --max_steps 10000 \
+    --save_steps 5000 \
+    --weight_decay 0.01 \
+    --warmup_ratio 0.01 \
+    --max_grad_norm 1.0 \
+    --logging_steps 20\
+    --dataloader_num_workers 1 \
+    --sharding "stage2" \
+    --eval_steps 1000 \
+    --report_to "visualdl" \
+    --disable_tqdm true \
+    --continue_training 1\
+    --recompute 1 \
+    --do_train \
+    --do_eval \
+    --device "gpu"
+```
+
+
 ## 微调
 
 ```shell

examples/language_model/llama/dataset.py

@@ -0,0 +1 @@
+../../..//model_zoo/gpt/dataset.py

examples/language_model/llama/modeling_pp.py

@@ -121,6 +121,7 @@ class PipelinePretrainedModel(PretrainedModel):
     _pipeline_name_mapping = None
 
     def __init__(self, config, *args, **kwargs):
+        raise ValueError()
         super().__init__(config, *args, **kwargs)
 
     def add_sequential_layer(self, layer_desc, name_prefix=""):
@@ -138,23 +139,39 @@ def _set_pipeline_name_mapping(self, mappings=None):
         else:
             mapping = {}
             state_dict_keys = list(super().state_dict().keys())
+            first_key = state_dict_keys[0].split(".")
+            # if use virtual pp_degree, the prefix is like 0.0.xxx
+            # else it will be like 0.xxx
+            use_virtual_pp_degree = first_key[0].isdigit() and first_key[1].isdigit()
+
             prefixs = self.get_sequential_name_prefixs()
             for k in state_dict_keys:
                 name_splited = k.split(".")
-                name_splited[0] = prefixs[name_splited[0]]
-                mapping[".".join(name_splited)] = k
+                if use_virtual_pp_degree:
+                    idx = str(int(name_splited[0]) + int(name_splited[1]))
+                    single_name = [prefixs[idx]]
+                    single_name.extend(name_splited[2:])
+                else:
+                    idx = name_splited[0]
+                    single_name = [prefixs[idx]]
+                    single_name.extend(name_splited[1:])
+                mapping[".".join(single_name)] = k
+
             self._pipeline_name_mapping = mapping
 
         return self._pipeline_name_mapping
 
     def state_dict(self, *args, **kwargs):
         state_dict = super().state_dict(*args, **kwargs)
-        prefixs = self.get_sequential_name_prefixs()
+
+        if self._pipeline_name_mapping is None:
+            self._set_pipeline_name_mapping()
+        assert len(self._pipeline_name_mapping) > 0, "The pipeline stage must have parameters!"
+        pp_to_single_mapping = {v: k for k, v in self._pipeline_name_mapping.items()}
+
         for k in list(state_dict.keys()):
             v = state_dict.pop(k)
-            name_splited = k.split(".")
-            name_splited[0] = prefixs[name_splited[0]]
-            state_dict[".".join(name_splited)] = v
+            state_dict[pp_to_single_mapping[k]] = v
 
         return state_dict
 
@@ -169,7 +186,8 @@ def set_state_dict(self, state_dict, *args, **kwargs):
                 continue
             state_dict[self._pipeline_name_mapping[k]] = v
 
-        return super().set_state_dict(state_dict, *args, **kwargs)
+        ret = super().set_state_dict(state_dict, *args, **kwargs)
+        return ret
 
 
 class LlamaForCausalLMPipe(PipelinePretrainedModel, PipelineLayer):
@@ -182,28 +200,25 @@ class LlamaForCausalLMPipe(PipelinePretrainedModel, PipelineLayer):
     config_class = LlamaConfig
 
     _get_tensor_parallel_mappings = LlamaPretrainedModel._get_tensor_parallel_mappings
+
     # NO base_model_prefix !!!!
 
     def __init__(
         self,
         config,
-        # num_partitions=1,
-        # topology=None,
-        use_recompute=None,
-        # fused_linear=False,
-        # fuse_attn_qkv=False,
+        # use_recompute=None,
         # scale_qk_by_layer_num=True,
-        recompute_granularity="full",
-        virtual_pp_degree=1,
+        # recompute_granularity="full",
+        # virtual_pp_degree=4,
         # sequence_parallel=False,
         # no_recompute_layers=None,
         pp_recompute_interval=1,
-        # use_flash_attn=False,
-        # fused_softmax_with_triangular=False,
     ):
         self.config = config
-        if use_recompute is None:
-            use_recompute = self.config.use_recompute
+
+        use_recompute = self.config.use_recompute
+        recompute_granularity = self.config.recompute_granularity
+        virtual_pp_degree = self.config.virtual_pp_degree
 
         hcg = get_hcg()
         tensor_parallel_degree = max(hcg.get_model_parallel_world_size(), 1)