fix ernie4_5 moe_layer to allgather (#2633)

Author: cheng221

examples/alignment/dpo/dpo_argument.py

@@ -102,6 +102,8 @@ class DPOConfig:
     ref_model_update_steps: int = field(default=-1, metadata={"help": "Update ref model state dict "})
     reference_free: bool = field(default=False, metadata={"help": "No reference model."})
     lora: bool = field(default=False, metadata={"help": "Use LoRA model."})
+    offset_alpha: float = field(default=0.0, metadata={"help": "offset alpha"})
+    normalize_logps: bool = field(default=True, metadata={"help": "normalize logps"})
 
 
 @dataclass

paddleformers/nn/criterion/dpo_loss.py

@@ -11,12 +11,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Optional
 
 import paddle
 import paddle.nn as nn
 import paddle.nn.functional as F
-from paddle.distributed.fleet.utils.sequence_parallel_utils import AllGatherOp
+from paddle.distributed.fleet.utils.sequence_parallel_utils import GatherOp
 
 from ...transformers.model_outputs import CausalLMOutputWithPast
 from ...transformers.sequence_parallel_utils import (
@@ -57,10 +56,12 @@ def dpo_logps(
     bias = lm_head_bias
     transpose_y = self.tie_word_embeddings
     labels = chosen_labels + rejected_labels
+    ignore_index = kwargs.pop("ignore_index", 0)  # default is 0
+
     # drop ignored index token
     if self.use_filtered_label_loss:
         if self.config.tensor_parallel_degree > 1 and self.config.sequence_parallel and logits is None:
-            labels, sparse_tgt_idx = sequence_parallel_sparse_mask_labels(labels, 0)
+            labels, sparse_tgt_idx = sequence_parallel_sparse_mask_labels(labels, ignore_index)
 
             if hidden_states is not None:
                 hidden_states = paddle.gather(hidden_states, sparse_tgt_idx, axis=0)
@@ -77,8 +78,15 @@ def dpo_logps(
             if logits is not None:
                 logits = paddle.gather(logits, sparse_tgt_idx, axis=1)
     else:
-        if hidden_states is not None:
-            hidden_states = AllGatherOp.apply(hidden_states)
+        if self.config.tensor_parallel_degree > 1 and self.config.sequence_parallel and hidden_states is not None:
+            hidden_states = GatherOp.apply(hidden_states)
+            hidden_states = hidden_states.reshape(
+                [
+                    -1,
+                    self.config.max_sequence_length,
+                    hidden_states.shape[-1],
+                ]
+            )
 
     #   bsz,seq_len,hidden_size or seq_len,hidden_size
     seq_len = labels.shape[1] if labels.ndim == 2 else labels.shape[0]
@@ -97,7 +105,7 @@ def dpo_logps(
             False,  # fused_linear
             self.loss_subbatch_sequence_length,
             return_token_loss=True,
-            ignore_index=0,
+            ignore_index=ignore_index,
         )
         per_token_logps = per_token_logps.reshape([1, per_token_logps.shape[-1], 1])
     else:
@@ -109,7 +117,6 @@ def dpo_logps(
                 transpose_y=transpose_y,
                 tensor_parallel_output=self.config.tensor_parallel_output,
             )
-
         if isinstance(logits, tuple):
             logits = logits[0]
         elif isinstance(logits, CausalLMOutputWithPast):
@@ -129,14 +136,15 @@ def dpo_logps(
                 1,
             )
 
-            per_token_logps = sb_loss_func(logits, labels.unsqueeze(-1))
+            per_token_logps = -sb_loss_func(logits, labels.unsqueeze(-1))
         else:
-            per_token_logps = self.loss_func(logits, labels.unsqueeze(-1))
+            per_token_logps = -self.loss_func(logits, labels.unsqueeze(-1))
 
     if len(response_indexs.shape) == 3:
         response_indexs = response_indexs[0]
 
     offset = 1 if self.ignore_eos_token else 0
+
     if self.use_filtered_label_loss:
         chosen_logps = paddle.stack(
             [
@@ -146,6 +154,8 @@ def dpo_logps(
                         paddle.arange(response_index[1], response_index[2], dtype=paddle.int32),
                         axis=0,
                     ).sum()
+                    if response_index[3] != 0
+                    else paddle.to_tensor(100.0)
                 )
                 for response_index in response_indexs
             ],
@@ -159,6 +169,8 @@ def dpo_logps(
                         paddle.arange(response_index[2] + offset, response_index[3], dtype=paddle.int32),
                         axis=0,
                     ).sum()
+                    if response_index[3] != 0
+                    else paddle.to_tensor(100.0)
                 )
                 for response_index in response_indexs
             ],
@@ -173,6 +185,8 @@ def dpo_logps(
                         paddle.arange(response_index[1], response_index[2], dtype=paddle.int32),
                         axis=0,
                     ).sum()
+                    if response_index[3] != 0
+                    else paddle.to_tensor(100.0)
                 )
                 for response_index in response_indexs
             ],
@@ -186,6 +200,8 @@ def dpo_logps(
                         paddle.arange(response_index[2] + offset, response_index[3], dtype=paddle.int32),
                         axis=0,
                     ).sum()
+                    if response_index[3] != 0
+                    else paddle.to_tensor(100.0)
                 )
                 for response_index in response_indexs
             ],
@@ -194,22 +210,36 @@ def dpo_logps(
 
     sft_loss = -chosen_logps.sum() / (chosen_labels != 0).sum()
     if average_log_prob:
-        chosen_response_length = response_indexs[:, 2] - response_indexs[:, 1] - offset
+        chosen_response_length = response_indexs[:, 2] - response_indexs[:, 1]
         rejected_response_length = response_indexs[:, 3] - response_indexs[:, 2]
         chosen_logps /= chosen_response_length.astype("float32")
         rejected_logps /= rejected_response_length.astype("float32")
+    elif self.dpo_config.normalize_logps:
+        avg_response_length = (response_indexs[:, 3] - response_indexs[:, 1]) / 2
+        chosen_response_length = response_indexs[:, 2] - response_indexs[:, 1]
+        rejected_response_length = response_indexs[:, 3] - response_indexs[:, 2]
+        chosen_logps *= avg_response_length / chosen_response_length.astype("float32")
+        rejected_logps *= avg_response_length / rejected_response_length.astype("float32")
     return chosen_logps, rejected_logps, sft_loss * self.dpo_config.sft_loss_ratio
 
 
 def cal_dpo_loss(
-    self, policy_chosen_logps, policy_rejected_logps, reference_chosen_logps, reference_rejected_logps, **kwargs
+    self,
+    policy_chosen_logps,
+    policy_rejected_logps,
+    reference_chosen_logps,
+    reference_rejected_logps,
+    score_deltas,
+    **kwargs
 ):
     """DPO Loss"""
     pi_logratios = policy_chosen_logps - policy_rejected_logps
     ref_logratios = reference_chosen_logps - reference_rejected_logps
     logits = pi_logratios - ref_logratios
 
     if self.dpo_config.loss_type == "sigmoid":
+        if self.dpo_config.offset_alpha > 0 and score_deltas is not None:
+            logits = logits - self.dpo_config.offset_alpha / self.dpo_config.beta * paddle.log(score_deltas + 1e-6)
         loss = (
             -F.log_sigmoid(self.dpo_config.beta * logits) * (1 - self.dpo_config.label_smoothing)
             - F.log_sigmoid(-self.dpo_config.beta * logits) * self.dpo_config.label_smoothing
@@ -282,21 +312,31 @@ def cal_dpo_loss(
 
 
 def dpo_loss_forward(
-    self: nn.Layer, logits: paddle.Tensor, labels: paddle.Tensor, loss_mask: Optional[paddle.Tensor] = None, **kwargs
+    self: nn.Layer, logits: paddle.Tensor, labels: paddle.Tensor, loss_mask: paddle.Tensor = None, **kwargs
 ):
     # unpack logtis and labels
     logits, labels, hidden_states, lm_head_weight, lm_head_bias, transpose_y = dpo_preprocess_inputs(
         self, logits, labels
     )
 
-    (
-        chosen_labels,
-        rejected_labels,
-        response_indexs,
-        score_deltas,
-        reference_chosen_logps,
-        reference_rejected_logps,
-    ) = labels
+    if self.dpo_config.offset_alpha > 0 or len(labels) == 6:
+        (
+            chosen_labels,
+            rejected_labels,
+            response_indexs,
+            score_deltas,
+            reference_chosen_logps,
+            reference_rejected_logps,
+        ) = labels
+    else:
+        (
+            chosen_labels,
+            rejected_labels,
+            response_indexs,
+            reference_chosen_logps,
+            reference_rejected_logps,
+        ) = labels
+        score_deltas = None
 
     average_log_prob = False
     if self.dpo_config.loss_type in ["ipo", "or", "simpo"]:
@@ -336,7 +376,12 @@ def dpo_loss_forward(
         **kwargs,
     )
     dpo_loss = cal_dpo_loss(
-        self, policy_chosen_logps, policy_rejected_logps, reference_chosen_logps, reference_rejected_logps
+        self,
+        policy_chosen_logps,
+        policy_rejected_logps,
+        reference_chosen_logps,
+        reference_rejected_logps,
+        score_deltas,
     )
 
     loss = dpo_loss + sft_loss

paddleformers/nn/criterion/sft_loss.py

@@ -15,6 +15,7 @@
 
 import paddle
 import paddle.nn as nn
+from paddle.distributed.fleet.utils import recompute
 from paddle.distributed.fleet.utils.sequence_parallel_utils import AllGatherOp
 
 from ...transformers.sequence_parallel_utils import (
@@ -51,41 +52,14 @@ def sft_postprocess_loss(self, masked_lm_loss, labels, loss_mask, **kwargs):
     return loss, loss_sum
 
 
-def sft_loss_forward(
-    self: nn.Layer,
-    logits: Union[paddle.Tensor, Tuple[paddle.Tensor]],
-    labels: Union[paddle.Tensor, Tuple[paddle.Tensor]],
-    loss_mask: paddle.Tensor = None,
-    **kwargs
-):
-    logits, labels, hidden_states, lm_head_weight, lm_head_bias, transpose_y = sft_preprocess_inputs(
-        self, logits, labels
-    )
-    if self.use_filtered_label_loss:
-        if self.tensor_parallel and self.sequence_parallel and logits is None:
-            masked_lm_labels, sparse_label_idx = sequence_parallel_sparse_mask_labels(labels, self.ignored_index)
-            sparse_label_idx = sparse_label_idx.reshape([-1, 1])
-            if hidden_states is not None:
-                hidden_states = paddle.gather(hidden_states, sparse_label_idx, axis=0)
-                hidden_states = AllGatherVarlenOp.apply(hidden_states)
-        else:
-            masked_lm_labels = labels.flatten()
-            sparse_label_idx = paddle.nonzero(masked_lm_labels != self.ignored_index).flatten()
-            masked_lm_labels = paddle.take_along_axis(masked_lm_labels, sparse_label_idx, axis=0)
-            if hidden_states is not None:
-                hidden_states = hidden_states.reshape([-1, hidden_states.shape[-1]])
-                hidden_states = paddle.take_along_axis(hidden_states, sparse_label_idx.reshape([-1, 1]), axis=0)
-            if logits is not None:
-                logits = paddle.gather(logits, sparse_label_idx, axis=1)
-        labels = masked_lm_labels
-    else:
-        if self.sequence_parallel:
-            if hidden_states is not None:
-                hidden_states = AllGatherOp.apply(hidden_states)
+def loss_impl(self, logits, labels):
+    logits = logits.cast("float32")
+    loss = self.loss_func(logits, labels)
+    return loss
 
-    masked_lm_labels = labels
-    # bsz,seq_len,hidden_size or seq_len,hidden_size
-    seq_len = masked_lm_labels.shape[1] if masked_lm_labels.ndim == 2 else masked_lm_labels.shape[0]
+
+def sft_calculate_loss(self, logits, hidden_states, lm_head_weight, lm_head_bias, labels, loss_mask, transpose_y):
+    seq_len = labels.shape[1] if labels.ndim == 2 else labels.shape[0]
     if self.use_fused_head_and_loss_fn and self.use_subbatch and seq_len > self.loss_subbatch_sequence_length:
         masked_lm_loss = fused_head_and_loss_fn(
             hidden_states,
@@ -123,7 +97,6 @@ def sft_loss_forward(
                 f" {logits.shape[-1]}, {self.config.vocab_size}"
             )
 
-        logits = logits.cast("float32")
         if logits.dim() == 2 and labels.dim() == 2:
             logits = logits.unsqueeze(0)
         elif logits.dim() == 3 and labels.dim() == 1:
@@ -133,16 +106,77 @@ def sft_loss_forward(
         # labels: bsz seq_len vocab_size
         if self.use_subbatch and seq_len > self.loss_subbatch_sequence_length:
             sb_loss_func = subbatch(
-                self.loss_func,
-                [0, 1],
-                [1, 1],
-                self.loss_subbatch_sequence_length,
-                1,
+                loss_impl,
+                arg_idx=[1, 2],
+                axis=[1, 1],
+                bs=self.loss_subbatch_sequence_length,
+                out_idx=1,
             )
-            masked_lm_loss = sb_loss_func(logits, labels.unsqueeze(-1))
+            masked_lm_loss = sb_loss_func(self, logits, labels.unsqueeze(-1))
         else:
-            masked_lm_loss = self.loss_func(logits, labels.unsqueeze(-1))
-    loss = sft_postprocess_loss(self, masked_lm_loss, labels, loss_mask, **kwargs)
+            masked_lm_loss = loss_impl(self, logits, labels.unsqueeze(-1))
+
+    masked_lm_loss = sft_postprocess_loss(self, masked_lm_loss, labels, loss_mask)
+    return masked_lm_loss
+
+
+def sft_loss_forward(
+    self: nn.Layer,
+    logits: Union[paddle.Tensor, Tuple[paddle.Tensor]],
+    labels: Union[paddle.Tensor, Tuple[paddle.Tensor]],
+    loss_mask: paddle.Tensor = None,
+    **kwargs
+):
+    logits, labels, hidden_states, lm_head_weight, lm_head_bias, transpose_y = sft_preprocess_inputs(
+        self, logits, labels
+    )
+    if self.use_filtered_label_loss:
+        if self.tensor_parallel and self.sequence_parallel and logits is None:
+            masked_lm_labels, sparse_label_idx = sequence_parallel_sparse_mask_labels(labels, self.ignored_index)
+            sparse_label_idx = sparse_label_idx.reshape([-1, 1])
+            if hidden_states is not None:
+                hidden_states = paddle.gather(hidden_states, sparse_label_idx, axis=0)
+                hidden_states = AllGatherVarlenOp.apply(hidden_states)
+        else:
+            masked_lm_labels = labels.flatten()
+            sparse_label_idx = paddle.nonzero(masked_lm_labels != self.ignored_index).flatten()
+            masked_lm_labels = paddle.take_along_axis(masked_lm_labels, sparse_label_idx, axis=0)
+            if hidden_states is not None:
+                hidden_states = hidden_states.reshape([-1, hidden_states.shape[-1]])
+                hidden_states = paddle.take_along_axis(hidden_states, sparse_label_idx.reshape([-1, 1]), axis=0)
+            if logits is not None:
+                logits = paddle.gather(logits, sparse_label_idx, axis=1)
+        labels = masked_lm_labels
+    else:
+        if self.sequence_parallel:
+            if hidden_states is not None:
+                hidden_states = AllGatherOp.apply(hidden_states)
+
+    masked_lm_labels = labels
+    # bsz,seq_len,hidden_size or seq_len,hidden_size
+    if self.config.recompute:
+        loss = recompute(
+            sft_calculate_loss,
+            self,
+            logits,
+            hidden_states,
+            lm_head_weight,
+            lm_head_bias,
+            labels,
+            loss_mask,
+            transpose_y,
+        )
+    else:
+        loss = sft_calculate_loss(
+            self,
+            logits,
+            hidden_states,
+            lm_head_weight,
+            lm_head_bias,
+            labels,
+            loss_mask,
+            transpose_y,
+        )
     return loss
 
 

paddleformers/nn/mlp.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import paddle
 import paddle.nn as nn
 from paddle.incubate.nn.functional import swiglu as fused_swiglu
 
@@ -117,5 +117,9 @@ def forward(self, x):
         else:
             gate = self.gate_proj(x)
             up = self.up_proj(x)
-            x = self.act_fn(gate) * up
+            if self.fuse_swiglu:
+                x = paddle.concat([gate, up], axis=-1)
+                x = fused_swiglu(x)
+            else:
+                x = self.act_fn(gate) * up
         return self.down_proj(x)

paddleformers/nn/moe/all_gather.py

@@ -308,7 +308,7 @@ def forward(
 
                 recv_mask_alltoall_out = paddle.cat(recv_mask_alltoall_out, 0)
                 distributed_input_to_alltoall_out = paddle.maximum(
-                    recv_mask_alltoall_out.cumsum() - 1,
+                    (recv_mask_alltoall_out.cumsum() - 1).astype(recv_mask_alltoall_out.dtype),
                     paddle.zeros([1], dtype=recv_mask_alltoall_out.dtype),
                 )
                 distributed_input_to_alltoall_out = distributed_input_to_alltoall_out.split(alltoall_shape)