models/dimamba.py

import math
from functools import partial
from typing import Optional, Tuple, Union

import huggingface_hub
import numpy as np
import omegaconf
import torch
import torch.nn as nn
import torch.nn.functional as F
from causal_conv1d import (
    causal_conv1d_fn,
    causal_conv1d_update,
)
from einops import rearrange, repeat
from mamba_ssm.ops.selective_scan_interface import (
    mamba_inner_fn,
    selective_scan_fn,
)
from torch import Tensor
from transformers import PretrainedConfig, PreTrainedModel
from transformers.modeling_outputs import (
    BaseModelOutputWithNoAttention,
    MaskedLMOutput,
)

try:
    from mamba_ssm.ops.triton.layernorm import (
        RMSNorm,
        layer_norm_fn,
        rms_norm_fn,
    )
except ImportError:
    RMSNorm, layer_norm_fn, rms_norm_fn = None, None, None
from mamba_ssm.ops.triton.selective_state_update import (
    selective_state_update,
)

from models.dit import (
    TimestepEmbedder,
    bias_dropout_add_scale_fused_inference,
    bias_dropout_add_scale_fused_train,
    modulate_fused,
)

# sys.path.append('mamba_wrappers/mamba2')
# from .mamba2.src.modules.ssd import SSD as Mamba


class Mamba(nn.Module):
    def __init__(
        self,
        d_model,
        d_state=16,
        d_conv=4,
        expand=2,
        dt_rank='auto',
        dt_min=0.001,
        dt_max=0.1,
        dt_init='random',
        dt_scale=1.0,
        dt_init_floor=1e-4,
        conv_bias=True,
        bias=False,
        use_fast_path=True,  # Fused kernel options
        layer_idx=None,
        device=None,
        dtype=None,
    ):
        factory_kwargs = {'device': device, 'dtype': dtype}
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.d_conv = d_conv
        self.expand = expand
        self.d_inner = int(self.expand * self.d_model)
        self.dt_rank = math.ceil(self.d_model / 16) if dt_rank == 'auto' else dt_rank
        self.use_fast_path = use_fast_path
        self.layer_idx = layer_idx

        self.in_proj = nn.Linear(
            self.d_model, self.d_inner * 2, bias=bias, **factory_kwargs
        )

        self.conv1d = nn.Conv1d(
            in_channels=self.d_inner,
            out_channels=self.d_inner,
            bias=conv_bias,
            kernel_size=d_conv,
            groups=self.d_inner,
            padding=d_conv - 1,
            **factory_kwargs,
        )

        self.activation = 'silu'
        self.act = nn.SiLU()

        self.x_proj = nn.Linear(
            self.d_inner, self.dt_rank + self.d_state * 2, bias=False, **factory_kwargs
        )
        self.dt_proj = nn.Linear(
            self.dt_rank, self.d_inner, bias=True, **factory_kwargs
        )

        # Initialize special dt projection to preserve variance at initialization
        dt_init_std = self.dt_rank**-0.5 * dt_scale
        if dt_init == 'constant':
            nn.init.constant_(self.dt_proj.weight, dt_init_std)
        elif dt_init == 'random':
            nn.init.uniform_(self.dt_proj.weight, -dt_init_std, dt_init_std)
        else:
            raise NotImplementedError

        # Initialize dt bias so that F.softplus(dt_bias) is between dt_min and dt_max
        dt = torch.exp(
            torch.rand(self.d_inner, **factory_kwargs)
            * (math.log(dt_max) - math.log(dt_min))
            + math.log(dt_min)
        ).clamp(min=dt_init_floor)
        # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
        inv_dt = dt + torch.log(-torch.expm1(-dt))
        with torch.no_grad():
            self.dt_proj.bias.copy_(inv_dt)
        # Our initialization would set all Linear.bias to zero, need to mark this one as _no_reinit
        self.dt_proj.bias._no_reinit = True

        # S4D real initialization
        A = repeat(
            torch.arange(1, self.d_state + 1, dtype=torch.float32, device=device),
            'n -> d n',
            d=self.d_inner,
        ).contiguous()
        A_log = torch.log(A)  # Keep A_log in fp32
        self.A_log = nn.Parameter(A_log)
        self.A_log._no_weight_decay = True

        # D 'skip' parameter
        self.D = nn.Parameter(torch.ones(self.d_inner, device=device))  # Keep in fp32
        self.D._no_weight_decay = True

        self.out_proj = nn.Linear(
            self.d_inner, self.d_model, bias=bias, **factory_kwargs
        )

    def forward(self, hidden_states, inference_params=None):
        """
        hidden_states: (B, L, D)
        Returns: same shape as hidden_states
        """
        batch, seqlen, dim = hidden_states.shape

        conv_state, ssm_state = None, None
        if inference_params is not None:
            conv_state, ssm_state = self._get_states_from_cache(inference_params, batch)
            if inference_params.seqlen_offset > 0:
                # The states are updated inplace
                out, _, _ = self.step(hidden_states, conv_state, ssm_state)
                return out

        # We do matmul and transpose BLH -> HBL at the same time
        xz = rearrange(
            self.in_proj.weight @ rearrange(hidden_states, 'b l d -> d (b l)'),
            'd (b l) -> b d l',
            l=seqlen,
        )
        if self.in_proj.bias is not None:
            xz = xz + rearrange(self.in_proj.bias.to(dtype=xz.dtype), 'd -> d 1')

        A = -torch.exp(self.A_log.float())  # (d_inner, d_state)
        # In the backward pass we write dx and dz next to each other to avoid torch.cat

        if (
            self.use_fast_path
            and causal_conv1d_fn is not None
            and inference_params is None
        ):  # Doesn't support outputting the states
            out = mamba_inner_fn(
                xz,
                self.conv1d.weight,
                self.conv1d.bias,
                self.x_proj.weight,
                self.dt_proj.weight,
                self.out_proj.weight,
                self.out_proj.bias,
                A,
                None,  # input-dependent B
                None,  # input-dependent C
                self.D.float(),
                delta_bias=self.dt_proj.bias.float(),
                delta_softplus=True,
            )

        else:
            x, z = xz.chunk(2, dim=1)
            # Compute short convolution
            if conv_state is not None:
                # If we just take x[:, :, -self.d_conv :], it will error if seqlen < self.d_conv
                # Instead F.pad will pad with zeros if seqlen < self.d_conv, and truncate otherwise.
                conv_state.copy_(
                    F.pad(x, (self.d_conv - x.shape[-1], 0))
                )  # Update state (B D W)
            if causal_conv1d_fn is None:
                x = self.act(self.conv1d(x)[..., :seqlen])
            else:
                assert self.activation in ['silu', 'swish']
                x = causal_conv1d_fn(
                    x=x,
                    weight=rearrange(self.conv1d.weight, 'd 1 w -> d w'),
                    bias=self.conv1d.bias,
                    activation=self.activation,
                )

            # We're careful here about the layout, to avoid extra transposes.
            # We want dt to have d as the slowest moving dimension
            # and L as the fastest moving dimension, since those are what the ssm_scan kernel expects.
            x_dbl = self.x_proj(rearrange(x, 'b d l -> (b l) d'))  # (bl d)
            dt, B, C = torch.split(
                x_dbl, [self.dt_rank, self.d_state, self.d_state], dim=-1
            )
            dt = self.dt_proj.weight @ dt.t()
            dt = rearrange(dt, 'd (b l) -> b d l', l=seqlen)
            B = rearrange(B, '(b l) dstate -> b dstate l', l=seqlen).contiguous()
            C = rearrange(C, '(b l) dstate -> b dstate l', l=seqlen).contiguous()

            assert self.activation in ['silu', 'swish']

            y = selective_scan_fn(
                x,
                dt,
                A,
                B,
                C,
                self.D.float(),
                z=z,
                delta_bias=self.dt_proj.bias.float(),
                delta_softplus=True,
                return_last_state=ssm_state is not None,
            )

            if ssm_state is not None:
                y, last_state = y
                ssm_state.copy_(last_state)
            y = rearrange(y, 'b d l -> b l d')

            out = self.out_proj(y)

        return out

    def step(self, hidden_states, conv_state, ssm_state):
        dtype = hidden_states.dtype
        assert (
            hidden_states.shape[1] == 1
        ), 'Only support decoding with 1 token at a time for now'
        xz = self.in_proj(hidden_states.squeeze(1))  # (B 2D)
        x, z = xz.chunk(2, dim=-1)  # (B D)

        # Conv step
        if causal_conv1d_update is None:
            conv_state.copy_(
                torch.roll(conv_state, shifts=-1, dims=-1)
            )  # Update state (B D W)
            conv_state[:, :, -1] = x
            x = torch.sum(
                conv_state * rearrange(self.conv1d.weight, 'd 1 w -> d w'), dim=-1
            )  # (B D)
            if self.conv1d.bias is not None:
                x = x + self.conv1d.bias
            x = self.act(x).to(dtype=dtype)
        else:
            x = causal_conv1d_update(
                x,
                conv_state,
                rearrange(self.conv1d.weight, 'd 1 w -> d w'),
                self.conv1d.bias,
                self.activation,
            )

        x_db = self.x_proj(x)  # (B dt_rank+2*d_state)
        dt, B, C = torch.split(x_db, [self.dt_rank, self.d_state, self.d_state], dim=-1)
        # Don't add dt_bias here
        dt = F.linear(dt, self.dt_proj.weight)  # (B d_inner)
        A = -torch.exp(self.A_log.float())  # (d_inner, d_state)

        # SSM step
        if selective_state_update is None:
            # Discretize A and B
            dt = F.softplus(dt + self.dt_proj.bias.to(dtype=dt.dtype))
            dA = torch.exp(torch.einsum('bd,dn->bdn', dt, A))
            dB = torch.einsum('bd,bn->bdn', dt, B)
            ssm_state.copy_(ssm_state * dA + rearrange(x, 'b d -> b d 1') * dB)
            y = torch.einsum('bdn,bn->bd', ssm_state.to(dtype), C)
            y = y + self.D.to(dtype) * x
            y = y * self.act(z)  # (B D)
        else:
            y = selective_state_update(
                ssm_state,
                x,
                dt,
                A,
                B,
                C,
                self.D,
                z=z,
                dt_bias=self.dt_proj.bias,
                dt_softplus=True,
            )

        out = self.out_proj(y)
        return out.unsqueeze(1), conv_state, ssm_state

    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
        device = self.out_proj.weight.device
        conv_dtype = self.conv1d.weight.dtype if dtype is None else dtype
        conv_state = torch.zeros(
            batch_size,
            self.d_model * self.expand,
            self.d_conv,
            device=device,
            dtype=conv_dtype,
        )
        ssm_dtype = self.dt_proj.weight.dtype if dtype is None else dtype
        # ssm_dtype = torch.float32
        ssm_state = torch.zeros(
            batch_size,
            self.d_model * self.expand,
            self.d_state,
            device=device,
            dtype=ssm_dtype,
        )
        return conv_state, ssm_state

    def _get_states_from_cache(
        self, inference_params, batch_size, initialize_states=False
    ):
        assert self.layer_idx is not None
        if self.layer_idx not in inference_params.key_value_memory_dict:
            batch_shape = (batch_size,)
            conv_state = torch.zeros(
                batch_size,
                self.d_model * self.expand,
                self.d_conv,
                device=self.conv1d.weight.device,
                dtype=self.conv1d.weight.dtype,
            )
            ssm_state = torch.zeros(
                batch_size,
                self.d_model * self.expand,
                self.d_state,
                device=self.dt_proj.weight.device,
                dtype=self.dt_proj.weight.dtype,
                # dtype=torch.float32,
            )
            inference_params.key_value_memory_dict[self.layer_idx] = (
                conv_state,
                ssm_state,
            )
        else:
            conv_state, ssm_state = inference_params.key_value_memory_dict[
                self.layer_idx
            ]
            # TODO: What if batch size changes between generation, and we reuse the same states?
            if initialize_states:
                conv_state.zero_()
                ssm_state.zero_()
        return conv_state, ssm_state


class Block(nn.Module):
    def __init__(
      self,
      dim,
      mixer_cls,
      norm_cls=nn.LayerNorm,
      fused_add_norm=False,
      residual_in_fp32=False,
      modulate=False,
      t_dim=0,
    ):
      """
      Simple block wrapping a mixer class with LayerNorm/RMSNorm and residual connection'

      This Block has a slightly different structure compared to a regular
      prenorm Transformer block.
      The standard block is: LN -> MHA/MLP -> Add.
      [Ref: https://arxiv.org/abs/2002.04745]
      Here we have: Add -> LN -> Mixer, returning both
      the hidden_states (output of the mixer) and the residual.
      This is purely for performance reasons, as we can fuse add and LayerNorm.
      The residual needs to be provided (except for the very first block).
      """
      super().__init__()
      self.residual_in_fp32 = residual_in_fp32
      self.fused_add_norm = fused_add_norm
      self.mixer = mixer_cls(dim)
      self.norm = norm_cls(dim)

      if self.fused_add_norm:
        assert RMSNorm is not None, 'RMSNorm import fails'
        assert isinstance(
            self.norm, (nn.LayerNorm, RMSNorm)
        ), 'Only LayerNorm and RMSNorm are supported for fused_add_norm'

      self.dropout = 0.1

      self.modulate = modulate
      self.t_dim = t_dim
      if modulate:
        self.adaLN_modulation = nn.Linear(t_dim,
                                          3 * dim,
                                          bias=True)
        self.adaLN_modulation.weight.data.zero_()
        self.adaLN_modulation.bias.data.zero_()

    def _get_bias_dropout_scale(self):
      return (
        bias_dropout_add_scale_fused_train
        if self.training
        else bias_dropout_add_scale_fused_inference
      )

    def forward(
      self,
      hidden_states: Tensor,
      residual: Optional[Tensor] = None,
      inference_params=None,
      time_embeds=None,
    ):
        r"""Pass the input through the encoder layer.

        Args:
            hidden_states: the sequence to the encoder layer (required).
            residual: hidden_states = Mixer(LN(residual))
        """
        if not self.fused_add_norm:
          residual = (
            (hidden_states + residual)
            if residual is not None
            else hidden_states
          )

          hidden_states = self.norm(
            residual.to(dtype=self.norm.weight.dtype))
          if self.residual_in_fp32:
            residual = residual.to(torch.float32)
        else:
          fused_add_norm_fn = (
            rms_norm_fn
            if isinstance(self.norm, RMSNorm)
            else layer_norm_fn
          )

          hidden_states, residual = fused_add_norm_fn(
            hidden_states,
            self.norm.weight,
            self.norm.bias,
            residual=residual,
            prenorm=True,
            residual_in_fp32=self.residual_in_fp32,
            eps=self.norm.eps)

        if self.modulate and time_embeds is not None:
          (shift_msa,
           scale_msa,
           gate_msa) = self.adaLN_modulation(
              time_embeds)[:, None].chunk(3, dim=-1)
          hidden_states = modulate_fused(hidden_states,
                                         shift_msa,
                                         scale_msa)

        mixer_out = self.mixer(hidden_states, inference_params=inference_params)

        hidden_states = mixer_out
        if self.modulate and time_embeds is not None:
          bias_dropout_scale_fn = self._get_bias_dropout_scale()
          hidden_states = bias_dropout_scale_fn(
            hidden_states,
            None,
            gate_msa,
            residual,
            self.dropout)

        return hidden_states, residual

    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
      return self.mixer.allocate_inference_cache(
        batch_size, max_seqlen, dtype=dtype, **kwargs)

class BiMambaConfig(PretrainedConfig):
    """Config that extends the original MambaConfig with params relevant to bi-directionality."""

    model_type = 'bimamba'

    def __init__(
        self,
        # From original MambaConfig
        d_model: int = 2560,
        n_layer: int = 64,
        vocab_size: int = 50277,
        ssm_cfg: Optional[dict] = None,
        rms_norm: bool = True,
        residual_in_fp32: bool = True,
        fused_add_norm: bool = True,
        pad_vocab_size_multiple: int = 8,
        tie_word_embeddings: bool = True,
        # Not in original MambaConfig, but default arg in create_block in mamba_ssm repo; used in layer norm
        norm_epsilon: float = 1e-5,
        # Used in init_weights
        initializer_cfg: Optional[dict] = None,
        # Caduceus-specific params
        bidirectional: bool = True,
        bidirectional_strategy: Union[str, None] = 'add',
        bidirectional_weight_tie: bool = True,
        temb_strategy: Union[str, None] = None,
        d_temb: int = 0,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.d_model = d_model
        self.n_layer = n_layer
        self.vocab_size = vocab_size
        self.ssm_cfg = ssm_cfg
        self.rms_norm = rms_norm
        self.residual_in_fp32 = residual_in_fp32
        self.fused_add_norm = fused_add_norm
        self.pad_vocab_size_multiple = pad_vocab_size_multiple
        self.tie_word_embeddings = tie_word_embeddings
        self.norm_epsilon = norm_epsilon
        self.initializer_cfg = initializer_cfg
        self.bidirectional = bidirectional
        self.bidirectional_strategy = bidirectional_strategy
        self.bidirectional_weight_tie = bidirectional_weight_tie

        self.temb_strategy = temb_strategy
        self.d_temb = d_temb


def create_block(
    d_model,
    ssm_cfg=None,
    norm_epsilon=1e-5,
    rms_norm=False,
    residual_in_fp32=False,
    fused_add_norm=False,
    layer_idx=None,
    bidirectional=True,
    bidirectional_strategy='add',
    bidirectional_weight_tie=True,
    device=None,
    dtype=None,
    modulate=False,
    d_temb=0,
):
    """Create BiMamba block.

    Adapted from: https://github.com/state-spaces/mamba/blob/main/mamba_ssm/models/mixer_seq_simple.py
    """
    if ssm_cfg is None:
        ssm_cfg = {}
    factory_kwargs = {'device': device, 'dtype': dtype}
    bidirectional_kwargs = {
        'bidirectional': bidirectional,
        'bidirectional_strategy': bidirectional_strategy,
        'bidirectional_weight_tie': bidirectional_weight_tie,
    }
    mixer_cls = partial(
        BiMambaWrapper,
        layer_idx=layer_idx,
        **ssm_cfg,
        **bidirectional_kwargs,
        **factory_kwargs,
    )
    norm_cls = partial(
        nn.LayerNorm if not rms_norm else RMSNorm, eps=norm_epsilon, **factory_kwargs
    )
    block_cls = Block
    block = block_cls(
        d_model,
        mixer_cls,
        norm_cls=norm_cls,
        fused_add_norm=fused_add_norm,
        residual_in_fp32=residual_in_fp32,
        t_dim=d_temb,
        modulate=modulate,
    )
    block.layer_idx = layer_idx

    return block


class BiMambaWrapper(nn.Module):
    """Thin wrapper around Mamba to support bi-directionality."""

    def __init__(
        self,
        d_model: int,
        bidirectional: bool = True,
        bidirectional_strategy: Optional[str] = 'add',
        bidirectional_weight_tie: bool = True,
        **mamba_kwargs,
    ):
        super().__init__()
        if bidirectional and bidirectional_strategy is None:
            bidirectional_strategy = 'add'  # Default strategy: `add`
        if bidirectional and bidirectional_strategy not in ['add', 'ew_multiply']:
            raise NotImplementedError(
                f'`{bidirectional_strategy}` strategy for bi-directionality is not implemented!'
            )
        self.bidirectional = bidirectional
        self.bidirectional_strategy = bidirectional_strategy

        self.mamba_fwd = Mamba(d_model=d_model, **mamba_kwargs)

        self.mamba_rev = None
        if bidirectional:
            self.mamba_rev = Mamba(d_model=d_model, **mamba_kwargs)
            if (
                bidirectional_weight_tie
            ):  # Tie in and out projections (where most of param count lies)
                self.mamba_rev.in_proj.weight = self.mamba_fwd.in_proj.weight
                self.mamba_rev.in_proj.bias = self.mamba_fwd.in_proj.bias
                self.mamba_rev.out_proj.weight = self.mamba_fwd.out_proj.weight
                self.mamba_rev.out_proj.bias = self.mamba_fwd.out_proj.bias
        else:
            self.mamba_rev = None

    def forward(self, hidden_states, inference_params=None):
        """Bidirectional-enabled forward pass

        hidden_states: (B, L, D)
        Returns: same shape as hidden_states
        """

        out = self.mamba_fwd(
            hidden_states,
            inference_params=inference_params,
        )

        if self.bidirectional:

            hidden_states_flipped = torch.flip(hidden_states, dims=(1,))

            out_rev = self.mamba_rev(
                hidden_states_flipped,  # Flip along the sequence length dimension
                inference_params=inference_params,
            )

            out_rev_flipped = torch.flip(out_rev, dims=(1,))
            if self.bidirectional_strategy == 'add':
                out = (
                    out + out_rev_flipped
                )  # Flip back for combining with forward hidden states
            elif self.bidirectional_strategy == 'ew_multiply':
                out = out * out_rev_flipped
            else:
                raise NotImplementedError(
                    f'`{self.bidirectional_strategy}` for bi-directionality not implemented!'
                )

        return out


class BiMambaEmbeddings(nn.Module):
    def __init__(
        self,
        config: BiMambaConfig,
        input_dim=None,
        device=None,
        dtype=None,
    ):
        super().__init__()
        factory_kwargs = {'device': device, 'dtype': dtype}
        if input_dim is None:
            input_dim = config.vocab_size
        self.word_embeddings = nn.Embedding(
            input_dim, config.d_model, **factory_kwargs
        )

    def forward(self, input_ids):
        """
        input_ids: (batch, seqlen)
        """
        return self.word_embeddings(input_ids)


class BiMambaMixerModel(nn.Module):
    def __init__(
        self,
        config: BiMambaConfig,
        device=None,
        dtype=None,
    ) -> None:
        super().__init__()
        factory_kwargs = {'device': device, 'dtype': dtype}
        self.temb_strategy = config.temb_strategy
        self.config = config
        input_dim = config.vocab_size
        d_model = config.d_model
        if self.temb_strategy and self.temb_strategy == 'concat':
            input_dim += config.d_temb
            d_model += config.d_temb
        if self.temb_strategy is None:
            config.d_temb = 0

        self.fused_add_norm = config.fused_add_norm
        self.residual_in_fp32 = config.residual_in_fp32

        self.embeddings = BiMambaEmbeddings(
            config,input_dim=input_dim, **factory_kwargs)

        # Mamba changes the order of residual and layer norm:
        # Instead of LN -> Attn / MLP -> Add, we do:
        # Add -> LN -> Attn / MLP / Mixer, returning both the residual branch (output of Add) and
        # the main branch (output of MLP / Mixer). The model definition is unchanged.
        # This is for performance reason: we can fuse add + layer_norm.
        if config.fused_add_norm:
            if layer_norm_fn is None or rms_norm_fn is None:
                raise ImportError('Failed to import Triton LayerNorm / RMSNorm kernels')

        self.layers = nn.ModuleList(
            [
                create_block(
                    d_model,
                    ssm_cfg=config.ssm_cfg,
                    norm_epsilon=config.norm_epsilon,
                    rms_norm=config.rms_norm,
                    residual_in_fp32=config.residual_in_fp32,
                    fused_add_norm=config.fused_add_norm,
                    layer_idx=i,
                    bidirectional=config.bidirectional,
                    bidirectional_strategy=config.bidirectional_strategy,
                    bidirectional_weight_tie=config.bidirectional_weight_tie,
                    modulate=True if config.temb_strategy and 'adaln' in config.temb_strategy else False,
                    d_temb=config.d_temb,
                    **factory_kwargs,
                )
                for i in range(config.n_layer)
            ]
        )

        if self.temb_strategy and 'adaln' in self.temb_strategy:
            self.adaLN_modulation_final = nn.Linear(
                config.d_temb, 2 * d_model, bias=True
            )
            self.adaLN_modulation_final.weight.data.zero_()
            self.adaLN_modulation_final.bias.data.zero_()

        norm_f = (nn.LayerNorm if not config.rms_norm else RMSNorm)(
            d_model, eps=config.norm_epsilon, **factory_kwargs
        )
        self.norm_f = norm_f

    def pre_apply_temb(self, input_embeds, time_embeds):
        """Prepend/add time embeddings to input embeddings at the start of the forward pass.
        
        Args:
            input_embeds: Input embeddings. (batch, seqlen, d_model)
            time_embeds: Timestep embeddings. (batch, d_temb)
        Returns:
            if self.temb_strategy == 'concat':
                input_embeds: (batch, seqlen, d_model + d_temb)
            if self.temb_strategy == 'add':
                input_embeds: (batch, seqlen, d_model)
        """
        if self.temb_strategy == 'concat':
            input_embeds = torch.cat([time_embeds.unsqueeze(1).tile(
                1, input_embeds.shape[1], 1), input_embeds], axis=-1)
        elif self.temb_strategy == 'add':
            input_embeds += time_embeds.unsqueeze(1).tile(1, input_embeds.shape[1], 1)
        return input_embeds

    def forward(
        self,
        input_ids,
        inputs_embeds=None,
        output_hidden_states=False,
        time_embeds=None,
    ):
        """Mixer forward."""
        all_hidden_states = []
        if inputs_embeds is not None:
            hidden_states = inputs_embeds
        else:
            hidden_states = self.embeddings(input_ids)
        if (
            time_embeds is not None
            and self.temb_strategy in ['concat', 'add']
        ):
            hidden_states = self.pre_apply_temb(hidden_states, time_embeds)

        residual = None

        for ind, layer in enumerate(self.layers):
            if output_hidden_states:
                all_hidden_states.append(hidden_states)
            # TODO: Add support for gradient checkpointing
            layer_out = layer(
                hidden_states, residual, inference_params=None, time_embeds=time_embeds
            )

            hidden_states, residuals = layer_out

        if not self.fused_add_norm:
            if self.temb_strategy and 'adaln' in self.temb_strategy:
                raise NotImplementedError('adaln only implemented for 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:
            if time_embeds is not None and self.temb_strategy and 'adaln' in self.temb_strategy:
                shift, scale = self.adaLN_modulation_final(time_embeds)[:, None].chunk(
                    2, dim=2
                )

            fused_add_norm_fn = (
                rms_norm_fn if isinstance(self.norm_f, RMSNorm) else layer_norm_fn
            )

            # Set prenorm=False here since we don't need the residual
            hidden_states = fused_add_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,
            )
            if time_embeds is not None and self.temb_strategy and 'adaln' in self.temb_strategy:
                hidden_states = modulate_fused(hidden_states, shift, scale)

        if output_hidden_states:
            all_hidden_states.append(hidden_states)

        return hidden_states, all_hidden_states


def cross_entropy(logits, y, ignore_index=-100):
    """Cross entropy loss."""
    logits = logits.view(-1, logits.shape[-1])
    y = y.view(-1)
    return F.cross_entropy(logits, y, ignore_index=ignore_index)


def weighted_cross_entropy(logits, y, loss_weights, ignore_index=-100):
    """Weighted cross entropy loss (discounts certain tokens)."""
    logits = logits.view(-1, logits.shape[-1])
    y = y.view(-1)
    ce = F.cross_entropy(logits, y, ignore_index=ignore_index, reduction='none')
    loss_weights = loss_weights.view(-1)
    loss_weights[y == ignore_index] = 0.0
    # TODO: Follows GPN implementation, but should we remove weight normalization?
    return (ce * (loss_weights / loss_weights.sum())).sum()


class BiMambaPreTrainedModel(PreTrainedModel):
    """PreTrainedModel wrapper for BiMamba backbone."""

    config_class = BiMambaConfig
    base_model_prefix = 'bimamba'
    supports_gradient_checkpointing = False
    _no_split_modules = ['BiMambaWrapper']

    def _init_weights(
        self,
        module,
        initializer_range=0.02,  # Now only used for embedding layer.
        **kwargs,
    ):
        """Adapted from: https://github.com/state-spaces/mamba/blob/main/mamba_ssm/models/mixer_seq_simple.py"""

        n_layer = self.config.n_layer
        initialized_cfg = (
            self.config.initializer_cfg
            if self.config.initializer_cfg is not None
            else {}
        )
        rescale_prenorm_residual = initialized_cfg.get('rescale_prenorm_residual', True)
        initializer_range = initialized_cfg.get('initializer_range', initializer_range)
        n_residuals_per_layer = initialized_cfg.get('n_residuals_per_layer', 1)

        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 * n_layer)


class BiMamba(BiMambaPreTrainedModel):
    """BiMamba model that can be instantiated using HF patterns."""

    def __init__(self, config: BiMambaConfig, device=None, dtype=None, **kwargs):
        super().__init__(config)

        # Adjust vocab size if vocab padding is set.
        if config.vocab_size % config.pad_vocab_size_multiple != 0:
            config.vocab_size += config.pad_vocab_size_multiple - (
                config.vocab_size % config.pad_vocab_size_multiple
            )

        self.config = config
        factory_kwargs = {'device': device, 'dtype': dtype}
        self.backbone = BiMambaMixerModel(config, **factory_kwargs, **kwargs)

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        time_embeds: Optional[bool] = None,
    ) -> Union[torch.Tensor, Tuple, BaseModelOutputWithNoAttention]:
        """HF-compatible forward method."""
        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            else self.config.output_hidden_states
        )
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        backbone_out = self.backbone(
            input_ids,
            inputs_embeds=inputs_embeds,
            output_hidden_states=output_hidden_states,
            time_embeds=time_embeds,
        )

        hidden_states, all_hidden_states = backbone_out

        if return_dict:
            return BaseModelOutputWithNoAttention(
                last_hidden_state=hidden_states,
                hidden_states=all_hidden_states if output_hidden_states else None,
            )
        elif output_hidden_states:
            return hidden_states, all_hidden_states
        else:
            return hidden_states


class BiMambaForMaskedLM(BiMambaPreTrainedModel):
    """HF-compatible BiMamba model for masked language modeling."""

    def __init__(self, config: BiMambaConfig, device=None, dtype=None, **kwargs):
        super().__init__(config, **kwargs)
        factory_kwargs = {'device': device, 'dtype': dtype}
        self.bimamba = BiMamba(config, **factory_kwargs, **kwargs)
        self.config = config
        self.temb_strategy = config.temb_strategy
        lm_head_in_dim = config.d_model
        # LM head may only take in concatenated timestep embeddings
        # if its weights are not tied to the vocab embedding
        if (
            not config.tie_word_embeddings
            and config.temb_strategy == 'concat'
        ):
            lm_head_in_dim += config.d_temb
        self.lm_head = nn.Linear(
            lm_head_in_dim,
            self.config.vocab_size,  # Use BiMamba config as it might have been updated
            bias=False,
            **factory_kwargs,
        )
        # Initialize weights and apply final processing
        self.post_init()
        if self.config.tie_word_embeddings:
            self.tie_weights()

    def init_weights(self):
        """
        If needed prunes and maybe initializes weights. If using a custom `PreTrainedModel`, you need to implement any
        initialization logic in `_init_weights`.
        """

        # Initialize weights
        self.apply(self._initialize_weights)

        # Tie weights should be skipped when not initializing all weights
        # since from_pretrained(...) calls tie weights anyways

    def post_init(self):
        """
        A method executed at the end of each Transformer model initialization, to execute code that needs the model's
        modules properly initialized (such as weight initialization).
        """
        self.init_weights()
        self._backward_compatibility_gradient_checkpointing()

    def get_input_embeddings(self):
        return self.bimamba.backbone.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.bimamba.backbone.embeddings.word_embeddings = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        """Overrides output embeddings."""
        self.lm_head = new_embeddings

    def tie_weights(self):
        """Tie weights."""
        super().tie_weights()

    def get_decoder(self):
        """Get decoder (backbone) for the model."""
        return self.bimamba

    def set_decoder(self, decoder):
        """Set decoder (backbone) for the model."""
        self.bimamba = decoder

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        loss_weights: Optional[torch.FloatTensor] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        time_embeds: Optional[torch.FloatTensor] = None,
    ) -> Union[Tuple, MaskedLMOutput]:
        """HF-compatible forward method."""

        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            else self.config.output_hidden_states
        )
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
        outputs = self.bimamba(
            input_ids=input_ids,
            inputs_embeds=inputs_embeds,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            time_embeds=time_embeds,
        )
        hidden_states = outputs[0]
        if (
            self.config.tie_word_embeddings
            and time_embeds is not None
            and self.temb_strategy is not None
            and self.temb_strategy == 'concat'
        ):
            hidden_states = hidden_states[:, :, self.config.d_temb:]

        logits = self.lm_head(hidden_states)

        loss = None
        if labels is not None:
            if loss_weights is not None:
                loss = weighted_cross_entropy(
                    logits, labels, loss_weights, ignore_index=self.config.pad_token_id
                )
            else:
                loss = cross_entropy(
                    logits, labels, ignore_index=self.config.pad_token_id
                )

        if not return_dict:
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output

        return MaskedLMOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
        )

class DiMamba(nn.Module, huggingface_hub.PyTorchModelHubMixin):
  def __init__(self, config, vocab_size: int, pad_token_id: int):
    super().__init__()
    if type(config) == dict:
      config = omegaconf.OmegaConf.create(config)

    self.temb_strategy = config.model.temb_strategy

    if self.temb_strategy == 'add':
      self.sigma_map = TimestepEmbedder(config.model.hidden_size)
    elif self.temb_strategy != 'none':
      self.sigma_map = TimestepEmbedder(config.model.cond_dim)

    mamba_config = BiMambaConfig(
      d_model=config.model.hidden_size,
      n_layer=config.model.n_blocks,
      pad_token_id=pad_token_id,
      vocab_size=vocab_size,
      pad_vocab_size_multiple=1,
      tie_word_embeddings=config.model.tie_word_embeddings,
      temb_strategy=self.temb_strategy,
      d_temb=config.model.cond_dim,
      bidirectional=True)

    self.model = BiMambaForMaskedLM(config=mamba_config)

  def _get_bias_dropout_scale(self):
    if self.training:
      return bias_dropout_add_scale_fused_train
    else:
      return bias_dropout_add_scale_fused_inference

  def forward(self, indices, sigma):
    c = None
    if self.temb_strategy is not None:
      c = F.silu(self.sigma_map(sigma))

    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
      x = self.model(indices, time_embeds=c).logits

    return x