update loss function

Author: yiwzhao

veomni/models/transformers/qwen3/modeling_qwen3.py

@@ -2,7 +2,6 @@
 
 import torch
 from torch import nn
-from torch.nn import CrossEntropyLoss
 from transformers.activations import ACT2FN
 from transformers.cache_utils import Cache, DynamicCache
 from transformers.generation import GenerationMixin
@@ -27,7 +26,7 @@
 
 from ....distributed.parallel_state import get_parallel_state
 from ....distributed.sequence_parallel import slice_position_embedding
-from ....ops.loss import causallm_loss_function, seqcls_token_loss_sp_aware
+from ....ops.loss import causallm_loss_function, seqcls_token_loss_function
 from ....utils import logging
 from ....utils.import_utils import is_liger_kernel_available
 from ...module_utils import GradientCheckpointingLayer
@@ -708,7 +707,6 @@ def forward(
             cache_position=cache_position,
             **kwargs,
         )
-
         hidden_states = outputs.last_hidden_state
         # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
         slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
@@ -743,7 +741,7 @@ def __init__(self, config):
         self.num_labels = config.num_labels
         self.model = Qwen3Model(config)
         self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
-        self.loss_fct = CrossEntropyLoss(ignore_index=-100, reduction="none")
+        self.loss_function = seqcls_token_loss_function
         # Initialize weights and apply final processing
         self.post_init()
 
@@ -779,23 +777,7 @@ def forward(
         logits = self.score(hidden_states)
 
         loss = None
-        if labels is not None:
-            # labels are token-level now, shape must match logits tokens
-            if logits.dim() == 3:
-                # [B, L, C] -> [B*L, C]
-                B, L, C = logits.shape
-                logits_2d = logits.view(B * L, C)
-                labels_1d = labels.view(B * L).to(logits.device)
-            elif logits.dim() == 2:
-                # [T, C] -> [T, C]
-                logits_2d = logits
-                labels_1d = labels.view(-1).to(logits.device)
-            else:
-                raise ValueError(f"Unexpected logits shape: {logits.shape}")
-
-            ps = get_parallel_state()
-            sp_group = ps.sp_group if ps.sp_enabled else None
-            loss = seqcls_token_loss_sp_aware(logits_2d, labels_1d, self.loss_fct, sp_group)
+        loss, _ = self.loss_function(hidden_states, self.score.weight, labels)
 
         return SequenceClassifierOutputWithPast(
             loss=loss,

veomni/ops/loss.py

@@ -1,7 +1,6 @@
 from typing import Optional
 
 import torch
-import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
 
@@ -93,21 +92,44 @@ def causallm_loss_function(
     return loss, logits
 
 
-def seqcls_token_loss_sp_aware(
-    logits: torch.Tensor,  # [N, C]
-    labels: torch.Tensor,  # [N]
-    loss_fct: nn.Module,
-    sp_group,
+def seqcls_token_loss_function(
+    hidden_states: torch.Tensor,
+    weight: torch.Tensor,
+    labels: torch.Tensor,
+    num_items_in_batch: Optional[int] = None,
     ignore_index: int = -100,
+    shift_labels: Optional[torch.Tensor] = None,
+    **kwargs,
 ) -> torch.Tensor:
-    # local sum loss
-    # CrossEntropyLoss(reduction="none") + mask + sum
-    per = loss_fct(logits, labels)  # [N] if reduction="none"
-    valid = labels != ignore_index
-    loss_sum = (per * valid).sum()
-    cnt = valid.sum().to(dtype=loss_sum.dtype)
-
-    if sp_group is not None:
-        dist.all_reduce(loss_sum, op=dist.ReduceOp.SUM, group=sp_group)
-        dist.all_reduce(cnt, op=dist.ReduceOp.SUM, group=sp_group)
-    return loss_sum / cnt.clamp_min(1.0)
+    # We don't use shift_labels
+    assert shift_labels is None
+
+    loss = None
+    logits = None
+
+    if labels is None:
+        logits = F.linear(hidden_states, weight)
+        return loss, logits
+
+    sp_enabled = get_parallel_state().sp_enabled
+
+    # Flatten the labels and hidden_states
+    labels = labels.view(-1)
+    hidden_states = hidden_states.view(-1, hidden_states.size(-1))
+
+    # Calculate loss
+    if fused_linear_cross_entropy is not None:  # use kernels
+        if is_seed_kernels_available():
+            loss = fused_linear_cross_entropy(hidden_states, weight, labels, ignore_index=ignore_index)
+        elif is_liger_kernel_available():
+            loss = fused_linear_cross_entropy(weight, hidden_states, labels)
+    else:
+        logits = F.linear(hidden_states, weight).float()
+        loss = fixed_cross_entropy(logits, labels, num_items_in_batch, ignore_index, **kwargs)
+
+    # Reduce loss when using sp
+    if sp_enabled:
+        num_valid_tokens = (labels != ignore_index).sum()
+        loss = reduce_sequence_parallel_loss(loss, num_valid_tokens)
+
+    return loss, logits