[Llama3] Fix: add pad token fallback and improve tensor reshaping (#3025)

Author: LittleHeroZZZX

paddleformers/cli/train/auto_parallel/workflow.py

@@ -208,6 +208,8 @@ def run_auto_parallel(model_args, data_args, generating_args, training_args):
 
     config = config_class.from_pretrained(model_args.model_name_or_path)
     tokenizer = AutoTokenizer.from_pretrained(model_args.tokenizer_name_or_path)
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token_id = tokenizer.eos_token_id
     # config = AutoConfig.from_pretrained(model_args.model_name_or_path)
     LlmMetaConfig.set_llm_config(config, training_args)
     config.use_fast_layer_norm = model_args.use_fast_layer_norm

paddleformers/cli/train/dpo/workflow.py

@@ -191,6 +191,8 @@ def run_dpo(
         tokenizer = AutoTokenizer.from_pretrained(model_args.tokenizer_name_or_path)
     else:
         tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path)
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token_id = tokenizer.eos_token_id
 
     logger.info("Loading model & tokenizer successfully !")
 

paddleformers/cli/train/pretrain/workflow.py

@@ -409,6 +409,8 @@ def run_dsv3_pretrain(model_args, data_args, generating_args, training_args):
             )
 
     tokenizer = AutoTokenizer.from_pretrained(model_args.tokenizer_name_or_path)
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token_id = tokenizer.eos_token_id
     config = DeepseekV2FastConfig.from_pretrained(model_args.model_name_or_path)
 
     # set all llm config

paddleformers/cli/train/sft/workflow.py

@@ -266,6 +266,8 @@ def neft_post_hook(module, input, output):
 
     # Load tokenizer & dataset
     tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path)
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token_id = tokenizer.eos_token_id
 
     # if using chat_template, data_args.eval_with_do_generation must be false
     if tokenizer.chat_template is not None:

paddleformers/transformers/llama/modeling.py

@@ -159,9 +159,9 @@ def forward(
         q_shape = (batch_size, seq_len, self.num_heads, self.head_dim)
         kv_shape = (batch_size, seq_len, self.num_key_value_heads, self.head_dim)
 
-        query_states = self.q_proj(hidden_states).view(q_shape).transpose(1, 2)
-        key_states = self.k_proj(hidden_states).view(kv_shape).transpose(1, 2)
-        value_states = self.v_proj(hidden_states).view(kv_shape).transpose(1, 2)
+        query_states = self.q_proj(hidden_states).reshape(q_shape).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).reshape(kv_shape).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).reshape(kv_shape).transpose(1, 2)
 
         cos, sin = position_embeddings
         query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

tests/transformers/llama/test_modeling.py

@@ -439,6 +439,7 @@ def test_inference_no_attention(self):
             "Paddleformers/tiny-random-llama3",
             download_hub="aistudio",
             convert_from_hf=True,
+            dtype="float32",
         )
         model.eval()
         input_ids = paddle.to_tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]])
@@ -452,9 +453,9 @@ def test_inference_no_attention(self):
         expected_slice = paddle.to_tensor(
             [
                 [
-                    [0.02366970, -0.42482421, 0.47202760],
-                    [-0.12180223, 0.00559035, 0.83846688],
-                    [0.45073321, 0.25703996, 1.36826384],
+                    [0.01802453, -0.42128855, 0.45844582],
+                    [-0.12787277, 0.00660499, 0.83033413],
+                    [0.44403678, 0.26123494, 1.36080980],
                 ]
             ],
             dtype=output.dtype,
@@ -467,6 +468,7 @@ def test_inference_with_attention(self):
             "Paddleformers/tiny-random-llama3",
             download_hub="aistudio",
             convert_from_hf=True,
+            dtype="float32",
         )
         model.eval()
         input_ids = paddle.to_tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]])
@@ -479,9 +481,9 @@ def test_inference_with_attention(self):
         expected_slice = paddle.to_tensor(
             [
                 [
-                    [0.02366970, -0.42482421, 0.47202760],
-                    [-0.12180223, 0.00559035, 0.83846688],
-                    [0.45073321, 0.25703996, 1.36826384],
+                    [0.01802453, -0.42128855, 0.45844582],
+                    [-0.12787277, 0.00660499, 0.83033413],
+                    [0.44403678, 0.26123494, 1.36080980],
                 ]
             ],
             dtype=output.dtype,

tests/transformers/test_shard_checkpoint.py

@@ -81,7 +81,7 @@ def test_from_pretrained_low_cpu_mem_usage_functional(self):
                 convert_from_hf=convert,
             )
             for p1, p2 in zip(m1.parameters(), m2.parameters()):
-                self.assertTrue(paddle.allclose(p1, p2))
+                self.assertTrue(paddle.allclose(p1.float(), p2.float()))
 
     @unittest.skipIf(not is_paddle_cuda_available(), "some op is missing in cpu mode")
     def test_keep_in_fp32_modules(self):