fix llama pretrain init. (#6116)

Author: ZHUI

examples/language_model/llama/modeling_pp.py

@@ -199,6 +199,7 @@ class LlamaForCausalLMPipe(PipelinePretrainedModel, PipelineLayer):
     config_class = LlamaConfig
 
     _get_tensor_parallel_mappings = LlamaPretrainedModel._get_tensor_parallel_mappings
+    _init_weights = LlamaPretrainedModel._init_weights
 
     # NO base_model_prefix !!!!
 
@@ -258,5 +259,6 @@ def __init__(
             },
             num_virtual_pipeline_stages=virtual_pp_degree,
         )
+        self.apply(self._init_weights)
         # DON'T init PipelinePretrainedModel
         # PipelinePretrainedModel.__init__(self.super(), config=config)

examples/language_model/llama/run_pretrain.py

@@ -111,19 +111,18 @@ class ModelArguments:
         default="llama", metadata={"help": "Only support for llama pre-training for now."}
     )
     model_name_or_path: str = field(
-        default="gpt2-meidum-en",
+        default="facebook/tiny-random-llama",
         metadata={
             "help": "Path to pretrained model or model identifier from https://paddlenlp.readthedocs.io/zh/latest/model_zoo/transformers.html"
         },
     )
-    hidden_dropout_prob: float = field(default=0.1, metadata={"help": "The hidden dropout prob."})
-    attention_probs_dropout_prob: float = field(default=0.1, metadata={"help": "The attention probs dropout prob."})
+    tokenizer_name_or_path: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
+    )
+
     config_name: Optional[str] = field(
         default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
     )
-    tokenizer_name_or_path: Optional[str] = field(
-        default="gpt2-meidum-en", metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
-    )
     use_flash_attention: bool = field(
         default=False,
         metadata={"help": "use_flash_attention"},
@@ -202,7 +201,7 @@ def print_dataset(data, mode="train"):
     def build_dataset(index, name):
         dataset = GPTDataset(
             file_prefix=input_prefix,
-            build_data_file=training_args.local_rank == 0,
+            build_data_file=training_args.local_process_index == 0,
             micro_batch_size=training_args.per_device_train_batch_size
             if name == "train"
             else training_args.per_device_eval_batch_size,

model_zoo/gpt/run_pretrain_trainer.py

@@ -195,7 +195,7 @@ def print_dataset(data, mode="train"):
     def build_dataset(index, name):
         dataset = GPTDataset(
             file_prefix=input_prefix,
-            build_data_file=training_args.local_rank == 0,
+            build_data_file=training_args.local_process_index == 0,
             micro_batch_size=training_args.per_device_train_batch_size
             if name == "train"
             else training_args.per_device_eval_batch_size,

paddlenlp/trainer/trainer.py

@@ -115,6 +115,11 @@
 except:
     mix_precision_utils = None
 
+try:
+    from paddle.io.dataloader.dataloader_iter import _DataLoaderIterBase
+except:
+    from paddle.fluid.dataloader.dataloader_iter import _DataLoaderIterBase
+
 
 def paddlenlp_load(path, return_numpy=False):
     if return_numpy:
@@ -752,9 +757,9 @@ def train(
                         for p in model._layers.parameters():
                             if hasattr(p, "main_grad") and p.main_grad is not None:
                                 assert p.grad is None
-                                p.main_grad = p.main_grad.scale(1.0 / model.accumulate_steps)
+                                p.main_grad = p.main_grad.scale(1.0 / self.args.gradient_accumulation_steps)
                             elif p.grad is not None:
-                                p.grad = p.grad.scale(1.0 / model.accumulate_steps)
+                                p.grad = p.grad.scale(1.0 / self.args.gradient_accumulation_steps)
 
                     # Optimizer step
                     optimizer_was_run = True
@@ -1930,7 +1935,7 @@ def evaluation_loop(
 
         if isinstance(dataloader, paddle.io.DataLoader):
             batch_size = dataloader.batch_sampler.batch_size
-        elif isinstance(dataloader, paddle.fluid.dataloader.dataloader_iter._DataLoaderIterBase):
+        elif isinstance(dataloader, _DataLoaderIterBase):
             # support for inner dataloader
             batch_size = dataloader._batch_sampler.batch_size
             # alias for inner dataloader
@@ -1942,7 +1947,7 @@ def evaluation_loop(
         if max_eval_iters > 0:
             # on eval limit steps
             num_samples = batch_size * self.args.dataset_world_size * max_eval_iters
-            if isinstance(dataloader, paddle.fluid.dataloader.dataloader_iter._DataLoaderIterBase) and isinstance(
+            if isinstance(dataloader, _DataLoaderIterBase) and isinstance(
                 dataloader._batch_sampler, NlpDistributedBatchSampler
             ):
                 consumed_samples = (

paddlenlp/transformers/llama/modeling.py

@@ -21,6 +21,7 @@
 
 import numpy as np
 import paddle
+import paddle.distributed.fleet.meta_parallel as mpu
 import paddle.nn.functional as F
 from paddle import Tensor, nn
 from paddle.distributed import fleet
@@ -296,19 +297,19 @@ def __init__(self, config):
         self.intermediate_size = config.intermediate_size
 
         if config.tensor_parallel_degree > 1:
-            self.gate_proj = fleet.meta_parallel.ColumnParallelLinear(
+            self.gate_proj = mpu.ColumnParallelLinear(
                 self.hidden_size,
                 self.intermediate_size,
                 gather_output=False,
                 has_bias=False,
             )
-            self.down_proj = fleet.meta_parallel.RowParallelLinear(
+            self.down_proj = mpu.RowParallelLinear(
                 self.intermediate_size,
                 self.hidden_size,
                 input_is_parallel=True,
                 has_bias=False,
             )
-            self.up_proj = fleet.meta_parallel.ColumnParallelLinear(
+            self.up_proj = mpu.ColumnParallelLinear(
                 self.hidden_size,
                 self.intermediate_size,
                 gather_output=False,
@@ -339,19 +340,19 @@ def __init__(self, config):
             self.num_heads = self.num_heads // config.tensor_parallel_degree
 
         if config.tensor_parallel_degree > 1:
-            self.q_proj = fleet.meta_parallel.ColumnParallelLinear(
+            self.q_proj = mpu.ColumnParallelLinear(
                 self.hidden_size,
                 self.hidden_size,
                 has_bias=False,
                 gather_output=False,
             )
-            self.k_proj = fleet.meta_parallel.ColumnParallelLinear(
+            self.k_proj = mpu.ColumnParallelLinear(
                 self.hidden_size,
                 self.hidden_size,
                 has_bias=False,
                 gather_output=False,
             )
-            self.v_proj = fleet.meta_parallel.ColumnParallelLinear(
+            self.v_proj = mpu.ColumnParallelLinear(
                 self.hidden_size,
                 self.hidden_size,
                 has_bias=False,
@@ -375,7 +376,7 @@ def __init__(self, config):
             )
 
         if config.tensor_parallel_degree > 1:
-            self.o_proj = fleet.meta_parallel.RowParallelLinear(
+            self.o_proj = mpu.RowParallelLinear(
                 self.hidden_size,
                 self.hidden_size,
                 has_bias=False,
@@ -581,7 +582,17 @@ def get_tensor_parallel_split_mappings(num_layers):
 
     def _init_weights(self, layer):
         """Initialization hook"""
-        if isinstance(layer, (nn.Linear, nn.Embedding)):
+        if isinstance(
+            layer,
+            (
+                nn.Linear,
+                nn.Embedding,
+                mpu.VocabParallelEmbedding,
+                mpu.ColumnParallelLinear,
+                mpu.RowParallelLinear,
+                LlamaLMHead,
+            ),
+        ):
             # In the dygraph mode, use the `set_value` to reset the parameter directly,
             # and reset the `state_dict` to update parameter in static mode.
             if isinstance(layer.weight, paddle.Tensor):
@@ -594,6 +605,16 @@ def _init_weights(self, layer):
                         shape=layer.weight.shape,
                     )
                 )
+        # Layer.apply is DFS https://github.com/PaddlePaddle/Paddle/blob/a6f5021fcc58b21f4414bae6bf4731ef6971582c/python/paddle/nn/layer/layers.py#L527-L530
+        # sublayer is init first
+        # scale RowParallelLinear weight
+        with paddle.no_grad():
+            if isinstance(layer, LlamaMLP):
+                factor = 1 / math.sqrt(2 * self.config.num_hidden_layers)
+                layer.down_proj.weight.scale_(factor)
+            if isinstance(layer, LlamaAttention):
+                factor = 1 / math.sqrt(2 * self.config.num_hidden_layers)
+                layer.o_proj.weight.scale_(factor)
 
 
 @register_base_model
@@ -610,7 +631,7 @@ def __init__(self, config: LlamaConfig):
         self.hidden_size = config.hidden_size
 
         if config.tensor_parallel_degree > 1:
-            self.embed_tokens = fleet.meta_parallel.VocabParallelEmbedding(
+            self.embed_tokens = mpu.VocabParallelEmbedding(
                 self.vocab_size,
                 self.hidden_size,
                 weight_attr=paddle.ParamAttr(initializer=nn.initializer.XavierNormal()),
@@ -800,7 +821,7 @@ def __init__(self, config):
         self.enable_parallel_cross_entropy = config.tensor_parallel_degree > 1 and config.tensor_parallel_output
 
         if self.enable_parallel_cross_entropy:  # and False: # and lm_head is distributed
-            self.loss_func = fleet.meta_parallel.ParallelCrossEntropy(ignore_index=self.ignore_index)
+            self.loss_func = mpu.ParallelCrossEntropy(ignore_index=self.ignore_index)
         else:
             self.loss_func = paddle.nn.CrossEntropyLoss(reduction="none", ignore_index=self.ignore_index)