import math from dataclasses import dataclass from typing import Optional, Tuple, Union import torch import torch.nn.functional as F from torch import nn from transformers.activations import ACT2FN from transformers.cache_utils import Cache, DynamicCache from transformers.generation.utils import GenerationMixin from transformers.modeling_outputs import MoeModelOutputWithPast from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update from transformers.modeling_utils import PreTrainedModel from transformers.utils import ModelOutput, add_start_docstrings from transformers.utils import logging as hf_logging from .configuration_maple import MapleConfig from .fa3 import flash_attention_forward logger = hf_logging.get_logger(__name__) @dataclass class MapleOutputWithPast(ModelOutput): loss: Optional[torch.FloatTensor] = None logits: Optional[torch.FloatTensor] = None past_key_values: Optional[Cache] = None hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None attentions: Optional[tuple[torch.FloatTensor, ...]] = None aux_loss: Optional[torch.FloatTensor] = None router_logits: Optional[tuple[torch.FloatTensor, ...]] = None class MapleModelOutputWithPast(MoeModelOutputWithPast): """Maple base-model output with an auxiliary router loss.""" def __init__(self, aux_loss=0.0, **kwargs): super().__init__(**kwargs) self.aux_loss = aux_loss class MapleRotaryEmbedding(nn.Module): def __init__(self, config: MapleConfig, device=None): super().__init__() if hasattr(config, "rope_scaling") and config.rope_scaling is not None: self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) else: self.rope_type = "default" self.max_seq_len_cached = config.max_position_embeddings self.original_max_seq_len = config.max_position_embeddings self.config = config self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) self.register_buffer("inv_freq", inv_freq, persistent=False) self.original_inv_freq = self.inv_freq @torch.no_grad() @dynamic_rope_update def forward(self, x, position_ids): inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) position_ids_expanded = position_ids[:, None, :].float() device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" with torch.autocast(device_type=device_type, enabled=False): freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) emb = torch.cat((freqs, freqs), dim=-1) cos = emb.cos() * self.attention_scaling sin = emb.sin() * self.attention_scaling freqs = torch.cat([freqs, freqs], dim=-1) return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype), freqs.float() def rotate_half(x): x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) rotary_dim = cos.shape[-1] q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin) k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin) q_embed = torch.cat([q_embed, q_pass], dim=-1) k_embed = torch.cat([k_embed, k_pass], dim=-1) return q_embed, k_embed class MapleMLP(nn.Module): def __init__(self, config: MapleConfig, intermediate_size: int): super().__init__() self.hidden_size = config.hidden_size self.intermediate_size = intermediate_size self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) self.act_fn = ACT2FN[config.hidden_act] def forward(self, x): gate_weight, up_weight, down_weight = self.gate_proj.weight, self.up_proj.weight, self.down_proj.weight return torch.nn.functional.linear( self.act_fn(torch.clamp(torch.nn.functional.linear(x, gate_weight), max=7.0)) * torch.clamp(torch.nn.functional.linear(x, up_weight), min=-7.0, max=7.0), down_weight, ) class MapleRMSNorm(nn.Module): def __init__(self, hidden_size, eps=1e-6): super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps def forward(self, hidden_states): input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) return self.weight * hidden_states.to(input_dtype) try: from liger_kernel.transformers.rms_norm import LigerRMSNorm MapleRMSNorm = LigerRMSNorm except ImportError: pass class MapleGate(nn.Module): def __init__(self, config: MapleConfig): super().__init__() self.top_k = config.num_experts_per_tok self.num_experts = config.num_experts self.gating_dim = config.hidden_size self.weight = nn.Parameter(torch.empty((self.num_experts, self.gating_dim))) self.reset_parameters() def reset_parameters(self) -> None: import torch.nn.init as init init.kaiming_uniform_(self.weight, a=math.sqrt(5)) def forward(self, hidden_states: torch.Tensor): hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32)) routing_weights = F.softmax(logits, dim=1, dtype=torch.float) scores, topk_idx = torch.topk(routing_weights, self.top_k, dim=-1) scores = scores.type_as(logits) topk_weight = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) return topk_idx, topk_weight, logits class MapleSparseMoeBlock(nn.Module): """Unfused Maple mixture-of-experts block.""" def __init__(self, config) -> None: super().__init__() self.config = config self.num_experts_per_tok = config.num_experts_per_tok self._setup_experts() self.gate = MapleGate(config) def _setup_experts(self): self.experts = nn.ModuleList( [ MapleMLP( config=self.config, intermediate_size=self.config.moe_intermediate_size, ) for _ in range(self.config.num_experts) ] ) def forward( self, hidden_states: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: bsz, seq_len, h = hidden_states.shape topk_idx, topk_weight, router_logits = self.gate(hidden_states) hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) flat_topk_idx = topk_idx.view(-1) if self.training: hidden_states = hidden_states.repeat_interleave(self.num_experts_per_tok, dim=0) y = torch.empty_like(hidden_states) for i, expert in enumerate(self.experts): y[flat_topk_idx == i] = expert(hidden_states[flat_topk_idx == i]) y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1) y = y.to(hidden_states.dtype).view(bsz, seq_len, h) else: y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(bsz, seq_len, h) return y, router_logits @torch.no_grad() def moe_infer(self, x, topk_ids, topk_weight): cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts))) cnts.scatter_(1, topk_ids, 1) tokens_per_expert = cnts.sum(dim=0) idxs = topk_ids.view(-1).argsort() sorted_tokens = x[idxs // topk_ids.shape[1]] tokens_per_expert = tokens_per_expert.cpu().numpy() outputs = [] start_idx = 0 for i, num_tokens in enumerate(tokens_per_expert): end_idx = start_idx + num_tokens if num_tokens == 0: continue expert = self.experts[i] tokens_for_this_expert = sorted_tokens[start_idx:end_idx] expert_out = expert(tokens_for_this_expert) outputs.append(expert_out.to(x.device)) start_idx = end_idx outs = torch.cat(outputs, dim=0) if outputs else sorted_tokens.new_empty(0) new_x = torch.empty_like(outs) new_x[idxs] = outs final_out = ( new_x.view(*topk_ids.shape, -1) .type(topk_weight.dtype) .mul_(topk_weight.unsqueeze(dim=-1)) .sum(dim=1) .type(new_x.dtype) ) return final_out class MapleAttention(nn.Module): """Maple grouped-query attention implemented with FlashAttention.""" def __init__(self, config: MapleConfig, layer_idx: Optional[int] = None): super().__init__() self.config = config self.layer_idx = layer_idx if layer_idx is None: logger.warning_once( f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will " "lead to errors during the forward call if caching is used. Please pass `layer_idx`." ) self.attention_dropout = config.attention_dropout self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = config.head_dim or self.hidden_size // self.num_heads self.scaling = self.head_dim**-0.5 self.num_key_value_heads = config.num_key_value_heads self.is_causal = True layer_type = config.layer_types[layer_idx] if hasattr(config, "layer_types") else None self.sliding_window = config.sliding_window if layer_type == "sliding_attention" else None self.q_proj = nn.Linear( config.hidden_size, config.num_attention_heads * self.head_dim, bias=False ) self.k_proj = nn.Linear( config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False ) self.v_proj = nn.Linear( config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False ) self.q_norm = MapleRMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = MapleRMSNorm(self.head_dim, eps=config.rms_norm_eps) self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.use_bias) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]: bsz, q_len, _ = hidden_states.size() qkv_weight = torch.cat([self.q_proj.weight, self.k_proj.weight, self.v_proj.weight], dim=0) out_qkv = torch.nn.functional.linear(hidden_states, qkv_weight) cos, sin, _freqs = position_embeddings qkv = out_qkv.view(bsz, q_len, self.num_heads + 2 * self.num_key_value_heads, self.head_dim) query_states, key_states, value_states = qkv.split( [self.num_heads, self.num_key_value_heads, self.num_key_value_heads], dim=-2 ) query_states = query_states.transpose(1, 2) key_states = key_states.transpose(1, 2) value_states = value_states.transpose(1, 2) query_states = self.q_norm(query_states) key_states = self.k_norm(key_states) if self.sliding_window is not None: query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) if use_cache and past_key_value is not None: cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) attn_output, attn_weights = flash_attention_forward( self, query_states, key_states, value_states, attention_mask, dropout=self.attention_dropout if self.training else 0.0, position_ids=position_ids, scaling=self.scaling, sliding_window=self.sliding_window, **kwargs, ) attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() attn_output = torch.nn.functional.linear(attn_output, self.o_proj.weight) if not output_attentions: attn_weights = None return attn_output, attn_weights, past_key_value class MapleDecoderLayer(nn.Module): def __init__(self, config: MapleConfig, layer_idx: int): super().__init__() self.self_attn = MapleAttention(config=config, layer_idx=layer_idx) self.mlp = MapleSparseMoeBlock(config) self.input_layernorm = MapleRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = MapleRMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: Optional[bool] = False, output_router_logits: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = None, **kwargs, ) -> Tuple[ torch.Tensor, Optional[torch.Tensor], Optional[Cache], torch.Tensor, Optional[torch.Tensor], ]: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) attn_out, 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=bool(output_attentions), use_cache=bool(use_cache), cache_position=cache_position, position_embeddings=position_embeddings, **kwargs, ) hidden_states = residual + attn_out residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states, router_logits = self.mlp(hidden_states) aux_loss = 0.0 hidden_states = residual + hidden_states.to(residual.device) return ( hidden_states, self_attn_weights, present_key_value, aux_loss, router_logits, ) @add_start_docstrings( "The bare Maple model, which outputs raw hidden states without a task-specific head.", ) class MaplePreTrainedModel(PreTrainedModel): config_class = MapleConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["MapleDecoderLayer"] _skip_keys_device_placement = "past_key_values" _supports_attention_backend = True _supports_flash_attn_2 = True _supports_sdpa = True _supports_cache_class = 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_() @add_start_docstrings( "The bare Maple model, which outputs raw hidden states without a task-specific head.", ) class MapleModel(MaplePreTrainedModel): def __init__(self, config: MapleConfig): super().__init__(config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) layers = [] for layer_idx in range(config.num_hidden_layers): layers.append(MapleDecoderLayer(config, layer_idx)) self.layers = nn.ModuleList(layers) self.config = config self.norm = MapleRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = MapleRotaryEmbedding(config=config) self.gradient_checkpointing = False self.post_init() def get_input_embeddings(self): return self.word_embeddings def set_input_embeddings(self, value): self.word_embeddings = value def prepare_fa2_from_position_ids(self, position_ids: torch.Tensor): position_ids = position_ids.flatten() total_tokens = position_ids.numel() indices_q = torch.arange(total_tokens, device=position_ids.device, dtype=torch.int32) starts = indices_q[position_ids == 0] # If no segment-start markers exist (common in decoding where pos ids are offset), # treat as a single sequence. if starts.numel() == 0: cu_seq_lens = torch.tensor([0, total_tokens], device=position_ids.device, dtype=torch.int32) else: if starts[0].item() != 0: starts = torch.cat([starts.new_zeros(1), starts], dim=0) if starts[-1].item() != total_tokens: starts = torch.cat([starts, starts.new_tensor([total_tokens])], dim=0) cu_seq_lens = starts max_length = (cu_seq_lens[1:] - cu_seq_lens[:-1]).max().item() return (indices_q, (cu_seq_lens, cu_seq_lens), (max_length, max_length)) def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, output_router_logits: Optional[bool] = None, return_dict: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, **kwargs, ) -> Union[Tuple, MapleModelOutputWithPast]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_router_logits = ( output_router_logits if output_router_logits is not None else self.config.output_router_logits ) output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) 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 if (input_ids is None) == (inputs_embeds is None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") 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`." ) use_cache = False if use_cache and past_key_values is None: past_key_values = DynamicCache() if inputs_embeds is None: inputs_embeds = self.word_embeddings(input_ids) forward_batch = kwargs.get("forward_batch", None) is_decode_step = False forward_mode = getattr(forward_batch, "forward_mode", None) if forward_batch is not None else None if forward_mode is not None: for mode_name in ( "is_decode", "is_decode_or_idle", "is_target_verify", "is_draft_decode", ): mode_fn = getattr(forward_mode, mode_name, None) if callable(mode_fn) and bool(mode_fn()): is_decode_step = True break past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 if cache_position is None: cache_position = torch.arange( past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device, ) if position_ids is not None: # Expand shared position IDs before preparing packed-sequence metadata. batch_size = input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0] if position_ids.shape[0] != batch_size: position_ids = position_ids.expand(batch_size, -1) # Decode does not need cu_seq_lens/max_length metadata and creating # them every step hurts CUDA graph capture stability. if (not is_decode_step) and inputs_embeds.shape[1] > 1: _, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = self.prepare_fa2_from_position_ids( position_ids ) kwargs["cu_seq_lens_q"] = cu_seq_lens_q kwargs["cu_seq_lens_k"] = cu_seq_lens_k kwargs["max_length_q"] = max_length_q kwargs["max_length_k"] = max_length_k if position_ids is None: position_ids = cache_position.unsqueeze(0) causal_mask = attention_mask hidden_states = inputs_embeds position_embeddings = self.rotary_emb(hidden_states, position_ids) all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None all_router_logits = () if output_router_logits else None aux_loss_sum = 0.0 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, output_attentions, output_router_logits, use_cache, cache_position, position_embeddings, **kwargs, ) else: layer_outputs = decoder_layer( hidden_states, attention_mask=causal_mask, position_ids=position_ids, past_key_value=past_key_values, output_attentions=output_attentions, output_router_logits=output_router_logits, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, **kwargs, ) hidden_states = layer_outputs[0] if output_attentions: all_self_attns += (layer_outputs[1],) aux_loss_sum = aux_loss_sum + layer_outputs[3] if output_router_logits: all_router_logits += (layer_outputs[4],) hidden_states = self.norm(hidden_states) if output_hidden_states: all_hidden_states += (hidden_states,) moe_layer_count = max(len(self.layers), 1) out = MapleModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values if use_cache else None, hidden_states=all_hidden_states, attentions=all_self_attns, router_logits=all_router_logits, aux_loss=aux_loss_sum / moe_layer_count, ) return ( out if return_dict else ( out.last_hidden_state, out.past_key_values, out.hidden_states, out.attentions, ) ) class MapleForCausalLM(MaplePreTrainedModel, GenerationMixin): _tied_weights_keys = ["lm_head.weight"] def __init__(self, config: MapleConfig): super().__init__(config) self.model = MapleModel(config) self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.post_init() def get_input_embeddings(self): return self.model.word_embeddings def set_input_embeddings(self, value): self.model.word_embeddings = 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: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.Tensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, output_router_logits: Optional[bool] = None, return_dict: Optional[bool] = None, logits_to_keep: Union[int, torch.Tensor] = 0, **kwargs, ) -> Union[Tuple, MapleOutputWithPast]: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) output_router_logits = ( output_router_logits if output_router_logits is not None else self.config.output_router_logits ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict 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_logits=output_router_logits, return_dict=True, **kwargs, ) hidden_states = outputs.last_hidden_state assert isinstance(hidden_states, torch.Tensor) loss = None logits = None if labels is not None: loss, logits = self.loss_function(hidden_states, self.lm_head.weight, labels) else: slice_indices = ( slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep ) logits = self.lm_head(hidden_states[:, slice_indices, :]) out = MapleOutputWithPast( loss=loss, aux_loss=getattr(outputs, "aux_loss", 0.0), logits=logits, past_key_values=outputs.past_key_values if hasattr(outputs, "past_key_values") else None, hidden_states=outputs.hidden_states if hasattr(outputs, "hidden_states") else None, attentions=outputs.attentions if hasattr(outputs, "attentions") else None, router_logits=outputs.router_logits if hasattr(outputs, "router_logits") else None, ) return out if return_dict else out.to_tuple()