from __future__ import annotations import math from typing import Any, ClassVar import torch import torch.nn.functional as functional from torch import nn from transformers.activations import ACT2FN from transformers.cache_utils import Cache, DynamicCache from transformers.generation import GenerationMixin from transformers.modeling_layers import GradientCheckpointingLayer from transformers.modeling_outputs import ( MoeCausalLMOutputWithPast, MoeModelOutputWithPast, ) from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from transformers.utils import can_return_tuple from transformers.utils.generic import merge_with_config_defaults from transformers.utils.output_capturing import capture_outputs from .configuration_alice_ai import AliceAIConfig def _rms_norm( hidden_states: torch.Tensor, weight: torch.Tensor, eps: float, *, zero_centered: bool = False, ) -> torch.Tensor: dtype = hidden_states.dtype normalized = hidden_states.float() * torch.rsqrt( hidden_states.float().square().mean(dim=-1, keepdim=True) + eps ) scale = 1.0 + weight.float() if zero_centered else weight.float() return (normalized * scale).to(dtype) class AliceAIRMSNorm(nn.Module): def __init__(self, hidden_size: int, eps: float, *, zero_centered: bool) -> None: super().__init__() initial_value = 0.0 if zero_centered else 1.0 self.weight = nn.Parameter(torch.full((hidden_size,), initial_value)) self.eps = eps self.zero_centered = zero_centered def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return _rms_norm( hidden_states, self.weight, self.eps, zero_centered=self.zero_centered, ) def _depth_softmax_mix( sources: list[torch.Tensor], query_weight: torch.Tensor, norm_weight: torch.Tensor, eps: float, ) -> torch.Tensor: stacked = torch.stack(sources, dim=0) keys = _rms_norm(stacked, norm_weight, eps) scores = functional.linear(keys.float(), query_weight.float()).squeeze(-1) weights = scores.softmax(dim=0).to(stacked.dtype).unsqueeze(-1) return (stacked * weights).sum(dim=0) class AliceAIFinalBlockAttnRes(nn.Module): def __init__(self, config: AliceAIConfig) -> None: super().__init__() self.res_proj = nn.Linear(config.hidden_size, 1, bias=False) self.res_norm_weight = nn.Parameter(torch.ones(config.hidden_size)) self.eps = config.rms_norm_eps def forward( self, completed: list[torch.Tensor], partial: torch.Tensor | None = None ) -> torch.Tensor: sources = completed if partial is None else [*completed, partial] return _depth_softmax_mix( sources, self.res_proj.weight, self.res_norm_weight, self.eps ) class AliceAIRotaryEmbedding(nn.Module): def __init__(self, config: AliceAIConfig) -> None: super().__init__() rotary_dim = int(config.head_dim * config.partial_rotary_factor) if rotary_dim % 2: raise ValueError("The rotary dimension must be even") inv_freq = 1.0 / ( config.rope_theta ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim) ) self.register_buffer("inv_freq", inv_freq, persistent=False) def forward( self, position_ids: torch.LongTensor, dtype: torch.dtype ) -> tuple[torch.Tensor, torch.Tensor]: frequencies = torch.einsum( "bi,j->bij", position_ids.float(), self.inv_freq.float() ) embeddings = torch.cat((frequencies, frequencies), dim=-1) return embeddings.cos().to(dtype), embeddings.sin().to(dtype) def _rotate_half(hidden_states: torch.Tensor) -> torch.Tensor: first, second = hidden_states.chunk(2, dim=-1) return torch.cat((-second, first), dim=-1) def _apply_rotary( hidden_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor ) -> torch.Tensor: rotary_dim = cos.shape[-1] rotary, remainder = hidden_states[..., :rotary_dim], hidden_states[..., rotary_dim:] cos = cos.unsqueeze(1) sin = sin.unsqueeze(1) rotary = rotary * cos + _rotate_half(rotary) * sin return torch.cat((rotary, remainder), dim=-1) def eager_attention_forward( module: nn.Module, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: torch.Tensor | None, dropout: float = 0.0, scaling: float | None = None, **_: Any, ) -> tuple[torch.Tensor, torch.Tensor]: key = key.repeat_interleave(module.num_key_value_groups, dim=1) value = value.repeat_interleave(module.num_key_value_groups, dim=1) scores = torch.matmul(query.float(), key.float().transpose(-1, -2)) scores = scores * (module.scaling if scaling is None else scaling) if attention_mask is not None: if attention_mask.dtype == torch.bool: scores = scores.masked_fill(~attention_mask, torch.finfo(scores.dtype).min) else: scores = scores + attention_mask probabilities = scores.softmax(dim=-1).to(query.dtype) probabilities = functional.dropout( probabilities, p=dropout, training=module.training ) output = torch.matmul(probabilities, value).transpose(1, 2).contiguous() return output, probabilities class AliceAIAttention(nn.Module): def __init__(self, config: AliceAIConfig, layer_idx: int) -> None: super().__init__() self.config = config self.layer_idx = layer_idx self.num_heads = config.num_attention_heads self.num_kv_heads = config.num_key_value_heads self.num_key_value_groups = self.num_heads // self.num_kv_heads self.head_dim = config.head_dim self.scaling = self.head_dim**-0.5 self.attention_dropout = config.attention_dropout self.is_causal = True self.q_proj = nn.Linear( config.hidden_size, 2 * self.num_heads * self.head_dim, bias=False ) self.k_proj = nn.Linear( config.hidden_size, self.num_kv_heads * self.head_dim, bias=False ) self.v_proj = nn.Linear( config.hidden_size, self.num_kv_heads * self.head_dim, bias=False ) self.o_proj = nn.Linear( self.num_heads * self.head_dim, config.hidden_size, bias=False ) self.q_norm = AliceAIRMSNorm( self.head_dim, config.rms_norm_eps, zero_centered=True ) self.k_norm = AliceAIRMSNorm( self.head_dim, config.rms_norm_eps, zero_centered=True ) def _backend_mask( self, attention_mask: torch.Tensor | None, cache_position: torch.LongTensor, key_length: int, implementation: str, ) -> torch.Tensor | None: if attention_mask is not None and attention_mask.ndim == 4: return attention_mask[..., :key_length] padding_mask = ( attention_mask[..., :key_length].to(torch.bool) if attention_mask is not None else None ) if padding_mask is not None and bool(padding_mask.all()): padding_mask = None if implementation.startswith("flash_attention"): return padding_mask query_length = cache_position.shape[0] is_plain_prefill = query_length == key_length and cache_position[0] == 0 is_single_token_decode = query_length == 1 if ( implementation == "sdpa" and padding_mask is None and (is_plain_prefill or is_single_token_decode) ): return None key_positions = torch.arange(key_length, device=cache_position.device) causal_mask = key_positions.view(1, 1, 1, -1) <= cache_position.view( 1, 1, -1, 1 ) if padding_mask is not None: causal_mask = causal_mask & padding_mask[:, None, None, :] return causal_mask def forward( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: torch.Tensor | None, cache: Cache | None, cache_position: torch.LongTensor, ) -> tuple[torch.Tensor, torch.Tensor | None]: batch_size, sequence_length, _ = hidden_states.shape query, output_gate = ( self.q_proj(hidden_states) .view(batch_size, sequence_length, self.num_heads, 2 * self.head_dim) .chunk(2, dim=-1) ) output_gate = output_gate.reshape( batch_size, sequence_length, self.num_heads * self.head_dim ) key = self.k_proj(hidden_states).view( batch_size, sequence_length, self.num_kv_heads, self.head_dim ) value = self.v_proj(hidden_states).view( batch_size, sequence_length, self.num_kv_heads, self.head_dim ) query = self.q_norm(query).transpose(1, 2) key = self.k_norm(key).transpose(1, 2) value = value.transpose(1, 2) cos, sin = position_embeddings query = _apply_rotary(query, cos, sin) key = _apply_rotary(key, cos, sin) if cache is not None: key, value = cache.update(key, value, self.layer_idx) implementation = self.config._attn_implementation backend_mask = self._backend_mask( attention_mask, cache_position, key.shape[-2], implementation ) attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface( implementation, eager_attention_forward ) output, attention_weights = attention_interface( self, query, key, value, backend_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, is_causal=True, ) output = output.reshape( batch_size, sequence_length, self.num_heads * self.head_dim ) output = output * torch.sigmoid(output_gate) return self.o_proj(output), attention_weights class AliceAIKDA(nn.Module): def __init__(self, config: AliceAIConfig, layer_idx: int) -> None: super().__init__() self.config = config self.layer_idx = layer_idx self.num_k_heads = config.linear_num_key_heads self.num_v_heads = config.linear_num_value_heads self.head_k_dim = config.linear_key_head_dim self.head_v_dim = config.linear_value_head_dim self.key_dim = self.num_k_heads * self.head_k_dim self.value_dim = self.num_v_heads * self.head_v_dim self.conv_kernel_size = config.linear_conv_kernel_dim self.act_fn = ACT2FN[config.hidden_act] conv_kwargs = { "kernel_size": self.conv_kernel_size, "padding": self.conv_kernel_size - 1, "bias": False, } self.q_conv1d = nn.Conv1d( self.key_dim, self.key_dim, groups=self.key_dim, **conv_kwargs ) self.k_conv1d = nn.Conv1d( self.key_dim, self.key_dim, groups=self.key_dim, **conv_kwargs ) self.v_conv1d = nn.Conv1d( self.value_dim, self.value_dim, groups=self.value_dim, **conv_kwargs ) self.q_proj = nn.Linear(config.hidden_size, self.key_dim, bias=False) self.k_proj = nn.Linear(config.hidden_size, self.key_dim, bias=False) self.v_proj = nn.Linear(config.hidden_size, self.value_dim, bias=False) self.f_a_proj = nn.Linear(config.hidden_size, self.head_v_dim, bias=False) self.f_b_proj = nn.Linear(self.head_v_dim, self.key_dim, bias=False) self.b_proj = nn.Linear(config.hidden_size, self.num_k_heads, bias=False) self.g_a_proj = nn.Linear(config.hidden_size, self.head_v_dim, bias=False) self.g_b_proj = nn.Linear(self.head_v_dim, self.value_dim, bias=False) self.o_norm = AliceAIRMSNorm( self.head_v_dim, config.rms_norm_eps, zero_centered=False ) self.o_proj = nn.Linear(self.value_dim, config.hidden_size, bias=False) self.a_log_bias = nn.Parameter(torch.empty(self.num_k_heads)) self.dt_bias = nn.Parameter(torch.empty(self.key_dim)) def _causal_conv( self, hidden_states: torch.Tensor, convolution: nn.Conv1d, cache: Cache | None, state_idx: int, ) -> torch.Tensor: channels_first = hidden_states.transpose(1, 2) sequence_length = channels_first.shape[-1] if cache is not None: channels_first = cache.update_conv_state( channels_first, self.layer_idx, state_idx, conv_kernel_size=self.conv_kernel_size, ) output = self.act_fn( convolution(channels_first)[..., : channels_first.shape[-1]] ) return output[..., -sequence_length:].transpose(1, 2) def _torch_kda( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, alpha: torch.Tensor, beta: torch.Tensor, initial_state: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor]: query = functional.normalize(query.float(), dim=-1) * self.head_k_dim**-0.5 key = functional.normalize(key.float(), dim=-1) gate = -self.a_log_bias.float().exp().view( 1, 1, self.num_k_heads, 1 ) * functional.softplus( alpha.float() + self.dt_bias.float().view(1, 1, self.num_k_heads, self.head_k_dim) ) beta = beta.float().sigmoid() if self.config.kda_allow_negative_eigenvalues: beta = beta * 2.0 state = ( initial_state.float() if initial_state is not None else query.new_zeros( query.shape[0], self.num_v_heads, self.head_k_dim, self.head_v_dim, ) ) repeat = self.num_v_heads // self.num_k_heads if repeat > 1: query = query.repeat_interleave(repeat, dim=2) key = key.repeat_interleave(repeat, dim=2) gate = gate.repeat_interleave(repeat, dim=2) beta = beta.repeat_interleave(repeat, dim=2) outputs = [] for token_idx in range(query.shape[1]): query_i = query[:, token_idx] key_i = key[:, token_idx] value_i = value[:, token_idx].float() state = state * gate[:, token_idx].exp().unsqueeze(-1) prediction = torch.einsum("bhk,bhkv->bhv", key_i, state) delta = (value_i - prediction) * beta[:, token_idx].unsqueeze(-1) state = state + torch.einsum("bhk,bhv->bhkv", key_i, delta) outputs.append(torch.einsum("bhk,bhkv->bhv", query_i, state)) return torch.stack(outputs, dim=1).to(value.dtype), state def _kda( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, alpha: torch.Tensor, beta: torch.Tensor, initial_state: torch.Tensor | None, use_cache: bool, ) -> tuple[torch.Tensor, torch.Tensor | None]: if query.is_cuda: try: from fla.ops.kda import chunk_kda, fused_recurrent_kda # noqa: PLC0415 kernel = ( fused_recurrent_kda if query.shape[1] == 1 and initial_state is not None else chunk_kda ) kernel_beta = beta.float().sigmoid() if self.config.kda_allow_negative_eigenvalues: kernel_beta = kernel_beta * 2.0 return kernel( q=query, k=key, v=value, g=alpha, beta=kernel_beta, A_log=self.a_log_bias.float(), dt_bias=self.dt_bias.float(), initial_state=( initial_state.float() if initial_state is not None else None ), output_final_state=use_cache, use_qk_l2norm_in_kernel=True, use_gate_in_kernel=True, ) except ImportError as error: raise ImportError( "CUDA execution requires flash-linear-attention with KDA support; " "install flash-linear-attention>=0.5.0" ) from error output, state = self._torch_kda(query, key, value, alpha, beta, initial_state) return output, state if use_cache else None def forward( self, hidden_states: torch.Tensor, cache: Cache | None, attention_mask: torch.Tensor | None, ) -> torch.Tensor: if attention_mask is not None: hidden_states = ( hidden_states * attention_mask[..., -hidden_states.shape[1] :, None] ) query = self.q_proj(hidden_states) key = self.k_proj(hidden_states) value = self.v_proj(hidden_states) output_gate = self.g_b_proj(self.g_a_proj(hidden_states)) alpha = self.f_b_proj(self.f_a_proj(hidden_states)) beta = self.b_proj(hidden_states) has_previous_state = cache is not None and cache.has_previous_state( self.layer_idx ) query = self._causal_conv(query, self.q_conv1d, cache, 0) key = self._causal_conv(key, self.k_conv1d, cache, 1) value = self._causal_conv(value, self.v_conv1d, cache, 2) query = query.view(*query.shape[:2], self.num_k_heads, self.head_k_dim) key = key.view(*key.shape[:2], self.num_k_heads, self.head_k_dim) value = value.view(*value.shape[:2], self.num_v_heads, self.head_v_dim) alpha = alpha.view(*alpha.shape[:2], self.num_k_heads, self.head_k_dim) initial_state = ( cache.layers[self.layer_idx].recurrent_states[0] if has_previous_state else None ) output, final_state = self._kda( query, key, value, alpha, beta, initial_state, cache is not None ) if cache is not None: cache.update_recurrent_state(final_state, self.layer_idx) output = self.o_norm(output) output = output * torch.sigmoid(output_gate.view_as(output)) return self.o_proj(output.flatten(-2)) class AliceAIMLP(nn.Module): def __init__( self, hidden_size: int, intermediate_size: int, hidden_act: str ) -> None: super().__init__() self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) self.act_fn = ACT2FN[hidden_act] def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return self.down_proj( self.act_fn(self.gate_proj(hidden_states)) * self.up_proj(hidden_states) ) class AliceAIExperts(nn.Module): def __init__(self, config: AliceAIConfig) -> None: super().__init__() self.num_experts = config.num_experts self.act_fn = ACT2FN[config.hidden_act] self.gate_up_proj = nn.Parameter( torch.empty( config.num_experts, 2 * config.moe_intermediate_size, config.hidden_size, ) ) self.down_proj = nn.Parameter( torch.empty( config.num_experts, config.hidden_size, config.moe_intermediate_size, ) ) def forward( self, hidden_states: torch.Tensor, expert_indices: torch.LongTensor, expert_weights: torch.Tensor, ) -> torch.Tensor: output = torch.zeros_like(hidden_states) for expert_idx in range(self.num_experts): token_indices, slots = torch.where(expert_indices == expert_idx) if token_indices.numel() == 0: continue expert_input = hidden_states[token_indices] gate_up = functional.linear(expert_input, self.gate_up_proj[expert_idx]) gate, up = gate_up.chunk(2, dim=-1) expert_output = functional.linear( self.act_fn(gate) * up, self.down_proj[expert_idx] ) expert_output = expert_output * expert_weights[token_indices, slots, None] output.index_add_(0, token_indices, expert_output.to(output.dtype)) return output class AliceAISigmoidTopKRouter(nn.Module): def __init__(self, config: AliceAIConfig) -> None: super().__init__() self.weight = nn.Parameter(torch.empty(config.num_experts, config.hidden_size)) self.top_k = config.num_experts_per_tok if config.router_bias_correction: self.register_buffer( "e_score_correction_bias", torch.zeros(config.num_experts) ) else: self.e_score_correction_bias = None def forward( self, hidden_states: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor, torch.LongTensor]: router_logits = functional.linear(hidden_states.float(), self.weight.float()) scores = router_logits.sigmoid() selection_scores = ( scores if self.e_score_correction_bias is None else scores + self.e_score_correction_bias ) indices = selection_scores.topk(self.top_k, dim=-1).indices weights = scores.gather(-1, indices) weights = weights / weights.sum(dim=-1, keepdim=True) return router_logits, weights.to(hidden_states.dtype), indices class AliceAISparseMoEBlock(nn.Module): def __init__(self, config: AliceAIConfig) -> None: super().__init__() self.gate = AliceAISigmoidTopKRouter(config) self.experts = AliceAIExperts(config) self.shared_expert = AliceAIMLP( config.hidden_size, config.shared_expert_intermediate_size, config.hidden_act, ) self.shared_expert_gate = nn.Linear(config.hidden_size, 1, bias=False) def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: original_shape = hidden_states.shape flattened = hidden_states.reshape(-1, original_shape[-1]) router_logits, weights, indices = self.gate(flattened) routed = self.experts(flattened, indices, weights) shared = self.shared_expert(flattened) * torch.sigmoid( self.shared_expert_gate(flattened) ) return (routed + shared).reshape(original_shape), router_logits class AliceAIDecoderLayer(GradientCheckpointingLayer): def __init__(self, config: AliceAIConfig, layer_idx: int) -> None: super().__init__() self.layer_idx = layer_idx self.layer_type = config.layer_types[layer_idx] if self.layer_type == "linear_attention": self.linear_attn = AliceAIKDA(config, layer_idx) else: self.self_attn = AliceAIAttention(config, layer_idx) self.input_layernorm = AliceAIRMSNorm( config.hidden_size, config.rms_norm_eps, zero_centered=True ) self.post_attention_layernorm = AliceAIRMSNorm( config.hidden_size, config.rms_norm_eps, zero_centered=True ) self.mlp = AliceAISparseMoEBlock(config) if layer_idx != 0: self.attn_res_proj = nn.Linear(config.hidden_size, 1, bias=False) self.attn_res_norm_weight = nn.Parameter(torch.ones(config.hidden_size)) self.mlp_res_proj = nn.Linear(config.hidden_size, 1, bias=False) self.mlp_res_norm_weight = nn.Parameter(torch.ones(config.hidden_size)) self.eps = config.rms_norm_eps def _mix( self, completed: list[torch.Tensor], partial: torch.Tensor | None, projection: nn.Linear, norm_weight: torch.Tensor, ) -> torch.Tensor: sources = completed if partial is None else [*completed, partial] return _depth_softmax_mix(sources, projection.weight, norm_weight, self.eps) def forward( self, embeddings: torch.Tensor, completed_blocks: tuple[torch.Tensor, ...], partial: torch.Tensor | None, position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: torch.Tensor | None, past_key_values: Cache | None, cache_position: torch.LongTensor, ) -> tuple[torch.Tensor, torch.Tensor]: if self.layer_idx == 0: mixed = embeddings else: mixed = self._mix( completed_blocks, partial, self.attn_res_proj, self.attn_res_norm_weight ) mixed = self.input_layernorm(mixed) if self.layer_type == "linear_attention": attention_output = self.linear_attn( mixed, past_key_values, attention_mask ) else: attention_output, _ = self.self_attn( mixed, position_embeddings, attention_mask, past_key_values, cache_position, ) partial = attention_output if partial is None else partial + attention_output mixed = self._mix( completed_blocks, partial, self.mlp_res_proj, self.mlp_res_norm_weight ) moe_output, router_logits = self.mlp(self.post_attention_layernorm(mixed)) return partial + moe_output, router_logits class AliceAIPreTrainedModel(PreTrainedModel): config_class = AliceAIConfig base_model_prefix = "model" supports_gradient_checkpointing = True _is_stateful = True _no_split_modules: ClassVar[list[str]] = ["AliceAIDecoderLayer"] _keys_to_ignore_on_load_unexpected: ClassVar[list[str]] = [r"^mtp\."] _supports_sdpa = True _supports_flash_attn = True _can_record_outputs: ClassVar[dict[str, type[nn.Module]]] = { "hidden_states": AliceAIDecoderLayer, "attentions": AliceAIAttention, } @staticmethod def _needs_initialization(parameter: torch.Tensor) -> bool: return not getattr(parameter, "_is_hf_initialized", False) @torch.no_grad() def _init_weights(self, module: nn.Module) -> None: # noqa: PLR0912 if isinstance(module, (nn.Linear, nn.Conv1d)): module.weight.normal_(mean=0.0, std=self.config.initializer_range) if module.bias is not None: module.bias.zero_() elif isinstance(module, nn.Embedding): module.weight.normal_(mean=0.0, std=self.config.initializer_range) if module.padding_idx is not None: module.weight[module.padding_idx].zero_() elif isinstance(module, AliceAIExperts): if self._needs_initialization(module.gate_up_proj): module.gate_up_proj.normal_(mean=0.0, std=self.config.initializer_range) if self._needs_initialization(module.down_proj): module.down_proj.normal_(mean=0.0, std=self.config.initializer_range) elif isinstance(module, AliceAISigmoidTopKRouter): if self._needs_initialization(module.weight): module.weight.normal_(mean=0.0, std=self.config.initializer_range) elif isinstance(module, AliceAIKDA): if self._needs_initialization(module.a_log_bias): module.a_log_bias.uniform_(0, 16).log_() if self._needs_initialization(module.dt_bias): dt = torch.exp( torch.rand_like(module.dt_bias) * (math.log(0.1) - math.log(0.001)) + math.log(0.001) ).clamp_min(1e-4) module.dt_bias.copy_(dt + torch.log(-torch.expm1(-dt))) elif isinstance(module, AliceAIFinalBlockAttnRes): if self._needs_initialization(module.res_proj.weight): module.res_proj.weight.zero_() elif isinstance(module, AliceAIDecoderLayer): if hasattr(module, "attn_res_proj"): if self._needs_initialization(module.attn_res_proj.weight): module.attn_res_proj.weight.zero_() if self._needs_initialization(module.mlp_res_proj.weight): module.mlp_res_proj.weight.zero_() class AliceAIModel(AliceAIPreTrainedModel): def __init__(self, config: AliceAIConfig) -> None: super().__init__(config) self.embed_tokens = nn.Embedding( config.vocab_size, config.hidden_size, config.pad_token_id ) self.layers = nn.ModuleList( [ AliceAIDecoderLayer(config, index) for index in range(config.num_hidden_layers) ] ) self.attnres_final = AliceAIFinalBlockAttnRes(config) self.norm = AliceAIRMSNorm( config.hidden_size, config.rms_norm_eps, zero_centered=True ) self.rotary_emb = AliceAIRotaryEmbedding(config) self.gradient_checkpointing = False self.post_init() @merge_with_config_defaults @capture_outputs def forward( self, input_ids: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, use_cache: bool | None = None, output_router_logits: bool | None = None, **_: Any, ) -> MoeModelOutputWithPast: if (input_ids is None) == (inputs_embeds is None): raise ValueError("Specify exactly one of input_ids or inputs_embeds") if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) use_cache = self.config.use_cache if use_cache is None else use_cache output_router_logits = ( self.config.output_router_logits if output_router_logits is None else output_router_logits ) if use_cache and past_key_values is None: past_key_values = DynamicCache(config=self.config) past_length = ( past_key_values.get_seq_length() if past_key_values is not None else 0 ) cache_position = torch.arange( past_length, past_length + inputs_embeds.shape[1], device=inputs_embeds.device, ) if position_ids is None: position_ids = cache_position.unsqueeze(0) position_embeddings = self.rotary_emb(position_ids, inputs_embeds.dtype) if isinstance(attention_mask, dict): linear_mask = attention_mask["linear_attention"] full_attention_mask = attention_mask["full_attention"] else: linear_mask = attention_mask full_attention_mask = attention_mask if cache_position[0] > 0 or ( attention_mask is not None and torch.all(attention_mask == 1) ): linear_mask = None completed_blocks = [inputs_embeds] partial = None router_logits = [] for layer_idx, layer in enumerate(self.layers): if layer_idx > 0 and layer_idx % self.config.block_attn_res_block_size == 0: completed_blocks.append(partial) partial = None layer_mask = ( linear_mask if layer.layer_type == "linear_attention" else full_attention_mask ) partial, layer_router_logits = layer( inputs_embeds, completed_blocks=tuple(completed_blocks), partial=partial, position_embeddings=position_embeddings, attention_mask=layer_mask, past_key_values=past_key_values, cache_position=cache_position, ) if output_router_logits: router_logits.append(layer_router_logits) hidden_states = self.norm(self.attnres_final(completed_blocks, partial)) return MoeModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values, router_logits=tuple(router_logits) if output_router_logits else None, ) class AliceAIForCausalLM(AliceAIPreTrainedModel, GenerationMixin): _tied_weights_keys: ClassVar[dict[str, str]] = { "lm_head.weight": "model.embed_tokens.weight" } def __init__(self, config: AliceAIConfig) -> None: super().__init__(config) self.model = AliceAIModel(config) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.vocab_size = config.vocab_size self.post_init() def get_input_embeddings(self) -> nn.Module: return self.model.embed_tokens def set_input_embeddings(self, value: nn.Module) -> None: self.model.embed_tokens = value def get_output_embeddings(self) -> nn.Module: return self.lm_head def set_output_embeddings(self, value: nn.Module) -> None: self.lm_head = value @can_return_tuple def forward( self, input_ids: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, output_router_logits: bool | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Any, ) -> MoeCausalLMOutputWithPast: 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_router_logits=output_router_logits, **kwargs, ) indices = ( slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep ) logits = self.lm_head(outputs.last_hidden_state[:, indices, :]) loss = ( self.loss_function(logits, labels, self.vocab_size, **kwargs) if labels is not None else None ) return MoeCausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, router_logits=outputs.router_logits, ) AliceAIConfig.register_for_auto_class() AliceAIModel.register_for_auto_class("AutoModel") AliceAIForCausalLM.register_for_auto_class("AutoModelForCausalLM")