Upload model

config.json

@@ -1,6 +1,10 @@
 {
+  "architectures": [
+    "MambaLMHeadModel"
+  ],
   "auto_map": {
-    "AutoConfig": "configuration_mamba.MambaConfig"
+    "AutoConfig": "configuration_mamba.MambaConfig",
+    "AutoModelForCausalLM": "modeling_mamba.MambaLMHeadModel"
   },
   "bias": false,
   "conv_bias": true,
@@ -14,6 +18,7 @@
   "model_type": "mamba",
   "n_layer": 24,
   "pad_vocab_size_multiple": 8,
+  "torch_dtype": "float32",
   "transformers_version": "4.37.2",
   "vocab_size": 50280
 }

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6504b24e9ba95e4a6bad94a346c849040623647d1a99a47f4f5e1cd32cbd9572
-size 259551392
+oid sha256:1bd3ca62665de4bfabff9d443f87a11090a10e505c0ccb56e6f9ca495b6e05bd
+size 671027808

modeling_mamba.py

@@ -1,82 +1,617 @@
-from typing import Optional, Tuple
+import json
+import math
+import os
+from collections import namedtuple
+from dataclasses import dataclass
+from functools import partial
+from typing import Dict, Optional, Tuple, Union
 
-from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
 import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import transformers
+from einops import einsum, rearrange, repeat
+from torch import FloatTensor, Tensor, nn
 from transformers import GenerationMixin, PreTrainedModel
-from transformers.generation import TextStreamer
+from transformers.modeling_outputs import (
+    BaseModelOutput,
+    BaseModelOutputWithPast,
+    CausalLMOutput,
+    ImageClassifierOutput,
+    QuestionAnsweringModelOutput,
+    SequenceClassifierOutput,
+)
+from trl import PreTrainedModelWrapper
 
 from .configuration_mamba import MambaConfig
 
-class MambaModel(PreTrainedModel):
-    config_class = MambaConfig
 
-    def __init__(
+# class SwiGLU(nn.Module):
+#     def forward(self, x, W, V, b, c, beta):
+#         return F.silu(x * W + b) * (x * V + c)
+
+
+# Inspired by:
+# - https://huggingface.co/Q-bert/Mamba-130M/blob/f0d00db98acaa62b1ee4304cd11643e69aa62a71/modeling_mamba.py#L31
+# - https://github.com/johnma2006/mamba-minimal/blob/03de542a36d873f6e6c4057ad687278cc6ae944d/model.py#L177
+# - https://github.com/state-spaces/mamba/blob/009bec5ee37f586844a3fc89c040a9c1a9d8badf/mamba_ssm/modules/mamba_simple.py#L31
+class MambaBlock(nn.Module):
+    def __init__(self, config: MambaConfig):
+        """A single Mamba block, as described in Figure 3 in Section 3.4 in the Mamba paper [1].
+        Furthermore, in section E.2.2 of the paper, the authors describe the Mamba block as:
+            "[T]he Mamba block is simply the standard SwiGLU block with an extra conv → SSM path added."
+        """
+        super().__init__()
+
+        self.config = config
+
+        self.in_proj = nn.Linear(config.d_model, config.d_inner * 2, bias=config.bias)
+
+        self.conv1d = nn.Conv1d(
+            in_channels=config.d_inner,
+            out_channels=config.d_inner,
+            bias=config.conv_bias,
+            kernel_size=config.d_conv,
+            groups=config.d_inner,
+            padding=config.d_conv - 1,
+        )
+
+        # x_proj takes in `x` and outputs the input-specific Δ, B, C
+        self.x_proj = nn.Linear(
+            config.d_inner, config.dt_rank + config.d_state * 2, bias=False
+        )
+
+        # dt_proj projects Δ from dt_rank to d_in
+        self.dt_proj = nn.Linear(config.dt_rank, config.d_inner, bias=True)
+
+        A = repeat(torch.arange(1, config.d_state + 1), "n -> d n", d=config.d_inner)
+        self.A_log = nn.Parameter(torch.log(A))
+        self.D = nn.Parameter(torch.ones(config.d_inner))
+
+        self.out_proj = nn.Linear(config.d_inner, config.d_model, bias=config.bias)
+        # self.norm = RMSNorm(config.d_model)
+
+
+
+    def forward(self, x):
+        """Mamba block forward. This looks the same as Figure 3 in Section 3.4 in the Mamba paper [1].
+
+        Args:
+            x: shape (b, l, d)    (See Glossary at top for definitions of b, l, d_in, n...)
+
+        Returns:
+            output: shape (b, l, d)
+
+        Official Implementation:
+            class Mamba, https://github.com/state-spaces/mamba/blob/main/mamba_ssm/modules/mamba_simple.py#L119
+            mamba_inner_ref(), https://github.com/state-spaces/mamba/blob/main/mamba_ssm/ops/selective_scan_interface.py#L311
+
+        """
+        (b, l, d) = x.shape
+        # x_copy = x # There was a separate class for residual, I deleted that part and added it here.
+        # x = self.norm(x)
+        x_and_res = self.in_proj(x)  # shape (b, l, 2 * d_in)
+        (x, res) = x_and_res.split(
+            split_size=[self.config.d_inner, self.config.d_inner], dim=-1
+        )
+
+        x = rearrange(x, "b l d_in -> b d_in l")
+        x = self.conv1d(x)[:, :, :l]
+        x = rearrange(x, "b d_in l -> b l d_in")
+
+        x = F.silu(x)
+
+        y = self.ssm(x)
+
+        y = y * F.silu(res) # SwiGLU: Swish_β(xW + b) ⊗ (xV + c) => torch.kron(F.silu(xW + b), xV + c) => torch.kron(F.silu(res), y)
+
+        output = self.out_proj(y)  # output = self.out_proj(y) + x_copy
+
+        # "the Mamba block is simply the standard SwiGLU block with an extra 𝖼𝗈𝗇𝗏 → 𝖲𝖲𝖬 path added"
+
+        return output
+
+    def ssm(self, x):
+        """Runs the SSM. See:
+            - Algorithm 2 in Section 3.2 in the Mamba paper [1]
+            - run_SSM(A, B, C, u) in The Annotated S4 [2]
+
+        Args:
+            x: shape (b, l, d_in)    (See Glossary at top for definitions of b, l, d_in, n...)
+
+        Returns:
+            output: shape (b, l, d_in)
+
+        Official Implementation:
+            mamba_inner_ref(), https://github.com/state-spaces/mamba/blob/main/mamba_ssm/ops/selective_scan_interface.py#L311
+
+        """
+        (d_in, n) = self.A_log.shape
+
+        # Compute ∆ A B C D, the state space parameters.
+        #     A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective)
+        #     ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4,
+        #                                  and is why Mamba is called **selective** state spaces)
+
+        A = -torch.exp(self.A_log.float())  # shape (d_in, n)
+        D = self.D.float()
+
+        x_dbl = self.x_proj(x)  # (b, l, dt_rank + 2*n)
+
+        (delta, B, C) = x_dbl.split(
+            split_size=[self.config.dt_rank, n, n], dim=-1
+        )  # delta: (b, l, dt_rank). B, C: (b, l, n)
+        delta = F.softplus(self.dt_proj(delta))  # (b, l, d_in)
+
+        y = self.selective_scan(
+            x, delta, A, B, C, D
+        )  # This is similar to run_SSM(A, B, C, u) in The Annotated S4 [2]
+
+        return y
+
+    def selective_scan(self, u, delta, A, B, C, D):
+        """Does selective scan algorithm. See:
+            - Section 2 State Space Models in the Mamba paper [1]
+            - Algorithm 2 in Section 3.2 in the Mamba paper [1]
+            - run_SSM(A, B, C, u) in The Annotated S4 [2]
+
+        This is the classic discrete state space formula:
+            x(t + 1) = Ax(t) + Bu(t)
+            y(t)     = Cx(t) + Du(t)
+        except B and C (and the step size delta, which is used for discretization) are dependent on the input x(t).
+
+        Args:
+            u: shape (b, l, d_in)    (See Glossary at top for definitions of b, l, d_in, n...)
+            delta: shape (b, l, d_in)
+            A: shape (d_in, n)
+            B: shape (b, l, n)
+            C: shape (b, l, n)
+            D: shape (d_in,)
+
+        Returns:
+            output: shape (b, l, d_in)
+
+        Official Implementation:
+            selective_scan_ref(), https://github.com/state-spaces/mamba/blob/main/mamba_ssm/ops/selective_scan_interface.py#L86
+            Note: I refactored some parts out of `selective_scan_ref` out, so the functionality doesn't match exactly.
+
+        """
+        (b, l, d_in) = u.shape
+        n = A.shape[1]
+
+        # Discretize continuous parameters (A, B)
+        # - A is discretized using zero-order hold (ZOH) discretization (see Section 2 Equation 4 in the Mamba paper [1])
+        # - B is discretized using a simplified Euler discretization instead of ZOH. From a discussion with authors:
+        #   "A is the more important term and the performance doesn't change much with the simplification on B"
+        deltaA = torch.exp(einsum(delta, A, "b l d_in, d_in n -> b l d_in n"))
+        deltaB_u = einsum(delta, B, u, "b l d_in, b l n, b l d_in -> b l d_in n")
+
+        # Perform selective scan (see scan_SSM() in The Annotated S4 [2])
+        # Note that the below is sequential, while the official implementation does a much faster parallel scan that
+        # is additionally hardware-aware (like FlashAttention).
+        x = torch.zeros((b, d_in, n), device=deltaA.device)
+        ys = []
+
+        for i in range(l):
+            x = deltaA[:, i] * x + deltaB_u[:, i]
+            y = einsum(x, C[:, i, :], "b d_in n, b n -> b d_in")
+            ys.append(y)
+
+        y = torch.stack(ys, dim=1)  # shape (b, l, d_in)
+
+        y = y + u * D
+
+        return y
+
+
+# Inspired by:
+# - https://huggingface.co/Q-bert/Mamba-130M/blob/f0d00db98acaa62b1ee4304cd11643e69aa62a71/modeling_mamba.py#L19
+# - https://github.com/johnma2006/mamba-minimal/blob/03de542a36d873f6e6c4057ad687278cc6ae944d/model.py#L328
+# - https://github.com/state-spaces/mamba/blob/009bec5ee37f586844a3fc89c040a9c1a9d8badf/mamba_ssm/ops/triton/layernorm.py#L481
+class RMSNorm(nn.Module):
+    def __init__(self, d_model: int, eps: float = 1e-5):
+        super().__init__()
+
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(d_model))
+
+    def forward(self, x):
+        output = (
+            x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
+        )
+
+        return output
+
+
+class ResidualBlock(
+    nn.Module
+):  # Copied and modified from https://github.com/johnma2006/mamba-minimal/blob/03de542a36d873f6e6c4057ad687278cc6ae944d/model.py#L143
+    def __init__(self, config: MambaConfig):
+        """Simple block wrapping Mamba block with normalization and residual connection."""
+        super().__init__()
+
+        # self.args = args
+        self.mixer = MambaBlock(config)
+        self.norm = RMSNorm(config.d_model)
+
+    def forward(self, x):
+        """
+        Args:
+            x: shape (b, l, d)    (See Glossary at top for definitions of b, l, d_in, n...)
+
+        Returns:
+            output: shape (b, l, d)
+
+        Official Implementation:
+            Block.forward(), https://github.com/state-spaces/mamba/blob/main/mamba_ssm/modules/mamba_simple.py#L297
+
+            Note: the official repo chains residual blocks that look like
+                [Add -> Norm -> Mamba] -> [Add -> Norm -> Mamba] -> [Add -> Norm -> Mamba] -> ...
+            where the first Add is a no-op. This is purely for performance reasons as this
+            allows them to fuse the Add->Norm.
+
+            We instead implement our blocks as the more familiar, simpler, and numerically equivalent
+                [Norm -> Mamba -> Add] -> [Norm -> Mamba -> Add] -> [Norm -> Mamba -> Add] -> ....
+
+        """
+        output = self.mixer(self.norm(x)) + x
+
+        return output
+
+# Inspired by:
+# - https://huggingface.co/Q-bert/Mamba-130M/blob/f0d00db98acaa62b1ee4304cd11643e69aa62a71/modeling_mamba.py#L181
+# class MambaPretrainedModel(PreTrainedModel, nn.Module):
+class MambaPretrainedModel(PreTrainedModel):
+    r"""
+    Base class for all models.
+
+    [`PreTrainedModel`] takes care of storing the configuration of the models and handles methods for loading,
+    downloading and saving models as well as a few methods common to all models to:
+
+        - resize the input embeddings,
+        - prune heads in the self-attention heads.
+
+    Class attributes (overridden by derived classes):
+
+        - **config_class** ([`PretrainedConfig`]) -- A subclass of [`PretrainedConfig`] to use as configuration class
+          for this model architecture.
+        - **load_tf_weights** (`Callable`) -- A python *method* for loading a TensorFlow checkpoint in a PyTorch model,
+          taking as arguments:
+
+            - **model** ([`PreTrainedModel`]) -- An instance of the model on which to load the TensorFlow checkpoint.
+            - **config** ([`PreTrainedConfig`]) -- An instance of the configuration associated to the model.
+            - **path** (`str`) -- A path to the TensorFlow checkpoint.
+
+        - **base_model_prefix** (`str`) -- A string indicating the attribute associated to the base model in derived
+          classes of the same architecture adding modules on top of the base model.
+        - **is_parallelizable** (`bool`) -- A flag indicating whether this model supports model parallelization.
+        - **main_input_name** (`str`) -- The name of the principal input to the model (often `input_ids` for NLP
+          models, `pixel_values` for vision models and `input_values` for speech models).
+    """
+
+    config_class = MambaConfig  # TODO: Build on top of MambaConfig?
+    # base_model_prefix = "backbone"
+    base_model_prefix = "mamba"
+    main_input_name = "input_ids"
+    model_tags = None
+
+    _auto_class = None
+    _no_split_modules = ["MambaBlock"]
+    _skip_keys_device_placement = None
+    _keep_in_fp32_modules = None
+
+    # a list of `re` patterns of `state_dict` keys that should be removed from the list of missing
+    # keys we find (keys inside the model but not in the checkpoint) and avoid unnecessary warnings.
+    _keys_to_ignore_on_load_missing = None
+    # a list of `re` patterns of `state_dict` keys that should be removed from the list of
+    # unexpected keys we find (keys inside the checkpoint but not the model) and avoid unnecessary
+    # warnings.
+    _keys_to_ignore_on_load_unexpected = None
+    # a list of `state_dict` keys to ignore when saving the model (useful for keys that aren't
+    # trained, but which are either deterministic or tied variables)
+    _keys_to_ignore_on_save = None
+    # a list of `state_dict` keys that are potentially tied to another key in the state_dict.
+    _tied_weights_keys = None
+
+    is_parallelizable = False
+    supports_gradient_checkpointing = True
+
+    # Flash Attention 2 support
+    _supports_flash_attn_2 = False
+
+    # SDPA support
+    _supports_sdpa = False
+
+    # Has support for a `Cache` instance as `past_key_values`
+    _supports_cache_class = False
+
+    def __init__(self, *inputs, **kwargs):
+        super().__init__(*inputs, **kwargs)
+
+    # https://github.com/state-spaces/mamba/blob/009bec5ee37f586844a3fc89c040a9c1a9d8badf/mamba_ssm/models/mixer_seq_simple.py#L54
+    def _init_weights(
         self,
-        config,
-        initializer_cfg=None,
-        device=None,
-        dtype=None,
-        **kwargs,
+        module,
+        initializer_range=0.02,  # Now only used for embedding layer.
+        rescale_prenorm_residual=True,
+        n_residuals_per_layer=1,  # Change to 2 if we have MLP
     ):
+        if isinstance(module, nn.Linear):
+            if module.bias is not None:
+                if not getattr(module.bias, "_no_reinit", False):
+                    nn.init.zeros_(module.bias)
+
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, std=initializer_range)
+
+        if rescale_prenorm_residual:
+            # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
+            #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
+            #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
+            #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/
+            #
+            # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
+            for name, p in module.named_parameters():
+                if name in ["out_proj.weight", "fc2.weight"]:
+                    # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
+                    # Following Pytorch init, except scale by 1/sqrt(2 * n_layer)
+                    # We need to reinit p since this code could be called multiple times
+                    # Having just p *= scale would repeatedly scale it down
+                    nn.init.kaiming_uniform_(p, a=math.sqrt(5))
+                    with torch.no_grad():
+                        p /= math.sqrt(n_residuals_per_layer * self.config.n_layer)
+
+    # def _set_gradient_checkpointing(self, module, value=False):
+    #     if isinstance(module, GPT2Model):
+    #         module.gradient_checkpointing = value
+
+
+class MambaModel(MambaPretrainedModel):
+    def __init__(
+        self, config: MambaConfig = MambaConfig(), **kwargs
+    ) -> None:
+        """Full Mamba model.
+        Mamba model decoder consisting of *config.n_layer* layers. Each layer is a [`MambaBlock`]
+        Args:
+            config: MambaConfig
+        """
         super().__init__(
             config,
             **kwargs,
         )
 
-        self.model = MambaLMHeadModel(
-            config,
-            initializer_cfg=initializer_cfg,
-            device=device,
-            dtype=dtype,
+        # self.embedding = nn.Embedding(
+        #     num_embeddings=config.vocab_size,
+        #     embedding_dim=config.d_model,
+        # )
+
+
+        self.embedding = nn.Embedding(
+            num_embeddings=config.vocab_size,
+            embedding_dim=config.d_model,
         )
 
-    def forward(
-            self,
-            input_ids,
-            position_ids=None,
-            inference_params=None,
-            num_last_tokens=0,
-            **kwargs,
-        ):
-        return self.model.forward(
-            input_ids,
-            position_ids,
-            inference_params,
-            num_last_tokens
+        self.layers = nn.ModuleList(
+            [ResidualBlock(config) for _ in range(self.config.n_layer)]
         )
+        # self.layers = nn.ModuleList([MambaBlock(config) for _ in range(config.n_layer)])
+        # # self.norm_f = RMSNorm(d_model=embedding_dim)
+        self.norm_f = RMSNorm(config.d_model)
+
+        # self.gradient_checkpointing = False
+        # # self.post_init()
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    # def _init_weights(self, module):
+    #     std = 0.02
+
+    #     if isinstance(module, (nn.Linear, nn.Conv1d)):
+    #         module.weight.data.normal_(mean=0.0, std=std)
+
+    #         if module.bias is not None:
+    #             module.bias.data.zero_()
+
+    #     elif isinstance(module, nn.Embedding):
+    #         module.weight.data.normal_(mean=0.0, std=std)
+
+    #         if module.padding_idx is not None:
+    #             module.weight.data[module.padding_idx].zero_()
+
+# Inspired by:
+# - https://huggingface.co/Q-bert/Mamba-130M/blob/f0d00db98acaa62b1ee4304cd11643e69aa62a71/modeling_mamba.py#L198
+# - https://github.com/state-spaces/mamba/blob/009bec5ee37f586844a3fc89c040a9c1a9d8badf/mamba_ssm/models/mixer_seq_simple.py#L86
+# class MambaModel(MambaPretrainedModel):
+    # def __init__(
+    #     self,
+    #     config: MambaConfig = MambaConfig(),
+    #     **kwargs,
+    # ) -> None:
+    #     super().__init__(
+    #         config,
+    #         **kwargs,
+    #     )
+
+    #     self.embedding = nn.Embedding(
+    #         num_embeddings=config.vocab_size,
+    #         embedding_dim=config.d_model,
+    #     )
+
+    #     # # self.layers = nn.ModuleList(
+    #     # #     [ResidualBlock(args=model_args) for _ in range(model_args.n_layer)]
+    #     # # )
+    #     self.layers = nn.ModuleList([MambaBlock(config) for _ in range(config.n_layer)])
+    #     # # self.norm_f = RMSNorm(d_model=embedding_dim)
+    #     self.norm_f = RMSNorm(config.d_model)
+
+    #     # self.gradient_checkpointing = False
+    #     # # self.post_init()
+
+    # def get_input_embeddings(self):
+    #     return self.embed_out
 
-class MambaModelForCausalLM(MambaModel, GenerationMixin):
-    def generate(
+    # def set_input_embeddings(self, value):
+    #     self.embed_out = value
+
+    # def forward(
+    #     self,
+    #     input_ids: torch.LongTensor = None,
+    #     output_hidden_states=False,
+    #     return_dict: Optional[bool] = None,
+    #     **kwargs,
+    #     # ) -> BaseModelOutput:
+    # ) -> Union[Tuple, BaseModelOutputWithPast]:
+    #     batch_size = input_ids.shape[0]
+    #     hidden_size = self.config.hidden_size
+    #     hidden_states: Tuple[Tensor[(batch_size, sequence_length, hidden_size)]] = ()
+    #     sequence_length = input_ids.shape[1]
+    #     output_hidden_states = output_hidden_states or self.config.output_hidden_states
+
+    #     last_hidden_state = self.embed_out(input_ids)
+    #     assert last_hidden_state.shape == (
+    #         batch_size,
+    #         sequence_length,
+    #         hidden_size,
+    #     ), f"{last_hidden_state.shape} != {(batch_size, sequence_length, hidden_size)}"
+    #     hidden_states += (last_hidden_state,)
+
+    #     for layer in self.layers:
+    #         last_hidden_state = layer(last_hidden_state)
+    #         assert last_hidden_state.shape == (
+    #             batch_size,
+    #             sequence_length,
+    #             hidden_size,
+    #         ), f"{last_hidden_state.shape} != {(batch_size, sequence_length, hidden_size)}"
+    #         hidden_states += (last_hidden_state,)
+
+    #     last_hidden_state = self.norm_f(last_hidden_state)
+    #     assert last_hidden_state.shape == (
+    #         batch_size,
+    #         sequence_length,
+    #         hidden_size,
+    #     ), f"{last_hidden_state.shape} != {(batch_size, sequence_length, hidden_size)}"
+    #     hidden_states += (last_hidden_state,)
+
+    #     assert (
+    #         len(hidden_states) == self.config.n_layer + 2
+    #     ), f"{len(hidden_states)} != {self.config.n_layer + 2}"
+
+    #     # return BaseModelOutput(
+    #     return BaseModelOutputWithPast(
+    #         hidden_states=hidden_states if output_hidden_states else None,
+    #         last_hidden_state=last_hidden_state,
+    #     )
+
+
+# Influences:
+# - https://huggingface.co/Q-bert/Mamba-130M/blob/f0d00db98acaa62b1ee4304cd11643e69aa62a71/modeling_mamba.py#L238
+# - https://github.com/state-spaces/mamba/blob/009bec5ee37f586844a3fc89c040a9c1a9d8badf/mamba_ssm/models/mixer_seq_simple.py#L176
+# class MambaModelForCausalLM(MambaModel, GenerationMixin):
+# class MambaModelForCausalLM(PreTrainedModel, GenerationMixin):
+# class MambaLMHeadModel(MambaPretrainedModel, GenerationMixin):
+class MambaLMHeadModel(MambaPretrainedModel):
+    # _tied_weights_keys = ["lm_head.weight",
+
+    def __init__(
         self,
-        input_ids,
-        max_length: int = 2048,
-        top_k: int = 1,
-        top_p: float = 0.0,
-        temperature: float = 1.0,
-        return_dict_in_generate: bool = False,
-        output_scores: bool = False,
-        repetition_penalty: float = 1.0,
-        eos_token_id: Optional[int] = None,
-        teacher_outputs: Optional[torch.Tensor] = None,
-        vocab_size: Optional[int] = None,
-        cg: bool = False,
-        enable_timing: bool = False,
-        streamer: Optional[TextStreamer] = None,
+        config: MambaConfig = MambaConfig(),
         **kwargs,
-    ):
-        return self.model.generate(
-            input_ids=input_ids,
-            max_length=max_length,
-            top_k=top_k,
-            top_p=top_p,
-            temperature=temperature,
-            return_dict_in_generate=return_dict_in_generate,
-            output_scores=output_scores,
-            repetition_penalty=repetition_penalty,
-            eos_token_id=eos_token_id,
-            teacher_outputs=teacher_outputs,
-            vocab_size=vocab_size,
-            cg=cg,
-            enable_timing=enable_timing,
-            streamer=streamer,
+    ) -> None:
+        super().__init__(
+            config,
+            **kwargs,
+        )
+
+        self.backbone = MambaModel(
+            config=self.config,
+        )
+
+        self.lm_head = nn.Linear(
+            in_features=self.config.hidden_size,
+            out_features=self.config.vocab_size,
+            bias=False,
         )
+
+#         # self.head.weight = self.backbone.embedding.weight # TODO: there's some logic in GenerationMix that does this
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+#     # def forward(
+#     #     self, input_ids, output_hidden_states=False, **kwargs
+#     # ) -> CausalLMOutput:
+#     #     batch_size = input_ids.shape[0]
+#     #     sequence_length = input_ids.shape[1]
+#     #     vocab_size = self.config.vocab_size
+#     #     output_hidden_states = output_hidden_states or self.config.output_hidden_states
+
+#     #     outputs = self.backbone(
+#     #         input_ids=input_ids,
+#     #         output_hidden_states=output_hidden_states,
+#     #     )
+
+#     #     last_hidden_state = outputs.last_hidden_state
+
+#     #     logits: torch.FloatTensor[batch_size, sequence_length, vocab_size] = (
+#     #         self.lm_head(
+#     #             last_hidden_state,
+#     #         )
+#     #     )
+
+#     #     return CausalLMOutput(
+#     #         hidden_states=outputs.hidden_states if output_hidden_states else None,
+#     #         logits=logits,
+#     #     )
+
+#     # def prepare_inputs_for_generation(
+#     #     self, input_ids, attention_mask=None, **model_kwargs
+#     # ):
+#     #     return {
+#     #         "input_ids": input_ids,
+#     #     }
+
+
+# class MultimodalMambaModelForCausalLMWithValueHead(PreTrainedModelWrapper):
+#     lm_head_namings: Tuple[str, str] = ("lm_head", "embed_out")
+#     transformers_parent_class: transformers.PreTrainedModel = transformers.AutoModelForCausalLM
+
+#     # def __init__(
+#     #     self,
+#     #     config: MultimodalMambaConfig = MultimodalMambaConfig(),
+#     #     **kwargs,
+#     # ) -> None:
+#     #     super().__init__(
+#     #         config,
+#     #         **kwargs,
+#     #     )
+
+#     #     self.model = MultimodalMambaModelForCausalLM(
+#     #         config=config,
+#     #     )
+
+#     #     self.value_head = nn.Linear(
+#     #         in_features=config.embedding_dim,
+#     #         out_features=1,
+#     #         bias=False,
+#     #     )
+
+#     # def forward(
+#     #     self, input_ids, output_hidden_states=False, **kwargs
+#     # ) -> CausalLMOutput:
+#     #     outputs = self.model(
+#     #         input_ids=input_ids,
+#     #         output_hidden_states=output_hidden_states,
+#     #     )
+
+#     #     last_hidden_state = outputs.last_hidden_state
+
+#     #     value: torch.FloatTensor[batch_size, sequence_length, 1] = self.value_head(
+#     #         last_hidden_state,
+#     #     )
+
+#     #     return CausalLMOutput(
+#     #         hidden_states=outputs.hidden_states if output_hidden_states else None,
+#     #         logits=outputs.logits,
+#     #         value=value,
+#     #     )