# Copyright © 2023-2025 Apple Inc. from dataclasses import dataclass from typing import Any, Dict, Optional, Union import mlx.core as mx import mlx.nn as nn from .activations import swiglu from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention @dataclass class ModelArgs(BaseModelArgs): model_type: str vocab_size: int hidden_size: int num_hidden_layers: int intermediate_size: int num_attention_heads: int num_key_value_heads: int rms_norm_eps: float rope_theta: float = 10000 max_position_embeddings: int = 32768 attention_bias: bool = False use_qk_norm: bool = True rope_scaling: Optional[Dict[str, Union[float, str]]] = None tie_word_embeddings: bool = False head_dim: Optional[int] = None def __post_init__(self): if self.rope_scaling: required_keys = {"alpha", "factor", "type"} if not all(key in self.rope_scaling for key in required_keys): raise ValueError(f"rope_scaling must contain keys {required_keys}") class DynamicNTKAlphaRoPE(nn.Module): def __init__( self, dims: int, base: float = 10000, scaling_alpha: float = 1.0, ): super().__init__() self.dims = dims base = base * scaling_alpha ** (dims / (dims - 2)) self._freqs = base ** (mx.arange(0, self.dims, 2) / self.dims) def __call__(self, x, offset: int = 0): return mx.fast.rope( x, self.dims, traditional=False, base=None, scale=1.0, offset=offset, freqs=self._freqs, ) class Attention(nn.Module): def __init__(self, args: ModelArgs): super().__init__() dim = args.hidden_size self.n_heads = n_heads = args.num_attention_heads self.n_kv_heads = n_kv_heads = args.num_key_value_heads head_dim = ( args.head_dim if args.head_dim is not None else args.hidden_size // n_heads ) self.head_dim = head_dim self.scale = head_dim**-0.5 self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=args.attention_bias) self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.attention_bias) self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.attention_bias) self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=args.attention_bias) self.use_qk_norm = args.use_qk_norm if self.use_qk_norm: self.query_layernorm = nn.RMSNorm(head_dim, args.rms_norm_eps) self.key_layernorm = nn.RMSNorm(head_dim, args.rms_norm_eps) scaling_alpha = 1.0 if args.rope_scaling and "alpha" in args.rope_scaling: scaling_alpha = args.rope_scaling["alpha"] self.rope = DynamicNTKAlphaRoPE( head_dim, base=args.rope_theta, scaling_alpha=scaling_alpha, ) def __call__( self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None, ) -> mx.array: B, L, D = x.shape queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x) queries = queries.reshape(B, L, self.n_heads, self.head_dim).transpose( 0, 2, 1, 3 ) keys = keys.reshape(B, L, self.n_kv_heads, self.head_dim).transpose(0, 2, 1, 3) values = values.reshape(B, L, self.n_kv_heads, self.head_dim).transpose( 0, 2, 1, 3 ) if cache is not None: queries = self.rope(queries, offset=cache.offset) keys = self.rope(keys, offset=cache.offset) else: queries = self.rope(queries) keys = self.rope(keys) if self.use_qk_norm: queries = self.query_layernorm(queries) keys = self.key_layernorm(keys) if cache is not None: keys, values = cache.update_and_fetch(keys, values) output = scaled_dot_product_attention( queries, keys, values, cache=cache, scale=self.scale, mask=mask ) output = output.transpose(0, 2, 1, 3).reshape(B, L, -1) return self.o_proj(output) class MLP(nn.Module): def __init__(self, args: ModelArgs): super().__init__() dim = args.hidden_size hidden_dim = args.intermediate_size self.gate_proj = nn.Linear(dim, hidden_dim, bias=False) self.down_proj = nn.Linear(hidden_dim, dim, bias=False) self.up_proj = nn.Linear(dim, hidden_dim, bias=False) def __call__(self, x) -> mx.array: return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x))) class TransformerBlock(nn.Module): def __init__(self, args: ModelArgs): super().__init__() self.num_attention_heads = args.num_attention_heads self.hidden_size = args.hidden_size self.self_attn = Attention(args) self.mlp = MLP(args) self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) self.post_attention_layernorm = nn.RMSNorm( args.hidden_size, eps=args.rms_norm_eps ) self.args = args def __call__( self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None, ) -> mx.array: r = self.self_attn(self.input_layernorm(x), mask, cache) h = x + r r = self.mlp(self.post_attention_layernorm(h)) out = h + r return out class HunyuanV1DenseModel(nn.Module): def __init__(self, args: ModelArgs): super().__init__() self.args = args self.vocab_size = args.vocab_size self.num_hidden_layers = args.num_hidden_layers self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size) self.layers = [TransformerBlock(args) for _ in range(args.num_hidden_layers)] self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) def __call__( self, inputs: mx.array, cache=None, ): h = self.embed_tokens(inputs) if cache is None: cache = [None] * len(self.layers) mask = create_attention_mask(h, cache[0]) for layer, c in zip(self.layers, cache): h = layer(h, mask, c) return self.norm(h) class Model(nn.Module): def __init__(self, args: ModelArgs): super().__init__() self.args = args self.model_type = args.model_type self.model = HunyuanV1DenseModel(args) if not args.tie_word_embeddings: self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False) def __call__( self, inputs: mx.array, cache=None, ): out = self.model(inputs, cache) if self.args.tie_word_embeddings: return self.model.embed_tokens.as_linear(out) else: return self.lm_head(out) @property def layers(self): return self.model.layers def sanitize(self, weights): if self.args.tie_word_embeddings: weights.pop("lm_head.weight", None) return weights