rene.py

from functools import partial

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
from flash_attn import flash_attn_with_kvcache
from mamba_ssm.models.mixer_seq_simple import _init_weights
from mamba_ssm.modules.mamba2 import Mamba2
from mamba_ssm.modules.mha import _update_kv_cache
from mamba_ssm.utils.generation import GenerationMixin as MambaGenerationMixin
from transformers.modeling_outputs import CausalLMOutput
from transformers.modeling_utils import PreTrainedModel

from .configuration_rene import ReneConfig


class ReneMLP(nn.Module):
 """One-hidden-layer network with GELU activation.

 Args:
 d_input: Block input dimension.
 d_output: Block output dimension.
 expand: Block expansion factor.
 bias: Use biases in linear layers.
 """

 def __init__(self, d_input, d_output=None, expand=3, bias=True, device=None, dtype=None):
 super().__init__()
 factory_kwargs = {"device": device, "dtype": dtype}
 self.d_input = d_input
 self.d_output = d_input if d_output is None else d_output
 self.d_inner = int(round(expand * d_input))
 self.in_proj = nn.Linear(self.d_input, self.d_inner, bias=bias, **factory_kwargs)
 self.activation = nn.GELU()
 self.out_proj = nn.Linear(self.d_inner, self.d_input, bias=bias, **factory_kwargs)

 def forward(self, x, inference_params=None):
 """Forward pass through the MLP module."""
 y = self.in_proj(x)
 y = self.activation(y)
 y = self.out_proj(y)
 return y

 def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
 """Allocate inference cache for ReneMLP. (There is nothing to cache for this module)."""
 return None


class ReneMHA(nn.Module):
 """Multi-head self-attention. Adapted from mamba_ssm MHA class."""

 def __init__(
 self,
 embed_dim,
 num_heads,
 num_heads_kv=None,
 head_dim=None, # If None, use embed_dim // num_heads
 qkv_proj_bias=True,
 out_proj_bias=True,
 softmax_scale=None,
 causal=True,
 sliding_window_length=None, # If None, infinite context
 layer_idx=None,
 device=None,
 dtype=None,
 ) -> None:
 """
 num_heads_kv: can be used to toggle MQA / GQA. If None, use num_heads.
 return_residual: whether to return the input x along with the output. This is for
 performance reason: for post-norm architecture, returning the input allows us
 to fuse the backward of nn.Linear with the residual connection.
 """
 super().__init__()
 factory_kwargs = {"device": device, "dtype": dtype}
 self.embed_dim = embed_dim
 self.layer_idx = layer_idx
 self.softmax_scale = softmax_scale
 self.causal = causal
 assert self.causal, "Rene does not yet support non-causal modeling"

 self.num_heads = num_heads
 self.num_heads_kv = num_heads_kv if num_heads_kv is not None else num_heads
 assert (
 self.num_heads % self.num_heads_kv == 0
 ), "num_heads must be divisible by num_heads_kv"
 if head_dim is None:
 assert self.embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
 self.head_dim = head_dim if head_dim is not None else self.embed_dim // num_heads
 qkv_dim = self.head_dim * (self.num_heads + 2 * self.num_heads_kv)
 out_dim = self.head_dim * self.num_heads

 self.sliding_window_length = sliding_window_length
 if self.sliding_window_length is None:
 self.window_size = (-1, -1)
 else:
 self.window_size = (self.sliding_window_length - 1, 0) # for flash_attn

 self.in_proj = nn.Linear(embed_dim, qkv_dim, bias=qkv_proj_bias, **factory_kwargs)
 self.out_proj = nn.Linear(out_dim, embed_dim, bias=out_proj_bias, **factory_kwargs)

 def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None):
 """Allocate inference cache for the multi-head self-attention module."""
 dtype = self.out_proj.weight.dtype if dtype is None else dtype
 device = self.out_proj.weight.device
 kv_cache = torch.empty(
 batch_size,
 max_seqlen,
 2,
 self.num_heads_kv,
 self.head_dim,
 dtype=dtype,
 device=device,
 )
 return kv_cache, None

 def _pytorch_attn(self, q, kv):
 k, v = kv.unbind(dim=-3)
 k = torch.repeat_interleave(k, dim=2, repeats=self.num_heads // self.num_heads_kv)
 v = torch.repeat_interleave(v, dim=2, repeats=self.num_heads // self.num_heads_kv)
 q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
 L, S = q.size(-2), k.size(-2)
 if S > self.sliding_window_length:
 attn_mask = (
 torch.ones(L, S, dtype=torch.bool)
 .tril(diagonal=0)
 .triu(-self.window_size[0])
 .to(device=q.device)
 )
 # Since we pass in an attn_mask explicitly, we need to pass is_causal=False to
 # `scaled_dot_product_attention` (even though the attn_mask itself is in fact causal).
 is_causal_arg = False
 else:
 # The previous branch would also handle this case correctly, but it is more efficient
 # to omit the attn_mask when we don't need it.
 attn_mask = None
 is_causal_arg = True
 return F.scaled_dot_product_attention(
 q, k, v, attn_mask=attn_mask, is_causal=is_causal_arg, scale=self.softmax_scale
 ).transpose(1, 2)

 def _update_kv_cache(self, kv, inference_params):
 """kv: (batch_size, seqlen, 2, nheads, head_dim) or (batch_size, 1, 2, nheads, head_dim)."""
 assert self.layer_idx is not None, "Generation requires layer_idx in the constructor"
 return _update_kv_cache(kv, inference_params, self.layer_idx)

 def _update_kvcache_attention(self, q, kv, inference_params):
 """Write kv to inference_params, then compute attention."""
 if inference_params.seqlen_offset == 0 or flash_attn_with_kvcache is None:
 # TODO: this only uses seqlen_offset and not lengths_per_sample.
 kv = self._update_kv_cache(kv, inference_params)
 return self._pytorch_attn(q, kv)
 else:
 batch = q.shape[0]
 kv_cache, _ = inference_params.key_value_memory_dict[self.layer_idx]
 kv_cache = kv_cache[:batch]
 cache_seqlens = (
 inference_params.lengths_per_sample[:batch]
 if inference_params.lengths_per_sample is not None
 else inference_params.seqlen_offset
 )
 return flash_attn_with_kvcache(
 q,
 kv_cache[:, :, 0],
 kv_cache[:, :, 1],
 kv[:, :, 0],
 kv[:, :, 1],
 cache_seqlens=cache_seqlens,
 softmax_scale=self.softmax_scale,
 causal=self.causal,
 window_size=self.window_size,
 )

 def forward(self, x, inference_params=None):
 """Forward pass through the multi-head self-attention module."""
 if (
 inference_params is not None
 and self.layer_idx not in inference_params.key_value_memory_dict
 ):
 inference_params.key_value_memory_dict[self.layer_idx] = self.allocate_inference_cache(
 x.shape[0], inference_params.max_seqlen, dtype=x.dtype
 )
 qkv = self.in_proj(x)
 q, kv = qkv.split(
 [self.num_heads * self.head_dim, self.num_heads_kv * 2 * self.head_dim], dim=-1
 )
 q = rearrange(q, "... (h d) -> ... h d", d=self.head_dim)
 kv = rearrange(kv, "... (two hkv d) -> ... two hkv d", two=2, d=self.head_dim)
 if inference_params is None:
 context = self._pytorch_attn(q, kv)
 else:
 context = self._update_kvcache_attention(q, kv, inference_params)
 context = rearrange(context, "... h d -> ... (h d)")
 out = self.out_proj(context)
 return out


class Block(nn.Module):
 """Simple residual block with normalization that wraps an inner "mixer" module."""

 def __init__(self, dim, mixer_cls, norm_cls=nn.LayerNorm, residual_in_fp32=False):
 """
 dim: The dimension of the input data.
 mixer_cls: The class of the mixer module.
 norm_cls: The class of the normalization module.
 residual_in_fp32: Whether to keep residuals in fp32.
 """
 super().__init__()
 self.residual_in_fp32 = residual_in_fp32
 self.norm = norm_cls(dim)
 self.mixer = mixer_cls(dim)

 def forward(self, x, inference_params=None, **mixer_kwargs):
 """Forward pass through the block."""
 y = self.norm(x.to(dtype=self.norm.weight.dtype))
 y = self.mixer(y, inference_params=inference_params, **mixer_kwargs)

 residual = x
 if self.residual_in_fp32:
 residual = residual.to(torch.float32)
 y = y + residual
 y = y.to(dtype=x.dtype)

 return y

 def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
 """Allocate inference cache for the mixer module."""
 return self.mixer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)


def _create_block(
 d_model,
 norm_cls,
 ssm_cfg=None,
 attn_layer_idx=None,
 attn_cfg=None,
 mlp_layer_idx=None,
 mlp_cfg=None,
 residual_in_fp32=False,
 layer_idx=None,
 device=None,
 dtype=None,
):
 factory_kwargs = {"device": device, "dtype": dtype}
 if ssm_cfg is None:
 ssm_cfg = {}
 if attn_layer_idx is None:
 attn_layer_idx = []
 if attn_cfg is None:
 attn_cfg = {}
 if mlp_layer_idx is None:
 mlp_layer_idx = []
 if mlp_cfg is None:
 mlp_cfg = {}
 if layer_idx in attn_layer_idx:
 mixer_cls = partial(ReneMHA, layer_idx=layer_idx, **attn_cfg, **factory_kwargs)
 elif layer_idx in mlp_layer_idx:
 mixer_cls = partial(ReneMLP, **mlp_cfg, **factory_kwargs)
 else:
 mixer_cls = partial(Mamba2, layer_idx=layer_idx, **ssm_cfg, **factory_kwargs)
 return Block(d_model, mixer_cls, norm_cls=norm_cls, residual_in_fp32=residual_in_fp32)


class MixerModel(nn.Module):
 """Adapted from mamba_ssm.models.mixer_seq_simple.MixerModel."""

 def __init__(
 self,
 d_model: int,
 n_layer: int,
 vocab_size: int,
 ssm_cfg=None,
 attn_layer_idx=None,
 attn_cfg=None,
 mlp_layer_idx=None,
 mlp_cfg=None,
 norm_epsilon: float = 1e-5,
 rms_norm: bool = False,
 initializer_cfg=None,
 residual_in_fp32=False,
 device=None,
 dtype=None,
 ) -> None:
 super().__init__()
 factory_kwargs = {"device": device, "dtype": dtype}
 self.residual_in_fp32 = residual_in_fp32

 if rms_norm:
 from mamba_ssm.ops.triton.layer_norm import RMSNorm as norm_cls_base
 else:
 norm_cls_base = nn.LayerNorm
 norm_cls = partial(norm_cls_base, eps=norm_epsilon, **factory_kwargs)

 self.embedding = nn.Embedding(vocab_size, d_model, **factory_kwargs)

 self.layers = nn.ModuleList(
 [
 _create_block(
 d_model,
 norm_cls=norm_cls,
 ssm_cfg=ssm_cfg,
 attn_layer_idx=attn_layer_idx,
 attn_cfg=attn_cfg,
 mlp_layer_idx=mlp_layer_idx,
 mlp_cfg=mlp_cfg,
 residual_in_fp32=residual_in_fp32,
 layer_idx=i,
 **factory_kwargs,
 )
 for i in range(n_layer)
 ]
 )

 self.norm_f = norm_cls(d_model)

 self.apply(
 partial(
 _init_weights,
 n_layer=n_layer,
 **(initializer_cfg if initializer_cfg is not None else {}),
 n_residuals_per_layer=1,
 )
 )

 def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
 """Allocate inference cache for all layers."""
 return {
 i: layer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)
 for i, layer in enumerate(self.layers)
 }

 def forward(self, input_ids, inference_params=None, **mixer_kwargs):
 """Forward pass through the model."""
 hidden_states = self.embedding(input_ids)
 for layer in self.layers:
 hidden_states = layer(hidden_states, inference_params=inference_params, **mixer_kwargs)
 hidden_states = self.norm_f(hidden_states.to(dtype=self.norm_f.weight.dtype))
 return hidden_states


class ReneLMHeadModel(PreTrainedModel, MambaGenerationMixin):
 """
 Rene language model architecture.
 Based on mamba_ssm.models.mixer_seq_simple.MambaLMHeadModel, with several adaptations.
 """

 config_class = ReneConfig
 base_model_prefix = "backbone"
 _no_split_modules = ["Block", "Mamba2"]
 supports_gradient_checkpointing = True
 _is_stateful = True
 _tied_weights_keys = ["lm_head.weight"]

 def __init__(
 self,
 config: ReneConfig,
 initializer_cfg=None,
 device=None,
 dtype=None,
 ) -> None:
 super().__init__(config)
 d_model = config.d_model
 n_layer = config.n_layer
 vocab_size = config.vocab_size
 ssm_cfg = config.ssm_cfg
 attn_layer_idx = config.attn_layer_idx
 attn_cfg = config.attn_cfg
 mlp_layer_idx = config.mlp_layer_idx
 mlp_cfg = config.mlp_cfg
 rms_norm = config.rms_norm
 residual_in_fp32 = config.residual_in_fp32
 pad_vocab_size_multiple = config.pad_vocab_size_multiple
 factory_kwargs = {"device": device, "dtype": dtype}

 if set(attn_layer_idx).intersection(mlp_layer_idx):
 raise ValueError(f"Conflicting {attn_layer_idx=} and {mlp_layer_idx=}")

 if vocab_size % pad_vocab_size_multiple != 0:
 vocab_size += pad_vocab_size_multiple - (vocab_size % pad_vocab_size_multiple)

 self.backbone = MixerModel(
 d_model=d_model,
 n_layer=n_layer,
 vocab_size=vocab_size,
 ssm_cfg=ssm_cfg,
 attn_layer_idx=attn_layer_idx,
 attn_cfg=attn_cfg,
 mlp_layer_idx=mlp_layer_idx,
 mlp_cfg=mlp_cfg,
 rms_norm=rms_norm,
 initializer_cfg=initializer_cfg,
 residual_in_fp32=residual_in_fp32,
 **factory_kwargs,
 )
 self.lm_head = nn.Linear(d_model, vocab_size, bias=False, **factory_kwargs)

 # Initialize weights
 self.apply(
 partial(
 _init_weights,
 n_layer=n_layer,
 **(initializer_cfg if initializer_cfg is not None else {}),
 )
 )
 self.tie_weights()

 def tie_weights(self):
 """Tie embeddings and softmax layer weights if specified by config."""
 if self.config.tie_word_embeddings:
 self.lm_head.weight = self.backbone.embedding.weight

 def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
 """Allocate inference cache."""
 return self.backbone.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)

 def forward(
 self, input_ids, position_ids=None, inference_params=None, num_last_tokens=0, **mixer_kwargs
 ):
 """
 "position_ids" is just to be compatible with Transformer generation. We don't use it.
 num_last_tokens: if > 0, only return the logits for the last n tokens.
 """
 hidden_states = self.backbone(input_ids, inference_params=inference_params, **mixer_kwargs)
 if num_last_tokens > 0:
 hidden_states = hidden_states[:, -num_last_tokens:]
 lm_logits = self.lm_head(hidden_states)

 return CausalLMOutput(logits=lm_logits)

 def generate(self, *args, **kwargs):
 """
 Calls the custom `generate` method from `mamba_ssm.utils.generation.GenerationMixin`.
 Refer to that method for argument names and defaults.
 """
 return MambaGenerationMixin.generate(self, *args, **kwargs)