# Copyright © 2025 Apple Inc. # # mlx-lm port of Upstage Solar Open 2 (250B-A15B hybrid MoE). # # Architecture (ground truth: upstageAI/transformers @ v5.14.1-solar-open2, # src/transformers/models/solar_open2/modeling_solar_open2.py): # * 48 decoder layers; explicit ``gqa_layers`` (0,4,8,...,44) use NoPE # grouped-query full attention, the other 36 use Kimi-Delta-Attention # (KDA, gated delta rule with a per-(head, dim) vector decay). # * Every layer past ``first_k_dense_replace`` is a MoE layer: # 320 routed experts (top-8, sigmoid router with grouped top-k and a # non-trainable e_score_correction_bias) + 1 shared expert. # # Reuse strategy: structure follows mlx_lm/models/kimi_linear.py (same KDA # family) with the Solar-specific gate math (lower-bounded log-decay gate, # optional beta x2 for negative eigenvalues) and a plain NoPE GQA layer # instead of MLA. The recurrent metal kernel / ops fallbacks are imported # from mlx_lm/models/gated_delta.py; the MoE feed-forward is # mlx_lm/models/switch_layers.py:SwitchGLU so QuantizedSwitchGLU conversion # and alis-dwq's SwitchGLU hooks keep working unchanged. from dataclasses import dataclass from functools import partial from typing import Any, Dict, List, Optional, Tuple import mlx.core as mx import mlx.nn as nn from .activations import swiglu from .base import ( BaseModelArgs, create_attention_mask, create_ssm_mask, scaled_dot_product_attention, ) from .cache import ArraysCache, KVCache from .gated_delta import gated_delta_kernel, gated_delta_ops from .rope_utils import initialize_rope from .switch_layers import SwitchGLU @dataclass class ModelArgs(BaseModelArgs): model_type: str vocab_size: int = 196608 hidden_size: int = 4096 num_hidden_layers: int = 48 num_attention_heads: int = 64 num_key_value_heads: int = 8 head_dim: int = 128 intermediate_size: int = 10240 moe_intermediate_size: int = 1280 rms_norm_eps: float = 1e-5 # Kept for checkpoint-config compatibility; unused unless use_rope=True. rope_theta: float = 10000.0 rope_scaling: Optional[Dict[str, Any]] = None partial_rotary_factor: float = 1.0 # was silently dropped by from_dict max_position_embeddings: int = 1048576 tie_word_embeddings: bool = False attention_bias: bool = False # Full-attention (GQA) layer options. use_rope: bool = False # NoPE by default use_qk_norm: bool = False use_gqa_gate: bool = True use_gqa_gate_bias: bool = False gqa_interval: int = 4 gqa_layers: Optional[List[int]] = None # takes priority over gqa_interval # KDA (linear attention) options. linear_attn_config: Optional[Dict[str, Any]] = None kda_use_full_proj: bool = False kda_gate_lower_bound: Optional[float] = -5.0 # HF SolarOpen2Config's *code* default is True (configuration:145; its # docstring wrongly says False); the shipped checkpoint config sets it # explicitly to true, so real-model behavior is unchanged either way. # Match the HF code default. kda_allow_neg_eigval: bool = True # MoE options. n_routed_experts: int = 320 n_shared_experts: int = 1 num_experts_per_tok: int = 8 norm_topk_prob: bool = True routed_scaling_factor: float = 1.0 n_group: int = 1 topk_group: int = 1 first_k_dense_replace: int = 0 hidden_act: str = "silu" # SwitchGLU hard-codes SwiGLU; silu-only port def __post_init__(self): if self.hidden_act != "silu": raise ValueError( f"hidden_act={self.hidden_act!r} is not supported: SwitchGLU " "and SolarMLP hard-code SwiGLU (silu-only port)." ) if self.linear_attn_config is None: self.linear_attn_config = { "short_conv_kernel_size": 4, "head_dim": self.head_dim, "num_heads": self.num_attention_heads, "num_kv_heads": None, } # Per-layer attention pattern, mirroring SolarOpen2Config.__post_init__. if self.gqa_layers is not None: full = set(self.gqa_layers) self.layer_types = [ "full_attention" if i in full else "linear_attention" for i in range(self.num_hidden_layers) ] else: self.layer_types = [ ( "full_attention" if (i + 1) % self.gqa_interval == 0 else "linear_attention" ) for i in range(self.num_hidden_layers) ] # -------------------------------------------------------------------------- # KDA helpers # -------------------------------------------------------------------------- @partial(mx.compile, shapeless=True) def _kda_decay_gate( g_raw: mx.array, A_log: mx.array, dt_bias: mx.array, lower_bound: Optional[float], ) -> mx.array: """Solar Open 2 KDA log-decay gate (reference: ``torch_kda_gate`` + ``fla.ops.kda.gate.fused_kda_gate``). ``g = -exp(A_log) * softplus(g_raw + dt_bias)`` computed in float32, clamped from below at ``lower_bound`` (-5.0 by default), then exponentiated once so the result is the multiplicative per-(head, dim) decay consumed by ``gated_delta_ops`` / ``gated_delta_kernel``. Differs from kimi_linear's ``compute_g``: that variant has no lower-bound clamp because Kimi Linear leaves the gate unclamped. """ g = g_raw.astype(mx.float32) + dt_bias.astype(mx.float32) g = -mx.exp(A_log.astype(mx.float32)) * nn.softplus(g) if lower_bound is not None: g = mx.maximum(g, mx.array(lower_bound, dtype=mx.float32)) return mx.exp(g) @partial(mx.compile, shapeless=True) def _l2norm(x: mx.array, eps: float) -> mx.array: """x * rsqrt(sum(x^2) + eps) computed in float32 (HF l2norm reference).""" xf = x.astype(mx.float32) return xf * mx.rsqrt((xf * xf).sum(axis=-1, keepdims=True) + eps) class SolarRMSNormGated(nn.Module): """Gated RMSNorm at the KDA output: ``RMSNorm(x) * sigmoid(gate)``, computed in float32 like ``SolarOpen2RMSNormGated`` / ``fla.modules.FusedRMSNormGated(activation="sigmoid")``.""" def __init__(self, hidden_size: int, eps: float = 1e-5): super().__init__() self.weight = mx.ones((hidden_size,)) self.eps = eps def __call__(self, x: mx.array, gate: mx.array) -> mx.array: return _gated_rms_norm(x, self.weight, gate, self.eps) @partial(mx.compile, shapeless=True) def _gated_rms_norm(x, weight, gate, eps): xf = x.astype(mx.float32) variance = (xf * xf).mean(axis=-1, keepdims=True) xf = xf * mx.rsqrt(variance + eps) xf = xf * weight.astype(mx.float32) * mx.sigmoid(gate.astype(mx.float32)) return xf.astype(x.dtype) class ShortConv1d(nn.Module): """Depthwise causal short conv (silu, no bias) with a (kernel - 1)-token recurrent state. Identical to kimi_linear.ShortConv1d; kept local so the port is self-contained.""" def __init__(self, channels: int, kernel_size: int): super().__init__() self.kernel_size = kernel_size self.conv = nn.Conv1d( in_channels=channels, out_channels=channels, kernel_size=kernel_size, bias=False, groups=channels, padding=0, ) def __call__( self, x: mx.array, state: Optional[mx.array], mask: Optional[mx.array], lengths: Optional[mx.array], ) -> Tuple[mx.array, mx.array]: if mask is not None: x = mx.where(mask[..., None], x, 0) if state is None: state = mx.zeros( (x.shape[0], self.kernel_size - 1, x.shape[-1]), dtype=x.dtype ) conv_input = mx.concatenate([state, x], axis=1) out = nn.silu(self.conv(conv_input)) n_keep = self.kernel_size - 1 if lengths is not None: ends = mx.clip(lengths, 0, x.shape[1]) positions = (ends[:, None] + mx.arange(n_keep))[..., None] new_state = mx.take_along_axis(conv_input, positions, axis=1) else: new_state = mx.contiguous(conv_input[:, -n_keep:, :]) return out, new_state class SolarDeltaAttention(nn.Module): """Kimi-Delta-Attention layer as configured for Solar Open 2. Differences vs kimi_linear.KimiDeltaAttention: * optional GQA-style kv heads (``linear_attn_config.num_kv_heads``; Solar Open 2 sets it to null -> equal to num_heads), * ``kda_use_full_proj`` switches between full-rank ``f_proj``/``g_proj`` and the factored low-rank ``f_a_proj``/``f_b_proj`` / ``g_a_proj``/``g_b_proj`` (Solar Open 2 uses the factored form), * the decay gate is lower-bounded (``kda_gate_lower_bound=-5.0``), * ``kda_allow_neg_eigval`` scales beta by 2 (beta in [0, 2]). """ def __init__(self, args: ModelArgs, layer_idx: int): super().__init__() cfg = args.linear_attn_config self.layer_idx = layer_idx self.num_heads = cfg["num_heads"] self.head_dim = cfg["head_dim"] self.num_kv_heads = cfg.get("num_kv_heads") or self.num_heads self.n_rep = self.num_heads // self.num_kv_heads self.conv_kernel = cfg.get("short_conv_kernel_size", 4) self.use_full_proj = args.kda_use_full_proj self.gate_lower_bound = args.kda_gate_lower_bound self.allow_neg_eigval = args.kda_allow_neg_eigval self.projection_dim = self.num_heads * self.head_dim self.kv_projection_dim = self.num_kv_heads * self.head_dim hidden = args.hidden_size self.scale = float(self.head_dim) ** -0.5 self.q_proj = nn.Linear(hidden, self.projection_dim, bias=False) self.k_proj = nn.Linear(hidden, self.kv_projection_dim, bias=False) self.v_proj = nn.Linear(hidden, self.kv_projection_dim, bias=False) self.b_proj = nn.Linear(hidden, self.num_heads, bias=False) self.q_conv = ShortConv1d(self.projection_dim, self.conv_kernel) self.k_conv = ShortConv1d(self.kv_projection_dim, self.conv_kernel) self.v_conv = ShortConv1d(self.kv_projection_dim, self.conv_kernel) if self.use_full_proj: self.f_proj = nn.Linear(hidden, self.projection_dim, bias=False) self.g_proj = nn.Linear(hidden, self.projection_dim, bias=False) else: self.f_a_proj = nn.Linear(hidden, self.head_dim, bias=False) self.f_b_proj = nn.Linear(self.head_dim, self.projection_dim, bias=False) self.g_a_proj = nn.Linear(hidden, self.head_dim, bias=False) self.g_b_proj = nn.Linear(self.head_dim, self.projection_dim, bias=False) # Shapes match the HF checkpoint: A_log (1, 1, H, 1), dt_bias (H*Dh,). self.A_log = mx.expand_dims( mx.log(mx.random.uniform(low=1.0, high=16.0, shape=(self.num_heads,))), (0, 1, 3), ) self.dt_bias = mx.ones((self.projection_dim,)) self.o_norm = SolarRMSNormGated(self.head_dim, eps=args.rms_norm_eps) self.o_proj = nn.Linear(self.projection_dim, hidden, bias=False) def __call__( self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None, ) -> mx.array: B, T, _ = x.shape dtype = x.dtype if cache is not None: q_state, k_state, v_state, ssm_state = cache lengths = cache.lengths else: q_state = k_state = v_state = ssm_state = None lengths = None q_conv, q_state = self.q_conv(self.q_proj(x), q_state, mask, lengths) k_conv, k_state = self.k_conv(self.k_proj(x), k_state, mask, lengths) v_conv, v_state = self.v_conv(self.v_proj(x), v_state, mask, lengths) if cache is not None: cache[0] = q_state cache[1] = k_state cache[2] = v_state q = q_conv.reshape(B, T, self.num_heads, self.head_dim) k = k_conv.reshape(B, T, self.num_kv_heads, self.head_dim) v = v_conv.reshape(B, T, self.num_kv_heads, self.head_dim) # l2-normalize q/k exactly like the reference # (use_qk_l2norm_in_kernel=True): x * rsqrt(sum(x^2) + eps) in f32. # Folding this through mx.fast.rms_norm is NOT equivalent: rms_norm # puts eps under the mean (i.e. D*eps on the sum), which measurably # shifts small-norm head vectors (validated on real weights, see # VALIDATION_RESULTS.md). q/k/v go into the recurrence in f32, # mirroring the f32 casts in torch_recurrent_kda. q = _l2norm(q, 1e-6) * self.scale k = _l2norm(k, 1e-6) v = v.astype(mx.float32) if self.use_full_proj: g_raw = self.f_proj(x) else: g_raw = self.f_b_proj(self.f_a_proj(x)) g_raw = g_raw.reshape(B, T, self.num_heads, self.head_dim) beta = mx.sigmoid(self.b_proj(x)) if self.allow_neg_eigval: beta = beta * 2.0 g = _kda_decay_gate( g_raw, self.A_log.reshape(self.num_heads, 1), self.dt_bias.reshape(self.num_heads, self.head_dim), self.gate_lower_bound, ) if ssm_state is None: ssm_state = mx.zeros( (B, self.num_kv_heads, self.head_dim, self.head_dim), dtype=mx.float32, ) if ( self.training or mx.default_device() != mx.gpu or not mx.metal.is_available() ): out, ssm_state = gated_delta_ops(q, k, v, g, beta, ssm_state, mask) else: out, ssm_state = gated_delta_kernel(q, k, v, g, beta, ssm_state, mask) if cache is not None: cache[3] = ssm_state cache.advance(T) if self.use_full_proj: g_out = self.g_proj(x) else: g_out = self.g_b_proj(self.g_a_proj(x)) g_out = g_out.reshape(B, T, self.num_heads, self.head_dim) out = self.o_norm(out.reshape(B, T, self.num_heads, self.head_dim), g_out) out = out.reshape(B, T, -1).astype(dtype) return self.o_proj(out) class SolarFullAttention(nn.Module): """Grouped-query full attention: NoPE by default (use_rope=False), optional per-head qk RMSNorm (off in Solar Open 2), and a sigmoid output gate (``use_gqa_gate=True``: ``attn_out * sigmoid(g_proj(x))``).""" def __init__(self, args: ModelArgs, layer_idx: int): super().__init__() self.layer_idx = layer_idx self.num_heads = args.num_attention_heads self.num_key_value_heads = args.num_key_value_heads self.head_dim = args.head_dim self.scale = float(self.head_dim) ** -0.5 self.use_gqa_gate = args.use_gqa_gate self.use_qk_norm = args.use_qk_norm self.use_rope = args.use_rope hidden = args.hidden_size self.q_proj = nn.Linear( hidden, self.num_heads * self.head_dim, bias=args.attention_bias ) self.k_proj = nn.Linear( hidden, self.num_key_value_heads * self.head_dim, bias=args.attention_bias ) self.v_proj = nn.Linear( hidden, self.num_key_value_heads * self.head_dim, bias=args.attention_bias ) self.o_proj = nn.Linear(self.num_heads * self.head_dim, hidden, bias=False) if self.use_qk_norm: self.q_norm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps) self.k_norm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps) if self.use_rope: self.rope = initialize_rope( int(self.head_dim * args.partial_rotary_factor), base=args.rope_theta, traditional=False, scaling_config=args.rope_scaling, max_position_embeddings=args.max_position_embeddings, ) else: self.rope = None if self.use_gqa_gate: self.g_proj = nn.Linear( hidden, self.num_heads * self.head_dim, bias=args.use_gqa_gate_bias, ) def __call__( self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[KVCache] = None, ) -> mx.array: B, L, _ = x.shape q = self.q_proj(x).reshape(B, L, self.num_heads, self.head_dim) k = self.k_proj(x).reshape(B, L, self.num_key_value_heads, self.head_dim) v = self.v_proj(x).reshape(B, L, self.num_key_value_heads, self.head_dim) if self.use_qk_norm: q = self.q_norm(q) k = self.k_norm(k) q = q.transpose(0, 2, 1, 3) k = k.transpose(0, 2, 1, 3) v = v.transpose(0, 2, 1, 3) if self.rope is not None: offset = cache.offset if cache is not None else 0 q = self.rope(q, offset=offset) k = self.rope(k, offset=offset) if cache is not None: k, v = cache.update_and_fetch(k, v) out = scaled_dot_product_attention(q, k, v, cache=cache, scale=self.scale, mask=mask) out = out.transpose(0, 2, 1, 3).reshape(B, L, -1) if self.use_gqa_gate: out = out * mx.sigmoid(self.g_proj(x)) return self.o_proj(out) # -------------------------------------------------------------------------- # MoE # -------------------------------------------------------------------------- class SolarMLP(nn.Module): def __init__( self, args: ModelArgs, hidden_size: Optional[int] = None, intermediate_size: Optional[int] = None, ): super().__init__() dim = hidden_size or args.hidden_size hidden = intermediate_size or args.intermediate_size self.gate_proj = nn.Linear(dim, hidden, bias=False) self.up_proj = nn.Linear(dim, hidden, bias=False) self.down_proj = nn.Linear(hidden, dim, bias=False) def __call__(self, x: mx.array) -> mx.array: return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x))) @mx.compile def _expert_select( gates: mx.array, bias: Optional[mx.array], top_k: int, n_group: int, topk_group: int, routed_scaling_factor: float, renormalize: bool, ) -> Tuple[mx.array, mx.array]: """Sigmoid router with optional grouped top-k (SolarOpen2TopkRouter). Selection uses bias-adjusted scores; the returned weights are the raw sigmoid scores at the selected experts (bias only affects *which* experts are picked), optionally renormalized (``norm_topk_prob``) and scaled by ``routed_scaling_factor``. """ scores = mx.sigmoid(gates) orig_scores = scores if bias is not None: scores = scores + bias.astype(scores.dtype) if n_group > 1: scores = mx.unflatten(scores, axis=-1, shape=(n_group, -1)) group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True) k = n_group - topk_group group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :] # HF masks dropped groups with -inf (masked_fill), not 0.0: with a # negative e_score_correction_bias an unmasked adjusted score can be # < 0, and a 0.0 mask would then pick a MASKED expert HF never picks # (attacks/attack6 A6a). scores = mx.put_along_axis( scores, mx.stop_gradient(group_idx), mx.array(-float("inf"), dtype=scores.dtype), axis=-2, ) scores = mx.flatten(scores, -2, -1) inds = mx.argpartition(-scores, kth=top_k - 1, axis=-1)[..., :top_k] weights = mx.take_along_axis(orig_scores, inds, axis=-1) if top_k > 1 and renormalize: weights = weights / (weights.sum(axis=-1, keepdims=True) + 1e-20) return inds, weights * routed_scaling_factor class SolarSparseMoE(nn.Module): """Routed MoE + shared expert. Naming contract (kept identical to kimi_linear so alis-dwq's hooks match): * ``gate`` -- router Linear; module path ``...mlp.gate`` matches alis-dwq's ``(?:^|\.)(?:gate|router)$`` router regex * ``switch_mlp`` -- SwitchGLU instance (expert_traffic hooks the class) * ``shared_experts`` * ``e_score_correction_bias`` lives on this module (not on ``gate``) so an 8-bit quantization of ``gate`` cannot drop it; ``sanitize`` moves the HF key ``...mlp.gate.e_score_correction_bias`` here. """ def __init__(self, args: ModelArgs): super().__init__() self.args = args hidden = args.hidden_size experts = args.n_routed_experts self.gate = nn.Linear(hidden, experts, bias=False) self.switch_mlp = SwitchGLU(hidden, args.moe_intermediate_size, experts) self.e_score_correction_bias = mx.zeros((experts,), dtype=mx.float32) if args.n_shared_experts: shared_hidden = args.moe_intermediate_size * args.n_shared_experts self.shared_experts = SolarMLP(args, intermediate_size=shared_hidden) else: self.shared_experts = None def __call__(self, x: mx.array) -> mx.array: # HF computes the router logits in float32: # F.linear(hidden_states.float(), gate.weight.float()) # Mirror that exactly — a bf16 matmul's reduction-order noise flips # razor-tie routing decisions (attacks/attack1b: 7/64 tokens flipped # vs HF f32 logits; both-sides-f32 gives 0/64). A quantized router # (packed uint32 weight) cannot upcast, so that path keeps the # activation-dtype matmul with an f32 cast of the logits. if hasattr(self.gate, "bits"): gates = self.gate(x).astype(mx.float32) else: gates = x.astype(mx.float32) @ self.gate.weight.astype(mx.float32).T inds, weights = _expert_select( gates, self.e_score_correction_bias, self.args.num_experts_per_tok, self.args.n_group, self.args.topk_group, self.args.routed_scaling_factor, self.args.norm_topk_prob, ) out = self.switch_mlp(x, inds) # HF casts the weighted expert sum back to the activation dtype. out = (out * weights[..., None]).sum(axis=-2).astype(x.dtype) if self.shared_experts is not None: out = out + self.shared_experts(x) return out class SolarDecoderLayer(nn.Module): def __init__(self, args: ModelArgs, layer_idx: int): super().__init__() self.is_linear = args.layer_types[layer_idx] == "linear_attention" if self.is_linear: self.self_attn = SolarDeltaAttention(args, layer_idx) else: self.self_attn = SolarFullAttention(args, layer_idx) if layer_idx >= args.first_k_dense_replace: self.mlp = SolarSparseMoE(args) else: self.mlp = SolarMLP(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 ) def __call__( self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None, ) -> mx.array: y = self.self_attn(self.input_layernorm(x), mask, cache) h = x + y z = self.mlp(self.post_attention_layernorm(h)) return h + z class SolarOpen2Model(nn.Module): def __init__(self, args: ModelArgs): super().__init__() self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size) self.layers = [SolarDecoderLayer(args, i) for i in range(args.num_hidden_layers)] self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) self.ssm_idx = None self.attn_idx = None for i, layer in enumerate(self.layers): if layer.is_linear: self.ssm_idx = i break for i, layer in enumerate(self.layers): if not layer.is_linear: self.attn_idx = i break def __call__( self, inputs: mx.array, cache: Optional[List[Any]] = None, ) -> mx.array: h = self.embed_tokens(inputs) if cache is None: cache = [None] * len(self.layers) ssm_mask = ( create_ssm_mask(h, cache[self.ssm_idx]) if self.ssm_idx is not None else None ) attn_mask = ( create_attention_mask(h, cache[self.attn_idx]) if self.attn_idx is not None else None ) for layer, layer_cache in zip(self.layers, cache): mask = ssm_mask if layer.is_linear else attn_mask h = layer(h, mask=mask, cache=layer_cache) return self.norm(h) def sanitize_layer_weights( weights: Dict[str, mx.array], prefix: str, layer: "SolarDecoderLayer", num_experts: int, ) -> Dict[str, mx.array]: """In-place HF -> MLX key conversion for one decoder layer. Handles both on-disk expert layouts: * per-expert tensors (the actual upstage/Solar-Open2-250B checkpoint): ``{prefix}.mlp.experts.{i}.{gate,up,down}_proj.weight`` -> stacked ``switch_mlp.{gate,up,down}_proj.weight``; * fused 3D tensors (module layout of the HF modeling code): ``experts.gate_up_proj`` (E, 2I, H) split into gate/up, and ``experts.down_proj`` renamed. """ if isinstance(layer.mlp, SolarSparseMoE): mlp_prefix = f"{prefix}.mlp" if f"{mlp_prefix}.experts.0.gate_proj.weight" in weights: # Per-expert layout: refuse silently-misconverted pre-quantized # tensors (scales/biases companions would be left behind and the # stacked uint32 "weights" would be nonsense). if any( k.startswith(f"{mlp_prefix}.experts.") and (k.endswith(".scales") or k.endswith(".biases")) for k in weights ): raise NotImplementedError( "Pre-quantized HF expert tensors are not supported by " "this port; load the BF16 checkpoint and quantize " "with mlx-lm instead." ) for name in ("gate_proj", "up_proj", "down_proj"): stacked = [ weights.pop(f"{mlp_prefix}.experts.{i}.{name}.weight") for i in range(num_experts) ] weights[f"{mlp_prefix}.switch_mlp.{name}.weight"] = mx.stack(stacked) else: if f"{mlp_prefix}.experts.gate_up_proj.scales" in weights: raise NotImplementedError( "Pre-quantized HF expert tensors are not supported by " "this port; load the BF16 checkpoint and quantize " "with mlx-lm instead." ) gu_key = f"{mlp_prefix}.experts.gate_up_proj" if gu_key in weights: gu = weights.pop(gu_key) # mlx#3836: mx.split silently corrupts >2**31-element tensors; # a fused HF-layout gate_up_proj (E, 2I, H) is 1.56x that for # this model. Use strided slices instead. half = gu.shape[1] // 2 gate_w, up_w = gu[:, :half], gu[:, half:] weights[f"{mlp_prefix}.switch_mlp.gate_proj.weight"] = mx.contiguous( gate_w ) weights[f"{mlp_prefix}.switch_mlp.up_proj.weight"] = mx.contiguous( up_w ) down_key = f"{mlp_prefix}.experts.down_proj" if down_key in weights: weights[f"{mlp_prefix}.switch_mlp.down_proj.weight"] = weights.pop( down_key ) # Router bias lives on the MoE block in MLX (see class note). bias_key = f"{mlp_prefix}.gate.e_score_correction_bias" if bias_key in weights: weights[f"{mlp_prefix}.e_score_correction_bias"] = weights.pop(bias_key) attn = layer.self_attn if isinstance(attn, SolarDeltaAttention): attn_prefix = f"{prefix}.self_attn" for src_name, dst_name in ( ("q_conv1d", "q_conv"), ("k_conv1d", "k_conv"), ("v_conv1d", "v_conv"), ): src_key = f"{attn_prefix}.{src_name}.weight" if src_key in weights: w = weights.pop(src_key) if w.ndim == 3: # HF depthwise conv weight (C, 1, K) -> # MLX Conv1d weight (C, K, 1) w = w.moveaxis(2, 1) weights[f"{attn_prefix}.{dst_name}.conv.weight"] = w dt_key = f"{attn_prefix}.dt_bias" if dt_key in weights and weights[dt_key].ndim > 1: weights[dt_key] = mx.reshape(weights[dt_key], (-1,)) return weights class Model(nn.Module): def __init__(self, args: ModelArgs): super().__init__() self.args = args self.model_type = args.model_type self.model = SolarOpen2Model(args) if args.tie_word_embeddings: self.lm_head = None else: self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False) def __call__( self, inputs: mx.array, cache: Optional[List[Any]] = None, ) -> mx.array: out = self.model(inputs, cache) if self.lm_head is None: return self.model.embed_tokens.as_linear(out) return self.lm_head(out) @property def layers(self): return self.model.layers def make_cache(self): caches: List[Any] = [] for layer in self.layers: if layer.is_linear: # [q_conv_state, k_conv_state, v_conv_state, recurrent_state] caches.append(ArraysCache(size=4)) else: caches.append(KVCache()) return caches def sanitize(self, weights: Dict[str, mx.array]) -> Dict[str, mx.array]: # Drop non-persistent HF rope buffers if a checkpoint carries them. weights = {k: v for k, v in weights.items() if "rotary_emb.inv_freq" not in k} if self.args.tie_word_embeddings: weights.pop("lm_head.weight", None) for layer_idx, layer in enumerate(self.layers): sanitize_layer_weights( weights, f"model.layers.{layer_idx}", layer, self.args.n_routed_experts, ) return weights @property def cast_predicate(self): def predicate(path: str): if "e_score_correction_bias" in path: return False if path.endswith("A_log") or path.endswith("dt_bias"): return False return True return predicate def quant_predicate_builder(self, group_size: int, bits: int): """Per-tensor bit assignment for Solar-Open2. The CLI global ``bits``/``group_size`` apply to the routed experts only (they dominate the parameter count, so they alone decide the model size). Everything else follows the runbook intent — sensitive paths run one tier up: - router (``mlp.gate``) → 8b/g64: routing is selection-sensitive and the tensor is tiny (320 x 4096 per layer). - KDA low-rank chokepoints (``f_a/f_b/g_a/g_b/b_proj``) → 8b/g64: these feed the recurrent state, so quantization error compounds across time steps. - ``embed_tokens`` / ``lm_head`` → 6b/g64 (8b when the global is 8). - attention q/k/v/o, GQA elementwise gate (``g_proj``), and the shared expert → global bits + 2 (capped at 8). At bits=8/g64 this reproduces the uniform 8-bit map, so the recipe is a no-op for the reference build. """ attn_bits = min(bits + 2, 8) embed_bits = max(6, min(bits, 8)) chokepoint = ("f_a_proj", "f_b_proj", "g_a_proj", "g_b_proj") upper = ("q_proj", "k_proj", "v_proj", "o_proj", "g_proj", "shared_experts") def predicate(path, _): if path.endswith("mlp.gate"): return {"group_size": 64, "bits": 8} if "switch_mlp" in path: return True # routed experts take the global CLI values if any(k in path for k in chokepoint) or path.endswith("b_proj"): return {"group_size": 64, "bits": 8} if "embed_tokens" in path or "lm_head" in path: return {"group_size": 64, "bits": embed_bits} if any(k in path for k in upper): return {"group_size": 64, "bits": attn_bits} return True return predicate