*That’s not much code. I could blog that much code.<\/p>\n
https:\/\/github.com\/ml-explore\/mlx-examples\/blob\/main\/llms\/mlx_lm\/models\/deepseek_v3.py<\/p>\n
# Copyright \u00a9 2024 Apple Inc.<\/p>\n
import math
\nfrom dataclasses import dataclass
\nfrom typing import Any, Dict, Optional, Tuple<\/p>\n
import mlx.core as mx
\nimport mlx.nn as nn<\/p>\n
from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
\nfrom .switch_layers import SwitchGLU<\/p>\n
@dataclass
\nclass ModelArgs(BaseModelArgs):
\n model_type: str = “deepseek_v3”
\n vocab_size: int = 102400
\n hidden_size: int = 4096
\n intermediate_size: int = 11008
\n moe_intermediate_size: int = 1407
\n num_hidden_layers: int = 30
\n num_attention_heads: int = 32
\n num_key_value_heads: int = 32
\n n_shared_experts: Optional[int] = None
\n n_routed_experts: Optional[int] = None
\n routed_scaling_factor: float = 1.0
\n kv_lora_rank: int = 512
\n q_lora_rank: int = 1536
\n qk_rope_head_dim: int = 64
\n v_head_dim: int = 128
\n qk_nope_head_dim: int = 128
\n topk_method: str = “noaux_tc”
\n scoring_func: str = “sigmoid”
\n norm_topk_prob: bool = True
\n n_group: Optional[int] = None
\n topk_group: Optional[int] = None
\n num_experts_per_tok: Optional[int] = None
\n moe_layer_freq: int = 1
\n first_k_dense_replace: int = 0
\n max_position_embeddings: int = 2048
\n rms_norm_eps: float = 1e-6
\n rope_theta: float = 10000.0
\n rope_scaling: Dict = None
\n attention_bias: bool = False<\/p>\n
def yarn_find_correction_dim(
\n num_rotations, dim, base=10000, max_position_embeddings=2048
\n):
\n return (dim * math.log(max_position_embeddings \/ (num_rotations * 2 * math.pi))) \/ (
\n 2 * math.log(base)
\n )<\/p>\n
def yarn_find_correction_range(
\n low_rot, high_rot, dim, base=10000, max_position_embeddings=2048
\n):
\n low = math.floor(
\n yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings)
\n )
\n high = math.ceil(
\n yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings)
\n )
\n return max(low, 0), min(high, dim – 1)<\/p>\n
def yarn_get_mscale(scale=1, mscale=1):
\n if scale <= 1:\n return 1.0\n return 0.1 * mscale * math.log(scale) + 1.0\n\n\ndef yarn_linear_ramp_mask(min_val, max_val, dim):\n if min_val == max_val:\n max_val += 0.001 # Prevent singularity\n\n linear_func = (mx.arange(dim, dtype=mx.float32) - min_val) \/ (max_val - min_val)\n return mx.clip(linear_func, 0, 1)\n\n\nclass DeepseekV3YarnRotaryEmbedding(nn.Module):\n def __init__(\n self,\n dim,\n max_position_embeddings=2048,\n base=10000,\n scaling_factor=1.0,\n original_max_position_embeddings=4096,\n beta_fast=32,\n beta_slow=1,\n mscale=1,\n mscale_all_dim=0,\n ):\n super().__init__()\n self.mscale = yarn_get_mscale(scaling_factor, mscale) \/ yarn_get_mscale(\n scaling_factor, mscale_all_dim\n )\n freq_extra = base ** (mx.arange(0, dim, 2, dtype=mx.float32) \/ dim)\n freq_inter = scaling_factor * base ** (\n mx.arange(0, dim, 2, dtype=mx.float32) \/ dim\n )\n low, high = yarn_find_correction_range(\n beta_fast,\n beta_slow,\n dim,\n base,\n original_max_position_embeddings,\n )\n freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dim \/\/ 2)\n self._freqs = (freq_inter * freq_extra) \/ (\n freq_inter * freq_mask + freq_extra * (1 - freq_mask)\n )\n\n def __call__(self, x, offset=0):\n if self.mscale != 1.0:\n x = self.mscale * x\n return mx.fast.rope(\n x,\n x.shape[-1],\n traditional=True,\n base=None,\n scale=1.0,\n offset=offset,\n freqs=self._freqs,\n )\n\n\nclass DeepseekV3Attention(nn.Module):\n def __init__(self, config: ModelArgs):\n super().__init__()\n self.config = config\n self.hidden_size = config.hidden_size\n self.num_heads = config.num_attention_heads\n self.max_position_embeddings = config.max_position_embeddings\n self.rope_theta = config.rope_theta\n self.q_lora_rank = config.q_lora_rank\n self.qk_rope_head_dim = config.qk_rope_head_dim\n self.kv_lora_rank = config.kv_lora_rank\n self.v_head_dim = config.v_head_dim\n self.qk_nope_head_dim = config.qk_nope_head_dim\n self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim\n\n self.scale = self.q_head_dim**-0.5\n\n if self.q_lora_rank is None:\n self.q_proj = nn.Linear(\n self.hidden_size, self.num_heads * self.q_head_dim, bias=False\n )\n else:\n self.q_a_proj = nn.Linear(\n self.hidden_size, self.q_lora_rank, bias=config.attention_bias\n )\n self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank)\n self.q_b_proj = nn.Linear(\n self.q_lora_rank, self.num_heads * self.q_head_dim, bias=False\n )\n\n self.kv_a_proj_with_mqa = nn.Linear(\n self.hidden_size,\n self.kv_lora_rank + self.qk_rope_head_dim,\n bias=config.attention_bias,\n )\n self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank)\n self.kv_b_proj = nn.Linear(\n self.kv_lora_rank,\n self.num_heads\n * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),\n bias=False,\n )\n\n self.o_proj = nn.Linear(\n self.num_heads * self.v_head_dim,\n self.hidden_size,\n bias=config.attention_bias,\n )\n\n mscale_all_dim = self.config.rope_scaling.get(\"mscale_all_dim\", 0)\n scaling_factor = self.config.rope_scaling[\"factor\"]\n if mscale_all_dim:\n mscale = yarn_get_mscale(scaling_factor, mscale_all_dim)\n self.scale = self.scale * mscale * mscale\n\n rope_kwargs = {\n key: self.config.rope_scaling[key]\n for key in [\n \"original_max_position_embeddings\",\n \"beta_fast\",\n \"beta_slow\",\n \"mscale\",\n \"mscale_all_dim\",\n ]\n if key in self.config.rope_scaling\n }\n self.rope = DeepseekV3YarnRotaryEmbedding(\n dim=self.qk_rope_head_dim,\n max_position_embeddings=self.max_position_embeddings,\n scaling_factor=scaling_factor,\n base=self.rope_theta,\n **rope_kwargs,\n )\n\n def __call__(\n self,\n x: mx.array,\n mask: Optional[mx.array] = None,\n cache: Optional[Any] = None,\n ) -> mx.array:
\n B, L, D = x.shape<\/p>\n
if self.q_lora_rank is None:
\n q = self.q_proj(x)
\n else:
\n q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x)))<\/p>\n
q = q.reshape(B, L, self.num_heads, self.q_head_dim).transpose(0, 2, 1, 3)
\n q_nope, q_pe = mx.split(q, [self.qk_nope_head_dim], axis=-1)
\n compressed_kv = self.kv_a_proj_with_mqa(x)
\n compressed_kv, k_pe = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
\n k_pe = k_pe.reshape(B, L, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3)
\n kv = self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
\n kv = kv.reshape(B, L, self.num_heads, -1).transpose(0, 2, 1, 3)<\/p>\n
k_nope, values = mx.split(kv, [self.qk_nope_head_dim], axis=-1)<\/p>\n
if cache is not None:
\n q_pe = self.rope(q_pe, cache.offset)
\n k_pe = self.rope(k_pe, cache.offset)
\n k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
\n keys, values = cache.update_and_fetch(
\n mx.concatenate([k_nope, k_pe], axis=-1), values
\n )
\n else:
\n q_pe = self.rope(q_pe)
\n k_pe = self.rope(k_pe)
\n k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
\n keys = mx.concatenate([k_nope, k_pe], axis=-1)<\/p>\n
queries = mx.concatenate([q_nope, q_pe], axis=-1)<\/p>\n
output = scaled_dot_product_attention(
\n queries, keys, values, cache=cache, scale=self.scale, mask=mask
\n )
\n output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
\n return self.o_proj(output)<\/p>\n
class DeepseekV3MLP(nn.Module):
\n def __init__(
\n self, config: ModelArgs, hidden_size: int = None, intermediate_size: int = None
\n ):
\n super().__init__()
\n self.config = config
\n self.hidden_size = config.hidden_size if hidden_size is None else hidden_size
\n self.intermediate_size = (
\n config.intermediate_size if intermediate_size is None else intermediate_size
\n )<\/p>\n
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
\n self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
\n self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)<\/p>\n
def __call__(self, x):
\n down_proj = self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
\n return down_proj<\/p>\n
class MoEGate(nn.Module):
\n def __init__(self, config: ModelArgs):
\n super().__init__()
\n self.config = config
\n self.top_k = config.num_experts_per_tok
\n self.norm_topk_prob = config.norm_topk_prob
\n self.n_routed_experts = config.n_routed_experts
\n self.routed_scaling_factor = config.routed_scaling_factor
\n self.topk_method = config.topk_method
\n self.n_group = config.n_group
\n self.topk_group = config.topk_group
\n self.weight = mx.zeros((self.n_routed_experts, config.hidden_size))
\n self.e_score_correction_bias = mx.zeros((self.n_routed_experts,))<\/p>\n
def __call__(self, x):
\n gates = x @ self.weight.T<\/p>\n
scores = mx.sigmoid(gates.astype(mx.float32))<\/p>\n
assert self.topk_method == “noaux_tc”, “Unsupported topk method.”
\n bsz, seq_len = x.shape[:2]
\n scores = scores + self.e_score_correction_bias
\n scores = scores.reshape(bsz, seq_len, self.n_group, -1)
\n group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1)
\n k = self.n_group – self.topk_group
\n group_idx = mx.argpartition(group_scores, kth=k – 1, axis=-1)[…, :k]
\n batch_idx = mx.expand_dims(mx.arange(bsz), (1, 2))
\n seq_idx = mx.expand_dims(mx.arange(seq_len), (0, 2))
\n scores[batch_idx, seq_idx, group_idx] = 0.0
\n scores = scores.reshape(bsz, seq_len, -1)<\/p>\n
k = self.top_k
\n inds = mx.argpartition(-scores, kth=k – 1, axis=-1)[…, :k]
\n scores = mx.take_along_axis(scores, inds, axis=-1)
\n if self.top_k > 1 and self.norm_topk_prob:
\n denominator = scores.sum(axis=-1, keepdims=True) + 1e-20
\n scores = scores \/ denominator
\n scores = scores * self.routed_scaling_factor<\/p>\n
return inds, scores<\/p>\n
class DeepseekV3MoE(nn.Module):
\n def __init__(self, config: ModelArgs):
\n super().__init__()
\n self.config = config
\n self.num_experts_per_tok = config.num_experts_per_tok
\n self.switch_mlp = SwitchGLU(
\n config.hidden_size, config.moe_intermediate_size, config.n_routed_experts
\n )<\/p>\n
self.gate = MoEGate(config)
\n if config.n_shared_experts is not None:
\n intermediate_size = config.moe_intermediate_size * config.n_shared_experts
\n self.shared_experts = DeepseekV3MLP(
\n config=config, intermediate_size=intermediate_size
\n )<\/p>\n
def __call__(self, x):
\n inds, scores = self.gate(x)
\n y = self.switch_mlp(x, inds)
\n y = (y * scores[…, None]).sum(axis=-2).astype(y.dtype)
\n if self.config.n_shared_experts is not None:
\n y = y + self.shared_experts(x)<\/p>\n
return y<\/p>\n
class DeepseekV3DecoderLayer(nn.Module):
\n def __init__(self, config: ModelArgs, layer_idx: int):
\n super().__init__()
\n self.self_attn = DeepseekV3Attention(config)
\n self.mlp = (
\n DeepseekV3MoE(config)
\n if (
\n config.n_routed_experts is not None
\n and layer_idx >= config.first_k_dense_replace
\n and layer_idx % config.moe_layer_freq == 0
\n )
\n else DeepseekV3MLP(config)
\n )
\n self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
\n self.post_attention_layernorm = nn.RMSNorm(
\n config.hidden_size, eps=config.rms_norm_eps
\n )<\/p>\n
def __call__(
\n self,
\n x: mx.array,
\n mask: Optional[mx.array] = None,
\n cache: Optional[Any] = None,
\n ) -> mx.array:
\n r = self.self_attn(self.input_layernorm(x), mask, cache)
\n h = x + r
\n r = self.mlp(self.post_attention_layernorm(h))
\n out = h + r
\n # Protect against overflow for fp16
\n if out.dtype == mx.float16:
\n out = mx.clip(out, a_min=None, a_max=mx.finfo(mx.float16).max – 1000)
\n return out<\/p>\n
class DeepseekV3Model(nn.Module):
\n def __init__(self, config: ModelArgs):
\n super().__init__()
\n self.vocab_size = config.vocab_size
\n self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
\n self.layers = [
\n DeepseekV3DecoderLayer(config, idx)
\n for idx in range(config.num_hidden_layers)
\n ]
\n self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
\n self.pipeline_rank = 0
\n self.pipeline_size = 1<\/p>\n
def pipeline(self, group):
\n # Split layers in reverse so rank=0 gets the last layers and
\n # rank=pipeline_size-1 gets the first
\n self.pipeline_rank = group.rank()
\n self.pipeline_size = group.size()
\n layers_per_rank = (
\n len(self.layers) + self.pipeline_size – 1
\n ) \/\/ self.pipeline_size
\n start = (self.pipeline_size – self.pipeline_rank – 1) * layers_per_rank
\n self.layers = self.layers[start : start + layers_per_rank]<\/p>\n
def __call__(
\n self,
\n x: mx.array,
\n cache: Optional[Any] = None,
\n mask: Optional[mx.array] = None,
\n ) -> mx.array:
\n h = self.embed_tokens(x)<\/p>\n
pipeline_rank = self.pipeline_rank
\n pipeline_size = self.pipeline_size
\n # Hack to avoid time-outs during prompt-processing
\n dist_stream = mx.cpu if h.shape[1] > 1 else mx.gpu
\n if mask is None:
\n mask = create_attention_mask(h, cache)<\/p>\n
if cache is None:
\n cache = [None] * len(self.layers)<\/p>\n
# Receive from the previous process in the pipeline<\/p>\n
if pipeline_rank < pipeline_size - 1:\n h = mx.distributed.recv_like(h, (pipeline_rank + 1), stream=dist_stream)\n\n for layer, c in zip(self.layers, cache):\n h = layer(h, mask, c)\n\n # Send to the next process in the pipeline\n if pipeline_rank != 0:\n h = mx.distributed.send(\n h, (pipeline_rank - 1) % pipeline_size, stream=dist_stream\n )\n\n # Broadcast h while keeping it in the graph\n h = mx.distributed.all_gather(h, stream=dist_stream)[: h.shape[0]]\n\n return self.norm(h)\n\n\nclass Model(nn.Module):\n def __init__(self, config: ModelArgs):\n super().__init__()\n self.args = config\n self.model_type = config.model_type\n self.model = DeepseekV3Model(config)\n self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)\n\n def __call__(\n self,\n inputs: mx.array,\n cache: Optional[Any] = None,\n mask: Optional[mx.array] = None,\n ):\n out = self.model(inputs, cache, mask)\n return self.lm_head(out)\n\n def sanitize(self, weights):\n for l in range(self.args.num_hidden_layers):\n prefix = f\"model.layers.{l}\"\n for n, m in [(\"w1\", \"gate_proj\"), (\"w2\", \"down_proj\"), (\"w3\", \"up_proj\")]:\n for k in [\"weight\", \"scales\", \"biases\"]:\n if f\"{prefix}.mlp.experts.0.{m}.{k}\" in weights:\n to_join = [\n weights.pop(f\"{prefix}.mlp.experts.{e}.{m}.{k}\")\n for e in range(self.args.n_routed_experts)\n ]\n weights[f\"{prefix}.mlp.switch_mlp.{m}.{k}\"] = mx.stack(to_join)\n\n # Remove multi-token prediction layer\n return {k: v for k, v in weights.items() if not k.startswith(\"model.layers.61\")}\n\n @property\n def layers(self):\n return self.model.layers\n<\/p>\n","protected":false},"excerpt":{"rendered":"
*That’s not much code. I could blog that much code. https:\/\/github.com\/ml-explore\/mlx-examples\/blob\/main\/llms\/mlx_lm\/models\/deepseek_v3.py # Copyright \u00a9 2024 Apple Inc. import math from dataclasses import dataclass from typing import Any, Dict, Optional, Tuple import mlx.core as mx import mlx.nn as nn from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention from .switch_layers import SwitchGLU @dataclass class ModelArgs(BaseModelArgs): model_type: str = “deepseek_v3” […]<\/p>\n","protected":false},"author":2,"featured_media":0,"comment_status":"open","ping_status":"open","sticky":false,"template":"","format":"standard","meta":{"footnotes":""},"categories":[1],"tags":[],"class_list":["post-5314","post","type-post","status-publish","format-standard","hentry","category-uncategorised"],"yoast_head":"\n