cherry pick from pr #2622 (#2625)

Author: Ace-To-HYB

paddleformers/nn/criterion/loss_utils.py

@@ -55,7 +55,7 @@ def calc_lm_head_logits(
         hidden_states,
         weight,
         bias=bias,
-        transpose_y=config.get("tie_word_embeddings", False),
+        transpose_y=True,
         tensor_parallel_degree=config.tensor_parallel_degree,
         tensor_parallel_output=tensor_parallel_output,
         fuse_linear=config.get("fuse_linear", False),

paddleformers/nn/lm_head.py

@@ -24,13 +24,9 @@
 
 class LMHead(nn.Layer):
     def __init__(self, config: PretrainedConfig):
-        """
-        transpose_y (bool): Whether to transpose the lm_head weight matrix before matrix multiplication.
-        """
         super().__init__()
         self.config = config
         self.use_bias = config.get("lm_head_bias", False)
-        self.transpose_y = config.get("tie_word_embeddings", False)
         self.vocab_parallel = False
 
         # apply vocab tensor parallel
@@ -45,21 +41,15 @@ def __init__(self, config: PretrainedConfig):
                     vocab_size,
                     config.tensor_parallel_degree,
                 )
-        self.lm_head_shape = (
-            [config.hidden_size, vocab_size] if not self.transpose_y else [vocab_size, config.hidden_size]
-        )
 
         self.weight = self.create_parameter(
-            shape=self.lm_head_shape,
+            shape=[vocab_size, config.hidden_size],
             dtype=paddle.get_default_dtype(),
             default_initializer=nn.initializer.XavierNormal(1.0),
         )
 
         # setting distributed attr for tensor parallel
-        self.weight.is_distributed = self.vocab_parallel
-
-        if self.weight.is_distributed:
-            self.weight.split_axis = 0 if self.transpose_y else 1
+        self._set_distributed_attr(self.weight)
 
         if self.use_bias:
             self.bias = self.create_parameter(
@@ -69,12 +59,15 @@ def __init__(self, config: PretrainedConfig):
             )
 
             # setting distributed attr for tensor parallel
-            self.bias.is_distributed = self.vocab_parallel
-            if self.bias.is_distributed:
-                self.bias.split_axis = 0
+            self._set_distributed_attr(self.bias)
         else:
             self.bias = None
 
+    def _set_distributed_attr(self, param):
+        param.is_distributed = self.vocab_parallel
+        if param.is_distributed:
+            param.split_axis = 0
+
     def forward(self, hidden_states, tensor_parallel_output=None):
         """Project hidden states to vocabulary logits.
 
@@ -114,5 +107,4 @@ def forward(self, hidden_states, tensor_parallel_output=None):
         )
 
     def extra_repr(self):
-        hidden_size, vocab_size = self.lm_head_shape if not self.transpose_y else self.lm_head_shape[::-1]
-        return f"hidden_size={hidden_size}, vocab_size={vocab_size}, dtype={self.weight.dtype}, vocab_parallel={self.vocab_parallel}"
+        return f"hidden_size={self.weight.shape[1]}, vocab_size={self.weight.shape[0]}, dtype={self.weight.dtype}, vocab_parallel={self.vocab_parallel}"

paddleformers/transformers/ernie4_5/modeling.py

@@ -439,7 +439,7 @@ def _get_tensor_parallel_mappings(cls, config, is_split=True):
 
         def make_base_actions():
             actions = {
-                "lm_head.weight": partial(fn, is_column=not config.tie_word_embeddings),
+                "lm_head.weight": partial(fn, is_column=False),
                 "embed_tokens.weight": partial(fn, is_column=False),
             }
             for layer_idx in range(config.num_hidden_layers):

paddleformers/transformers/ernie4_5_moe/modeling.py

@@ -503,7 +503,7 @@ def _get_tensor_parallel_mappings(cls, config, is_split=True):
 
         def make_base_actions():
             actions = {
-                "lm_head.weight": partial(fn, is_column=not config.tie_word_embeddings),
+                "lm_head.weight": partial(fn, is_column=False),
                 "embed_tokens.weight": partial(fn, is_column=False),
             }
             for layer_idx in range(config.num_hidden_layers):

paddleformers/transformers/gpt_oss/modeling.py

@@ -650,7 +650,7 @@ class GptOssPreTrainedModel(PretrainedModel):
     config_class = GptOssConfig
     base_model_prefix = "model"
     keys_to_ignore_on_load_unexpected = [r"self_attn.rotary_emb.inv_freq"]
-    transpose_weight_keys = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "lm_head"]
+    transpose_weight_keys = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
 
     @classmethod
     def _get_tensor_parallel_mappings(cls, config: GptOssConfig, is_split=True):
@@ -667,7 +667,7 @@ def get_tensor_parallel_split_mappings(num_layers, num_experts):
             final_actions = {}
 
             base_actions = {
-                "lm_head.weight": partial(fn, is_column=True),
+                "lm_head.weight": partial(fn, is_column=False),
                 # Row Linear
                 "embed_tokens.weight": partial(fn, is_column=False),
                 "layers.0.self_attn.o_proj.weight": partial(fn, is_column=False),

paddleformers/transformers/qwen2/modeling.py

@@ -348,7 +348,7 @@ class Qwen2PretrainedModel(PretrainedModel):
     config_class = Qwen2Config
     base_model_prefix = "model"
     _keys_to_ignore_on_load_unexpected = [r"self_attn.rotary_emb.inv_freq"]
-    transpose_weight_keys = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "lm_head"]
+    transpose_weight_keys = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
 
     @classmethod
     def _get_tensor_parallel_mappings(cls, config: Qwen2Config, is_split=True):
@@ -380,7 +380,7 @@ def _get_tensor_parallel_mappings(cls, config: Qwen2Config, is_split=True):
 
         def make_base_actions():
             actions = {
-                "lm_head.weight": partial(fn, is_column=not config.tie_word_embeddings),
+                "lm_head.weight": partial(fn, is_column=False),
                 "embed_tokens.weight": partial(fn, is_column=False),
             }
             for layer_idx in range(config.num_hidden_layers):

paddleformers/transformers/qwen2_moe/modeling.py

@@ -388,7 +388,6 @@ class Qwen2MoePretrainedModel(PretrainedModel):
         "down_proj",
         "gate",
         "shared_expert_gate",
-        "lm_head",
     ]
 
     @classmethod
@@ -433,7 +432,7 @@ def _get_tensor_parallel_mappings(cls, config: Qwen2MoeConfig, is_split=True):
 
         def make_base_actions():
             actions = {
-                "lm_head.weight": partial(fn, is_column=not config.tie_word_embeddings),
+                "lm_head.weight": partial(fn, is_column=False),
                 "embed_tokens.weight": partial(fn, is_column=False),
             }
             for layer_idx in range(config.num_hidden_layers):

paddleformers/transformers/qwen3/modeling.py

@@ -247,7 +247,7 @@ class Qwen3PretrainedModel(PretrainedModel):
     config_class = Qwen3Config
     base_model_prefix = "model"
     _keys_to_ignore_on_load_unexpected = [r"self_attn.rotary_emb.inv_freq"]
-    transpose_weight_keys = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "lm_head"]
+    transpose_weight_keys = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
 
     @classmethod
     def _get_tensor_parallel_mappings(cls, config: Qwen3Config, is_split=True):
@@ -279,7 +279,7 @@ def _get_tensor_parallel_mappings(cls, config: Qwen3Config, is_split=True):
 
         def make_base_actions():
             actions = {
-                "lm_head.weight": partial(fn, is_column=not config.tie_word_embeddings),
+                "lm_head.weight": partial(fn, is_column=False),
                 "embed_tokens.weight": partial(fn, is_column=False),
             }
             for layer_idx in range(config.num_hidden_layers):

paddleformers/transformers/qwen3_moe/modeling.py

@@ -273,7 +273,6 @@ class Qwen3MoePretrainedModel(PretrainedModel):
         "up_proj",
         "down_proj",
         "gate",
-        "lm_head",
     ]
 
     @classmethod
@@ -311,7 +310,7 @@ def _get_tensor_parallel_mappings(cls, config: Qwen3MoeConfig, is_split=True):
 
         def make_base_actions():
             actions = {
-                "lm_head.weight": partial(fn, is_column=not config.tie_word_embeddings),
+                "lm_head.weight": partial(fn, is_column=False),
                 "embed_tokens.weight": partial(fn, is_column=False),
             }
             for layer_idx in range(config.num_hidden_layers):

tests/nn/test_lm_head.py

@@ -27,11 +27,10 @@ def test_initialization_default(self):
         lm_head = LMHead(config)
 
         # Check weight shape and attributes
-        self.assertEqual(lm_head.weight.shape, [config.hidden_size, config.vocab_size])
+        self.assertEqual(lm_head.weight.shape, [config.vocab_size, config.hidden_size])
         self.assertFalse(lm_head.weight.is_distributed)
         self.assertIsNone(lm_head.bias)
         self.assertFalse(lm_head.vocab_parallel)
-        self.assertFalse(lm_head.transpose_y)
 
     def test_initialization_with_tie_word_embeddings(self):
         # Test initialization with tied embeddings
@@ -40,7 +39,6 @@ def test_initialization_with_tie_word_embeddings(self):
         lm_head = LMHead(config)
 
         self.assertEqual(lm_head.weight.shape, [config.vocab_size, config.hidden_size])
-        self.assertTrue(lm_head.transpose_y)
 
     def test_forward_normal(self):
         # Test normal forward pass