Upload modeling_hybridna.py with huggingface_hub

modeling_hybridna.py

@@ -35,8 +35,8 @@ from transformers.modeling_attn_mask_utils import (
     AttentionMaskConverter,
 )
 from transformers.modeling_outputs import (
-    MoeCausalLMOutputWithPast,
-    MoeModelOutputWithPast,
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
     SequenceClassifierOutputWithPast,
 )
 from transformers.modeling_utils import PreTrainedModel
@@ -90,85 +90,6 @@ logger = logging.get_logger(__name__)
 _CONFIG_FOR_DOC = "HybriDNAConfig"
 
 
-# Copied from transformers.models.mixtral.modeling_mixtral.load_balancing_loss_func with gate->router
-def load_balancing_loss_func(
-        router_logits: torch.Tensor,
-        num_experts: torch.Tensor = None,
-        top_k=2,
-        attention_mask: Optional[torch.Tensor] = None,
-) -> float:
-    r"""
-    Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
-    See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. This function implements the loss
-    function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
-    experts is too unbalanced.
-    Args:
-        router_logits (Union[`torch.Tensor`, Tuple[torch.Tensor]):
-            Logits from the `router`, should be a tuple of model.config.num_hidden_layers tensors of
-            shape [batch_size X sequence_length, num_experts].
-        attention_mask (`torch.Tensor`, None):
-            The attention_mask used in forward function
-            shape [batch_size X sequence_length] if not None.
-        num_experts (`int`, *optional*):
-            Number of experts
-    Returns:
-        The auxiliary loss.
-    """
-    if router_logits is None or not isinstance(router_logits, tuple):
-        return 0
-
-    if isinstance(router_logits, tuple):
-        compute_device = router_logits[0].device
-        concatenated_router_logits = torch.cat(
-            [layer_router.to(compute_device) for layer_router in router_logits], dim=0
-        )
-
-    routing_weights = torch.nn.functional.softmax(concatenated_router_logits, dim=-1)
-
-    _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)
-
-    expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
-
-    if attention_mask is None:
-        # Compute the percentage of tokens routed to each experts
-        tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
-
-        # Compute the average probability of routing to these experts
-        router_prob_per_expert = torch.mean(routing_weights, dim=0)
-    else:
-        batch_size, sequence_length = attention_mask.shape
-        num_hidden_layers = concatenated_router_logits.shape[0] // (batch_size * sequence_length)
-
-        # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask
-        expert_attention_mask = (
-            attention_mask[None, :, :, None, None]
-                .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))
-                .reshape(-1, top_k, num_experts)
-                .to(compute_device)
-        )
-
-        # Compute the percentage of tokens routed to each experts
-        tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(
-            expert_attention_mask, dim=0
-        )
-
-        # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert
-        router_per_expert_attention_mask = (
-            attention_mask[None, :, :, None]
-                .expand((num_hidden_layers, batch_size, sequence_length, num_experts))
-                .reshape(-1, num_experts)
-                .to(compute_device)
-        )
-
-        # Compute the average probability of routing to these experts
-        router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(
-            router_per_expert_attention_mask, dim=0
-        )
-
-    overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
-    return overall_loss * num_experts
-
-
 # Copied from transformers.models.llama.modeling_llama._get_unpad_data
 def _get_unpad_data(attention_mask):
     seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
@@ -1221,9 +1142,7 @@ class HybriDNAMLP(nn.Module):
 class HybriDNAAttentionDecoderLayer(nn.Module):
     def __init__(self, config: HybriDNAConfig, layer_idx: int):
         super().__init__()
-        # Remove MoE support: always use vanilla MLP
         self.self_attn = HYBRIDNA_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx)
-        # Previously: ffn_layer_class = HybriDNASparseMoeBlock if num_experts > 1 else HybriDNAMLP
         self.feed_forward = HybriDNAMLP(config)
         self.input_layernorm = HybriDNARMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.pre_ff_layernorm = HybriDNARMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -1231,10 +1150,9 @@ class HybriDNAAttentionDecoderLayer(nn.Module):
     def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None,
                 position_ids: Optional[torch.LongTensor] = None,
                 past_key_value: Optional[HybridMambaAttentionDynamicCache] = None,
-                output_attentions: Optional[bool] = False, output_router_logits: Optional[bool] = False,
+                output_attentions: Optional[bool] = False,
                 use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None
                 ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
-        # ...existing code before feed-forward...
         residual = hidden_states
         hidden_states = self.input_layernorm(hidden_states)
 
@@ -1251,28 +1169,24 @@ class HybriDNAAttentionDecoderLayer(nn.Module):
         # residual connection after attention
         hidden_states = residual + hidden_states
 
-        # feed-forward (experts/MLP)
+        # feed-forward
         residual = hidden_states
         hidden_states = self.pre_ff_layernorm(hidden_states)
-        # Remove MoE tuple-check: directly compute feed-forward output
         hidden_states = self.feed_forward(hidden_states)
         hidden_states = residual + hidden_states
 
         outputs = (hidden_states,)
         if output_attentions:
-            outputs += (self_attn_weights,)  # remains unchanged from attention branch
+            outputs += (self_attn_weights,)
         if use_cache:
             outputs += (present_key_value,)
-        # Remove router_logits branch completely
         return outputs
 
 
 class HybriDNAMambaDecoderLayer(nn.Module):
     def __init__(self, config: HybriDNAConfig, layer_idx: int):
         super().__init__()
-        # Remove MoE support: always use vanilla MLP
         self.mamba = HybriDNAMamba2Mixer(config=config, layer_idx=layer_idx)
-        # Previously: ffn_layer_class = HybriDNASparseMoeBlock if num_experts > 1 else HybriDNAMLP
         self.feed_forward = HybriDNAMLP(config)
         self.input_layernorm = HybriDNARMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.pre_ff_layernorm = HybriDNARMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -1280,10 +1194,9 @@ class HybriDNAMambaDecoderLayer(nn.Module):
     def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None,
                 position_ids: Optional[torch.LongTensor] = None,
                 past_key_value: Optional[HybridMambaAttentionDynamicCache] = None,
-                output_attentions: Optional[bool] = False, output_router_logits: Optional[bool] = False,
+                output_attentions: Optional[bool] = False,
                 use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None
                 ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
-        # ...existing code before feed-forward...
         residual = hidden_states
 
         hidden_states = self.input_layernorm(hidden_states)
@@ -1297,10 +1210,9 @@ class HybriDNAMambaDecoderLayer(nn.Module):
         # residual connection after mamba
         hidden_states = residual + hidden_states
 
-        # feed-forward (experts/MLP)
+        # feed-forward
         residual = hidden_states
         hidden_states = self.pre_ff_layernorm(hidden_states)
-        # Remove MoE tuple-check: directly compute feed-forward output
         hidden_states = self.feed_forward(hidden_states)
         hidden_states = residual + hidden_states
 
@@ -1309,11 +1221,10 @@ class HybriDNAMambaDecoderLayer(nn.Module):
             outputs += (self_attn_weights,)
         if use_cache:
             outputs += (past_key_value,)
-        # Remove router_logits branch completely
         return outputs
 
 
-JAMBA_START_DOCSTRING = r"""
+HYBRIDNA_START_DOCSTRING = r"""
     This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
     library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
     etc.)
@@ -1330,7 +1241,7 @@ JAMBA_START_DOCSTRING = r"""
 
 @add_start_docstrings(
     "The bare HybriDNA Model outputting raw hidden-states without any specific head on top.",
-    JAMBA_START_DOCSTRING,
+    HYBRIDNA_START_DOCSTRING,
 )
 class HybriDNAPreTrainedModel(PreTrainedModel):
     config_class = HybriDNAConfig
@@ -1354,7 +1265,7 @@ class HybriDNAPreTrainedModel(PreTrainedModel):
                 module.weight.data[module.padding_idx].zero_()
 
 
-JAMBA_INPUTS_DOCSTRING = r"""
+HYBRIDNA_INPUTS_DOCSTRING = r"""
     Args:
         input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
             Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
@@ -1404,9 +1315,6 @@ JAMBA_INPUTS_DOCSTRING = r"""
         output_hidden_states (`bool`, *optional*):
             Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
             more detail.
-        output_router_logits (`bool`, *optional*):
-            Whether or not to return the logits of all the routers. They are useful for computing the router loss, and
-            should not be returned during inference.
         return_dict (`bool`, *optional*):
             Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
         cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
@@ -1420,9 +1328,8 @@ ALL_DECODER_LAYER_TYPES = {"attention": HybriDNAAttentionDecoderLayer, "mamba":
 
 @add_start_docstrings(
     "The bare HybriDNA Model outputting raw hidden-states without any specific head on top.",
-    JAMBA_START_DOCSTRING,
+    HYBRIDNA_START_DOCSTRING,
 )
-# Adapted from transformers.models.mistral.modeling_mistral.MistralModel with MISTRAL->JAMBA, Mistral->HybriDNA
 class HybriDNAModel(HybriDNAPreTrainedModel):
     """
     Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`HybriDNADecoderLayer`]
@@ -1455,7 +1362,7 @@ class HybriDNAModel(HybriDNAPreTrainedModel):
     def set_input_embeddings(self, value):
         self.embed_tokens = value
 
-    @add_start_docstrings_to_model_forward(JAMBA_INPUTS_DOCSTRING)
+    @add_start_docstrings_to_model_forward(HYBRIDNA_INPUTS_DOCSTRING)
     def forward(
             self,
             input_ids: torch.LongTensor = None,
@@ -1466,14 +1373,10 @@ class HybriDNAModel(HybriDNAPreTrainedModel):
             use_cache: Optional[bool] = None,
             output_attentions: Optional[bool] = None,
             output_hidden_states: Optional[bool] = None,
-            output_router_logits: Optional[bool] = None,
             return_dict: Optional[bool] = None,
             cache_position: Optional[torch.LongTensor] = None,
-    ) -> Union[Tuple, MoeModelOutputWithPast]:
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_router_logits = (
-            output_router_logits if output_router_logits is not None else self.config.output_router_logits
-        )
         output_hidden_states = (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
@@ -1512,7 +1415,6 @@ class HybriDNAModel(HybriDNAPreTrainedModel):
 
         all_hidden_states = () if output_hidden_states else None
         all_self_attns = () if output_attentions else None
-        all_router_logits = () if output_router_logits else None
 
         for decoder_layer in self.layers:
             if output_hidden_states:
@@ -1526,7 +1428,6 @@ class HybriDNAModel(HybriDNAPreTrainedModel):
                     position_ids,
                     past_key_values,
                     output_attentions,
-                    output_router_logits,
                     use_cache,
                     cache_position,
                 )
@@ -1537,7 +1438,6 @@ class HybriDNAModel(HybriDNAPreTrainedModel):
                     position_ids=position_ids,
                     past_key_value=past_key_values,
                     output_attentions=output_attentions,
-                    output_router_logits=output_router_logits,
                     use_cache=use_cache,
                     cache_position=cache_position,
                 )
@@ -1549,11 +1449,6 @@ class HybriDNAModel(HybriDNAPreTrainedModel):
                     # append attentions only of attention layers. Mamba layers return `None` as the attention weights
                     all_self_attns += (layer_outputs[1],)
 
-            if output_router_logits:
-                if layer_outputs[-1] is not None:
-                    # append router logits only of expert layers. Regular MLP layers return `None` as the router logits
-                    all_router_logits += (layer_outputs[-1],)
-
         hidden_states = self.final_layernorm(hidden_states)
 
         # add hidden states from the last decoder layer
@@ -1570,15 +1465,14 @@ class HybriDNAModel(HybriDNAPreTrainedModel):
         if not return_dict:
             return tuple(
                 v
-                for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits]
+                for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]
                 if v is not None
             )
-        return MoeModelOutputWithPast(
+        return BaseModelOutputWithPast(
             last_hidden_state=hidden_states,
             past_key_values=next_cache,
             hidden_states=all_hidden_states,
             attentions=all_self_attns,
-            router_logits=all_router_logits,
         )
 
     def _update_causal_mask(self, attention_mask, input_tensor, cache_position):
@@ -1617,7 +1511,6 @@ class HybriDNAModel(HybriDNAPreTrainedModel):
         return causal_mask
 
 
-# Adapted from transformers.models.mixtral.modeling_mixtral.MixtralForCausalLM with MIXTRAL->JAMBA, Mixtral->HybriDNA
 class HybriDNAForCausalLM(HybriDNAPreTrainedModel, GenerationMixin):
     _tied_weights_keys = ["lm_head.weight"]
 
@@ -1626,9 +1519,6 @@ class HybriDNAForCausalLM(HybriDNAPreTrainedModel, GenerationMixin):
         self.model = HybriDNAModel(config)
         self.vocab_size = config.vocab_size
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-        self.router_aux_loss_coef = config.router_aux_loss_coef
-        self.num_experts = config.num_experts
-        self.num_experts_per_tok = config.num_experts_per_tok
         # Initialize weights and apply final processing
         self.post_init()
 
@@ -1650,9 +1540,8 @@ class HybriDNAForCausalLM(HybriDNAPreTrainedModel, GenerationMixin):
     def get_decoder(self):
         return self.model
 
-    @add_start_docstrings_to_model_forward(JAMBA_INPUTS_DOCSTRING)
-    @replace_return_docstrings(output_type=MoeCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
-    # Ignore copy
+    @add_start_docstrings_to_model_forward(HYBRIDNA_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
     def forward(
             self,
             input_ids: torch.LongTensor = None,
@@ -1664,11 +1553,10 @@ class HybriDNAForCausalLM(HybriDNAPreTrainedModel, GenerationMixin):
             use_cache: Optional[bool] = None,
             output_attentions: Optional[bool] = None,
             output_hidden_states: Optional[bool] = None,
-            output_router_logits: Optional[bool] = None,
             return_dict: Optional[bool] = None,
             cache_position: Optional[torch.LongTensor] = None,
             num_logits_to_keep: Optional[Union[int, None]] = None,
-    ) -> Union[Tuple, MoeCausalLMOutputWithPast]:
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
         r"""
         Args:
             labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
@@ -1682,22 +1570,17 @@ class HybriDNAForCausalLM(HybriDNAPreTrainedModel, GenerationMixin):
         Returns:
         Example:
         ```python
-        >>> from transformers import AutoTokenizer, HybriDNAForCausalLM
-        >>> model = HybriDNAForCausalLM.from_pretrained("ai21labs/HybriDNA-v0.1")
-        >>> tokenizer = AutoTokenizer.from_pretrained("ai21labs/HybriDNA-v0.1")
-        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> from transformers import AutoTokenizer, AutoModelForCausalLM
+        >>> model = AutoModelForCausalLM.from_pretrained("Mishamq/HybriDNA-300M", trust_remote_code=True)
+        >>> tokenizer = AutoTokenizer.from_pretrained("Mishamq/HybriDNA-300M", trust_remote_code=True)
+        >>> prompt = "ACGTACGTACGTACGT"
         >>> inputs = tokenizer(prompt, return_tensors="pt")
         >>> # Generate
         >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
         >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
-        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
         ```"""
 
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_router_logits = (
-            output_router_logits if output_router_logits is not None else self.config.output_router_logits
-        )
-
         output_hidden_states = (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
@@ -1713,7 +1596,6 @@ class HybriDNAForCausalLM(HybriDNAPreTrainedModel, GenerationMixin):
             use_cache=use_cache,
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
-            output_router_logits=output_router_logits,
             cache_position=cache_position,
             return_dict=return_dict,
         )
@@ -1738,31 +1620,16 @@ class HybriDNAForCausalLM(HybriDNAPreTrainedModel, GenerationMixin):
             shift_labels = shift_labels.to(shift_logits.device)
             loss = loss_fct(shift_logits, shift_labels)
 
-        aux_loss = None
-        if output_router_logits:
-            aux_loss = load_balancing_loss_func(
-                outputs.router_logits if return_dict else outputs[-1],
-                self.num_experts,
-                self.num_experts_per_tok,
-                attention_mask,
-            )
-            if labels is not None:
-                loss += self.router_aux_loss_coef * aux_loss.to(loss.device)  # make sure to reside in the same device
-
         if not return_dict:
             output = (logits,) + outputs[1:]
-            if output_router_logits:
-                output = (aux_loss,) + output
             return (loss,) + output if loss is not None else output
 
-        return MoeCausalLMOutputWithPast(
+        return CausalLMOutputWithPast(
             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(
@@ -1771,7 +1638,6 @@ class HybriDNAForCausalLM(HybriDNAPreTrainedModel, GenerationMixin):
             past_key_values=None,
             attention_mask=None,
             inputs_embeds=None,
-            output_router_logits=False,
             cache_position=None,
             **kwargs,
     ):
@@ -1827,7 +1693,6 @@ class HybriDNAForCausalLM(HybriDNAPreTrainedModel, GenerationMixin):
                 "past_key_values": past_key_values,
                 "use_cache": kwargs.get("use_cache"),
                 "attention_mask": attention_mask,
-                "output_router_logits": output_router_logits,
                 "num_logits_to_keep": self.config.num_logits_to_keep,
                 "cache_position": cache_position,
             }
@@ -1846,9 +1711,8 @@ class HybriDNAForCausalLM(HybriDNAPreTrainedModel, GenerationMixin):
     padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
     each row of the batch).
     """,
-    JAMBA_START_DOCSTRING,
+    HYBRIDNA_START_DOCSTRING,
 )
-# Copied from transformers.models.mixtral.modeling_mixtral.MixtralForSequenceClassification with Mixtral->HybriDNA, MIXTRAL->JAMBA
 class HybriDNAForSequenceClassification(HybriDNAPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
@@ -1865,7 +1729,7 @@ class HybriDNAForSequenceClassification(HybriDNAPreTrainedModel):
     def set_input_embeddings(self, value):
         self.model.embed_tokens = value
 
-    @add_start_docstrings_to_model_forward(JAMBA_INPUTS_DOCSTRING)
+    @add_start_docstrings_to_model_forward(HYBRIDNA_INPUTS_DOCSTRING)
     def forward(
             self,
             input_ids: torch.LongTensor = None,
@@ -1963,7 +1827,7 @@ class HybriDNAForSequenceClassification(HybriDNAPreTrainedModel):
     The input sequence is concatenated with its reverse complement before being processed by the model.
     [`HybriDNAForSequenceClassificationRCEcho`]
     """,
-    JAMBA_START_DOCSTRING,
+    HYBRIDNA_START_DOCSTRING,
 )
 class HybriDNAForSequenceClassificationRCEcho(HybriDNAPreTrainedModel):
     def __init__(self, config):
@@ -2005,7 +1869,7 @@ class HybriDNAForSequenceClassificationRCEcho(HybriDNAPreTrainedModel):
         rc = torch.flip(rc, dims=[1])
         return rc
 
-    @add_start_docstrings_to_model_forward(JAMBA_INPUTS_DOCSTRING)
+    @add_start_docstrings_to_model_forward(HYBRIDNA_INPUTS_DOCSTRING)
     def forward(
         self,
         input_ids: torch.LongTensor = None,