192 lines
8.7 KiB
Python
192 lines
8.7 KiB
Python
import math
|
||
from typing import Optional, Literal, Tuple
|
||
|
||
import torch
|
||
import torch.nn.functional as F
|
||
from torch import nn
|
||
|
||
from .LMConfig import LMConfig
|
||
|
||
|
||
block_size = 128
|
||
attn_impl: Literal["naive", "absorb"] = "absorb"
|
||
|
||
|
||
class Linear(nn.Module):
|
||
"""
|
||
这是deepseek的线性层实现，暂时没有使用
|
||
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, in_features: int, out_features: int, bias: bool = False, dtype = None):
|
||
super().__init__()
|
||
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))
|
||
nn.init.kaiming_normal_(self.weight, mode='fan_out', nonlinearity='linear')
|
||
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 F.linear(x, self.weight, self.bias)
|
||
|
||
|
||
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.
|
||
"""
|
||
dtype = x.dtype
|
||
x = torch.view_as_complex(x.float().view(*x.shape[:-1], -1, 2))
|
||
freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1))
|
||
y = torch.view_as_real(x * freqs_cis).flatten(3)
|
||
return y.to(dtype)
|
||
|
||
|
||
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, args: LMConfig):
|
||
super().__init__()
|
||
self.dim = args.dim
|
||
self.n_heads = args.n_heads
|
||
self.n_local_heads = args.n_heads
|
||
self.q_lora_rank = args.q_lora_rank
|
||
self.kv_lora_rank = args.kv_lora_rank
|
||
self.qk_nope_head_dim = args.qk_nope_head_dim
|
||
self.qk_rope_head_dim = args.qk_rope_head_dim
|
||
self.qk_head_dim = args.qk_nope_head_dim + args.qk_rope_head_dim
|
||
self.v_head_dim = args.v_head_dim
|
||
|
||
|
||
dtype = args.torch_dtype
|
||
if self.q_lora_rank == 0:
|
||
self.wq = nn.Linear(self.dim, self.n_heads * self.qk_head_dim, bias=False, dtype=dtype)
|
||
else:
|
||
self.wq_a = nn.Linear(self.dim, self.q_lora_rank, bias=False, dtype=dtype)
|
||
self.q_norm = nn.RMSNorm(self.q_lora_rank, dtype=dtype)
|
||
self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.qk_head_dim, bias=False, dtype=dtype)
|
||
self.wkv_a = nn.Linear(self.dim, self.kv_lora_rank + self.qk_rope_head_dim, bias=False, dtype=dtype)
|
||
self.kv_norm = nn.RMSNorm(self.kv_lora_rank, dtype=dtype)
|
||
self.wkv_b = nn.Linear(self.kv_lora_rank, self.n_heads * (self.qk_nope_head_dim + self.v_head_dim), bias=False, dtype=dtype)
|
||
self.wo = nn.Linear(self.n_heads * self.v_head_dim, self.dim, bias=False, dtype=dtype)
|
||
self.softmax_scale = self.qk_head_dim ** -0.5
|
||
if args.max_seq_len > args.original_seq_len:
|
||
mscale = 0.1 * args.mscale * math.log(args.rope_factor) + 1.0
|
||
self.softmax_scale = self.softmax_scale * mscale * mscale
|
||
|
||
self.use_cache = args.use_cache
|
||
if self.use_cache:
|
||
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, dtype=dtype), 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, dtype=dtype), persistent=False)
|
||
else:
|
||
self.register_buffer("kv_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.kv_lora_rank, dtype=dtype), persistent=False)
|
||
self.register_buffer("pe_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.qk_rope_head_dim, dtype=dtype), persistent=False)
|
||
|
||
def forward(self,
|
||
x: torch.Tensor,
|
||
start_pos: int,
|
||
freqs_cis: torch.Tensor,
|
||
past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
|
||
mask: Optional[torch.Tensor] = None):
|
||
"""
|
||
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.
|
||
"""
|
||
bsz, seqlen, _ = x.size()
|
||
end_pos = start_pos + seqlen
|
||
if self.q_lora_rank == 0:
|
||
q = self.wq(x)
|
||
else:
|
||
q = self.wq_b(self.q_norm(self.wq_a(x)))
|
||
q = q.view(bsz, seqlen, self.n_local_heads, self.qk_head_dim)
|
||
q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
|
||
q_pe = apply_rotary_emb(q_pe, freqs_cis)
|
||
kv = self.wkv_a(x)
|
||
kv, k_pe = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
|
||
k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis)
|
||
v = None
|
||
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)
|
||
if self.use_cache:
|
||
self.k_cache[:bsz, start_pos:end_pos] = k
|
||
self.v_cache[:bsz, start_pos:end_pos] = v
|
||
k = self.k_cache[:bsz, :end_pos]
|
||
v = self.k_cache[:bsz, :end_pos]
|
||
scores = torch.einsum("bshd,bthd->bsht", q, k) * self.softmax_scale
|
||
else:
|
||
# 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)
|
||
wkv_b = self.wkv_b.weight
|
||
wkv_b = wkv_b.view(self.n_local_heads, -1, self.kv_lora_rank)
|
||
q_nope = torch.einsum("bshd,hdc->bshc", q_nope, wkv_b[:, :self.qk_nope_head_dim])
|
||
kv = self.kv_norm(kv)
|
||
k_pe = k_pe.squeeze(2)
|
||
if self.use_cache:
|
||
self.kv_cache[:bsz, start_pos:end_pos] = kv
|
||
self.pe_cache[:bsz, start_pos:end_pos] = k_pe
|
||
kv = self.kv_cache[:bsz, :end_pos]
|
||
k_pe = self.pe_cache[:bsz, :end_pos]
|
||
scores = (torch.einsum("bshc,btc->bsht", q_nope, kv) +
|
||
torch.einsum("bshr,btr->bsht", q_pe, k_pe)) * self.softmax_scale
|
||
if mask is not None:
|
||
scores += mask.unsqueeze(1)
|
||
scores = scores.softmax(dim=-1, dtype=torch.float32).type_as(x)
|
||
if attn_impl == "naive":
|
||
x = torch.einsum("bsht,bthd->bshd", scores, v)
|
||
else:
|
||
x = torch.einsum("bsht,btc->bshc", scores, kv)
|
||
x = torch.einsum("bshc,hdc->bshd", x, wkv_b[:, -self.v_head_dim:])
|
||
x = self.wo(x.flatten(2))
|
||
return x, None
|