Unified FuseLoss (#11057)

Author: KB-Ding

paddlenlp/rl/models/ppo_model_utils.py

@@ -30,6 +30,12 @@
 from paddle.distributed.fleet.layers.mpu import mp_ops
 from paddle.distributed.fleet.meta_parallel import ParallelCrossEntropy
 
+try:
+    from paddle.distributed.fleet.utils.sequence_parallel_utils import GatherOp
+except:
+    pass
+
+
 from ...transformers.llama.modeling import (
     LlamaPretrainingCriterion as PretrainingCriterion,
 )
@@ -446,6 +452,16 @@ def forward(
 
         else:
             hidden_states, weight, bias, transpose_y = logits
+            if self.config.tensor_parallel_degree > 1 and self.config.sequence_parallel:
+                hidden_states = GatherOp.apply(hidden_states)
+                hidden_states = hidden_states.reshape(
+                    [
+                        input_ids.shape[0],
+                        -1,
+                        hidden_states.shape[-1],
+                    ]
+                )
+
             if use_remove_padding:
                 input_ids = raw_input_ids
                 if pad_size > 0:

paddlenlp/transformers/qwen2/modeling.py

@@ -1480,6 +1480,10 @@ def __init__(self, config: Qwen2Config, embedding_weights=None, transpose_y=Fals
             self.weight.split_axis = 0 if self.transpose_y else 1
 
     def forward(self, hidden_states, tensor_parallel_output=None, batch_size=None):
+        # add this for fused_head_and_loss_fn
+        if self.config.use_fused_head_and_loss_fn:
+            return hidden_states, self.weight, None, self.transpose_y
+
         if self.config.sequence_parallel:
             hidden_states = GatherOp.apply(hidden_states)
             hidden_states = paddle.reshape_(hidden_states, [batch_size, -1, self.config.hidden_size])
@@ -1667,19 +1671,6 @@ def forward(
 
         hidden_states = outputs[0]
 
-        # add this for fused_head_and_loss_fn
-        if self.config.use_fused_head_and_loss_fn and self.training:
-            if self.config.tensor_parallel_degree > 1 and self.config.sequence_parallel:
-                hidden_states = GatherOp.apply(hidden_states)
-                hidden_states = hidden_states.reshape(
-                    [
-                        batch_size,
-                        -1,
-                        hidden_states.shape[-1],
-                    ]
-                )
-            return hidden_states, self.lm_head.weight, None, self.lm_head.transpose_y
-
         # if labels is None，means we need full output, instead of tensor_parallel_output
         # tensor_parallel_output is together with ParallelCrossEntropy
         tensor_parallel_output = self.config.tensor_parallel_output and self.config.tensor_parallel_degree > 1

paddlenlp/transformers/qwen3/modeling.py

@@ -70,7 +70,6 @@
 
 try:
     from paddle.distributed.fleet.utils.sequence_parallel_utils import (
-        GatherOp,
         ScatterOp,
         mark_as_sequence_parallel_parameter,
     )
@@ -1139,19 +1138,6 @@ def forward(
 
         hidden_states = outputs[0]
 
-        # add this for fused_head_and_loss_fn
-        if self.config.use_fused_head_and_loss_fn and self.training:
-            if self.config.tensor_parallel_degree > 1 and self.config.sequence_parallel:
-                hidden_states = GatherOp.apply(hidden_states)
-                hidden_states = hidden_states.reshape(
-                    [
-                        batch_size,
-                        -1,
-                        hidden_states.shape[-1],
-                    ]
-                )
-            return hidden_states, self.lm_head.weight, None, self.lm_head.transpose_y
-
         # if labels is None，means we need full output, instead of tensor_parallel_output
         # tensor_parallel_output is together with ParallelCrossEntropy
         tensor_parallel_output = self.config.tensor_parallel_output and self.config.tensor_parallel_degree > 1