Upload modeling_BiMambaForMaskedLM.py

modeling_BiMambaForMaskedLM.py

@@ -0,0 +1,152 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, AutoConfig
+from transformers.modeling_outputs import MaskedLMOutput   
+from mamba_ssm.modules.mamba_simple import Mamba
+from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
+from mamba_ssm.models.config_mamba import MambaConfig
+from mamba_ssm.ops.selective_scan_interface import selective_scan_fn, mamba_inner_fn
+try:
+    from mamba_ssm.ops.triton.layer_norm import RMSNorm, layer_norm_fn, rms_norm_fn
+except ImportError:
+    RMSNorm, layer_norm_fn, rms_norm_fn = None, None, None
+from einops import rearrange
+
+def convert_hf_config_to_mamba(hf_config) -> MambaConfig:
+  return MambaConfig(
+      d_model=hf_config.d_model,
+      d_intermediate=getattr(hf_config, "intermediate_size", 4 * hf_config.d_model),
+      n_layer=getattr(hf_config, "n_layer", getattr(hf_config, "num_hidden_layers", 12)),
+      vocab_size=hf_config.vocab_size,
+      ssm_cfg=getattr(hf_config, "ssm_cfg", {}),
+      attn_layer_idx=getattr(hf_config, "attn_layer_idx", []),
+      attn_cfg=getattr(hf_config, "attn_cfg", {}),
+      rms_norm=getattr(hf_config, "rms_norm", True),
+      residual_in_fp32=getattr(hf_config, "residual_in_fp32", True),
+      fused_add_norm=getattr(hf_config, "fused_add_norm", False),
+      pad_vocab_size_multiple=getattr(hf_config, "pad_vocab_size_multiple", 8),
+      tie_embeddings=getattr(hf_config, "tie_embeddings", False),
+  )
+
+def patch_mixer_forward_to_accept_embeddings(model):
+    """
+    Injects a new forward method into a MixerModel instance,
+    allowing it to accept either input_ids or inputs_embeds.
+    """
+
+    def new_forward(self, input_ids=None, inputs_embeds=None, inference_params=None, attention_mask=None, **mixer_kwargs):
+        if inputs_embeds is not None:
+            hidden_states = inputs_embeds
+        elif input_ids is not None:
+            hidden_states = self.embedding(input_ids)
+        else:
+            raise ValueError("You must provide either input_ids or inputs_embeds.")
+
+        residual = None
+
+        # hiddens: (batch_size, seq_len, d_model)
+        # attention_mask: (batch_size, seq_len) -- 1 for real tokens, 0 for padding
+        mask = attention_mask.unsqueeze(-1)  # (batch_size, seq_len, 1)
+
+        for layer in self.layers:
+            hidden_states, residual = layer(
+                hidden_states, residual, inference_params=inference_params, **mixer_kwargs
+            )
+
+            # Add attention mask
+            hidden_states = hidden_states * mask
+            residual = residual * mask
+
+        if not self.fused_add_norm:
+            residual = (hidden_states + residual) if residual is not None else hidden_states
+            hidden_states = self.norm_f(residual.to(dtype=self.norm_f.weight.dtype))
+        else:
+            # Set prenorm=False here since we don't need the residual
+            hidden_states = layer_norm_fn(
+                hidden_states,
+                self.norm_f.weight,
+                self.norm_f.bias,
+                eps=self.norm_f.eps,
+                residual=residual,
+                prenorm=False,
+                residual_in_fp32=self.residual_in_fp32,
+                is_rms_norm=isinstance(self.norm_f, RMSNorm)
+            )
+        return hidden_states
+
+    # Bind the new forward method to the instance
+    model.backbone.forward = new_forward.__get__(model.backbone, model.backbone.__class__)
+
+class BiMambaForMaskedLM(PreTrainedModel):
+    config_class    = AutoConfig
+    base_model_prefix = "bimamba"
+
+    def __init__(self, config):
+        super().__init__(config)                    # <-- HF init
+        mamba_cfg = convert_hf_config_to_mamba(config)
+
+        # your embedding + two Mamba directions + proj
+        self.token_embedding = nn.Embedding(config.vocab_size, config.d_model, padding_idx=config.pad_token_id)
+        self.mamba_forward   = MambaLMHeadModel(mamba_cfg)
+        self.mamba_backward  = MambaLMHeadModel(mamba_cfg)
+        self.lm_head_proj    = nn.Linear(config.d_model * 2, config.d_model, bias=False)
+
+        # Patch mixer_forward_to accept embeddings
+        patch_mixer_forward_to_accept_embeddings(self.mamba_forward)
+        patch_mixer_forward_to_accept_embeddings(self.mamba_backward)
+
+        # self.post_init()  # wires up HF weight-tying & save/load
+
+    #### Added:
+    def get_input_embeddings(self):
+        return self.token_embedding
+
+    def set_input_embeddings(self, new_emb):
+        self.token_embedding = new_emb
+
+    def get_output_embeddings(self):
+        return self.lm_head_proj
+
+    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
+        for backbone in (self.mamba_forward.backbone,
+                         self.mamba_backward.backbone):
+            for block in backbone.layers:
+                block.gradient_checkpointing = True
+
+    def forward(
+        self,
+        input_ids=None,
+        inputs_embeds=None,
+        attention_mask=None,
+        labels=None,
+        return_dict=True,
+    ):
+        if inputs_embeds is None:
+            input_ids = input_ids.long()
+            inputs_embeds = self.token_embedding(input_ids)
+
+        hid_fwd = self.mamba_forward.backbone(inputs_embeds=inputs_embeds, attention_mask=attention_mask)
+        rev_emb = torch.flip(inputs_embeds, dims=[1])
+        rev_mask = torch.flip(attention_mask, dims=[1])
+        hid_bwd = self.mamba_backward.backbone(inputs_embeds=rev_emb, attention_mask=rev_mask)
+        hid_bwd = torch.flip(hid_bwd, dims=[1])
+
+        combined = torch.cat([hid_fwd, hid_bwd], dim=-1)
+        projected = self.lm_head_proj(combined)
+        logits    = F.linear(projected, self.token_embedding.weight)
+
+        loss = None
+        if labels is not None:
+            loss_fn = nn.CrossEntropyLoss(ignore_index=-100)
+            loss    = loss_fn(logits.view(-1, logits.size(-1)), labels.view(-1))
+
+        if not return_dict:
+            out = (logits, combined)
+            return (loss,) + out if loss is not None else out
+
+        return MaskedLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=projected,
+        )