import os import math import datetime from dataclasses import dataclass from typing import Tuple, Optional, Literal import torch from torch import nn import torch.nn.functional as F import torch.distributed as dist from safetensors.torch import load_model from kernel import act_quant, weight_dequant, fp8_gemm, int64_bmm_broadcast, \ complex_int64_mul_broadcast, einsum_bshd_hdc_bshc, einsum_bshc_btc_bsht, softmax_init_q21, softmax_q21, einsum_bsht_btc_bshc, einsum_bshc_hdc_bshd, \ silu_init_q25, silu_q25, sigmoid_q25, softmax_init_q19, softmax_q19, silu_init_q23, silu_q23, sigmoid_q23, RMS_Norm_int64 world_size = 1 rank = 0 block_size = 128 gemm_impl: Literal["bf16", "fp8"] = "bf16" attn_impl: Literal["naive", "absorb"] = "absorb" snark = True zkDataDir = '../zkdata' @dataclass class ModelArgs: """ Data class for defining model arguments and hyperparameters. Attributes: max_batch_size (int): Maximum batch size. max_seq_len (int): Maximum sequence length. dtype (Literal["bf16", "fp8"]): Data type for computations. vocab_size (int): Vocabulary size. dim (int): Model dimension. inter_dim (int): Intermediate dimension for MLP layers. moe_inter_dim (int): Intermediate dimension for MoE layers. n_layers (int): Number of transformer layers. n_dense_layers (int): Number of dense layers in the model. n_heads (int): Number of attention heads. n_routed_experts (int): Number of routed experts for MoE layers. n_shared_experts (int): Number of shared experts for MoE layers. n_activated_experts (int): Number of activated experts in MoE layers. n_expert_groups (int): Number of expert groups. n_limited_groups (int): Number of limited groups for MoE routing. score_func (Literal["softmax", "sigmoid"]): Scoring function for MoE routing. route_scale (float): Scaling factor for routing scores. q_lora_rank (int): LoRA rank for query projections. kv_lora_rank (int): LoRA rank for key-value projections. qk_nope_head_dim (int): Dimension for query-key projections without positional embeddings. qk_rope_head_dim (int): Dimension for query-key projections with rotary embeddings. v_head_dim (int): Dimension for value projections. original_seq_len (int): Original sequence length. rope_theta (float): Base for rotary positional encoding. rope_factor (float): Scaling factor for extended sequence lengths. beta_fast (int): Fast beta correction factor. beta_slow (int): Slow beta correction factor. mscale (float): Scaling factor for extended attention. """ max_batch_size: int = 8 max_seq_len: int = 4096 * 4 dtype: Literal["bf16", "fp8"] = "bf16" vocab_size: int = 102400 dim: int = 2048 inter_dim: int = 10944 moe_inter_dim: int = 1408 n_layers: int = 27 n_dense_layers: int = 1 n_heads: int = 16 # moe n_routed_experts: int = 64 n_shared_experts: int = 2 n_activated_experts: int = 6 n_expert_groups: int = 1 n_limited_groups: int = 1 score_func: Literal["softmax", "sigmoid"] = "softmax" route_scale: float = 1. # mla q_lora_rank: int = 0 kv_lora_rank: int = 512 qk_nope_head_dim: int = 128 qk_rope_head_dim: int = 64 v_head_dim: int = 128 # yarn original_seq_len: int = 4096 rope_theta: float = 10000.0 rope_factor: float = 40 beta_fast: int = 32 beta_slow: int = 1 mscale: float = 1. def saveTensor(fileName, t): with open(fileName, "w", encoding="utf-8") as f: t = t.detach() if t.device.type != "cpu": t = t.cpu() t = t.contiguous() with open(fileName, "wb") as f: # .numpy() -> bytes（C-order） f.write(t.numpy().tobytes(order="C")) class ParallelEmbedding(nn.Module): """ Embedding layer with parallelism support across distributed processes. Args: vocab_size (int): Vocabulary size. dim (int): Embedding dimension. """ def __init__(self, vocab_size: int, dim: int): super().__init__() self.vocab_size = vocab_size self.dim = dim assert vocab_size % world_size == 0, f"Vocabulary size must be divisible by world size (world_size={world_size})" self.part_vocab_size = (vocab_size // world_size) self.vocab_start_idx = rank * self.part_vocab_size self.vocab_end_idx = self.vocab_start_idx + self.part_vocab_size # weight 的 shape: [129280, 7168] self.register_buffer("weight", torch.empty(self.part_vocab_size, self.dim, dtype=torch.int64)) def forward(self, x: torch.Tensor) -> torch.Tensor: """ Forward pass for parallel embedding layer. Args: x (torch.Tensor): Input tensor containing token indices. Returns: torch.Tensor: Embedded representations. Raises: ValueError: If `world_size` is not defined. """ # print('aaab ' + str(self.weight[0][0].type())) if world_size > 1: # 找出 x 中 的值不在 [vocab_start_idx, vocab_end_idx) 范围内的下标 mask = (x < self.vocab_start_idx) | (x >= self.vocab_end_idx) # x 中所有的值都减去 vocab_start_idx x = x - self.vocab_start_idx # 之前找出的标记为 mask 下标的值设置为0 x[mask] = 0 y = F.embedding(x, self.weight) if world_size > 1: y[mask] = 0 dist.all_reduce(y) # print(f'ParallelEmbedding x: {x}', flush=True) return y def linear(x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] = None) -> torch.Tensor: """ Applies a linear transformation to the incoming data: y = xA^T + b. This function supports specialized implementations based on quantization and tensor formats. Args: x (torch.Tensor): The input tensor. weight (torch.Tensor): The weight tensor. It may be quantized and requires dequantization for certain cases. bias (Optional[torch.Tensor]): The bias tensor to be added. Default is None. Returns: torch.Tensor: The result of the linear transformation, which may involve quantization-aware computations depending on the input parameters. Notes: - If `weight` is quantized (e.g., `element_size() == 1`), a dequantized version is used for computation. - If `gemm_impl == "bf16"`, dequantization and a `bf16` GEMM operation are applied. - For other cases, the function applies quantization to `x` and uses `fp8_gemm` for computation. """ element_size = weight.element_size() typ = weight.type() # print(f'linear weight element_size {element_size}, type: {typ}', flush=True) if weight.element_size() > 1: # print('linear weight.element_size > 1, element_size=' + str(weight.element_size()), flush=True) return F.linear(x, weight, bias) elif gemm_impl == "bf16": weight = weight_dequant(weight, weight.scale) return F.linear(x, weight, bias) else: # print('linear act_quant', flush=True) x, scale = act_quant(x, block_size) y = fp8_gemm(x, scale, weight, weight.scale) if bias is not None: y += bias return y def linear_int(x: torch.Tensor, weight: torch.Tensor, x_rescale, weight_rescale, res_rescale, bias: Optional[torch.Tensor] = None) -> tuple[torch.Tensor]: if weight.element_size() > 1: (q, r) = int64_bmm_broadcast(x, weight, x_rescale, weight_rescale, res_rescale) return (q, r) elif gemm_impl == "bf16": weight = weight_dequant(weight, weight.scale) return (F.linear(x, weight, bias), torch.tensor(0, dtype=torch.int64)) else: print('linear act_quant', flush=True) x, scale = act_quant(x, block_size) y = fp8_gemm(x, scale, weight, weight.scale) if bias is not None: y += bias return (y, torch.tensor(0, dtype=torch.int64)) class Linear_int(nn.Module): """ Custom linear layer with support for quantized weights and optional bias. Args: in_features (int): Number of input features. out_features (int): Number of output features. bias (bool): Whether to include a bias term. Defaults to False. dtype (optional): Data type for the layer. Defaults to `torch.bfloat16`. """ dtype = torch.int64 def __init__(self, layer_id, in_features: int, out_features: int, x_rescale, weight_rescale, res_rescale, dtype, bias: bool = False): super().__init__() self.layer_id = layer_id self.in_features = in_features self.out_features = out_features self.x_rescale = x_rescale self.weight_rescale = weight_rescale self.res_rescale = res_rescale self.register_buffer("weight", torch.empty(out_features, in_features, dtype=dtype)) if bias: self.bias = nn.Parameter(torch.empty(out_features)) else: self.register_parameter("bias", None) def forward(self, x: torch.Tensor) -> tuple[torch.Tensor]: q, r = linear_int(x, self.weight, self.x_rescale, self.weight_rescale, self.res_rescale, self.bias) return q, r class Linear_rescale_int(nn.Module): """ Custom linear layer with support for quantized weights and optional bias. Args: in_features (int): Number of input features. out_features (int): Number of output features. bias (bool): Whether to include a bias term. Defaults to False. dtype (optional): Data type for the layer. Defaults to `torch.bfloat16`. """ dtype = torch.int64 def __init__(self, layer_id, in_features: int, out_features: int, x_rescale, weight_rescale, dtype, bias: bool = False): super().__init__() self.layer_id = layer_id self.in_features = in_features self.out_features = out_features self.x_rescale = x_rescale self.weight_rescale = weight_rescale self.register_buffer("weight", torch.empty(out_features, in_features, dtype=dtype)) self.register_buffer("scale", torch.tensor(0, dtype=torch.int32)) if bias: self.bias = nn.Parameter(torch.empty(out_features)) else: self.register_parameter("bias", None) def forward(self, x: torch.Tensor) -> torch.Tensor: rescale = self.scale.item() y, _r = linear_int(x, self.weight, self.x_rescale, self.weight_rescale, rescale, self.bias) return y class Linear(nn.Module): """ Custom linear layer with support for quantized weights and optional bias. Args: in_features (int): Number of input features. out_features (int): Number of output features. bias (bool): Whether to include a bias term. Defaults to False. dtype (optional): Data type for the layer. Defaults to `torch.bfloat16`. """ dtype = torch.bfloat16 def __init__(self, layer_id, in_features: int, out_features: int, bias: bool = False, dtype = None): super().__init__() self.layer_id = layer_id self.in_features = in_features self.out_features = out_features self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype or Linear.dtype)) # print('Linear.weight.element_size: ' + str(self.weight.element_size())) # nn.Parameter.element_size() 返回的是 每个元素在内存中占用的字节数 # torch.float32 -> 4 字节 # torch.float64 -> 8 字节 # torch.int64 -> 8 字节 # torch.bfloat16 -> 2 字节 # torch.float8_e4m3fn -> 1 字节 if self.weight.element_size() == 1: scale_out_features = (out_features + block_size - 1) // block_size scale_in_features = (in_features + block_size - 1) // block_size self.weight.scale = self.scale = nn.Parameter(torch.empty(scale_out_features, scale_in_features, dtype=torch.float32)) else: self.register_parameter("scale", None) if bias: self.bias = nn.Parameter(torch.empty(out_features)) else: self.register_parameter("bias", None) def forward(self, x: torch.Tensor) -> torch.Tensor: """ Forward pass for the custom linear layer. Args: x (torch.Tensor): Input tensor. Returns: torch.Tensor: Transformed tensor after linear computation. """ return linear(x, self.weight, self.bias) class ColumnParallelLinear(Linear): """ Linear layer with column parallelism, splitting output features across distributed processes. Args: in_features (int): Number of input features. out_features (int): Total number of output features. bias (bool): Whether to include a bias term. Defaults to False. dtype (optional): Data type for the layer. Defaults to `torch.bfloat16`. """ def __init__(self, layer_id, in_features: int, out_features: int, bias: bool = False, dtype = None): assert out_features % world_size == 0, f"Output features must be divisible by world size (world_size={world_size})" self.part_out_features = out_features // world_size super().__init__(layer_id, in_features, self.part_out_features, bias, dtype) def forward(self, x: torch.Tensor) -> torch.Tensor: """ Forward pass for column parallel linear layer. Args: x (torch.Tensor): Input tensor. Returns: torch.Tensor: Transformed tensor with column-parallel computation. """ y = linear(x, self.weight, self.bias) return y class ColumnParallelLinear_int(Linear_int): def __init__(self, layer_id, in_features: int, out_features: int, x_rescale, weight_rescale, res_rescale, dtype, bias: bool = False): assert out_features % world_size == 0, f"Output features must be divisible by world size (world_size={world_size})" self.part_out_features = out_features // world_size super().__init__(layer_id, in_features, self.part_out_features, x_rescale, weight_rescale, res_rescale, dtype, bias) def forward(self, x: torch.Tensor) -> torch.Tensor: y, _r = linear_int(x, self.weight, self.x_rescale, self.weight_rescale, self.res_rescale, self.bias) return y class ColumnParallelLinear_rescale_int(Linear_int): def __init__(self, layer_id, in_features: int, out_features: int, x_rescale, weight_rescale, dtype, bias: bool = False): assert out_features % world_size == 0, f"Output features must be divisible by world size (world_size={world_size})" self.part_out_features = out_features // world_size super().__init__(layer_id, in_features, self.part_out_features, x_rescale, weight_rescale, 1, dtype, bias) self.register_buffer("scale", torch.tensor(0, dtype=torch.int32)) # self.res_rescale = self.scale def forward(self, x: torch.Tensor) -> torch.Tensor: rescale = self.scale.item() y, _r = linear_int(x, self.weight, self.x_rescale, self.weight_rescale, rescale, self.bias) return y class RowParallelLinear(Linear): """ Linear layer with row parallelism, splitting input features across distributed processes. Args: in_features (int): Total number of input features. out_features (int): Number of output features. bias (bool): Whether to include a bias term. Defaults to False. dtype (optional): Data type for the layer. Defaults to `torch.bfloat16`. """ def __init__(self, layer_id, in_features: int, out_features: int, bias: bool = False, dtype = None): assert in_features % world_size == 0, f"Input features must be divisible by world size (world_size={world_size})" self.part_in_features = in_features // world_size super().__init__(layer_id, self.part_in_features, out_features, bias, dtype) def forward(self, x: torch.Tensor) -> torch.Tensor: """ Forward pass for row parallel linear layer. Args: x (torch.Tensor): Input tensor. Returns: torch.Tensor: Transformed tensor with row-parallel computation. """ y = linear(x, self.weight) if world_size > 1: dist.all_reduce(y) if self.bias is not None: y += self.bias return y class RowParallelLinear_rescale_int(Linear_int): """ Linear layer with row parallelism, splitting input features across distributed processes. Args: in_features (int): Total number of input features. out_features (int): Number of output features. bias (bool): Whether to include a bias term. Defaults to False. dtype (optional): Data type for the layer. Defaults to `torch.bfloat16`. """ def __init__(self, layer_id, in_features: int, out_features: int, x_rescale, weight_rescale, res_rescale, dtype, bias: bool = False): assert in_features % world_size == 0, f"Input features must be divisible by world size (world_size={world_size})" self.part_in_features = in_features // world_size super().__init__(layer_id, self.part_in_features, out_features, x_rescale, weight_rescale, res_rescale, dtype, bias) self.register_buffer("scale", torch.tensor(0, dtype=torch.int32)) self.res_rescale = self.scale # useless def forward(self, x: torch.Tensor) -> torch.Tensor: """ Forward pass for row parallel linear layer. Args: x (torch.Tensor): Input tensor. Returns: torch.Tensor: Transformed tensor with row-parallel computation. """ # rescale = 2 ** self.scale.item() rescale = self.scale.item() # print(f'RowParallelLinear_rescale_int forward scale: {self.scale} ' + str(rescale), flush=True) y, _ = linear_int(x, self.weight, self.x_rescale, self.weight_rescale, rescale, self.bias) if world_size > 1: dist.all_reduce(y) if self.bias is not None: y += self.bias return y class RMSNorm(nn.Module): """ Root Mean Square Layer Normalization (RMSNorm). Args: dim (int): Dimension of the input tensor. eps (float): Epsilon value for numerical stability. Defaults to 1e-6. """ def __init__(self, dim: int, eps: float = 1e-6): super().__init__() self.dim = dim self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x: torch.Tensor): """ Forward pass for RMSNorm. Args: x (torch.Tensor): Input tensor. Returns: torch.Tensor: Normalized tensor with the same shape as input. """ return F.rms_norm(x, (self.dim,), self.weight, self.eps) class RMSNorm_int(nn.Module): def __init__(self, dim: int, dtype, eps: float = 1e-6): super().__init__() self.dim = dim self.eps = eps self.register_buffer( "weight", torch.ones(dim, dtype=dtype)) def forward(self, x: torch.Tensor): # x 的 scale 为 2 ** 31 # weight的scale 为 2 ** 15, 范围为 2^7 - 2^14 # rms 的 scale 为 2 ** 28 # 返回的结果 scale 为 2 ** 16,因为中间计算的时候 除以了 (1 << 15)，44 + 15 - 28 - 15 = 16 (c, rms) = RMS_Norm_int64(x[0], self.weight, 1, self.dim) return (c[None, :], rms) def precompute_freqs_cis(args: ModelArgs) -> torch.Tensor: """ Precomputes frequency-based complex exponential values for rotary positional embeddings. Args: args (ModelArgs): Model arguments containing positional embedding parameters. Returns: torch.Tensor: Precomputed complex exponential values for positional embeddings. """ # dim = 64 dim = args.qk_rope_head_dim seqlen = args.max_seq_len beta_fast = args.beta_fast beta_slow = args.beta_slow base = args.rope_theta factor = args.rope_factor def find_correction_dim(num_rotations, dim, base, max_seq_len): """ Computes the correction dimension for a given number of rotations in the rotary positional embedding. Args: num_rotations (float): Number of rotations to compute the correction for. dim (int): Dimensionality of the embedding space. base (float): Base value for the exponential computation. max_seq_len (int): Maximum sequence length. Returns: float: The correction dimension based on the input parameters. """ return dim * math.log(max_seq_len / (num_rotations * 2 * math.pi)) / (2 * math.log(base)) def find_correction_range(low_rot, high_rot, dim, base, max_seq_len): """ Computes the range of correction dimensions for rotary positional embeddings. Args: low_rot (float): Lower bound for the number of rotations. high_rot (float): Upper bound for the number of rotations. dim (int): Dimensionality of the embedding space. base (float): Base value for the exponential computation. max_seq_len (int): Maximum sequence length. Returns: Tuple[int, int]: The range of correction dimensions (low, high), clamped to valid indices. """ low = math.floor(find_correction_dim(low_rot, dim, base, max_seq_len)) high = math.ceil(find_correction_dim(high_rot, dim, base, max_seq_len)) return max(low, 0), min(high, dim-1) def linear_ramp_factor(min, max, dim): """ Computes a linear ramp function used to smooth values between a minimum and maximum range. Args: min (float): Minimum value for the ramp function. max (float): Maximum value for the ramp function. dim (int): Dimensionality of the ramp tensor. Returns: torch.Tensor: A tensor of shape (dim,) with values linearly interpolated between 0 and 1, clamped to the range [0, 1]. """ if min == max: max += 0.001 linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min) ramp_func = torch.clamp(linear_func, 0, 1) return ramp_func # torch.arange(0, dim, 2, dtype=torch.float32) 的作用是： 生成从 0 开始、步长为 2、到 dim 之前（不含 dim）的一维张量，数据类型为 float32 # 1/10000^(2k/d_model) # freqs shape: 一维向量，长度为 dim /2 freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) # original_seq_len=4096 if seqlen > args.original_seq_len: low, high = find_correction_range(beta_fast, beta_slow, dim, base, args.original_seq_len) smooth = 1 - linear_ramp_factor(low, high, dim // 2) freqs = freqs / factor * (1 - smooth) + freqs * smooth t = torch.arange(seqlen) # torch.outer 的作用是计算两个向量的 外积 (outer product)，比如： # t = torch.tensor([1, 2, 3]) # shape = [3] # freqs = torch.tensor([10, 20]) # shape = [2] # out = torch.outer(t, freqs) # tensor([[10, 20], # [20, 40], # [30, 60]]) # freqs shape为 [seqlen, dim/2] freqs = torch.outer(t, freqs) # torch.polar(abs, angle) 的作用: 把 极坐标 (r, θ) 转换成 复数 (x + iy) 的函数 # freqs_cis_0 shape为 [seqlen, dim/2] freqs_cis_0 = torch.polar(torch.ones_like(freqs), freqs) # return freqs_cis_0 # 复数转换成实数, freqs_cis_1 shape为 [seqlen, dim] freqs_cis_1 = torch.view_as_real(freqs_cis_0) # freqs_cis = torch.empty_like(freqs_cis_1, dtype=torch.int64, device='cuda') # cols 为 2 * freqs_cis_1.shape[1] 是因为 复数的实部 和 虚部 # rescale 参数为 19 = 42 - 23, ex 部分加 +19，总的rescale为 2^42 freqs_cis = (freqs_cis_1 * (2 ** 42)).round().to(torch.int64) freqs_cis_abs = freqs_cis.abs() min1 = freqs_cis_abs.min() max1 = freqs_cis_abs.max() print(f'freqs_cis min {min1}, max: {max1}', flush=True) # print(f'freqs_cis: {freqs_cis}') # freqs_cis 的 rescale 为 2^42 return freqs_cis # x(q_pe) 的维度 [batch, seqLen, 128, 64] def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: """ Applies rotary positional embeddings to the input tensor. Args: x (torch.Tensor): Input tensor with positional embeddings to be applied. freqs_cis (torch.Tensor): Precomputed complex exponential values for positional embeddings. Returns: torch.Tensor: Tensor with rotary embeddings applied. """ # if x.dtype == torch.int64: # x 的维度 变为 [batch, seqLen, 128, 32, 2] ### important!!! 调用 so lib库之前，必须确保内存连续 x = x.contiguous().view(*x.shape[:-1], -1, 2) # freqs_cis 的维度为 [1, seqLen, 1, 32, 2] freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-2), 2) # freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1)) # 4194304 = 1 << (64 - 42), 42是 rescale, int64 * int64 结果的高 64位 乘以 4194304 # 4398046511104 = 1 << 42 # print(x) # print(f'x shape: {x.shape}, freqs_cis shape: {freqs_cis.shape}') # y = complex_int64_mul_broadcast(x, freqs_cis, 4194304, 4398046511104) y = complex_int64_mul_broadcast(x, freqs_cis) y2 = y.flatten(3) return y2 def getBF16PrintStr(ele): v = int(ele.cpu().view(torch.uint16).item()) ex = v >> 7 & 0xFF r = '(1+' + str(v & 0x7F) + '/128)' rraw = v & 0x7F if v & 0x8000: vstr = '-' + r + '*2^' + str(ex - 127) else: vstr = r + '*2^' + str(ex - 127) return vstr class MLA(nn.Module): """ Multi-Headed Attention Layer (MLA). Attributes: dim (int): Dimensionality of the input features. n_heads (int): Number of attention heads. n_local_heads (int): Number of local attention heads for distributed systems. q_lora_rank (int): Rank for low-rank query projection. kv_lora_rank (int): Rank for low-rank key/value projection. qk_nope_head_dim (int): Dimensionality of non-positional query/key projections. qk_rope_head_dim (int): Dimensionality of rotary-positional query/key projections. qk_head_dim (int): Total dimensionality of query/key projections. v_head_dim (int): Dimensionality of value projections. softmax_scale (float): Scaling factor for softmax in attention computation. """ def __init__(self, layer_id, args: ModelArgs): super().__init__() # RowParallelLinear和ColumnParallelLinear是将Linear层按照行和列划分为多个子线性层并分配到各个设备上，每个设备维护一个子线性层， # 如线性层的shape为[in_features, out_features]，RowParallelLinear的shape为[in_features/world_size, out_features]， # ColumnParallelLinear的shape为[in_features，out_features/world_size]，world_size是设备数 self.layer_id = layer_id # 7168 self.dim = args.dim # 128 self.n_heads = args.n_heads # 当前进程跑的header数目 self.n_local_heads = args.n_heads // world_size # query向下投影矩阵维度，默认为0表示不压缩，实际使用过程为 1536 self.q_lora_rank = args.q_lora_rank # key和value向下投影矩阵维度，实际使用过程为 512; self.kv_lora_rank = args.kv_lora_rank # query/key不包含位置信息的隐藏层维度, 实际使用过程为 128 self.qk_nope_head_dim = args.qk_nope_head_dim # query/key包含rope位置信息的隐藏层维度, 实际使用过程为 64 self.qk_rope_head_dim = args.qk_rope_head_dim # 192 self.qk_head_dim = args.qk_nope_head_dim + args.qk_rope_head_dim # value隐藏层维度, 实际使用过程为 128 self.v_head_dim = args.v_head_dim # query向下投影矩阵维度，默认为0表示不压缩，实际使用过程为 1536 if self.q_lora_rank == 0: self.wq = ColumnParallelLinear(layer_id, self.dim, self.n_heads * self.qk_head_dim) else: # query向下投影矩阵, shape [7168, 1536], Float8_e4m3fnTensor self.wq_a = Linear_int(layer_id, self.dim, self.q_lora_rank, 1, 1, 30, torch.int32) self.q_norm = RMSNorm_int(self.q_lora_rank, torch.int32) # query向上投影矩阵的列并行线性层, shape [1536, 24576(128 * 192)], Float8_e4m3fnTensor # self.wq_b = ColumnParallelLinear_int(layer_id, self.q_lora_rank, self.n_heads * self.qk_head_dim, 1, 1, (1 << 30), torch.int32) self.wq_b1 = ColumnParallelLinear_int(layer_id, self.q_lora_rank, self.n_heads * args.qk_nope_head_dim, 1, 1, 30, torch.int32) self.wq_b2 = ColumnParallelLinear_int(layer_id, self.q_lora_rank, self.n_heads * args.qk_rope_head_dim, 1, 1, 30, torch.int32) # key和value的向下投影矩阵, shape [576, 7168], Float8_e4m3fnTensor, kv_lora_rank=512, qk_rope_head_dim=64 # self.wkv_a = Linear_int(layer_id, self.dim, self.kv_lora_rank + self.qk_rope_head_dim, 1, 1, (1 << 29), torch.int32) self.wkv_a1 = Linear_int(layer_id, self.dim, self.kv_lora_rank, 1, 1, 29, torch.int32) self.wkv_a2 = Linear_int(layer_id, self.dim, self.qk_rope_head_dim, 1, 1, 29, torch.int32) # self.kv_norm = RMSNorm(self.kv_lora_rank) self.kv_norm = RMSNorm_int(self.kv_lora_rank, torch.int32) # key和value向上投影矩阵的列并行线性层, shape [32768, 512], Float8_e4m3fnTensor # kv_lora_rank=512, n_heads = 128, qk_nope_head_dim = 128, v_head_dim = 128 # self.wkv_b = ColumnParallelLinear(layer_id, self.kv_lora_rank, self.n_heads * (self.qk_nope_head_dim + self.v_head_dim)) self.wkv_b_1 = ColumnParallelLinear_rescale_int(layer_id, self.kv_lora_rank, self.n_heads * self.qk_nope_head_dim, 1, 1, torch.int32) self.wkv_b_2 = ColumnParallelLinear_rescale_int(layer_id, self.kv_lora_rank, self.n_heads * self.v_head_dim, 1, 1, torch.int32) # 输出投影行并行线性层, shape [7168, 16384], Float8_e4m3fnTensor self.wo = RowParallelLinear_rescale_int(layer_id, self.n_heads * self.v_head_dim, self.dim, 1, 1, 1, torch.int32) # softmax缩放系数, qk_head_dim = 192 # self.softmax_scale = self.qk_head_dim ** -0.5 # # max_seq_len = 4096 * 4, original_seq_len = 4096 # if args.max_seq_len > args.original_seq_len: # # mscale = 1.0, rope_factor = 40, math.log = ln 自然对数 # mscale = 0.1 * args.mscale * math.log(args.rope_factor) + 1.0 # self.softmax_scale = self.softmax_scale * mscale * mscale self.softmax_scale1 = 94 self.softmax_scale2 = 695 if attn_impl == "naive": 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) 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) else: # 缓存key和value向下投影表示 # self.register_buffer("kv_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.kv_lora_rank), persistent=False) # self.register_buffer("kv_cache", torch.zeros(1, args.max_seq_len, self.kv_lora_rank), persistent=False) self.register_buffer("kv_cache", torch.zeros(1, args.max_seq_len, self.kv_lora_rank, dtype=torch.int64), persistent=False) # 缓存key执行rope操作后的表示 # self.register_buffer("pe_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.qk_rope_head_dim), persistent=False) # self.register_buffer("pe_cache", torch.zeros(1, args.max_seq_len, self.qk_rope_head_dim), persistent=False) self.register_buffer("pe_cache", torch.zeros(1, args.max_seq_len, self.qk_rope_head_dim, dtype=torch.int64), persistent=False) # x shape [1, seqLen, 7168], x 的resacle 为 2^21 def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]): """ Forward pass for the Multi-Headed Attention Layer (MLA). Args: x (torch.Tensor): Input tensor of shape (batch_size, seq_len, dim). start_pos (int): Starting position in the sequence for caching. freqs_cis (torch.Tensor): Precomputed complex exponential values for rotary embeddings. mask (Optional[torch.Tensor]): Mask tensor to exclude certain positions from attention. Returns: torch.Tensor: Output tensor with the same shape as the input. """ # 从输入获取batch size和序列长度seqlen，并根据输入序列的起始位置计算输入序列的结束位置end_pos=start_pos+seqlen； bsz, seqlen, _ = x.size() end_pos = start_pos + seqlen # 获取query的投影表示：如果对query投影矩阵进行压缩(即q_lora_rank不为0)，则将输入乘以query的向下投影矩阵wq_a，然后经过归一化层q_norm， # 再乘以向上投影矩阵wq_b，否则直接乘以原始投影矩阵wq；将其维度调整为[batchsize, n_local_threads, qk_head_dim]； if self.q_lora_rank == 0: q = self.wq(x) else: # query向下投影矩阵, shape [7168, 1536], Float8_e4m3fnTensor # x(也就是 attn_normed) 的 scale 为 2^21, wq_a weight 的scale 为 2^30, q_down 的 scale 为 2^21 q_down, q_down_rem = self.wq_a(x) # q_down = self.wq_a(x) if snark: dirStr = f'{zkDataDir}/pos_{start_pos}/layer_{self.layer_id}' os.makedirs(dirStr, exist_ok=True) saveTensor(f'{dirStr}/wq_a_x.bin', x.cpu()) saveTensor(f'{dirStr}/wq_a_w.bin', self.wq_a.weight.view(torch.uint32).cpu()) saveTensor(f'{dirStr}/wq_a_y.bin', q_down.cpu()) saveTensor(f'{dirStr}/q_norm_r.bin', q_down_rem.cpu()) # q_down = (q_down.detach().to(torch.float32) * (2 ** -23)).to(torch.bfloat16) # q_norm 的 rescale 为 2^19 (q_normed, rms) = self.q_norm(q_down) if snark: dirStr = f'{zkDataDir}/pos_{start_pos}/layer_{self.layer_id}' os.makedirs(dirStr, exist_ok=True) saveTensor(f'{dirStr}/q_norm_x.bin', q_down.cpu()) saveTensor(f'{dirStr}/q_norm_weight.bin', self.q_norm.weight.view(torch.uint32).cpu()) saveTensor(f'{dirStr}/q_norm_rms.bin', rms.cpu()) saveTensor(f'{dirStr}/q_norm_y.bin', q_normed.cpu()) # q 的 rescale 为 2^19 # q = self.wq_b(q_normed) q_nope = self.wq_b1(q_normed) q_pe = self.wq_b2(q_normed) # 在pytorch中view函数的作用为重构张量的维度 # q = q.view(bsz, seqlen, self.n_local_heads, self.qk_head_dim) q_nope = q_nope.view(bsz, seqlen, self.n_local_heads, self.qk_nope_head_dim) q_pe = q_pe.view(bsz, seqlen, self.n_local_heads, self.qk_rope_head_dim) # 将query的投影表示按照最后一个维度拆分，前面qk_nope_head_dim维(128)作为query不包含位置信息的表示q_nope，后面qk_rope_head_dim维(64)添加rope位置信息 # (调用apply_rotary_emb函数，参考秀才经商：DeepSeek源码解析之RoPE)作为query包含位置信息的表示q_pe(即公式39)； # q_nope 的维度[batch, seqLen, 128, 128], q_pe 的维度 [batch, seqLen, 128, 64] # q_nope, q_pe 的 rescale 为 2^19 # q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) # freqs_cis 的 rescale 为 2^42, 计算之后 q_pe 的 rescale 为 2^19 if snark: saveTensor(f'{zkDataDir}/pos_{start_pos}/layer_{self.layer_id}/q_pe_x.bin', q_pe.cpu()) saveTensor(f'{zkDataDir}/freqs_cis.bin', freqs_cis.cpu()) q_pe = apply_rotary_emb(q_pe, freqs_cis) if snark: saveTensor(f'{zkDataDir}/pos_{start_pos}/layer_{self.layer_id}/q_pe_y.bin', self.q_norm.weight.view(torch.uint32).cpu()) # 获取key和value的联合表示kv(即公式41中的)和包含位置信息的key表示k_pe(即公式43中的)：输入乘以向下投影矩阵wkv_a后，按照最后一个维度拆分， # 前面kv_lora_rank维作为key和value的联合表示，后面qk_rope_head_dim维添加rope位置信息(调用apply_rotary_emb)后得到包含rope位置信息的key表示； # x 的resacle 为 2^21, kv shape [batch, seqLen, 512], kv 的resacle 为 2^21 kv, kv_rem = self.wkv_a1(x) if snark: dirStr = f'{zkDataDir}/pos_{start_pos}/layer_{self.layer_id}' os.makedirs(dirStr, exist_ok=True) saveTensor(f'{dirStr}/wkv_a1_x.bin', x.cpu()) saveTensor(f'{dirStr}/wkv_a1_w.bin', self.wkv_a1.weight.view(torch.uint32).cpu()) saveTensor(f'{dirStr}/wkv_a1_y.bin', kv.cpu()) saveTensor(f'{dirStr}/wkv_a1_r.bin', kv_rem.cpu()) k_pe, _ = self.wkv_a2(x) # print(f'k_pe 1 shape: {k_pe.shape}', flush=True) # unsqueeze()用于增加一个维度, k_pe.unsqueeze(2) 把 k_pe reshape 成 [batch, seqLen, 1, dim] # # kv, k_pe 的resacle 为 2^21 k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis) # print(f'k_pe 2 shape: {k_pe.shape}', flush=True) if attn_impl == "naive": q = torch.cat([q_nope, q_pe], dim=-1) kv = self.wkv_b(self.kv_norm(kv)) kv = kv.view(bsz, seqlen, self.n_local_heads, self.qk_nope_head_dim + self.v_head_dim) k_nope, v = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) k = torch.cat([k_nope, k_pe.expand(-1, -1, self.n_local_heads, -1)], dim=-1) self.k_cache[:bsz, start_pos:end_pos] = k self.v_cache[:bsz, start_pos:end_pos] = v scores = torch.einsum("bshd,bthd->bsht", q, self.k_cache[:bsz, :end_pos]) * self.softmax_scale else: # 计算query和key的注意力： # query中不包含位置信息的q_nope(乘以了key的向上投影矩阵后)与缓存kv_cache中的key表示求内积； # query中包含位置信息的q_pe与缓存pe_cache中的key表示求内积； # 两者相加后乘以softmax缩放系数softmax_scale # q_nope 的维度[batch, seqLen, 128, 128], wkv_b_1 shape: [128, 128, 512] # q_nope rescale 2^19, wkv_b_1 rescale 2 ** 32 # q_nope = torch.einsum("bshd,hdc->bshc", q_nope, wkv_b_1) # 调用 einsum_bshd_hdc_bshc 之后, q_nope维度 [batch, seqLen, 128, 512] wkv_b_1 = self.wkv_b_1.weight.view(self.n_local_heads, -1, self.kv_lora_rank) q_nope = einsum_bshd_hdc_bshc(q_nope.contiguous(), wkv_b_1.contiguous(), self.wkv_b_1.scale.item()) # print('q_nope type: ' + str(q_nope.type())) # print('q_nope shape: ' + str(q_nope.shape)) # kv_normed 的 rescale 为 2^23 (kv_normed, rms) = self.kv_norm(kv) # kv_cache 的 rescale 为 2^23, shape [batch, seqLen, 512], self.kv_cache[:bsz, start_pos:end_pos] = kv_normed # self.kv_cache[:bsz, start_pos:end_pos] = kv2 # kv = (kv.detach().to(torch.float32) * (2 ** -23)).to(torch.bfloat16) # pe_cache 的 rescale 为 2^21 self.pe_cache[:bsz, start_pos:end_pos] = k_pe.squeeze(2) # q_nope rescale: 2^19, kv_cache rescale: 2^23 # q_nope 的维度 [batch, seqLen, 128, 512], kv_cache 维度 (batch, args.max_seq_len, 512) # score1 = torch.einsum("bshc,btc->bsht", q_nope, self.kv_cache[:bsz, :end_pos]) kv_cache1 = self.kv_cache[:bsz, :end_pos] # score1 = einsum_bshc_btc_bsht(q_nope.contiguous(), kv_cache1.contiguous(), 25) # score1 的 rescale 为 2^19 score1 = einsum_bshc_btc_bsht(q_nope.contiguous(), kv_cache1.contiguous(), 23) # print(f'kv_cache1 type: {kv_cache1.type()}, shape: {kv_cache1.shape}', flush=True) # score1 = (score1.detach().to(torch.float32) * (2 ** -21)).to(torch.bfloat16) # score2 = torch.einsum("bshr,btr->bsht", q_pe, self.pe_cache[:bsz, :end_pos]) pe_cache1 = self.pe_cache[:bsz, :end_pos] # score2 = einsum_bshc_btc_bsht(q_pe.contiguous(), pe_cache1.contiguous(), 23) # q_pe 的 rescale 为 2^19, scores2 的rescale 为 2^19 score2 = einsum_bshc_btc_bsht(q_pe.contiguous(), pe_cache1.contiguous(), 21) # score2 = (score2.detach().to(torch.float32) * (2 ** -21)).to(torch.bfloat16) # scores = (score1 + score2) * self.softmax_scale # scores 的 rescale 为 2 ** 19 scores = (score1 + score2) * self.softmax_scale1 // self.softmax_scale2 # scores = torch.round(((score1 + score2) * self.softmax_scale1).to(torch.float32) / self.softmax_scale2).to(torch.int64) # mask 在 unsqueeze(1) 之后的 shape 为 [seqLen, 1, senLen], scores 的shape 为 [batch, seqLen, heads , t] if mask is not None: # print('mask type: ' + str(mask.type())) # print('mask shape: ' + str(mask.shape)) scores += mask.unsqueeze(1) # query和key的内积按照最后一个维度计算softmax值； # scores = scores.softmax(dim=-1, dtype=torch.float32).type_as(x) scores_new = torch.empty_like(scores, dtype=torch.int64, device='cuda') # scores 和 scores_new 的 rescale 为 2 ** 19, shape: [bsz, seqLen, headCount, seqLen] # # softmax_q19 会破坏 scores 的原始数据，先拷贝一份数据 if snark: saveTensor(f'{zkDataDir}/pos_{start_pos}/layer_{self.layer_id}/scores_softmax_x.bin', scores.contiguous().cpu()) softmax_q19(scores.contiguous(), scores_new) if snark: saveTensor(f'{zkDataDir}/pos_{start_pos}/layer_{self.layer_id}/scores_softmax_y.bin', scores_new.cpu()) if attn_impl == "naive": x = torch.einsum("bsht,bthd->bshd", scores, self.v_cache[:bsz, :end_pos]) else: kv_cache2 = self.kv_cache[:bsz, :end_pos] # kv_cache2 = (kv_cache2.detach().to(torch.float32) * (2 ** -25)).to(torch.bfloat16) # x = (x.detach().to(torch.float32) * (2 ** -23)).to(torch.bfloat16) # 计算最终输出： # 注意力分数乘以kv缓存后，再乘以value的向上投影矩阵wkv_b(实现公式45和46)； # 乘以输出投影矩阵wo(公式47)； # x = torch.einsum("bsht,btc->bshc", scores_new, kv_cache2) # scores_new 的 rescale 为 2^19, kv_cache2 的 rescale 为 2^23, bshc 的 rescale 为 2^19 # scores_new shape: [1, 8, 128, 8], bshc shape: [1, 8, 128, 512] # bshc = einsum_bsht_btc_bshc(scores_new.contiguous(), kv_cache2.contiguous(), 25) bshc = einsum_bsht_btc_bshc(scores_new.contiguous(), kv_cache2.contiguous(), 23) # # v_head_dim = 128, kv_lora_rank = 512, n_local_heads = 128 # wkv_b_2 = wkv_b[:, -self.v_head_dim:] # # print('wkv_b 2 type: ' + str(wkv_b_2.type())) # # print('wkv_b 2 shape: ' + str(wkv_b_2.shape)) wkv_b_2 = self.wkv_b_2.weight wkv_b_2 = wkv_b_2.view(self.n_local_heads, -1, self.kv_lora_rank) # wkv_b_2 = (wkv_b_2.detach().to(torch.float32) * (2 ** -self.wkv_b_2.scale.item())).to(torch.bfloat16) # x = torch.einsum("bshc,hdc->bshd", x, wkv_b_2) # bshc 的 rescale 为 2^19, wkv_b_2 的 rescale 为 self.wkv_b_2.scale # x 的 rescale 为 2 ** 19 # bshc shape: [1, seqLen, 128, 512], wkv_b_2 shape: [128, 128, 512] x = einsum_bshc_hdc_bshd(bshc.contiguous(), wkv_b_2.contiguous(), self.wkv_b_2.scale.item()) # x = (x.detach().to(torch.float32) * (2 ** -21)).to(torch.bfloat16) # x 返回的的 shape [1, seqLen, 7168] x = self.wo(x.flatten(2)) return x class MLP(nn.Module): """ Multi-Layer Perceptron (MLP) used as a feed-forward layer. Attributes: w1 (nn.Module): Linear layer for input-to-hidden transformation. w2 (nn.Module): Linear layer for hidden-to-output transformation. w3 (nn.Module): Additional linear layer for feature transformation. """ def __init__(self, layer_id, dim: int, inter_dim: int): """ Initializes the MLP layer. Args: dim (int): Input and output dimensionality. inter_dim (int): Hidden layer dimensionality. """ super().__init__() self.w1 = ColumnParallelLinear(layer_id, dim, inter_dim) self.w2 = RowParallelLinear(layer_id, inter_dim, dim) self.w3 = ColumnParallelLinear(layer_id, dim, inter_dim) def forward(self, x: torch.Tensor) -> torch.Tensor: """ Forward pass for the MLP layer. Args: x (torch.Tensor): Input tensor. Returns: torch.Tensor: Output tensor after MLP computation. """ return self.w2(F.silu(self.w1(x)) * self.w3(x)) class MLP_int(nn.Module): """ Multi-Layer Perceptron (MLP) used as a feed-forward layer. Attributes: w1 (nn.Module): Linear layer for input-to-hidden transformation. w2 (nn.Module): Linear layer for hidden-to-output transformation. w3 (nn.Module): Additional linear layer for feature transformation. """ def __init__(self, layer_id, dim: int, inter_dim: int): """ Initializes the MLP layer. Args: dim (int): Input and output dimensionality. inter_dim (int): Hidden layer dimensionality. """ super().__init__() self.layer_id = layer_id self.w1 = ColumnParallelLinear_rescale_int(layer_id, dim, inter_dim, 1, 1, torch.int32) self.w2 = RowParallelLinear_rescale_int(layer_id, inter_dim, dim, 1, 1, 1, torch.int32) self.w3 = ColumnParallelLinear_rescale_int(layer_id, dim, inter_dim, 1, 1, torch.int32) # 输入的 x 的rescale 为 2^23, [bsz, seqLen, 7168] def forward(self, start_pos: int, x: torch.Tensor) -> torch.Tensor: """ Forward pass for the MLP layer. Args: x (torch.Tensor): Input tensor. Returns: torch.Tensor: Output tensor after MLP computation. """ # r1 shape: [bsz, seqLen, inter_dim], r1 rescale: 2^23 r1 = self.w1(x) # s1 = F.silu(r1) # s1 shape: [bsz, seqLen, inter_dim], s1 rescale: 2^23 s1 = torch.empty_like(r1, dtype=torch.int64, device='cuda') # silu_q25(r1, s1) if snark: saveTensor(f'{zkDataDir}/pos_{start_pos}/layer_{self.layer_id}/mlp_silu_x.bin', r1.contiguous().cpu()) silu_q23(r1, s1) if snark: saveTensor(f'{zkDataDir}/pos_{start_pos}/layer_{self.layer_id}/mlp_silu_y.bin', s1.cpu()) # r2 rescale: 2^23, shape: [1, seqLen, inter_dim] r2 = self.w3(x) # 返回的 shape [bsz, seqLen, dim] q = self.w2(s1 * r2 // (1 << 23)) return q class Gate(nn.Module): """ Gating mechanism for routing inputs in a mixture-of-experts (MoE) model. Attributes: dim (int): Dimensionality of input features. topk (int): Number of top experts activated for each input. n_groups (int): Number of groups for routing. topk_groups (int): Number of groups to route inputs to. score_func (str): Scoring function ('softmax' or 'sigmoid'). route_scale (float): Scaling factor for routing weights. weight (torch.nn.Parameter): Learnable weights for the gate. bias (Optional[torch.nn.Parameter]): Optional bias term for the gate. """ def __init__(self, layer_id: int, args: ModelArgs): """ Initializes the Gate module. Args: args (ModelArgs): Model arguments containing gating parameters. """ super().__init__() self.layer_id = layer_id self.dim = args.dim # n_activated_experts = 8 self.topk = args.n_activated_experts # n_expert_groups = 8 self.n_groups = args.n_expert_groups # n_limited_groups = 4 self.topk_groups = args.n_limited_groups # score_func = 'sigmoid' self.score_func = args.score_func # route_scale = 2.5 self.route_scale = args.route_scale # n_routed_experts = 256 # self.weight = nn.Parameter(torch.empty(args.n_routed_experts, args.dim)) self.register_buffer("weight", torch.empty(args.n_routed_experts, args.dim, dtype=torch.int32)) self.register_buffer("scale", torch.tensor(0, dtype=torch.int32)) # self.bias = nn.Parameter(torch.empty(args.n_routed_experts, dtype=torch.int32)) if self.dim == 7168 else None if self.dim == 7168: self.register_buffer("bias", torch.empty(args.n_routed_experts, dtype=torch.int32)) else: self.bias = None # x 的 rescale 为 2^23 def forward(self, start_pos: int, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """ Forward pass for the gating mechanism. Args: x (torch.Tensor): Input tensor. Returns: Tuple[torch.Tensor, torch.Tensor]: Routing weights and selected expert indices. """ x = x.view(1, -1, self.dim) # scores = linear(x, self.weight) # self.weight shape: [256, 7168] # 当前 scores shape: [1, seqLen, 256] # rescale = 2 ** self.scale.item() rescale = self.scale.item() # scores 的 rescale 为 2^23 scores, scores_rem = linear_int(x, self.weight, 1, 1, rescale) # scores = int64_bmm_with_bias(x, self.weight, bias, 1, 1, self.scale) # x shape: [seqLen, 7168] x = x.view(-1, self.dim) if self.score_func == "softmax": scores = scores.softmax(dim=-1, dtype=torch.float32) else: # scores = scores.sigmoid() C = torch.empty_like(scores, dtype=torch.int64, device='cuda') if snark: saveTensor(f'{zkDataDir}/pos_{start_pos}/layer_{self.layer_id}/sigmoid_gate_x.bin', scores.cpu()) saveTensor(f'{zkDataDir}/pos_{start_pos}/layer_{self.layer_id}/sigmoid_gate_r.bin', scores_rem.cpu()) sigmoid_q23(scores, C) if snark: saveTensor(f'{zkDataDir}/pos_{start_pos}/layer_{self.layer_id}/sigmoid_gate_y.bin', C.cpu()) # 当前 scores shape: [seqLen, 256] scores = C.squeeze(0) # bias的rescale为2^23 original_scores = scores if self.bias is not None: # scores = scores + self.bias # 当前 scores shape: [seqLen, 256] scores = scores + self.bias if snark: saveTensor(f'{zkDataDir}/pos_{start_pos}/layer_{self.layer_id}/gate_original_scores.bin', original_scores.contiguous().cpu()) saveTensor(f'{zkDataDir}/pos_{start_pos}/layer_{self.layer_id}/gate_bias.bin', self.bias.view(torch.uint32).cpu()) # n_groups = 8 if self.n_groups > 1: # x.size(0) = 8，当前 scores shape: [seqLen, 8, 32] scores = scores.view(x.size(0), self.n_groups, -1) # print(f'scores shape 111: {scores.shape}', flush=True) if self.bias is None: group_scores = scores.amax(dim=-1) else: # topk 返回 -1维度上 最大的 前 2 个值，同时返回值和索引，[0] 表示 取值，sum(-1) 再把最大的两个值相加. # 256维，分成8个组，每个组挑最大的两个数相加，得到 [seqLen, 8] 的结果，代表 8 个组的 最大两个值的和。 # group_scores 的 shape: [8, 8] group_scores = scores.topk(2, dim=-1)[0].sum(dim=-1) # print(group_scores[0], flush=True) # print(f'group_scores shape: {group_scores.shape}') # topk_groups = 4, 从 8 个group中选择最大的 4个，返回其索引，比如返回 [[0, 2, 4, 6], ...] # indices shape: [seqLen, 4] indices = group_scores.topk(self.topk_groups, dim=-1)[1] # print(indices[0], flush=True) # mask shape: [seqLen, 8] # scatter_: 按照给定索引，把某个源张量的值写入到目标张量对应位置。 Tensor.scatter_(dim, index, src, reduce=None) # 比如 mask 为[[False, True, False, True, False, True, False, True], ...] # mask: 每一行最大的4个值相对应的 mask 为 False mask = scores.new_ones(x.size(0), self.n_groups, dtype=bool).scatter_(1, indices, False) # print(mask[0], flush=True) # 把满足布尔 mask 的位置替换成 "-inf", mask.unsqueeze(-1) shape: [8, 8, 1] # 把 scores 中 淘汰掉的4个group中的每一个值设置为 "-inf",总共设置 128个 "-inf"，占每一行中的一半 # scores shape: [seqLen, 256] # scores = scores.masked_fill_(mask.unsqueeze(-1), float("-inf")).flatten(1) scores = scores.masked_fill_(mask.unsqueeze(-1), -(1 << 42)).flatten(1) # 没有淘汰掉的group中的 128个值中，选择最大的8个值，返回其下标 # self.topk = 8, indices shape: [8, 8] indices = torch.topk(scores, self.topk, dim=-1)[1] # print(indices[0], flush=True) # gather 用来按照索引从一个张量中取值，按照8个最大值的下标，获取其值 # weights shape: [8, 8] weights = original_scores.gather(1, indices) if snark: saveTensor(f'{zkDataDir}/pos_{start_pos}/layer_{self.layer_id}/gate_indices.bin', indices.contiguous().cpu()) saveTensor(f'{zkDataDir}/pos_{start_pos}/layer_{self.layer_id}/gate_weights.bin', weights.contiguous().cpu()) # print(f'weights shape: {weights.shape}') if self.score_func == "sigmoid": sum1 = weights.sum(dim=-1, keepdim=True) # weights = (weights * (2 ** 25) + sum1 // 2) // sum1 weights = (weights * (2 ** 23)) // sum1 # weights /= weights.sum(dim=-1, keepdim=True) #self.route_scale = 2.5 # weights *= self.route_scale weights = weights * 5 // 2 # weights = (weights.to(torch.float32) * (2 ** -23)).to(torch.bfloat16) # return weights.type_as(x), indices return weights, indices class Expert_int(nn.Module): """ Expert layer for Mixture-of-Experts (MoE) models. Attributes: w1 (nn.Module): Linear layer for input-to-hidden transformation. w2 (nn.Module): Linear layer for hidden-to-output transformation. w3 (nn.Module): Additional linear layer for feature transformation. """ def __init__(self, layer_id, idx, dim: int, inter_dim: int): """ Initializes the Expert layer. Args: dim (int): Input and output dimensionality. inter_dim (int): Hidden layer dimensionality. """ super().__init__() # # w1 shape: [2048, 7168] # self.w1 = Linear(layer_id, dim, inter_dim) # # w2 shape: [7168, 2048] # self.w2 = Linear(layer_id, inter_dim, dim) # # w3 shape: [2048, 7168] # self.w3 = Linear(layer_id, dim, inter_dim) self.layer_id = layer_id self.idx = idx self.w1 = Linear_rescale_int(layer_id, dim, inter_dim, 1, 1, torch.int32) self.w2 = Linear_rescale_int(layer_id, inter_dim, dim, 1, 1, torch.int32) self.w3 = Linear_rescale_int(layer_id, dim, inter_dim, 1, 1, torch.int32) # todo: add row id in the forward function def forward(self, start_pos: int, x: torch.Tensor) -> torch.Tensor: """ Forward pass for the Expert layer. Args: x (torch.Tensor): Input tensor. Returns: torch.Tensor: Output tensor after expert computation. """ # 返回的 shape [bsz, seqLen, 7168] # return self.w2(F.silu(self.w1(x)) * self.w3(x)) # r1 shape: [bsz, seqLen, 18432], r1 rescale: 2^23 r1 = self.w1(x) # s1 = F.silu(r1) # s1 shape: [bsz, seqLen, 18432], s1 rescale: 2^23 s1 = torch.empty_like(r1, dtype=torch.int64, device='cuda') # silu_q25(r1, s1) if snark: saveTensor(f'{zkDataDir}/pos_{start_pos}/layer_{self.layer_id}/expert_{self.idx}_silu_x.bin', r1.contiguous().cpu()) silu_q23(r1, s1) if snark: saveTensor(f'{zkDataDir}/pos_{start_pos}/layer_{self.layer_id}/expert_{self.idx}_silu_y.bin', s1.cpu()) # r2 rescale: 2^23 r2 = self.w3(x) # 返回的 shape [bsz, seqLen, 7168] q = self.w2((s1 * r2) >> 23) return q class MoE(nn.Module): """ Mixture-of-Experts (MoE) module. Attributes: dim (int): Dimensionality of input features. n_routed_experts (int): Total number of experts in the model. n_local_experts (int): Number of experts handled locally in distributed systems. n_activated_experts (int): Number of experts activated for each input. gate (nn.Module): Gating mechanism to route inputs to experts. experts (nn.ModuleList): List of expert modules. shared_experts (nn.Module): Shared experts applied to all inputs. """ def __init__(self, layer_id, args: ModelArgs, ckpt_path): """ Initializes the MoE module. Args: args (ModelArgs): Model arguments containing MoE parameters. """ super().__init__() self.layer_id = layer_id self.ckpt_path = ckpt_path self.dim = args.dim self.moe_inter_dim = args.moe_inter_dim assert args.n_routed_experts % world_size == 0, f"Number of experts must be divisible by world size (world_size={world_size})" self.n_routed_experts = args.n_routed_experts self.n_local_experts = args.n_routed_experts // world_size self.n_activated_experts = args.n_activated_experts self.experts_start_idx = rank * self.n_local_experts self.experts_end_idx = self.experts_start_idx + self.n_local_experts self.gate = Gate(layer_id, args) # moe_inter_dim = 2048 # self.experts = nn.ModuleList([Expert(layer_id, args.dim, args.moe_inter_dim) if self.experts_start_idx <= i < self.experts_end_idx else None # for i in range(self.n_routed_experts)]) # self.experts = torch.nn.ModuleList() # dim = 7168, n_shared_experts = 1, moe_inter_dim = 2048 self.shared_experts = MLP_int(layer_id, args.dim, args.n_shared_experts * args.moe_inter_dim) # x 的 rescale 为 2^23, shape: [1, seqLen, 7168] def forward(self, start_pos: int, x: torch.Tensor) -> torch.Tensor: """ Forward pass for the MoE module. Args: x (torch.Tensor): Input tensor. Returns: torch.Tensor: Output tensor after expert routing and computation. """ # ffn_normed 的 rescale 为 2^23 # x = (x.to(torch.float32) * (2 ** -23)).to(torch.bfloat16) # z rescale: 2^23, z 的 shape [seqLen, 7168] z = self.shared_experts(start_pos, x) # x shape 之前为: [bsz, seqLen, 7168], 之后为 [8, 7168] shape = x.size() x = x.view(-1, self.dim) # weights shape: [seqLen, 8], indices shape: [seqLen, 8] # weights 的 rescale 为 2^23 weights, indices = self.gate(start_pos, x) # y shape: [seqLen, 7168] y = torch.zeros_like(x) # torch.bincount 用来统计非负整数张量中各个数值出现的次数，类似于直方图计数 # torch.bincount(input, weights=None, minlength=0) -> Tensor, weights: 可选的一维浮点张量，和 input 形状一致。若提供，就不是“次数统计”，而是“权重和” # 统计 256 个 专家 出现的次数 counts = torch.bincount(indices.flatten(), minlength=self.n_routed_experts).tolist() for i in range(self.experts_start_idx, self.experts_end_idx): if counts[i] == 0: continue # expert = self.experts[i] with torch.device("cuda"): expert = Expert_int(self.layer_id, i, self.dim, self.moe_inter_dim) # load_model(expert, f'/data3/DeepSeek-V3-Demo1/experts-{self.layer_id}/{i}.safetensors') expertModelPath = os.path.join(self.ckpt_path, f"experts-{self.layer_id}/{i}.safetensors") load_model(expert, expertModelPath) # 第 idx 个 token, 专家 i 出现的编号是 top # 比如 # [0, 1, 3, 2, 5, 4, 6, 9] # [7, 8, 3, 12, 5, 11, 6, 1] # [16, 10, 3, 2, 15, 4, 6, 9] # [10, 21, 3, 2, 5, 4, 1, 9] # torch.where(indices == 1) 返回的结果是 ([0, 1, 3], [1, 7, 6]) idx, top = torch.where(indices == i) # expert(x[idx]) 返回的 shape [seqLen, 2048], weights[idx, top, None] 的 shape 为 [seqLen, 1], 包含一个 weight 值 # y[idx] += expert(x[idx]) * weights[idx, top, None] x2 = x[idx].unsqueeze(0) y2 = expert(start_pos, x2) y2 = y2.view(-1, self.dim) # y[idx] += y2 * weights[idx, top, None] // (1 << 25) y[idx] += y2 * weights[idx, top, None] // (1 << 23) # z = self.shared_experts(x) if world_size > 1: dist.all_reduce(y) return (y + z).view(shape) def getBF8PrintStr(ele): v = int(ele.cpu().view(torch.uint8).item()) ex = v >> 3 & 0xF r = v & 0x7 if ex == 15 and r == 7: print(f'BF8 Nan: {ex} {r} !!!', flush=True) elif ex == 0: print(f'BF8 subnormal: {ex} {r} !!!', flush=True) if v & 0x80: vstr = f'-{ex} {r}' else: vstr = f'{ex} {r}' return vstr class Block(nn.Module): """ Transformer block combining attention and feed-forward layers. Attributes: attn (nn.Module): Attention layer (MLA). ffn (nn.Module): Feed-forward network (MLP or MoE). attn_norm (nn.Module): Layer normalization for attention. ffn_norm (nn.Module): Layer normalization for feed-forward network. """ def __init__(self, layer_id: int, args: ModelArgs, ckpt_path): """ Initializes the Transformer block. Args: layer_id (int): Layer index in the transformer. args (ModelArgs): Model arguments containing block parameters. """ super().__init__() self.layer_id = layer_id self.ckpt_path = ckpt_path self.attn = MLA(layer_id, args) self.ffn = MLP_int(layer_id, args.dim, args.inter_dim) if layer_id < args.n_dense_layers else MoE(layer_id, args, ckpt_path) # print('args.dim: ' + str(args.dim)) # args.dim = 7168 self.attn_norm = RMSNorm_int(args.dim, torch.int32) self.ffn_norm = RMSNorm_int(args.dim, torch.int32) # self.ffn_norm = RMSNorm(args.dim) def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]) -> torch.Tensor: """ Forward pass for the Transformer block. Args: x (torch.Tensor): Input tensor. start_pos (int): Starting position in the sequence. freqs_cis (torch.Tensor): Precomputed complex exponential values for rotary embeddings. mask (Optional[torch.Tensor]): Mask tensor to exclude certain positions from attention. Returns: torch.Tensor: Output tensor after block computation. """ x_abs = x.abs() x_abs_min = x_abs.min().item() x_abs_max = x_abs.max().item() print(f'x abs min: {x_abs_min}, max: {x_abs_max}', flush=True) # self.attn_norm(x): 在进行attention之前，先将7168维的embeding 进行 归一化 # attn_norm 的 scale 为 2^21, x 的 scale 为 2^31 (atten_normed, rms) = self.attn_norm(x) if snark: os.makedirs(f'{zkDataDir}/pos_{start_pos}/layer_{self.layer_id}', exist_ok=True) saveTensor(f'{zkDataDir}/pos_{start_pos}/layer_{self.layer_id}/attn_norm_x.bin', x.cpu()) saveTensor(f'{zkDataDir}/pos_{start_pos}/layer_{self.layer_id}/attn_norm_weight.bin', self.attn_norm.weight.view(torch.uint32).cpu()) saveTensor(f'{zkDataDir}/pos_{start_pos}/layer_{self.layer_id}/attn_norm_y.bin', atten_normed.cpu()) saveTensor(f'{zkDataDir}/pos_{start_pos}/layer_{self.layer_id}/attn_norm_rms.bin', rms.cpu()) # attned 的 rescale 是 2^19, shape: [1, seqLen, 7168] attned = self.attn(atten_normed, start_pos, freqs_cis, mask) # 调整 rescale，因为 x 的 rescale 是 2^31, attned 的 rescale 是 2^19，因此要乘以 2^12 # x = x + attned * (2 ** 10) x = x + attned * (2 ** 12) # ffn_normed 的 rescale 为 2^23 (ffn_normed, rms) = self.ffn_norm(x) ffned = self.ffn(start_pos, ffn_normed) # x = x + ffned * (2 ** 6) x = x + ffned * (2 ** 8) # 返回的 x 的rescale 为 2^31 return x # Transformer 类在初始化中就已经明确好了自己的进程（rank），并且可以发现它是由比较经典的transformer组件构成的： # embedding层（self.embed）、堆叠的decoding block（self.layers），标准的RMSnorm层（self.norm）与最后将隐藏状态投射到词表分布的output层（self.head） # 根据前面提及的初始化的参数来看，词表大小为129280，模型的hidden dim为7168，堆叠的decode block一共有61个。维度变换会在下面举例说明。 # Transformer 由61个Block组成，每个Block有 attn 和 ffd # Transformer类在初始化中就已经明确好了自己的进程（rank），并且可以发现它是由比较经典的transformer组件构成的 # embedding层（self.embed）、堆叠的decoding block（self.layers），标准的RMSnorm层（self.norm）与最后将隐藏状态投射到词表分布的output层（self.head）。 class Transformer(nn.Module): """ Transformer model with positional embeddings, multiple layers, and output projection. Attributes: max_seq_len (int): Maximum sequence length for the transformer. embed (nn.Module): Embedding layer for input tokens. layers (torch.nn.ModuleList): List of transformer blocks. norm (nn.Module): Layer normalization applied after all blocks. head (nn.Module): Output projection layer mapping to vocabulary size. freqs_cis (torch.Tensor): Precomputed complex exponential values for rotary(旋转的) embeddings. """ def __init__(self, args: ModelArgs): """ Initializes the Transformer model. Args: args (ModelArgs): Model arguments containing transformer parameters. """ global world_size, rank world_size = dist.get_world_size() if dist.is_initialized() else 1 rank = dist.get_rank() if dist.is_initialized() else 0 Linear.dtype = torch.float8_e4m3fn if args.dtype == "fp8" else torch.bfloat16 super().__init__() self.args = args self.max_seq_len = args.max_seq_len self.embed = ParallelEmbedding(args.vocab_size, args.dim) self.layers = torch.nn.ModuleList() for layer_id in range(args.n_layers): # self.layers.append(Block(layer_id, args)) self.layers.append(nn.Module()) self.norm = RMSNorm_int(args.dim, torch.int64) # self.head = ColumnParallelLinear(-1, args.dim, args.vocab_size, dtype=torch.get_default_dtype()) # 模型中的 head 的 rescale 为 2^43, 使用的过程中的rescale为 2^35, head 输入的 rescale为 2^15, 输出的 rescale为 2^21 # self.head = ColumnParallelLinear_int(-1, args.dim, args.vocab_size, 1, (1 << 8), (1 << 29), torch.int64) self.head = ColumnParallelLinear_int(-1, args.dim, args.vocab_size, 1, (1 << 8), 29, torch.int64) # self.head = ColumnParallelLinear_int(-1, args.dim, args.vocab_size, 1, (1 << 8), (1 << 31), torch.int64) # self.head = ColumnParallelLinear_int(-1, args.dim, args.vocab_size, (1 << 5), (1 << 11), (1 << 21), torch.int64) # register_buffer()注册了名为 "freqs_cis" 的缓冲区，缓冲区的值由 precompute_freqs_cis(args) 提供，并且由于设置了 persistent=False， # 该缓冲区不会被保存到模型的状态字典中。缓冲区注册的张量是该Transformer类的位置编码。 # register_buffer 用于注册一个非参数张量（tensor），这个张量虽然不是模型的可学习参数，但仍然是模型状态的一部分。 # 与参数不同，缓冲区不会在反向传播中计算梯度，也不会被优化器更新，但它会随模型一起移动到相应的设备（如 GPU）上。 # persistent=False表示这个参数表示该缓冲区不属于持久状态（persistent state）。也就是说，当你调用 model.state_dict() 保存模型时， # 这个缓冲区不会被包含进去。位置编码可以在模型加载后重新计算，不需要存储。 self.register_buffer("freqs_cis", precompute_freqs_cis(args), persistent=False) @torch.inference_mode() def prep_inference(self, tokens: torch.Tensor, start_pos: int = 0): # softmax_init() softmax_init_q19() softmax_init_q21() silu_init_q23() seqlen = tokens.size(1) # h 是经过embed之后的结果，embed将文本表达转化为词嵌入，h的形状为 (batch_size, seq_len, 7168) h = self.embed(tokens) # h = h.to(torch.bfloat16) * (1.0 / (1 << 44)) return (h, start_pos, seqlen) @torch.inference_mode() def layer_inference(self, layer_id, h, start_pos, seqlen): freqs_cis = self.freqs_cis[start_pos:start_pos+seqlen] mask = None # triu = triangle up # 返回上三角矩阵 # 参数 k=0 代表主对角线，k 为正数则从主对角线开始向上数第 k 条，k 为负数则从主对角线开始向下数第 k 条 if seqlen > 1: # mask = torch.full((seqlen, seqlen), float("-inf"), device="cuda").triu_(1) mask = torch.full((seqlen, seqlen), -(64 << 36), dtype=torch.int64, device="cuda").triu_(1) h = self.layers[layer_id](h, start_pos, freqs_cis, mask) h_abs = (h.to(torch.float32) * (2 ** -31)).to(torch.bfloat16).abs() h_abs_max = h_abs.max() h_abs[h_abs < (2 ** -125)] = h_abs_max h_abs_min = h_abs.min() h_abs_min_str = getBF16PrintStr(h_abs_min) h_abs_max_str = getBF16PrintStr(h_abs_max) print(f'h_abs min: {h_abs_min_str}, max: {h_abs_max_str}') # 返回的 h 的rescale 为 2^31 return h @torch.inference_mode() def finish_inference(self, h): # norm的结果的scale = 2^15, h 的 scale = 2^15 h = self.norm(h)[0][:, -1] # logits 的rescale 为 2^21 logits = self.head(h[None, :]) if world_size > 1: all_logits = [torch.empty_like(logits) for _ in range(world_size)] dist.all_gather(all_logits, logits) logits = torch.cat(all_logits, dim=-1) # logits 的 scale = 2^21 return logits # # 这里开始推理了，torch.inference_mode 这句话 关闭梯度计算 并 禁止 autograd 构建计算图，同时比 torch.no_grad() 还高效，专门为推理场景优化 # @torch.inference_mode() # def forward(self, tokens: torch.Tensor, start_pos: int = 0): # """ # Forward pass for the Transformer model. # Args: # tokens (torch.Tensor): Input tensor of token IDs with shape (batch_size, seq_len). # start_pos (int, optional): Starting position in the sequence for rotary(旋转的) embeddings. Defaults to 0. # Returns: # torch.Tensor: Logits tensor of shape (batch_size, vocab_size). # """ # seqlen = tokens.size(1) # # h 是经过embed之后的结果，embed将文本表达转化为词嵌入，h的形状为 (batch_size, seq_len, 7168) # h = self.embed(tokens) # freqs_cis = self.freqs_cis[start_pos:start_pos+seqlen] # print('freqs_cis: ' + str(freqs_cis.tolist())) # mask = None # # triu = triangle up # # 返回上三角矩阵 # # 参数 k=0 代表主对角线，k 为正数则从主对角线开始向上数第 k 条，k 为负数则从主对角线开始向下数第 k 条 # if seqlen > 1: # mask = torch.full((seqlen, seqlen), float("-inf"), device=tokens.device).triu_(1) # for layer in self.layers: # h = layer(h, start_pos, freqs_cis, mask) # # 只取最后一个 token # h = self.norm(h)[:, -1] # logits = self.head(h) # if world_size > 1: # all_logits = [torch.empty_like(logits) for _ in range(world_size)] # dist.all_gather(all_logits, logits) # logits = torch.cat(all_logits, dim=-1) # return logits if __name__ == "__main__": torch.set_default_dtype(torch.bfloat16) torch.set_default_device("cuda") torch.manual_seed(0) args = ModelArgs() x = torch.randint(0, args.vocab_size, (2, 128)) model = Transformer(0, args) print(model(x).size())