Update modeling_bailing_moe.py

modeling_bailing_moe.py

@@ -20,17 +20,14 @@
 """ PyTorch BailingMoE model."""
 import math
 import warnings
-from dataclasses import dataclass
 from typing import List, Optional, Tuple, Union
 
 import torch
-import torch.distributed as dist
 import torch.nn.functional as F
 import torch.utils.checkpoint
-import transformers
-from packaging import version
 from torch import nn
 from torch.nn import CrossEntropyLoss
+
 from transformers.activations import ACT2FN
 from transformers.cache_utils import Cache, DynamicCache
 from transformers.modeling_attn_mask_utils import (
@@ -40,10 +37,8 @@ from transformers.modeling_attn_mask_utils import (
     _prepare_4d_causal_attention_mask_for_sdpa,
 )
 from transformers.modeling_outputs import (
-    ModelOutput,
-    MoeCausalLMOutputWithPast,
     MoeModelOutputWithPast,
-    SequenceClassifierOutputWithPast,
+    MoeCausalLMOutputWithPast,
 )
 from transformers.modeling_utils import PreTrainedModel
 from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS, is_torch_greater_or_equal_than_1_13
@@ -56,9 +51,9 @@ from transformers.utils import (
     replace_return_docstrings,
 )
 from transformers.utils.import_utils import is_torch_fx_available
-
 from .configuration_bailing_moe import BailingMoeConfig
 
+
 if is_flash_attn_2_available():
     from flash_attn import flash_attn_func, flash_attn_varlen_func
     from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
@@ -108,220 +103,6 @@ def _make_causal_mask(
     )
 
 
-def _unpack_router_logits(router_outputs):
-    """
-    Unpack the router tuple for blance loss calculation.
-    """
-    total_router_logits = []
-    total_expert_indexes = []
-    for router_output in router_outputs:
-        if router_output[0] is not None:
-            router_logits, expert_indexes = router_output
-            total_router_logits.append(router_logits.unsqueeze(0))
-            total_expert_indexes.append(expert_indexes.unsqueeze(0))
-    return torch.cat(total_router_logits, dim=0), total_expert_indexes
-
-
-def load_balancing_loss_func(router_probs: torch.Tensor, expert_indices: torch.Tensor, labels: torch.Tensor) -> float:
-    num_layers, _, seq_len, num_experts = router_probs.shape
-    num_experts = router_probs.shape[-1]
-    new_labels = labels.clone().detach()
-    ##
-    for batch_tensor in new_labels:
-        neg_mask = batch_tensor == -100
-        diff_neg_ones = torch.diff(neg_mask.float())
-        start_pos = torch.where(diff_neg_ones == 1.0)[0]  # 找到-1序列开始的位置
-        if start_pos.nelement() == 0:  # 如果没有找到开始位置，可能需要根据实际情况调整
-            pass
-        else:
-            last_start = start_pos[-1]  # 需要修改的最后一串-1的开始位置
-            batch_tensor[:last_start] = 0  # 将这部分-1全部改为0
-    new_labels = new_labels.to(torch.int64)
-
-    # cast the expert indices to int64, otherwise one-hot encoding will fail
-
-    if expert_indices.dtype != torch.int64:
-        expert_indices = expert_indices.to(torch.int64)
-
-    if len(expert_indices.shape) == 3:
-        expert_indices = expert_indices.unsqueeze(3)
-
-    expert_mask = torch.nn.functional.one_hot(expert_indices, num_experts)
-
-    # For a given token, determine if it was routed to a given expert.
-    expert_mask = torch.max(expert_mask, axis=-2).values
-
-    # cast to float32 otherwise mean will fail
-    expert_mask = expert_mask.to(torch.float32)
-    labels_mask = (new_labels[None, ..., None].expand_as(expert_mask) != -100).long()
-
-    # sample level balance loss
-    tokens_per_group_and_expert = torch.sum(expert_mask * labels_mask, dim=-2) / torch.sum(labels_mask, dim=-2)
-    router_prob_per_group_and_expert = torch.sum(router_probs * labels_mask, dim=-2) / torch.sum(labels_mask, dim=-2)
-    return torch.mean(tokens_per_group_and_expert * router_prob_per_group_and_expert) * (num_experts**2)
-
-
-def router_z_loss_func(router_logits: torch.Tensor, labels: torch.Tensor) -> float:
-    r"""
-    Compute the router z-loss implemented in PyTorch.
-
-    The router z-loss was introduced in [Designing Effective Sparse Expert Models](https://arxiv.org/abs/2202.08906).
-    It encourages router logits to remain small in an effort to improve stability.
-
-    Args:
-        router_logits (`float`):
-            Input logits of shape [num_layers, batch_size, sequence_length, num_experts]
-
-    Returns:
-        Scalar router z-loss.
-    """
-    num_layers, num_groups, tokens_per_group, _ = router_logits.shape
-    labels_mask = (labels[None, ..., None].expand_as(router_logits) != -100).long()
-
-    ori_dtype = router_logits.dtype
-    if ori_dtype == torch.bfloat16:
-        loss_func_inputs = (router_logits * labels_mask).to(torch.float32)
-    else:
-        loss_func_inputs = router_logits * labels_mask
-    log_z = torch.logsumexp(loss_func_inputs, dim=-1).to(ori_dtype)
-    z_loss = log_z**2
-
-    return torch.sum(z_loss) / (num_layers * num_groups * tokens_per_group)
-
-
-def auxiliary_loss(router_tuple, lm_logits, labels, config: BailingMoeConfig):
-    balance_loss, z_loss, last_logits_l2_loss = 0.0, 0.0, 0.0
-
-    loss = 0
-    if router_tuple is not None:
-        router_logits, layer_router_index = _unpack_router_logits(router_tuple)
-        top1_expert_index = torch.cat(layer_router_index, dim=0)
-        z_loss = router_z_loss_func(router_logits, labels)
-        router_probs = torch.nn.Softmax(dim=-1)(router_logits)
-        balance_loss = load_balancing_loss_func(router_probs, top1_expert_index, labels)
-
-        num_layers = router_probs.shape[0]
-        num_experts = router_probs.shape[-1]
-        router_probs_log = router_probs.detach().view(num_layers, -1, num_experts)
-        router_probs_mean = router_probs_log.mean(1)
-        router_probs_sort_mean = router_probs_log.sort(-1, descending=True)[0].mean(1)
-        router_probs_log = torch.stack([router_probs_mean, router_probs_sort_mean], dim=1)
-        dist.all_reduce(router_probs_log, dist.ReduceOp.SUM)
-        router_probs_log = router_probs_log / torch.distributed.get_world_size()
-        if dist.get_rank() == 0:
-            router_probs_log = router_probs_log.float()
-            router_probs_log /= router_probs_log.sum(-1, keepdim=True)
-
-        loss = float(config.router_z_loss_alpha) * z_loss + float(config.router_balance_loss_alpha) * balance_loss
-
-    last_logits_l2_loss = 0.0
-    if float(config.last_logits_l2_alpha) >= 0:
-        shift_logits = lm_logits[..., :-1, :].contiguous()
-        shift_labels = labels[..., 1:].contiguous()
-
-        shift_logits = lm_logits.view(-1, lm_logits.size(-1))
-        labels_mask = (shift_labels.view(-1) != -100).long()
-
-        last_logits_l2_loss = torch.sum(torch.linalg.norm(shift_logits.float(), 2.0, dim=-1) * labels_mask) / torch.sum(
-            labels_mask
-        )
-        loss += float(config.last_logits_l2_alpha) * last_logits_l2_loss
-        last_logits_l2_loss = last_logits_l2_loss.item()
-
-    return loss, balance_loss, z_loss, last_logits_l2_loss
-
-
-def local_token_level_cross_entropy(logits, labels, **kwargs):
-    # 在每个batch内部做token-level的平均,然后在所有batch间做平均
-    if isinstance(logits, ModelOutput):
-        logits = logits.logits
-    elif isinstance(logits, Tuple):
-        logits = logits[0]
-
-    logits = logits.float()
-    shift_logits = logits[..., :-1, :].contiguous()
-    shift_labels = labels[..., 1:].contiguous()
-    loss_fct = torch.nn.CrossEntropyLoss(ignore_index=-100)
-    loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
-    return loss
-
-
-def sample_level_cross_entropy(logits, labels, **kwargs):
-    # 先对所有样本字token-level的平均,然后计算所有sample的平均值
-    if isinstance(logits, ModelOutput):
-        logits = logits.logits
-    elif isinstance(logits, Tuple):
-        logits = logits[0]
-
-    logits = logits.float()
-    shift_logits = logits[..., :-1, :].contiguous()
-    shift_labels = labels[..., 1:].contiguous()
-    loss_fct = CrossEntropyLoss(ignore_index=-100, reduction='none')
-    loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)).reshape(
-        shift_labels.shape[0], -1
-    )
-    loss = loss.sum(-1) / (shift_labels != -100).sum(-1)
-    loss = loss.mean()
-    return loss
-
-
-def global_token_level_cross_entropy(logits, labels, **kwargs):
-    # 对所有样本一起做token-level的平均
-    if isinstance(logits, ModelOutput):
-        logits = logits.logits
-    elif isinstance(logits, Tuple):
-        logits = logits[0]
-
-    logits = logits.float()
-    shift_logits = logits[..., :-1, :].contiguous()
-    shift_labels = labels[..., 1:].contiguous()
-    loss_fct = CrossEntropyLoss(ignore_index=-100, reduction='none')
-    loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)).reshape(
-        shift_labels.shape[0], -1
-    )
-    num_tokens = (shift_labels != -100).sum()
-    loss = loss.sum()
-
-    num_tokens_tensor = torch.zeros([1], device=loss.device, dtype=loss.dtype)
-    num_tokens_tensor[0] = num_tokens.item()
-
-    torch.distributed.all_reduce(num_tokens_tensor)
-
-    global_num_tokens = num_tokens_tensor.sum()
-
-    torch.distributed.barrier()
-    # global_num_tokens是全局的token数，因为在梯度更新的时候回自动对所有卡求mean
-    # 所有这里要乘一个world_size
-    loss = loss.sum() / global_num_tokens * torch.distributed.get_world_size()
-
-    return loss
-
-
-BAILING_LOSS_MAPPING = {
-    'local_token_level_cross_entropy': local_token_level_cross_entropy,
-    'sample_level_cross_entropy': sample_level_cross_entropy,
-    'global_token_level_cross_entropy': global_token_level_cross_entropy,
-}
-
-
-@dataclass
-class CustomMoeOutput(ModelOutput):
-    """完全自定义的输出类，包含所有需要的字段"""
-
-    loss: Optional[torch.FloatTensor] = None
-    aux_loss: Optional[torch.FloatTensor] = None
-    logits: torch.FloatTensor = None
-    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
-    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
-    attentions: Optional[Tuple[torch.FloatTensor]] = None
-    router_logits: Optional[Tuple[torch.FloatTensor]] = None
-    # 额外的损失组件
-    lm_loss: Optional[torch.FloatTensor] = None
-    balance_loss: Optional[torch.FloatTensor] = None
-    z_loss: Optional[torch.FloatTensor] = None
-    last_logits_l2_loss: Optional[torch.FloatTensor] = None
-
-
 class BailingMoeRMSNorm(nn.Module):
     def __init__(self, hidden_size, eps=1e-6):
         """
@@ -696,7 +477,6 @@ class BailingMoeAttention(nn.Module):
         value_states = value_states.transpose(1, 2)
 
         kv_seq_len = key_states.shape[-2]
-
         if past_key_value is not None:
             if self.layer_idx is None:
                 raise ValueError(
@@ -705,7 +485,6 @@ class BailingMoeAttention(nn.Module):
                     "with a layer index."
                 )
             kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
-
         cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
         query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
 
@@ -1564,67 +1343,36 @@ class BailingMoeForCausalLM(BailingMoePreTrainedModel):
 
         logits = logits.float()
 
-        lm_loss = None
+        loss = None
         aux_loss = None
 
         if labels is not None:
-            built_in_loss_mapping = {}
-            if version.parse(transformers.__version__) >= version.parse("4.46.0"):
-                from transformers.loss.loss_utils import LOSS_MAPPING
-
-                built_in_loss_mapping = dict(LOSS_MAPPING)
-            built_in_loss_mapping.update(BAILING_LOSS_MAPPING)
-
-            loss_type = getattr(self.config, "loss_type", None)
-            if loss_type is None or loss_type not in built_in_loss_mapping:
-                logger.warning_once(
-                    f"`loss_type={loss_type}` was set in the config but it is unrecognised. "
-                    f"Using the default loss: `global_token_level_cross_entropy`."
-                )
-                loss_type = "global_token_level_cross_entropy"
-
-            loss_fct = built_in_loss_mapping[loss_type]
-            lm_loss = loss_fct(logits, labels)
-
-        loss = lm_loss
-        if output_router_logits and labels is not None:
-            aux_loss, balance_loss, z_loss, last_logits_l2_loss = auxiliary_loss(
-                outputs.router_logits, logits, labels, self.config
-            )
-            loss = lm_loss + self.config.router_aux_loss_coef * aux_loss
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
 
         if not return_dict:
             output = (logits,) + outputs[1:]
-            if output_router_logits and labels is not None:
-                output = (aux_loss, balance_loss, z_loss, last_logits_l2_loss) + output
+            if output_router_logits:
+                output = (aux_loss,) + output
             return (loss,) + output if loss is not None else output
 
-        if output_router_logits and labels is not None:
-            moe_output = CustomMoeOutput(
-                loss=loss,
-                aux_loss=aux_loss,
-                logits=logits,
-                past_key_values=outputs.past_key_values,
-                hidden_states=outputs.hidden_states,
-                attentions=outputs.attentions,
-                router_logits=outputs.router_logits,
-                lm_loss=lm_loss,
-                balance_loss=balance_loss,
-                z_loss=z_loss,
-                last_logits_l2_loss=last_logits_l2_loss,
-            )
-
-            return moe_output
-        else:
-            return MoeCausalLMOutputWithPast(
-                loss=loss,
-                aux_loss=aux_loss,
-                logits=logits,
-                past_key_values=outputs.past_key_values,
-                hidden_states=outputs.hidden_states,
-                attentions=outputs.attentions,
-                router_logits=outputs.router_logits,
-            )
+        return MoeCausalLMOutputWithPast(
+            loss=loss,
+            aux_loss=aux_loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            router_logits=outputs.router_logits,
+        )
 
     def prepare_inputs_for_generation(
         self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, token_type_ids=None, **kwargs
@@ -1693,97 +1441,3 @@ class BailingMoeForCausalLM(BailingMoePreTrainedModel):
                 tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
             )
         return reordered_past
-
-
-class BailingMoeForRewardModel(BailingMoePreTrainedModel):
-    def __init__(self, config: BailingMoeConfig, model: BailingMoeModel = None):
-        super().__init__(config)
-        self.num_labels = 1  # config.num_labels
-        if model:
-            self.model = model
-        else:
-            self.model = BailingMoeModel(config)
-        self.value_head = nn.Sequential(
-            nn.Linear(config.hidden_size, config.hidden_size), nn.ReLU(), nn.Linear(config.hidden_size, self.num_labels)
-        )
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-    def get_input_embeddings(self):
-        return self.model.word_embeddings
-
-    def set_input_embeddings(self, value):
-        self.model.word_embeddings = value
-
-    @add_start_docstrings_to_model_forward(BAILINGMOE_INPUTS_DOCSTRING)
-    def forward(
-        self,
-        input_ids: torch.LongTensor = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[List[torch.FloatTensor]] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-    ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
-        r"""
-        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
-            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
-            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
-            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
-        """
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        transformer_outputs = self.model(
-            input_ids,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds,
-            use_cache=use_cache,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-
-        if return_dict:
-            last_hidden_state = transformer_outputs.last_hidden_state
-        else:
-            last_hidden_state = transformer_outputs[0]
-
-        logits = self.value_head(last_hidden_state)
-
-        if input_ids is not None:
-            batch_size = input_ids.shape[0]
-        else:
-            batch_size = inputs_embeds.shape[0]
-
-        if self.config.pad_token_id is None and batch_size != 1:
-            raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
-        if self.config.pad_token_id is None:
-            sequence_lengths = -1
-        else:
-            if input_ids is not None:
-                # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
-                sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
-                sequence_lengths = sequence_lengths % input_ids.shape[-1]
-                sequence_lengths = sequence_lengths.to(logits.device)
-            else:
-                sequence_lengths = -1
-
-        if isinstance(sequence_lengths, int) and sequence_lengths == -1:
-            sequence_lengths = (attention_mask.sum(dim=-1, keepdim=True) - 1).squeeze()
-
-        pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]  # logits of last token
-        pooled_logits = pooled_logits.squeeze()
-
-        return SequenceClassifierOutputWithPast(
-            logits=pooled_logits,
-            past_key_values=transformer_outputs.past_key_values,
-            hidden_states=transformer_outputs.hidden_states,
-            attentions=transformer_outputs.hidden_states,
-        )