# fmt: off
from __future__ import annotations
from dataclasses import dataclass
import torch
import torch.utils.checkpoint
from scipy.stats import norm
from torch import nn
from torch.nn import CrossEntropyLoss
from transformers.activations import ACT2FN
from transformers.cache_utils import Cache, DynamicCache, StaticCache
from transformers.modeling_attn_mask_utils import AttentionMaskConverter
from transformers.modeling_utils import PreTrainedModel
from transformers.models.llama.configuration_llama import LlamaConfig
from transformers.models.llama.modeling_llama import (
LLAMA_ATTENTION_CLASSES,
LlamaRMSNorm,
)
from transformers.utils import ModelOutput, logging
logger = logging.get_logger(__name__)
@dataclass
class MonetModelOutputWithPast(ModelOutput):
last_hidden_state: torch.FloatTensor = None
past_key_values: tuple[tuple[torch.FloatTensor]] | None = None
hidden_states: tuple[torch.FloatTensor, ...] | None = None
attentions: tuple[torch.FloatTensor, ...] | None = None
router_probs: tuple[tuple[torch.FloatTensor, ...], ...] | None = None
@dataclass
class MonetCausalLMOutputWithPast(ModelOutput):
loss: torch.FloatTensor | None = None
aux_loss: torch.FloatTensor | None = None
logits: torch.FloatTensor = None
past_key_values: tuple[tuple[torch.FloatTensor]] | None = None
hidden_states: tuple[torch.FloatTensor, ...] | None = None
attentions: tuple[torch.FloatTensor, ...] | None = None
router_probs: tuple[tuple[torch.FloatTensor, ...], ...] | None = None
class MonetConfig(LlamaConfig):
model_type = "monet"
keys_to_ignore_at_inference = ["past_key_values"]
def __init__(
self,
vocab_size=32000,
hidden_size=4096,
intermediate_size=None,
num_hidden_layers=32,
num_attention_heads=32,
num_key_value_heads=None,
hidden_act="relu2",
max_position_embeddings=2048,
initializer_range=0.02,
rms_norm_eps=1e-6,
use_cache=True,
pad_token_id=None,
bos_token_id=1,
eos_token_id=2,
pretraining_tp=1,
tie_word_embeddings=False,
rope_theta=10000.0,
rope_scaling=None,
attention_bias=False,
attention_dropout=0.0,
mlp_bias=None,
moe_dim=8,
moe_heads=8,
moe_experts=512,
moe_topk=32,
moe_groups=4,
moe_decompose="vertical",
output_router_probs=False,
**kwargs,
):
self.moe_dim = moe_dim
self.moe_heads = moe_heads
self.moe_experts = moe_experts
self.moe_topk = moe_topk
self.moe_groups = moe_groups
self.moe_decompose = moe_decompose
self.output_router_probs = output_router_probs
super().__init__(
vocab_size=vocab_size,
hidden_size=hidden_size,
intermediate_size=intermediate_size,
num_hidden_layers=num_hidden_layers,
num_attention_heads=num_attention_heads,
num_key_value_heads=num_key_value_heads,
hidden_act=hidden_act,
max_position_embeddings=max_position_embeddings,
initializer_range=initializer_range,
rms_norm_eps=rms_norm_eps,
use_cache=use_cache,
pad_token_id=pad_token_id,
bos_token_id=bos_token_id,
eos_token_id=eos_token_id,
pretraining_tp=pretraining_tp,
tie_word_embeddings=tie_word_embeddings,
rope_theta=rope_theta,
rope_scaling=rope_scaling,
attention_bias=attention_bias,
attention_dropout=attention_dropout,
mlp_bias=mlp_bias,
**kwargs,
)
class MonetRouter(nn.Module):
def __init__(self, config: MonetConfig):
super().__init__()
self.config = config
flatten_shape = config.moe_heads * config.moe_experts
self.w1 = nn.Linear(config.hidden_size, flatten_shape, bias=False)
self.w2 = nn.Linear(config.hidden_size, flatten_shape, bias=False)
self.norm1 = nn.BatchNorm1d(config.moe_heads, affine=False)
self.norm2 = nn.BatchNorm1d(config.moe_heads, affine=False)
def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
g1z = self.w1(x).unflatten(-1, (self.config.moe_heads, -1)).float()
g2z = self.w2(x).unflatten(-1, (self.config.moe_heads, -1)).float()
g1n = self.norm1(g1z.transpose(2, 3).flatten(0, -2))
g2n = self.norm2(g2z.transpose(2, 3).flatten(0, -2))
g1n = g1n.view(g1z.size(0), g1z.size(1), g1z.size(3), -1).transpose(2, 3)
g2n = g2n.view(g2z.size(0), g2z.size(1), g2z.size(3), -1).transpose(2, 3)
sigma = float(norm.ppf(1 - self.config.moe_topk / self.config.moe_experts))
g1s = g1n.amax(-1, keepdim=True).clamp_max_(sigma)
g2s = g2n.amax(-1, keepdim=True).clamp_max_(sigma)
g1 = nn.functional.softmax(torch.where(g1n >= g1s, g1z, -1e10), dim=-1)
g2 = nn.functional.softmax(torch.where(g2n >= g2s, g2z, -1e10), dim=-1)
return g1, g2
class MonetMoVDE(nn.Module):
def __init__(self, config: MonetConfig):
super().__init__()
self.config = config
self.act_fn = ACT2FN[config.hidden_act]
flatten_shape = config.moe_experts * config.moe_dim // 2
self.u1 = nn.Linear(config.hidden_size, flatten_shape)
self.u2 = nn.Linear(config.hidden_size, flatten_shape)
self.v11 = nn.Linear(flatten_shape, config.hidden_size // 2, bias=False)
self.v12 = nn.Linear(flatten_shape, config.hidden_size // 2, bias=False)
self.v21 = nn.Linear(flatten_shape, config.hidden_size // 2, bias=False)
self.v22 = nn.Linear(flatten_shape, config.hidden_size // 2, bias=False)
self.b1 = nn.Parameter(torch.zeros(config.moe_experts, config.hidden_size // 2))
self.b2 = nn.Parameter(torch.zeros(config.moe_experts, config.hidden_size // 2))
def forward(
self, x: torch.Tensor, g1: torch.Tensor, g2: torch.Tensor
) -> torch.Tensor:
g1, g2 = g1.type_as(x), g2.type_as(x)
x1 = self.act_fn(self.u1(x).unflatten(-1, (self.config.moe_experts, -1)))
x2 = self.act_fn(self.u2(x).unflatten(-1, (self.config.moe_experts, -1)))
x11 = self.v11(torch.einsum("btim,bthi->btim", x1, g1).flatten(-2))
x12 = self.v12(torch.einsum("btjm,bthj,bthi->btim", x2, g2, g1).flatten(-2))
x13 = torch.einsum("bthi,id->btd", g1, self.b1.type_as(x))
x21 = self.v21(torch.einsum("btim,bthi,bthj->btjm", x1, g1, g2).flatten(-2))
x22 = self.v22(torch.einsum("btjm,bthj->btjm", x2, g2).flatten(-2))
x23 = torch.einsum("bthj,jd->btd", g2, self.b2.type_as(x))
return torch.cat((x11 + x12 + x13, x21 + x22 + x23), dim=-1)
class MonetMoHDE(nn.Module):
def __init__(self, config: MonetConfig):
super().__init__()
self.config = config
self.act_fn = ACT2FN[config.hidden_act]
flatten_shape = config.moe_experts * config.moe_dim
self.u = nn.Linear(config.hidden_size, flatten_shape)
self.v = nn.Linear(flatten_shape, config.hidden_size, bias=False)
self.b = nn.Parameter(torch.zeros(config.moe_experts, config.hidden_size))
def forward(
self, x: torch.Tensor, g1: torch.Tensor, g2: torch.Tensor
) -> torch.Tensor:
g1, g2 = g1.type_as(x), g2.type_as(x)
x = self.act_fn(self.u(x).unflatten(-1, (self.config.moe_experts, -1)))
x = self.v(torch.einsum("btim,bthi,bthj->btjm", x, g1, g2).flatten(-2))
return x + torch.einsum("bthj,jd->btd", g2, self.b)
class MonetDecoderLayer(nn.Module):
def __init__(self, config: MonetConfig, layer_idx: int):
super().__init__()
self.hidden_size = config.hidden_size
self.self_attn = LLAMA_ATTENTION_CLASSES[config._attn_implementation](
config=config, layer_idx=layer_idx
)
self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.post_attention_layernorm = LlamaRMSNorm(
config.hidden_size, eps=config.rms_norm_eps
)
if config.moe_decompose == "vertical":
self.moe = MonetMoVDE(config)
elif config.moe_decompose == "horizontal":
self.moe = MonetMoHDE(config)
if layer_idx % config.moe_groups == 0:
self.router = MonetRouter(config).requires_grad_(False)
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor | None = None,
past_key_value: Cache | None = None,
previous_router_probs: tuple[torch.Tensor, torch.Tensor] | None = None,
output_attentions: bool | None = False,
use_cache: bool | None = False,
cache_position: torch.LongTensor | None = None,
**kwargs,
) -> tuple[torch.FloatTensor, ...]:
residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)
# Self Attention
hidden_states, self_attn_weights, present_key_value = self.self_attn(
hidden_states=hidden_states,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_value=past_key_value,
output_attentions=output_attentions,
use_cache=use_cache,
cache_position=cache_position,
)
hidden_states = residual + hidden_states
# Fully Connected
residual = hidden_states
hidden_states = self.post_attention_layernorm(hidden_states)
g1, g2 = (
self.router(hidden_states)
if hasattr(self, "router")
else previous_router_probs
)
hidden_states = self.moe(hidden_states, g1, g2)
hidden_states = residual + hidden_states
outputs = (hidden_states,)
if output_attentions:
outputs += (self_attn_weights,)
if use_cache:
outputs += (present_key_value,)
return outputs + ((g1, g2) if hasattr(self, "router") else None,)
class MonetPreTrainedModel(PreTrainedModel):
config_class = MonetConfig
base_model_prefix = "model"
supports_gradient_checkpointing = True
_no_split_modules = ["MonetDecoderLayer"]
_skip_keys_device_placement = ["past_key_values"]
_supports_flash_attn_2 = True
_supports_sdpa = True
_supports_cache_class = True
_supports_quantized_cache = True
_supports_static_cache = True
def _init_weights(self, module):
std = self.config.initializer_range
if isinstance(module, nn.Linear):
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_()
class MonetModel(MonetPreTrainedModel):
def __init__(self, config: MonetConfig):
super().__init__(config)
self.padding_idx = config.pad_token_id
self.vocab_size = config.vocab_size
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) # noqa
self.layers = nn.ModuleList([MonetDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) # noqa
self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.gradient_checkpointing = False
# Initialize weights and apply final processing
self.post_init()
def get_input_embeddings(self):
return self.embed_tokens
def set_input_embeddings(self, value):
self.embed_tokens = value
def forward(
self,
input_ids: torch.LongTensor = None,
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor | None = None,
past_key_values: Cache | list[torch.FloatTensor] | None = None,
inputs_embeds: torch.FloatTensor | None = None,
use_cache: bool | None = None,
output_attentions: bool | None = None,
output_hidden_states: bool | None = None,
output_router_probs: bool | None = None,
return_dict: bool | None = None,
cache_position: torch.LongTensor | None = None,
) -> tuple[torch.Tensor, ...] | MonetModelOutputWithPast:
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions # noqa
output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states # noqa
output_router_probs = output_router_probs if output_router_probs is not None else self.config.output_router_probs # noqa
use_cache = use_cache if use_cache is not None else self.config.use_cache
return_dict = return_dict if return_dict is not None else self.config.use_return_dict # noqa
if (input_ids is None) ^ (inputs_embeds is not None):
raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one") # noqa
if self.gradient_checkpointing and self.training and use_cache:
logger.warning_once("`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`.") # noqa
use_cache = False
if inputs_embeds is None:
inputs_embeds = self.embed_tokens(input_ids)
return_legacy_cache = False
if use_cache and not isinstance(past_key_values, Cache): # kept for BC (non `Cache` `past_key_values` inputs) # noqa
return_legacy_cache = True
past_key_values = DynamicCache.from_legacy_cache(past_key_values)
logger.warning_once(
"We detected that you are passing `past_key_values` as a tuple and this is deprecated and will be removed in v4.43. " # noqa
"Please use an appropriate `Cache` class (https://huggingface.co/docs/transformers/v4.41.3/en/internal/generation_utils#transformers.Cache)" # noqa
)
if cache_position is None:
past_seen_tokens = (
past_key_values.get_seq_length() if past_key_values is not None else 0
)
cache_position = torch.arange(past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device) # noqa
if position_ids is None:
position_ids = cache_position.unsqueeze(0)
causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions) # noqa
# embed positions
hidden_states = inputs_embeds
# decoder layers
all_hidden_states = () if output_hidden_states else None
all_self_attns = () if output_attentions else None
all_router_probs = () if output_router_probs else None
previous_router_probs, next_decoder_cache = None, None
for decoder_layer in self.layers:
if output_hidden_states:
all_hidden_states += (hidden_states,)
if self.gradient_checkpointing and self.training:
layer_outputs = self._gradient_checkpointing_func(
decoder_layer.__call__,
hidden_states,
causal_mask,
position_ids,
past_key_values,
previous_router_probs,
output_attentions,
use_cache,
cache_position,
)
else:
layer_outputs = decoder_layer(
hidden_states,
attention_mask=causal_mask,
position_ids=position_ids,
past_key_value=past_key_values,
previous_router_probs=previous_router_probs,
output_attentions=output_attentions,
use_cache=use_cache,
cache_position=cache_position,
)
hidden_states = layer_outputs[0]
if use_cache:
next_decoder_cache = layer_outputs[2 if output_attentions else 1]
if output_attentions:
all_self_attns += (layer_outputs[1],)
if output_router_probs:
all_router_probs += (layer_outputs[-1],)
previous_router_probs = (
layer_outputs[-1]
if layer_outputs[-1] is not None
else previous_router_probs
)
hidden_states = self.norm(hidden_states)
# add hidden states from the last decoder layer
if output_hidden_states:
all_hidden_states += (hidden_states,)
next_cache = next_decoder_cache if use_cache else None
if return_legacy_cache:
next_cache = next_cache.to_legacy_cache()
if not return_dict:
return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_probs] if v is not None) # noqa
return MonetModelOutputWithPast(
last_hidden_state=hidden_states,
past_key_values=next_cache,
hidden_states=all_hidden_states,
attentions=all_self_attns,
router_probs=all_router_probs,
)
def _update_causal_mask(
self,
attention_mask: torch.Tensor,
input_tensor: torch.Tensor,
cache_position: torch.Tensor,
past_key_values: Cache,
output_attentions: bool,
):
if self.config._attn_implementation == "flash_attention_2":
if attention_mask is not None and 0.0 in attention_mask:
return attention_mask
return None
past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 # noqa
using_static_cache = isinstance(past_key_values, StaticCache)
if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions: # noqa
if AttentionMaskConverter._ignore_causal_mask_sdpa(
attention_mask,
inputs_embeds=input_tensor,
past_key_values_length=past_seen_tokens,
is_training=self.training,
):
return None
dtype, device = input_tensor.dtype, input_tensor.device
min_dtype = torch.finfo(dtype).min
sequence_length = input_tensor.shape[1]
if using_static_cache:
target_length = past_key_values.get_max_length()
else:
target_length = (
attention_mask.shape[-1]
if isinstance(attention_mask, torch.Tensor)
else past_seen_tokens + sequence_length + 1
)
if attention_mask is not None and attention_mask.dim() == 4:
if attention_mask.max() != 0:
raise ValueError("Custom 4D attention mask should be passed in inverted form with max==0`") # noqa
causal_mask = attention_mask
else:
causal_mask = torch.full(
(sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device # noqa
)
if sequence_length != 1:
causal_mask = torch.triu(causal_mask, diagonal=1)
causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) # noqa
causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1) # noqa
if attention_mask is not None:
causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit # noqa
mask_length = attention_mask.shape[-1]
padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :] # noqa
padding_mask = padding_mask == 0
causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(padding_mask, min_dtype) # noqa
if (
self.config._attn_implementation == "sdpa"
and attention_mask is not None
and attention_mask.device.type == "cuda"
and not output_attentions
):
causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) # noqa
return causal_mask
class MonetForCausalLM(MonetPreTrainedModel):
_tied_weights_keys = ["lm_head.weight"]
def __init__(self, config):
super().__init__(config)
self.model = MonetModel(config)
self.vocab_size = config.vocab_size
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
# Initialize weights and apply final processing
self.post_init()
def get_input_embeddings(self):
return self.model.embed_tokens
def set_input_embeddings(self, value):
self.model.embed_tokens = value
def get_output_embeddings(self):
return self.lm_head
def set_output_embeddings(self, new_embeddings):
self.lm_head = new_embeddings
def set_decoder(self, decoder):
self.model = decoder
def get_decoder(self):
return self.model
def forward(
self,
input_ids: torch.LongTensor = None,
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor | None = None,
past_key_values: Cache | list[torch.FloatTensor] | None = None,
inputs_embeds: torch.FloatTensor | None = None,
labels: torch.LongTensor | None = None,
use_cache: bool | None = None,
output_attentions: bool | None = None,
output_hidden_states: bool | None = None,
output_router_probs: bool | None = None,
return_dict: bool | None = None,
cache_position: torch.LongTensor | None = None,
) -> tuple[torch.Tensor, ...] | MonetCausalLMOutputWithPast:
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions # noqa
output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states # noqa
output_router_probs = output_router_probs if output_router_probs is not None else self.config.output_router_probs # noqa
return_dict = return_dict if return_dict is not None else self.config.use_return_dict # noqa
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
outputs = self.model(
input_ids=input_ids,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
inputs_embeds=inputs_embeds,
use_cache=use_cache,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
output_router_probs=output_router_probs,
return_dict=return_dict,
cache_position=cache_position,
)
hidden_states = outputs[0]
logits = self.lm_head(hidden_states)
logits = logits.float()
loss = None
if labels is not None:
# Shift so that tokens < n predict n
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
# Flatten the tokens
loss_fct = CrossEntropyLoss()
shift_logits = shift_logits.view(-1, self.config.vocab_size)
shift_labels = shift_labels.view(-1)
# Enable model parallelism
shift_labels = shift_labels.to(shift_logits.device)
loss = loss_fct(shift_logits, shift_labels)
if not return_dict:
output = (logits,) + outputs[1:]
return (loss,) + output if loss is not None else output
return MonetCausalLMOutputWithPast(
loss=loss,
logits=logits,
past_key_values=outputs.past_key_values,
hidden_states=outputs.hidden_states,
attentions=outputs.attentions,
router_probs=outputs.router_probs,
)
def prepare_inputs_for_generation(
self,
input_ids,
past_key_values=None,
attention_mask=None,
inputs_embeds=None,
cache_position=None,
use_cache=True,
**kwargs,
):
past_length = 0
if past_key_values is not None:
past_length = cache_position[0] if cache_position is not None else past_key_values.get_seq_length() # noqa
max_cache_length = (
torch.tensor(past_key_values.get_max_length(), device=input_ids.device)
if past_key_values.get_max_length() is not None
else None
)
cache_length = past_length if max_cache_length is None else torch.min(max_cache_length, past_length) # noqa
# Keep only the unprocessed tokens:
if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]: # noqa
input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
# input_ids based on the past_length.
elif past_length < input_ids.shape[1]:
input_ids = input_ids[:, past_length:]
if (
max_cache_length is not None
and attention_mask is not None
and cache_length + input_ids.shape[1] > max_cache_length
):
attention_mask = attention_mask[:, -max_cache_length:]
position_ids = kwargs.get("position_ids", None)
if attention_mask is not None and position_ids is None:
# create position_ids on the fly for batch generation
position_ids = attention_mask.long().cumsum(-1) - 1
position_ids.masked_fill_(attention_mask == 0, 1)
if past_key_values:
position_ids = position_ids[:, -input_ids.shape[1] :]
if inputs_embeds is not None and past_length == 0:
model_inputs = {"inputs_embeds": inputs_embeds}
else:
model_inputs = {"input_ids": input_ids.contiguous()}
input_length = position_ids.shape[-1] if position_ids is not None else input_ids.shape[-1] # noqa
if cache_position is None:
cache_position = torch.arange(past_length, past_length + input_length, device=input_ids.device) # noqa
elif use_cache:
cache_position = cache_position[-input_length:]
model_inputs.update(
{
"position_ids": position_ids,
"cache_position": cache_position,
"past_key_values": past_key_values,
"use_cache": use_cache,
"attention_mask": attention_mask,
}
)
return model_inputs
@staticmethod
def _reorder_cache(past_key_values, beam_idx):
reordered_past = ()
for layer_past in past_key_values:
reordered_past += (
tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past), # noqa
)
return reordered_past