Python\uc73c\ub85c \uc791\uc131\ub41c \ub9ce\uc740 \ud504\ub85c\uc81d\ud2b8\ub4e4\uc774 \ube44\uad50\uc801 \uc9e7\uc740 \ucf54\ub4dc\ub85c\ub3c4 \uac15\ub825\ud55c \uae30\ub2a5\uc744 \uad6c\ud604\ud560 \uc218 \uc788\ub294 \uc774\uc720\ub294 \uc5ec\ub7ec \uac00\uc9c0\uac00 \uc788\uc2b5\ub2c8\ub2e4.<\/p>\n\n\n\n
\uc989, Deep Seek\uc758 \uc18c\uc2a4 \ucf54\ub4dc\uac00 \uc0c1\ub300\uc801\uc73c\ub85c \uc791\uc544 \ubcf4\uc774\ub294 \uac83\uc740 Python\uc758 \uac04\uacb0\ud55c \ubb38\ubc95\uacfc \uc678\ubd80 \ub77c\uc774\ube0c\ub7ec\ub9ac\uc758 \uac15\ub825\ud55c \uae30\ub2a5 \ub355\ubd84\uc774\ub77c\uace0 \ubcfc \uc218 \uc788\uc2b5\ub2c8\ub2e4.<\/p>\n\n\n\n
import math\nfrom dataclasses import dataclass\nfrom typing import Tuple, Optional, Literal\n\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\nimport torch.distributed as dist\n\nfrom kernel import act_quant, weight_dequant, fp8_gemm\n\n\nworld_size = 1\nrank = 0\nblock_size = 128\ngemm_impl: Literal[\"bf16\", \"fp8\"] = \"bf16\"\nattn_impl: Literal[\"naive\", \"absorb\"] = \"absorb\"\n\n@dataclass\nclass ModelArgs:\n \"\"\"\n Data class for defining model arguments and hyperparameters.\n\n Attributes:\n max_batch_size (int): Maximum batch size.\n max_seq_len (int): Maximum sequence length.\n dtype (Literal[\"bf16\", \"fp8\"]): Data type for computations.\n vocab_size (int): Vocabulary size.\n dim (int): Model dimension.\n inter_dim (int): Intermediate dimension for MLP layers.\n moe_inter_dim (int): Intermediate dimension for MoE layers.\n n_layers (int): Number of transformer layers.\n n_dense_layers (int): Number of dense layers in the model.\n n_heads (int): Number of attention heads.\n n_routed_experts (int): Number of routed experts for MoE layers.\n n_shared_experts (int): Number of shared experts for MoE layers.\n n_activated_experts (int): Number of activated experts in MoE layers.\n n_expert_groups (int): Number of expert groups.\n n_limited_groups (int): Number of limited groups for MoE routing.\n score_func (Literal[\"softmax\", \"sigmoid\"]): Scoring function for MoE routing.\n route_scale (float): Scaling factor for routing scores.\n q_lora_rank (int): LoRA rank for query projections.\n kv_lora_rank (int): LoRA rank for key-value projections.\n qk_nope_head_dim (int): Dimension for query-key projections without positional embeddings.\n qk_rope_head_dim (int): Dimension for query-key projections with rotary embeddings.\n v_head_dim (int): Dimension for value projections.\n original_seq_len (int): Original sequence length.\n rope_theta (float): Base for rotary positional encoding.\n rope_factor (float): Scaling factor for extended sequence lengths.\n beta_fast (int): Fast beta correction factor.\n beta_slow (int): Slow beta correction factor.\n mscale (float): Scaling factor for extended attention.\n \"\"\"\n max_batch_size: int = 8\n max_seq_len: int = 4096 * 4\n dtype: Literal[\"bf16\", \"fp8\"] = \"bf16\"\n vocab_size: int = 102400\n dim: int = 2048\n inter_dim: int = 10944\n moe_inter_dim: int = 1408\n n_layers: int = 27\n n_dense_layers: int = 1\n n_heads: int = 16\n # moe\n n_routed_experts: int = 64\n n_shared_experts: int = 2\n n_activated_experts: int = 6\n n_expert_groups: int = 1\n n_limited_groups: int = 1\n score_func: Literal[\"softmax\", \"sigmoid\"] = \"softmax\"\n route_scale: float = 1.\n # mla\n q_lora_rank: int = 0\n kv_lora_rank: int = 512\n qk_nope_head_dim: int = 128\n qk_rope_head_dim: int = 64\n v_head_dim: int = 128\n # yarn\n original_seq_len: int = 4096\n rope_theta: float = 10000.0\n rope_factor: float = 40\n beta_fast: int = 32\n beta_slow: int = 1\n mscale: float = 1.\n\n\nclass ParallelEmbedding(nn.Module):\n \"\"\"\n Embedding layer with parallelism support across distributed processes.\n\n Args:\n vocab_size (int): Vocabulary size.\n dim (int): Embedding dimension.\n \"\"\"\n def __init__(self, vocab_size: int, dim: int):\n super().__init__()\n self.vocab_size = vocab_size\n self.dim = dim\n assert vocab_size % world_size == 0, f\"Vocabulary size must be divisible by world size (world_size={world_size})\"\n self.part_vocab_size = (vocab_size \/\/ world_size)\n self.vocab_start_idx = rank * self.part_vocab_size\n self.vocab_end_idx = self.vocab_start_idx + self.part_vocab_size\n self.weight = nn.Parameter(torch.empty(self.part_vocab_size, self.dim))\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Forward pass for parallel embedding layer.\n\n Args:\n x (torch.Tensor): Input tensor containing token indices.\n\n Returns:\n torch.Tensor: Embedded representations.\n\n Raises:\n ValueError: If `world_size` is not defined.\n \"\"\"\n if world_size > 1:\n mask = (x < self.vocab_start_idx) | (x >= self.vocab_end_idx)\n x = x - self.vocab_start_idx\n x[mask] = 0\n y = F.embedding(x, self.weight)\n if world_size > 1:\n y[mask] = 0\n dist.all_reduce(y)\n return y\n\n\ndef linear(x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] = None) -> torch.Tensor:\n \"\"\"\n Applies a linear transformation to the incoming data: y = xA^T + b.\n This function supports specialized implementations based on quantization\n and tensor formats.\n\n Args:\n x (torch.Tensor): The input tensor.\n weight (torch.Tensor): The weight tensor. It may be quantized and \n requires dequantization for certain cases.\n bias (Optional[torch.Tensor]): The bias tensor to be added. Default is None.\n\n Returns:\n torch.Tensor: The result of the linear transformation, which may involve \n quantization-aware computations depending on the input parameters.\n\n Notes:\n - If `weight` is quantized (e.g., `element_size() == 1`), a dequantized version \n is used for computation.\n - If `gemm_impl == \"bf16\"`, dequantization and a `bf16` GEMM operation are applied.\n - For other cases, the function applies quantization to `x` and uses `fp8_gemm` for computation.\n \"\"\"\n if weight.element_size() > 1:\n return F.linear(x, weight, bias)\n elif gemm_impl == \"bf16\":\n weight = weight_dequant(weight, weight.scale)\n return F.linear(x, weight, bias)\n else:\n x, scale = act_quant(x, block_size)\n y = fp8_gemm(x, scale, weight, weight.scale)\n if bias is not None:\n y += bias\n return y\n\n\nclass Linear(nn.Module):\n \"\"\"\n Custom linear layer with support for quantized weights and optional bias.\n\n Args:\n in_features (int): Number of input features.\n out_features (int): Number of output features.\n bias (bool): Whether to include a bias term. Defaults to False.\n dtype (optional): Data type for the layer. Defaults to `torch.bfloat16`.\n \"\"\"\n dtype = torch.bfloat16\n\n def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None):\n super().__init__()\n self.in_features = in_features\n self.out_features = out_features\n self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype or Linear.dtype))\n if self.weight.element_size() == 1:\n scale_out_features = (out_features + block_size - 1) \/\/ block_size\n scale_in_features = (in_features + block_size - 1) \/\/ block_size\n self.weight.scale = self.scale = nn.Parameter(torch.empty(scale_out_features, scale_in_features, dtype=torch.float32))\n else:\n self.register_parameter(\"scale\", None)\n if bias:\n self.bias = nn.Parameter(torch.empty(out_features))\n else:\n self.register_parameter(\"bias\", None)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Forward pass for the custom linear layer.\n\n Args:\n x (torch.Tensor): Input tensor.\n\n Returns:\n torch.Tensor: Transformed tensor after linear computation.\n \"\"\"\n return linear(x, self.weight, self.bias)\n\n\nclass ColumnParallelLinear(Linear):\n \"\"\"\n Linear layer with column parallelism, splitting output features across distributed processes.\n\n Args:\n in_features (int): Number of input features.\n out_features (int): Total number of output features.\n bias (bool): Whether to include a bias term. Defaults to False.\n dtype (optional): Data type for the layer. Defaults to `torch.bfloat16`.\n \"\"\"\n def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None):\n assert out_features % world_size == 0, f\"Output features must be divisible by world size (world_size={world_size})\"\n self.part_out_features = out_features \/\/ world_size\n super().__init__(in_features, self.part_out_features, bias, dtype)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Forward pass for column parallel linear layer.\n\n Args:\n x (torch.Tensor): Input tensor.\n\n Returns:\n torch.Tensor: Transformed tensor with column-parallel computation.\n \"\"\"\n y = linear(x, self.weight, self.bias)\n return y\n\n\nclass RowParallelLinear(Linear):\n \"\"\"\n Linear layer with row parallelism, splitting input features across distributed processes.\n\n Args:\n in_features (int): Total number of input features.\n out_features (int): Number of output features.\n bias (bool): Whether to include a bias term. Defaults to False.\n dtype (optional): Data type for the layer. Defaults to `torch.bfloat16`.\n \"\"\"\n def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None):\n assert in_features % world_size == 0, f\"Input features must be divisible by world size (world_size={world_size})\"\n self.part_in_features = in_features \/\/ world_size\n super().__init__(self.part_in_features, out_features, bias, dtype)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Forward pass for row parallel linear layer.\n\n Args:\n x (torch.Tensor): Input tensor.\n\n Returns:\n torch.Tensor: Transformed tensor with row-parallel computation.\n \"\"\"\n y = linear(x, self.weight)\n if world_size > 1:\n dist.all_reduce(y)\n if self.bias is not None:\n y += self.bias\n return y\n\n\nclass RMSNorm(nn.Module):\n \"\"\"\n Root Mean Square Layer Normalization (RMSNorm).\n\n Args:\n dim (int): Dimension of the input tensor.\n eps (float): Epsilon value for numerical stability. Defaults to 1e-6.\n \"\"\"\n def __init__(self, dim: int, eps: float = 1e-6):\n super().__init__()\n self.dim = dim\n self.eps = eps\n self.weight = nn.Parameter(torch.ones(dim))\n\n def forward(self, x: torch.Tensor):\n \"\"\"\n Forward pass for RMSNorm.\n\n Args:\n x (torch.Tensor): Input tensor.\n\n Returns:\n torch.Tensor: Normalized tensor with the same shape as input.\n \"\"\"\n return F.rms_norm(x, (self.dim,), self.weight, self.eps)\n\n\ndef precompute_freqs_cis(args: ModelArgs) -> torch.Tensor:\n \"\"\"\n Precomputes frequency-based complex exponential values for rotary positional embeddings.\n\n Args:\n args (ModelArgs): Model arguments containing positional embedding parameters.\n\n Returns:\n torch.Tensor: Precomputed complex exponential values for positional embeddings.\n \"\"\"\n dim = args.qk_rope_head_dim\n seqlen = args.max_seq_len\n beta_fast = args.beta_fast\n beta_slow = args.beta_slow\n base = args.rope_theta\n factor = args.rope_factor\n\n def find_correction_dim(num_rotations, dim, base, max_seq_len):\n \"\"\"\n Computes the correction dimension for a given number of rotations in the rotary positional embedding.\n\n Args:\n num_rotations (float): Number of rotations to compute the correction for.\n dim (int): Dimensionality of the embedding space.\n base (float): Base value for the exponential computation.\n max_seq_len (int): Maximum sequence length.\n\n Returns:\n float: The correction dimension based on the input parameters.\n \"\"\"\n return dim * math.log(max_seq_len \/ (num_rotations * 2 * math.pi)) \/ (2 * math.log(base))\n\n def find_correction_range(low_rot, high_rot, dim, base, max_seq_len):\n \"\"\"\n Computes the range of correction dimensions for rotary positional embeddings.\n\n Args:\n low_rot (float): Lower bound for the number of rotations.\n high_rot (float): Upper bound for the number of rotations.\n dim (int): Dimensionality of the embedding space.\n base (float): Base value for the exponential computation.\n max_seq_len (int): Maximum sequence length.\n\n Returns:\n Tuple[int, int]: The range of correction dimensions (low, high), clamped to valid indices.\n \"\"\"\n low = math.floor(find_correction_dim(low_rot, dim, base, max_seq_len))\n high = math.ceil(find_correction_dim(high_rot, dim, base, max_seq_len))\n return max(low, 0), min(high, dim-1)\n\n def linear_ramp_factor(min, max, dim):\n \"\"\"\n Computes a linear ramp function used to smooth values between a minimum and maximum range.\n\n Args:\n min (float): Minimum value for the ramp function.\n max (float): Maximum value for the ramp function.\n dim (int): Dimensionality of the ramp tensor.\n\n Returns:\n torch.Tensor: A tensor of shape (dim,) with values linearly interpolated between 0 and 1,\n clamped to the range [0, 1].\n \"\"\"\n if min == max:\n max += 0.001\n linear_func = (torch.arange(dim, dtype=torch.float32) - min) \/ (max - min)\n ramp_func = torch.clamp(linear_func, 0, 1)\n return ramp_func\n\n freqs = 1.0 \/ (base ** (torch.arange(0, dim, 2, dtype=torch.float32) \/ dim))\n if seqlen > args.original_seq_len:\n low, high = find_correction_range(beta_fast, beta_slow, dim, base, args.original_seq_len)\n smooth = 1 - linear_ramp_factor(low, high, dim \/\/ 2)\n freqs = freqs \/ factor * (1 - smooth) + freqs * smooth\n\n t = torch.arange(seqlen)\n freqs = torch.outer(t, freqs)\n freqs_cis = torch.polar(torch.ones_like(freqs), freqs)\n return freqs_cis\n\n\ndef apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Applies rotary positional embeddings to the input tensor.\n\n Args:\n x (torch.Tensor): Input tensor with positional embeddings to be applied.\n freqs_cis (torch.Tensor): Precomputed complex exponential values for positional embeddings.\n\n Returns:\n torch.Tensor: Tensor with rotary embeddings applied.\n \"\"\"\n dtype = x.dtype\n x = torch.view_as_complex(x.float().view(*x.shape[:-1], -1, 2))\n freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1))\n y = torch.view_as_real(x * freqs_cis).flatten(3)\n return y.to(dtype)\n\n\nclass MLA(nn.Module):\n \"\"\"\n Multi-Headed Attention Layer (MLA).\n\n Attributes:\n dim (int): Dimensionality of the input features.\n n_heads (int): Number of attention heads.\n n_local_heads (int): Number of local attention heads for distributed systems.\n q_lora_rank (int): Rank for low-rank query projection.\n kv_lora_rank (int): Rank for low-rank key\/value projection.\n qk_nope_head_dim (int): Dimensionality of non-positional query\/key projections.\n qk_rope_head_dim (int): Dimensionality of rotary-positional query\/key projections.\n qk_head_dim (int): Total dimensionality of query\/key projections.\n v_head_dim (int): Dimensionality of value projections.\n softmax_scale (float): Scaling factor for softmax in attention computation.\n \"\"\"\n def __init__(self, args: ModelArgs):\n super().__init__()\n self.dim = args.dim\n self.n_heads = args.n_heads\n self.n_local_heads = args.n_heads \/\/ world_size\n self.q_lora_rank = args.q_lora_rank\n self.kv_lora_rank = args.kv_lora_rank\n self.qk_nope_head_dim = args.qk_nope_head_dim\n self.qk_rope_head_dim = args.qk_rope_head_dim\n self.qk_head_dim = args.qk_nope_head_dim + args.qk_rope_head_dim\n self.v_head_dim = args.v_head_dim\n\n if self.q_lora_rank == 0:\n self.wq = ColumnParallelLinear(self.dim, self.n_heads * self.qk_head_dim)\n else:\n self.wq_a = Linear(self.dim, self.q_lora_rank)\n self.q_norm = RMSNorm(self.q_lora_rank)\n self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.qk_head_dim)\n self.wkv_a = Linear(self.dim, self.kv_lora_rank + self.qk_rope_head_dim)\n self.kv_norm = RMSNorm(self.kv_lora_rank)\n self.wkv_b = ColumnParallelLinear(self.kv_lora_rank, self.n_heads * (self.qk_nope_head_dim + self.v_head_dim))\n self.wo = RowParallelLinear(self.n_heads * self.v_head_dim, self.dim)\n self.softmax_scale = self.qk_head_dim ** -0.5\n if args.max_seq_len > args.original_seq_len:\n mscale = 0.1 * args.mscale * math.log(args.rope_factor) + 1.0\n self.softmax_scale = self.softmax_scale * mscale * mscale\n\n if attn_impl == \"naive\":\n self.register_buffer(\"k_cache\", torch.zeros(args.max_batch_size, args.max_seq_len, self.n_local_heads, self.qk_head_dim), persistent=False)\n self.register_buffer(\"v_cache\", torch.zeros(args.max_batch_size, args.max_seq_len, self.n_local_heads, self.v_head_dim), persistent=False)\n else:\n self.register_buffer(\"kv_cache\", torch.zeros(args.max_batch_size, args.max_seq_len, self.kv_lora_rank), persistent=False)\n self.register_buffer(\"pe_cache\", torch.zeros(args.max_batch_size, args.max_seq_len, self.qk_rope_head_dim), persistent=False)\n\n def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]):\n \"\"\"\n Forward pass for the Multi-Headed Attention Layer (MLA).\n\n Args:\n x (torch.Tensor): Input tensor of shape (batch_size, seq_len, dim).\n start_pos (int): Starting position in the sequence for caching.\n freqs_cis (torch.Tensor): Precomputed complex exponential values for rotary embeddings.\n mask (Optional[torch.Tensor]): Mask tensor to exclude certain positions from attention.\n\n Returns:\n torch.Tensor: Output tensor with the same shape as the input.\n \"\"\"\n bsz, seqlen, _ = x.size()\n end_pos = start_pos + seqlen\n if self.q_lora_rank == 0:\n q = self.wq(x)\n else:\n q = self.wq_b(self.q_norm(self.wq_a(x)))\n q = q.view(bsz, seqlen, self.n_local_heads, self.qk_head_dim)\n q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)\n q_pe = apply_rotary_emb(q_pe, freqs_cis)\n kv = self.wkv_a(x)\n kv, k_pe = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)\n k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis)\n if attn_impl == \"naive\":\n q = torch.cat([q_nope, q_pe], dim=-1)\n kv = self.wkv_b(self.kv_norm(kv))\n kv = kv.view(bsz, seqlen, self.n_local_heads, self.qk_nope_head_dim + self.v_head_dim)\n k_nope, v = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)\n k = torch.cat([k_nope, k_pe.expand(-1, -1, self.n_local_heads, -1)], dim=-1)\n self.k_cache[:bsz, start_pos:end_pos] = k\n self.v_cache[:bsz, start_pos:end_pos] = v\n scores = torch.einsum(\"bshd,bthd->bsht\", q, self.k_cache[:bsz, :end_pos]) * self.softmax_scale\n else:\n wkv_b = self.wkv_b.weight if self.wkv_b.scale is None else weight_dequant(self.wkv_b.weight, self.wkv_b.scale, block_size) \n wkv_b = wkv_b.view(self.n_local_heads, -1, self.kv_lora_rank)\n q_nope = torch.einsum(\"bshd,hdc->bshc\", q_nope, wkv_b[:, :self.qk_nope_head_dim])\n self.kv_cache[:bsz, start_pos:end_pos] = self.kv_norm(kv)\n self.pe_cache[:bsz, start_pos:end_pos] = k_pe.squeeze(2)\n scores = (torch.einsum(\"bshc,btc->bsht\", q_nope, self.kv_cache[:bsz, :end_pos]) +\n torch.einsum(\"bshr,btr->bsht\", q_pe, self.pe_cache[:bsz, :end_pos])) * self.softmax_scale\n if mask is not None:\n scores += mask.unsqueeze(1)\n scores = scores.softmax(dim=-1, dtype=torch.float32).type_as(x)\n if attn_impl == \"naive\":\n x = torch.einsum(\"bsht,bthd->bshd\", scores, self.v_cache[:bsz, :end_pos])\n else:\n x = torch.einsum(\"bsht,btc->bshc\", scores, self.kv_cache[:bsz, :end_pos])\n x = torch.einsum(\"bshc,hdc->bshd\", x, wkv_b[:, -self.v_head_dim:])\n x = self.wo(x.flatten(2))\n return x\n\n\nclass MLP(nn.Module):\n \"\"\"\n Multi-Layer Perceptron (MLP) used as a feed-forward layer.\n\n Attributes:\n w1 (nn.Module): Linear layer for input-to-hidden transformation.\n w2 (nn.Module): Linear layer for hidden-to-output transformation.\n w3 (nn.Module): Additional linear layer for feature transformation.\n \"\"\"\n def __init__(self, dim: int, inter_dim: int):\n \"\"\"\n Initializes the MLP layer.\n\n Args:\n dim (int): Input and output dimensionality.\n inter_dim (int): Hidden layer dimensionality.\n \"\"\"\n super().__init__()\n self.w1 = ColumnParallelLinear(dim, inter_dim)\n self.w2 = RowParallelLinear(inter_dim, dim)\n self.w3 = ColumnParallelLinear(dim, inter_dim)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Forward pass for the MLP layer.\n\n Args:\n x (torch.Tensor): Input tensor.\n\n Returns:\n torch.Tensor: Output tensor after MLP computation.\n \"\"\"\n return self.w2(F.silu(self.w1(x)) * self.w3(x))\n\n\nclass Gate(nn.Module):\n \"\"\"\n Gating mechanism for routing inputs in a mixture-of-experts (MoE) model.\n\n Attributes:\n dim (int): Dimensionality of input features.\n topk (int): Number of top experts activated for each input.\n n_groups (int): Number of groups for routing.\n topk_groups (int): Number of groups to route inputs to.\n score_func (str): Scoring function ('softmax' or 'sigmoid').\n route_scale (float): Scaling factor for routing weights.\n weight (torch.nn.Parameter): Learnable weights for the gate.\n bias (Optional[torch.nn.Parameter]): Optional bias term for the gate.\n \"\"\"\n def __init__(self, args: ModelArgs):\n \"\"\"\n Initializes the Gate module.\n\n Args:\n args (ModelArgs): Model arguments containing gating parameters.\n \"\"\"\n super().__init__()\n self.dim = args.dim\n self.topk = args.n_activated_experts\n self.n_groups = args.n_expert_groups\n self.topk_groups = args.n_limited_groups\n self.score_func = args.score_func\n self.route_scale = args.route_scale\n self.weight = nn.Parameter(torch.empty(args.n_routed_experts, args.dim))\n self.bias = nn.Parameter(torch.empty(args.n_routed_experts)) if self.dim == 7168 else None\n\n def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:\n \"\"\"\n Forward pass for the gating mechanism.\n\n Args:\n x (torch.Tensor): Input tensor.\n\n Returns:\n Tuple[torch.Tensor, torch.Tensor]: Routing weights and selected expert indices.\n \"\"\"\n scores = linear(x, self.weight)\n if self.score_func == \"softmax\":\n scores = scores.softmax(dim=-1, dtype=torch.float32)\n else:\n scores = scores.sigmoid()\n original_scores = scores\n if self.bias is not None:\n scores = scores + self.bias\n if self.n_groups > 1:\n scores = scores.view(x.size(0), self.n_groups, -1)\n if self.bias is None:\n group_scores = scores.amax(dim=-1)\n else:\n group_scores = scores.topk(2, dim=-1)[0].sum(dim=-1)\n indices = group_scores.topk(self.topk_groups, dim=-1)[1]\n mask = scores.new_ones(x.size(0), self.n_groups, dtype=bool).scatter_(1, indices, False)\n scores = scores.masked_fill_(mask.unsqueeze(-1), float(\"-inf\")).flatten(1)\n indices = torch.topk(scores, self.topk, dim=-1)[1]\n weights = original_scores.gather(1, indices)\n if self.score_func == \"sigmoid\":\n weights \/= weights.sum(dim=-1, keepdim=True)\n weights *= self.route_scale\n return weights.type_as(x), indices\n\n\nclass Expert(nn.Module):\n \"\"\"\n Expert layer for Mixture-of-Experts (MoE) models.\n\n Attributes:\n w1 (nn.Module): Linear layer for input-to-hidden transformation.\n w2 (nn.Module): Linear layer for hidden-to-output transformation.\n w3 (nn.Module): Additional linear layer for feature transformation.\n \"\"\"\n def __init__(self, dim: int, inter_dim: int):\n \"\"\"\n Initializes the Expert layer.\n\n Args:\n dim (int): Input and output dimensionality.\n inter_dim (int): Hidden layer dimensionality.\n \"\"\"\n super().__init__()\n self.w1 = Linear(dim, inter_dim)\n self.w2 = Linear(inter_dim, dim)\n self.w3 = Linear(dim, inter_dim)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Forward pass for the Expert layer.\n\n Args:\n x (torch.Tensor): Input tensor.\n\n Returns:\n torch.Tensor: Output tensor after expert computation.\n \"\"\"\n return self.w2(F.silu(self.w1(x)) * self.w3(x))\n\n\nclass MoE(nn.Module):\n \"\"\"\n Mixture-of-Experts (MoE) module.\n\n Attributes:\n dim (int): Dimensionality of input features.\n n_routed_experts (int): Total number of experts in the model.\n n_local_experts (int): Number of experts handled locally in distributed systems.\n n_activated_experts (int): Number of experts activated for each input.\n gate (nn.Module): Gating mechanism to route inputs to experts.\n experts (nn.ModuleList): List of expert modules.\n shared_experts (nn.Module): Shared experts applied to all inputs.\n \"\"\"\n def __init__(self, args: ModelArgs):\n \"\"\"\n Initializes the MoE module.\n\n Args:\n args (ModelArgs): Model arguments containing MoE parameters.\n \"\"\"\n super().__init__()\n self.dim = args.dim\n assert args.n_routed_experts % world_size == 0, f\"Number of experts must be divisible by world size (world_size={world_size})\"\n self.n_routed_experts = args.n_routed_experts\n self.n_local_experts = args.n_routed_experts \/\/ world_size\n self.n_activated_experts = args.n_activated_experts\n self.experts_start_idx = rank * self.n_local_experts\n self.experts_end_idx = self.experts_start_idx + self.n_local_experts\n self.gate = Gate(args)\n self.experts = nn.ModuleList([Expert(args.dim, args.moe_inter_dim) if self.experts_start_idx <= i < self.experts_end_idx else None\n for i in range(self.n_routed_experts)])\n self.shared_experts = MLP(args.dim, args.n_shared_experts * args.moe_inter_dim)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n \"\"\"\n Forward pass for the MoE module.\n\n Args:\n x (torch.Tensor): Input tensor.\n\n Returns:\n torch.Tensor: Output tensor after expert routing and computation.\n \"\"\"\n shape = x.size()\n x = x.view(-1, self.dim)\n weights, indices = self.gate(x)\n y = torch.zeros_like(x)\n counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist()\n for i in range(self.experts_start_idx, self.experts_end_idx):\n if counts[i] == 0:\n continue\n expert = self.experts[i]\n idx, top = torch.where(indices == i)\n y[idx] += expert(x[idx]) * weights[idx, top, None]\n z = self.shared_experts(x)\n if world_size > 1:\n dist.all_reduce(y)\n return (y + z).view(shape)\n\n\nclass Block(nn.Module):\n \"\"\"\n Transformer block combining attention and feed-forward layers.\n\n Attributes:\n attn (nn.Module): Attention layer (MLA).\n ffn (nn.Module): Feed-forward network (MLP or MoE).\n attn_norm (nn.Module): Layer normalization for attention.\n ffn_norm (nn.Module): Layer normalization for feed-forward network.\n \"\"\"\n def __init__(self, layer_id: int, args: ModelArgs):\n \"\"\"\n Initializes the Transformer block.\n\n Args:\n layer_id (int): Layer index in the transformer.\n args (ModelArgs): Model arguments containing block parameters.\n \"\"\"\n super().__init__()\n self.attn = MLA(args)\n self.ffn = MLP(args.dim, args.inter_dim) if layer_id < args.n_dense_layers else MoE(args)\n self.attn_norm = RMSNorm(args.dim)\n self.ffn_norm = RMSNorm(args.dim)\n\n def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]) -> torch.Tensor:\n \"\"\"\n Forward pass for the Transformer block.\n\n Args:\n x (torch.Tensor): Input tensor.\n start_pos (int): Starting position in the sequence.\n freqs_cis (torch.Tensor): Precomputed complex exponential values for rotary embeddings.\n mask (Optional[torch.Tensor]): Mask tensor to exclude certain positions from attention.\n\n Returns:\n torch.Tensor: Output tensor after block computation.\n \"\"\"\n x = x + self.attn(self.attn_norm(x), start_pos, freqs_cis, mask)\n x = x + self.ffn(self.ffn_norm(x))\n return x\n\n\nclass Transformer(nn.Module):\n \"\"\"\n Transformer model with positional embeddings, multiple layers, and output projection.\n\n Attributes:\n max_seq_len (int): Maximum sequence length for the transformer.\n embed (nn.Module): Embedding layer for input tokens.\n layers (torch.nn.ModuleList): List of transformer blocks.\n norm (nn.Module): Layer normalization applied after all blocks.\n head (nn.Module): Output projection layer mapping to vocabulary size.\n freqs_cis (torch.Tensor): Precomputed complex exponential values for rotary embeddings.\n \"\"\"\n def __init__(self, args: ModelArgs):\n \"\"\"\n Initializes the Transformer model.\n\n Args:\n args (ModelArgs): Model arguments containing transformer parameters.\n \"\"\"\n global world_size, rank\n world_size = dist.get_world_size() if dist.is_initialized() else 1\n rank = dist.get_rank() if dist.is_initialized() else 0\n Linear.dtype = torch.float8_e4m3fn if args.dtype == \"fp8\" else torch.bfloat16\n super().__init__()\n self.max_seq_len = args.max_seq_len\n self.embed = ParallelEmbedding(args.vocab_size, args.dim)\n self.layers = torch.nn.ModuleList()\n for layer_id in range(args.n_layers):\n self.layers.append(Block(layer_id, args))\n self.norm = RMSNorm(args.dim)\n self.head = ColumnParallelLinear(args.dim, args.vocab_size, dtype=torch.get_default_dtype())\n self.register_buffer(\"freqs_cis\", precompute_freqs_cis(args), persistent=False)\n\n @torch.inference_mode()\n def forward(self, tokens: torch.Tensor, start_pos: int = 0):\n \"\"\"\n Forward pass for the Transformer model.\n\n Args:\n tokens (torch.Tensor): Input tensor of token IDs with shape (batch_size, seq_len).\n start_pos (int, optional): Starting position in the sequence for rotary embeddings. Defaults to 0.\n\n Returns:\n torch.Tensor: Logits tensor of shape (batch_size, vocab_size).\n \"\"\"\n seqlen = tokens.size(1)\n h = self.embed(tokens)\n freqs_cis = self.freqs_cis[start_pos:start_pos+seqlen]\n mask = None\n if seqlen > 1:\n mask = torch.full((seqlen, seqlen), float(\"-inf\"), device=tokens.device).triu_(1)\n for layer in self.layers:\n h = layer(h, start_pos, freqs_cis, mask)\n h = self.norm(h)[:, -1]\n logits = self.head(h)\n if world_size > 1:\n all_logits = [torch.empty_like(logits) for _ in range(world_size)]\n dist.all_gather(all_logits, logits)\n logits = torch.cat(all_logits, dim=-1)\n return logits\n\n\nif __name__ == \"__main__\":\n torch.set_default_dtype(torch.bfloat16)\n torch.set_default_device(\"cuda\")\n torch.manual_seed(0)\n args = ModelArgs()\n x = torch.randint(0, args.vocab_size, (2, 128))\n model = Transformer(args)\n print(model(x).size())\n<\/pre><\/div>\n","protected":false},"excerpt":{"rendered":"