\u6a21\u578b\u7ed3\u6784<\/h3>\n
\u4e0eGPT\u7b49\u751f\u6210\u6a21\u578b\u7c7b\u4f3c，LLaMA\u4e5f\u53ea\u4f7f\u7528\u4e86Transformer\u7684\u89e3\u7801\u5668，\u4f46\u57fa\u4e8eTransformer\u8fdb\u884c\u4e86\u4e09\u4e2a\u6539\u8fdb：<\/p>\n
\u4e0b\u9762\u9010\u4e00\u4ecb\u7ecd\u8fd9\u4e09\u4e2a\u6539\u8fdb：<\/p>\n
RMSNorm<\/h4>\n
RMSNorm（Root Mean Square Normalization）\u662f\u4e00\u79cd\u5f52\u4e00\u5316\u6280\u672f，\u7528\u4e8e\u7a33\u5b9a\u548c\u52a0\u901f\u795e\u7ecf\u7f51\u7edc\u7684\u8bad\u7ec3\u8fc7\u7a0b\u3002\u4e0e\u5176\u4ed6\u5f52\u4e00\u5316\u65b9\u6cd5（\u5982BatchNorm\u548cLayerNorm）\u4e0d\u540c，RMSNorm\u901a\u8fc7\u8ba1\u7b97\u8f93\u5165\u5f20\u91cf\u7684\u5747\u65b9\u6839（RMS）\u6765\u8fdb\u884c\u5f52\u4e00\u5316\u3002RMSNorm\u516c\u5f0f\u5982\u4e0b： RMSNorm ( x ) = x 1 d \u2211 i = 1 d x i 2 + \u03f5 \u22c5 \u03b3 \\\\text{RMSNorm}(x) = \\\\frac{x}{\\\\sqrt{\\\\frac{1}{d} \\\\sum_{i=1}^{d} x_i^2 + \\\\epsilon}} \\\\cdot \\\\gamma RMSNorm(x)=d1\u200b\u2211i=1d\u200bxi2\u200b+\u03f5 \u200bx\u200b\u22c5\u03b3 \u5176\u4e2d x x x\u662f\u8f93\u5165\u5411\u91cf， d d d \u662f\u8f93\u5165\u5411\u91cf\u7684\u7ef4\u5ea6， \u03f5 \\\\epsilon \u03f5\u662f\u4e00\u4e2a\u5c0f\u5e38\u6570，\u7528\u4e8e\u907f\u514d\u9664\u96f6\u9519\u8bef， \u03b3 \\\\gamma \u03b3\u662f\u4e00\u4e2a\u53ef\u5b66\u4e60\u7684\u7f29\u653e\u53c2\u6570\u3002<\/p>\n
LLaMa\u4e2d\u7684\u5b9e\u73b0\u5982\u4e0b：<\/p>\n
class RMSNorm(torch.nn.Module):
\n def __init__(self, dim: int, eps: float = 1e-6):
\n super().__init__()
\n self.eps = eps
\n self.weight = nn.Parameter(torch.ones(dim)) <\/p>\n
def _norm(self, x):
\n return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) <\/p>\n
def forward(self, x):
\n output = self._norm(x.float()).type_as(x)
\n return output * self.weight<\/p>\n
SwiGLU\u6fc0\u6d3b\u51fd\u6570<\/h4>\n
SwiGLU (Swish-Gated Linear Unit) \u662f\u4e00\u79cd\u7528\u4e8e\u795e\u7ecf\u7f51\u7edc\u7684\u6fc0\u6d3b\u51fd\u6570，\u5b83\u7ed3\u5408\u4e86Swish\u6fc0\u6d3b\u51fd\u6570\u548c\u95e8\u63a7\u673a\u5236，\u80fd\u591f\u6709\u6548\u5730\u589e\u5f3a\u6a21\u578b\u7684\u8868\u8fbe\u80fd\u529b\u548c\u6027\u80fd\u3002\u516c\u5f0f\u5982\u4e0b： SwiGLU ( x ) = Swish ( x ) \u22c5 ( Gated Linear Unit ( x ) ) \\\\text{SwiGLU}(x) = \\\\text{Swish}(x) \\\\cdot (\\\\text{Gated Linear Unit}(x)) SwiGLU(x)=Swish(x)\u22c5(Gated Linear Unit(x)) Swish ( x ) = x \u22c5 \u03c3 ( x ) \\\\text{Swish}(x) = x \\\\cdot \\\\sigma(x) Swish(x)=x\u22c5\u03c3(x) Gated Linear Unit ( x ) = Linear 1 ( x ) \u22c5 \u03c3 ( Linear 2 ( x ) ) \\\\text{Gated Linear Unit}(x) = \\\\text{Linear}_1(x) \\\\cdot \\\\sigma(\\\\text{Linear}_2(x)) Gated Linear Unit(x)=Linear1\u200b(x)\u22c5\u03c3(Linear2\u200b(x)) \u03c3 ( x ) = 1 1 + e \u2212 x \\\\sigma(x) = \\\\frac{1}{1 + e^{-x}} \u03c3(x)=1+e\u2212x1\u200b<\/p>\n
Linear 1 \\\\text{Linear}_1 Linear1\u200b\u548c Linear 2 \\\\text{Linear}_2 Linear2\u200b\u662f\u4e24\u4e2a\u5355\u72ec\u7684\u7ebf\u6027\u53d8\u6362\u3002<\/p>\n
LLaMa\u4ee3\u7801\u4e2d\u4f7f\u7528 F . s i l u ( x ) F.silu(x) F.silu(x)\u6dfb\u52a0SwiGLU\u6fc0\u6d3b\u51fd\u6570<\/p>\n
RoPE<\/h4>\n
\u65cb\u8f6c\u4f4d\u7f6e\u5d4c\u5165（Rotary Position Embedding, RoPE）\u662f\u4e00\u79cd\u4e3a\u5e8f\u5217\u6a21\u578b（\u5982Transformer）\u63d0\u4f9b\u4f4d\u7f6e\u7f16\u7801\u7684\u65b9\u6cd5\u3002RoPE\u901a\u8fc7\u5c06\u8f93\u5165\u5411\u91cf\u5728\u590d\u6570\u57df\u8fdb\u884c\u65cb\u8f6c\u53d8\u6362，\u6765\u7f16\u7801\u5e8f\u5217\u4e2d\u4f4d\u7f6e\u7684\u4fe1\u606f\u3002\u4e0e\u4f20\u7edf\u7684\u4f4d\u7f6e\u7f16\u7801\u65b9\u6cd5（\u5982\u6b63\u5f26-\u4f59\u5f26\u4f4d\u7f6e\u7f16\u7801）\u76f8\u6bd4，RoPE\u80fd\u591f\u66f4\u597d\u5730\u6355\u6349\u5e8f\u5217\u4e2d\u7684\u76f8\u5bf9\u4f4d\u7f6e\u4fe1\u606f，\u63d0\u9ad8\u6a21\u578b\u7684\u8868\u73b0\u529b\u3002<\/p>\n
\u65cb\u8f6c\u4f4d\u7f6e\u5d4c\u5165（RoPE）\u662f\u4e00\u79cd\u4e3a\u5e8f\u5217\u6a21\u578b\u63d0\u4f9b\u4f4d\u7f6e\u7f16\u7801\u7684\u65b9\u6cd5\u3002\u5176\u901a\u8fc7\u5c06\u8f93\u5165\u5411\u91cf\u5728\u590d\u6570\u57df\u8fdb\u884c\u65cb\u8f6c\u53d8\u6362\u6765\u7f16\u7801\u4f4d\u7f6e\u4fe1\u606f\u3002\u4ee5\u4e0b\u662fRoPE\u7684\u5177\u4f53\u5b9e\u73b0\u6b65\u9aa4：<\/p>\n
\u9891\u7387\u5411\u91cf\u7684\u8ba1\u7b97: f i = 1 \u03b8 2 i d f_i = \\\\frac{1}{\\\\theta^{\\\\frac{2i}{d}}} fi\u200b=\u03b8d2i\u200b1\u200b \u5176\u4e2d \u03b8 \\\\theta \u03b8\u662f\u4e00\u4e2a\u5e38\u6570（\u901a\u5e38\u53d6 10000）， i i i\u662f\u5411\u91cf\u7ef4\u5ea6\u7684\u7d22\u5f15\u3002<\/p>\n<\/li>\n
\u65cb\u8f6c\u89d2\u5ea6\u7684\u8ba1\u7b97: angle ( t ) = t \u22c5 f i \\\\text{angle}(t) = t \\\\cdot f_i angle(t)=t\u22c5fi\u200b \u5176\u4e2d t t t\u662f\u4f4d\u7f6e\u7d22\u5f15\u3002<\/p>\n<\/li>\n
\u5e94\u7528\u65cb\u8f6c\u53d8\u6362: \u5bf9\u6bcf\u4e2a\u4f4d\u7f6e t t t\u7684\u8f93\u5165\u5411\u91cf x t x_t xt\u200b，\u5728\u590d\u6570\u57df\u8fdb\u884c\u65cb\u8f6c\u53d8\u6362： x t \u2032 = x t \u22c5 e j \u22c5 angle ( t ) x_t\u2019 = x_t \\\\cdot e^{j \\\\cdot \\\\text{angle}(t)} xt\u2032\u200b=xt\u200b\u22c5ej\u22c5angle(t) \u5bf9\u4e8e\u4f4d\u7f6e\u7f16\u7801，\u5e38\u89c4\u7684\u505a\u6cd5\u662f\u5728\u8ba1\u7b97 query，key \u548c value \u5411\u91cf\u4e4b\u524d，\u4f1a\u8ba1\u7b97\u4e00\u4e2a\u4f4d\u7f6e\u7f16\u7801\u5411\u91cf \u52a0\u5230\u8bcd\u5d4c\u5165\u4e0a，\u4f4d\u7f6e\u7f16\u7801\u5411\u91cf\u540c\u6837\u4e5f\u662f\u7ef4\u5411\u91cf，\u7136\u540e\u518d\u4e58\u4ee5\u5bf9\u5e94\u7684\u53d8\u6362\u77e9\u9635\u3002<\/p>\n<\/li>\n
RoPE \u7684 self-attention \u64cd\u4f5c\u7684\u6d41\u7a0b\u662f：\u5bf9\u4e8e token \u5e8f\u5217\u4e2d\u7684\u6bcf\u4e2a\u8bcd\u5d4c\u5165\u5411\u91cf，\u9996\u5148\u8ba1\u7b97\u5176\u5bf9\u5e94\u7684 query \u548c key \u5411\u91cf，\u7136\u540e\u5bf9\u6bcf\u4e2a token \u4f4d\u7f6e\u90fd\u8ba1\u7b97\u5bf9\u5e94\u7684\u65cb\u8f6c\u4f4d\u7f6e\u7f16\u7801，\u63a5\u7740\u5bf9\u6bcf\u4e2a token \u4f4d\u7f6e\u7684 query \u548c key \u5411\u91cf\u7684\u5143\u7d20\u6309\u7167\u4e24\u4e24\u4e00\u7ec4\u5e94\u7528\u65cb\u8f6c\u53d8\u6362，\u6700\u540e\u518d\u8ba1\u7b97 query \u548c key \u4e4b\u95f4\u7684\u5185\u79ef\u5f97\u5230 self-attention \u7684\u8ba1\u7b97\u7ed3\u679c\u3002<\/p>\n
\u4e0b\u56fe\u5f88\u76f4\u89c2\u7684\u5c55\u793a\u4e86\u65cb\u8f6c\u53d8\u6362\u7684\u8fc7\u7a0b： ![\"image.png\"](\"https:\/\/www.wsisp.com\/helps\/wp-content\/uploads\/2025\/05\/20250507010932-681ab2ccd23ff.png\")
<\/p>\n
\u65cb\u8f6c\u7f16\u7801 RoPE \u53ef\u4ee5\u6709\u6548\u5730\u4fdd\u6301\u4f4d\u7f6e\u4fe1\u606f\u7684\u76f8\u5bf9\u5173\u7cfb，\u5373\u76f8\u90bb\u4f4d\u7f6e\u7684\u7f16\u7801\u4e4b\u95f4\u6709\u4e00\u5b9a\u7684\u76f8\u4f3c\u6027，\u800c\u8fdc\u79bb\u4f4d\u7f6e\u7684\u7f16\u7801\u4e4b\u95f4\u6709\u4e00\u5b9a\u7684\u5dee\u5f02\u6027\u3002 \u8fd9\u6837\u53ef\u4ee5\u589e\u5f3a\u6a21\u578b\u5bf9\u4f4d\u7f6e\u4fe1\u606f\u7684\u611f\u77e5\u548c\u5229\u7528\u3002\u8fd9\u4e00\u70b9\u662f\u5176\u4ed6\u7edd\u5bf9\u4f4d\u7f6e\u7f16\u7801\u65b9\u5f0f（\u5982\u6b63\u5f26\u4f4d\u7f6e\u7f16\u7801\u3001\u5b66\u4e60\u7684\u4f4d\u7f6e\u7f16\u7801\u7b49）\u6240\u4e0d\u5177\u5907\u7684，\u56e0\u4e3a\u5b83\u4eec\u53ea\u80fd\u8868\u793a\u7edd\u5bf9\u4f4d\u7f6e，\u800c\u4e0d\u80fd\u8868\u793a\u76f8\u5bf9\u4f4d\u7f6e\u3002<\/p>\n
\u4e3a\u4ec0\u4e48\u65cb\u8f6c\u4f4d\u7f6e\u5d4c\u5165\u6709\u6548？<\/p>\n
\u4e0b\u9762\u8fd9\u7bc7\u6587\u7ae0\u7ed9\u51fa\u4e86\u516c\u5f0f\u539f\u7406\u548c\u63a8\u5bfc，\u8bb2\u89e3\u5341\u5206\u8be6\u7ec6：\u70b9\u51fb\u6b64\u5904<\/p>\n
\u5728LLaMA\u4e2d，RoPE\u4f7f\u7528\u4e0b\u9762\u7684\u65b9\u5f0f\u5b9e\u73b0：<\/p>\n
def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
\n freqs = 1.0 \/ (theta ** (torch.arange(0, dim, 2)[: (dim \/\/ 2)].float() \/ dim))
\n t = torch.arange(end, device=freqs.device) # type: ignore
\n freqs = torch.outer(t, freqs).float() # type: ignore
\n freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64
\n return freqs_cis <\/p>\n
def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
\n ndim = x.ndim
\n assert 0 <= 1 < ndim
\n assert freqs_cis.shape == (x.shape[1], x.shape[-1])
\n shape = [d if i == 1 or i == ndim – 1 else 1 for i, d in enumerate(x.shape)]
\n return freqs_cis.view(*shape) <\/p>\n
def apply_rotary_emb(
\n xq: torch.Tensor,
\n xk: torch.Tensor,
\n freqs_cis: torch.Tensor,
\n) -> Tuple[torch.Tensor, torch.Tensor]:
\n xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
\n xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
\n freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
\n xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
\n xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
\n return xq_out.type_as(xq), xk_out.type_as(xk)<\/p>\n
\u4e0b\u9762\u7684\u4ee3\u7801\u7ed9\u51fa\u4e86\u52a0\u5165\u65cb\u8f6c\u4f4d\u7f6e\u5d4c\u5165\u7684\u6ce8\u610f\u529b\u673a\u5236：<\/p>\n
class Attention(nn.Module):
\n def __init__(self, args: ModelArgs):
\n super().__init__() <\/p>\n
self.n_local_heads = args.n_heads \/\/ fs_init.get_model_parallel_world_size()
\n self.head_dim = args.dim \/\/ args.n_heads <\/p>\n
self.wq = ColumnParallelLinear(
\n args.dim,
\n args.n_heads * self.head_dim,
\n bias=False,
\n gather_output=False,
\n init_method=lambda x: x,
\n )
\n self.wk = ColumnParallelLinear(
\n args.dim,
\n args.n_heads * self.head_dim,
\n bias=False,
\n gather_output=False,
\n init_method=lambda x: x,
\n )
\n self.wv = ColumnParallelLinear(
\n args.dim,
\n args.n_heads * self.head_dim,
\n bias=False,
\n gather_output=False,
\n init_method=lambda x: x,
\n )
\n self.wo = RowParallelLinear(
\n args.n_heads * self.head_dim,
\n args.dim,
\n bias=False,
\n input_is_parallel=True,
\n init_method=lambda x: x,
\n ) <\/p>\n
self.cache_k = torch.zeros(
\n (args.max_batch_size, args.max_seq_len, self.n_local_heads, self.head_dim)
\n ).cuda()
\n self.cache_v = torch.zeros(
\n (args.max_batch_size, args.max_seq_len, self.n_local_heads, self.head_dim)
\n ).cuda() <\/p>\n
def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]):
\n bsz, seqlen, _ = x.shape
\n xq, xk, xv = self.wq(x), self.wk(x), self.wv(x) <\/p>\n
xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
\n xk = xk.view(bsz, seqlen, self.n_local_heads, self.head_dim)
\n xv = xv.view(bsz, seqlen, self.n_local_heads, self.head_dim) <\/p>\n
xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis) <\/p>\n
self.cache_k = self.cache_k.to(xq)
\n self.cache_v = self.cache_v.to(xq) <\/p>\n
self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk
\n self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv <\/p>\n
keys = self.cache_k[:bsz, : start_pos + seqlen]
\n values = self.cache_v[:bsz, : start_pos + seqlen] <\/p>\n
xq = xq.transpose(1, 2)
\n keys = keys.transpose(1, 2)
\n values = values.transpose(1, 2)
\n scores = torch.matmul(xq, keys.transpose(2, 3)) \/ math.sqrt(self.head_dim)
\n if mask is not None:
\n scores = scores + mask # (bs, n_local_heads, slen, cache_len + slen)
\n scores = F.softmax(scores.float(), dim=-1).type_as(xq)
\n output = torch.matmul(scores, values) # (bs, n_local_heads, slen, head_dim)
\n output = output.transpose(
\n 1, 2
\n ).contiguous().view(bsz, seqlen, -1) <\/p>\n
return self.wo(output)<\/p>\n
\u63a5\u4e0b\u6765\u7ed9\u51faLLaMA\u5b9e\u73b0\u7684\u5168\u90e8\u4ee3\u7801：<\/p>\n
# Copyright (c) Meta Platforms, Inc. and affiliates.
\n# This software may be used and distributed according to the terms of the GNU General Public License version 3. <\/p>\n
from typing import Optional, Tuple
\nfrom dataclasses import dataclass
\nimport math <\/p>\n
import torch
\nfrom torch import nn
\nimport torch.nn.functional as F <\/p>\n
import fairscale.nn.model_parallel.initialize as fs_init
\nfrom fairscale.nn.model_parallel.layers import (
\n ParallelEmbedding,
\n RowParallelLinear,
\n ColumnParallelLinear,
\n) <\/p>\n
@dataclass
\nclass ModelArgs:
\n dim: int = 512
\n n_layers: int = 8
\n n_heads: int = 8
\n vocab_size: int = -1 # defined later by tokenizer
\n multiple_of: int = 256 # make SwiGLU hidden layer size multiple of large power of 2
\n norm_eps: float = 1e-5 <\/p>\n
max_batch_size: int = 32
\n max_seq_len: int = 2048 <\/p>\n
class RMSNorm(torch.nn.Module):
\n def __init__(self, dim: int, eps: float = 1e-6):
\n super().__init__()
\n self.eps = eps
\n self.weight = nn.Parameter(torch.ones(dim)) <\/p>\n
def _norm(self, x):
\n return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) <\/p>\n
def forward(self, x):
\n output = self._norm(x.float()).type_as(x)
\n return output * self.weight <\/p>\n
def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
\n freqs = 1.0 \/ (theta ** (torch.arange(0, dim, 2)[: (dim \/\/ 2)].float() \/ dim))
\n t = torch.arange(end, device=freqs.device) # type: ignore
\n freqs = torch.outer(t, freqs).float() # type: ignore
\n freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64
\n return freqs_cis <\/p>\n
def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
\n ndim = x.ndim
\n assert 0 <= 1 < ndim
\n assert freqs_cis.shape == (x.shape[1], x.shape[-1])
\n shape = [d if i == 1 or i == ndim – 1 else 1 for i, d in enumerate(x.shape)]
\n return freqs_cis.view(*shape) <\/p>\n
def apply_rotary_emb(
\n xq: torch.Tensor,
\n xk: torch.Tensor,
\n freqs_cis: torch.Tensor,
\n) -> Tuple[torch.Tensor, torch.Tensor]:
\n xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
\n xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
\n freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
\n xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
\n xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
\n return xq_out.type_as(xq), xk_out.type_as(xk) <\/p>\n
class Attention(nn.Module):
\n def __init__(self, args: ModelArgs):
\n super().__init__() <\/p>\n
self.n_local_heads = args.n_heads \/\/ fs_init.get_model_parallel_world_size()
\n self.head_dim = args.dim \/\/ args.n_heads <\/p>\n
self.wq = ColumnParallelLinear(
\n args.dim,
\n args.n_heads * self.head_dim,
\n bias=False,
\n gather_output=False,
\n init_method=lambda x: x,
\n )
\n self.wk = ColumnParallelLinear(
\n args.dim,
\n args.n_heads * self.head_dim,
\n bias=False,
\n gather_output=False,
\n init_method=lambda x: x,
\n )
\n self.wv = ColumnParallelLinear(
\n args.dim,
\n args.n_heads * self.head_dim,
\n bias=False,
\n gather_output=False,
\n init_method=lambda x: x,
\n )
\n self.wo = RowParallelLinear(
\n args.n_heads * self.head_dim,
\n args.dim,
\n bias=False,
\n input_is_parallel=True,
\n init_method=lambda x: x,
\n ) <\/p>\n
self.cache_k = torch.zeros(
\n (args.max_batch_size, args.max_seq_len, self.n_local_heads, self.head_dim)
\n ).cuda()
\n self.cache_v = torch.zeros(
\n (args.max_batch_size, args.max_seq_len, self.n_local_heads, self.head_dim)
\n ).cuda() <\/p>\n
def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]):
\n bsz, seqlen, _ = x.shape
\n xq, xk, xv = self.wq(x), self.wk(x), self.wv(x) <\/p>\n
xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
\n xk = xk.view(bsz, seqlen, self.n_local_heads, self.head_dim)
\n xv = xv.view(bsz, seqlen, self.n_local_heads, self.head_dim) <\/p>\n
xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis) <\/p>\n
self.cache_k = self.cache_k.to(xq)
\n self.cache_v = self.cache_v.to(xq) <\/p>\n
self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk
\n self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv <\/p>\n
keys = self.cache_k[:bsz, : start_pos + seqlen]
\n values = self.cache_v[:bsz, : start_pos + seqlen] <\/p>\n
xq = xq.transpose(1, 2)
\n keys = keys.transpose(1, 2)
\n values = values.transpose(1, 2)
\n scores = torch.matmul(xq, keys.transpose(2, 3)) \/ math.sqrt(self.head_dim)
\n if mask is not None:
\n scores = scores + mask # (bs, n_local_heads, slen, cache_len + slen)
\n scores = F.softmax(scores.float(), dim=-1).type_as(xq)
\n output = torch.matmul(scores, values) # (bs, n_local_heads, slen, head_dim)
\n output = output.transpose(
\n 1, 2
\n ).contiguous().view(bsz, seqlen, -1) <\/p>\n
return self.wo(output) <\/p>\n
class FeedForward(nn.Module):
\n def __init__(
\n self,
\n dim: int,
\n hidden_dim: int,
\n multiple_of: int,
\n ):
\n super().__init__()
\n hidden_dim = int(2 * hidden_dim \/ 3)
\n hidden_dim = multiple_of * ((hidden_dim + multiple_of – 1) \/\/ multiple_of) <\/p>\n
self.w1 = ColumnParallelLinear(
\n dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x
\n )
\n self.w2 = RowParallelLinear(
\n hidden_dim, dim, bias=False, input_is_parallel=True, init_method=lambda x: x
\n )
\n self.w3 = ColumnParallelLinear(
\n dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x
\n ) <\/p>\n
def forward(self, x):
\n return self.w2(F.silu(self.w1(x)) * self.w3(x)) <\/p>\n
class TransformerBlock(nn.Module):
\n def __init__(self, layer_id: int, args: ModelArgs):
\n super().__init__()
\n self.n_heads = args.n_heads
\n self.dim = args.dim
\n self.head_dim = args.dim \/\/ args.n_heads
\n self.attention = Attention(args)
\n self.feed_forward = FeedForward(
\n dim=args.dim, hidden_dim=4 * args.dim, multiple_of=args.multiple_of
\n )
\n self.layer_id = layer_id
\n self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
\n self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps) <\/p>\n
def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, mask: Optional[torch.Tensor]):
\n h = x + self.attention.forward(self.attention_norm(x), start_pos, freqs_cis, mask)
\n out = h + self.feed_forward.forward(self.ffn_norm(h))
\n return out <\/p>\n
class Transformer(nn.Module):
\n def __init__(self, params: ModelArgs):
\n super().__init__()
\n self.params = params
\n self.vocab_size = params.vocab_size
\n self.n_layers = params.n_layers <\/p>\n
self.tok_embeddings = ParallelEmbedding(
\n params.vocab_size, params.dim, init_method=lambda x: x
\n ) <\/p>\n
self.layers = torch.nn.ModuleList()
\n for layer_id in range(params.n_layers):
\n self.layers.append(TransformerBlock(layer_id, params)) <\/p>\n
self.norm = RMSNorm(params.dim, eps=params.norm_eps)
\n self.output = ColumnParallelLinear(
\n params.dim, params.vocab_size, bias=False, init_method=lambda x: x
\n ) <\/p>\n
self.freqs_cis = precompute_freqs_cis(
\n self.params.dim \/\/ self.params.n_heads, self.params.max_seq_len * 2
\n ) <\/p>\n
@torch.inference_mode()
\n def forward(self, tokens: torch.Tensor, start_pos: int):
\n _bsz, seqlen = tokens.shape
\n h = self.tok_embeddings(tokens)
\n self.freqs_cis = self.freqs_cis.to(h.device)
\n freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen] <\/p>\n
mask = None
\n if seqlen > 1:
\n mask = torch.full((1, 1, seqlen, seqlen), float("-inf"), device=tokens.device)
\n mask = torch.triu(mask, diagonal=start_pos + 1).type_as(h) <\/p>\n
for layer in self.layers:
\n h = layer(h, start_pos, freqs_cis, mask)
\n h = self.norm(h)
\n output = self.output(h[:, -1, :]) # only compute last logits
\n return output.float()<\/p>\n
\n
CSDN\u72ec\u5bb6\u798f\u5229<\/font><\/h2>\n\u6700\u540e，\u611f\u8c22\u6bcf\u4e00\u4e2a\u8ba4\u771f\u9605\u8bfb\u6211\u6587\u7ae0\u7684\u4eba，\u793c\u5c1a\u5f80\u6765\u603b\u662f\u8981\u6709\u7684，\u4e0b\u9762\u8d44\u6599\u867d\u7136\u4e0d\u662f\u4ec0\u4e48\u5f88\u503c\u94b1\u7684\u4e1c\u897f，\u5982\u679c\u4f60\u7528\u5f97\u5230\u7684\u8bdd\u53ef\u4ee5\u76f4\u63a5\u62ff\u8d70： ![\"\"](\"https:\/\/www.wsisp.com\/helps\/wp-content\/uploads\/2025\/05\/20250507010933-681ab2cd27cc3.png\")
<\/p>\n","protected":false},"excerpt":{"rendered":"
\u6587\u7ae0\u6d4f\u89c8\u9605\u8bfb897\u6b21\uff0c\u70b9\u8d5e18\u6b21\uff0c\u6536\u85cf21\u6b21\u3002\u4e0b\u56fe\u5f88\u76f4\u89c2\u7684\u5c55\u793a\u4e86\u65cb\u8f6c\u53d8\u6362\u7684\u8fc7\u7a0b\uff1a\u65cb\u8f6c\u7f16\u7801 RoPE \u53ef\u4ee5\u6709\u6548\u5730\u4fdd\u6301\u4f4d\u7f6e\u4fe1\u606f\u7684\u76f8\u5bf9\u5173\u7cfb\uff0c_swigu layer<\/p>\n","protected":false},"author":2,"featured_media":35906,"comment_status":"open","ping_status":"open","sticky":false,"template":"","format":"standard","meta":{"footnotes":""},"categories":[1],"tags":[1699,66,347,50,2215,224,51],"topic":[],"class_list":["post-35908","post","type-post","status-publish","format-standard","has-post-thumbnail","hentry","category-server","tag-agi","tag-ai","tag-llama","tag-50","tag-2215","tag-224","tag-51"],"yoast_head":"\n
\u4e00\u6587\u8be6\u89e3LLaMa\u7cfb\u5217\u6a21\u578b\uff1a\u539f\u7406\u4ecb\u7ecd\u3001\u4ee3\u7801\u89e3\u8bfb - \u7f51\u7855\u4e92\u8054\u5e2e\u52a9\u4e2d\u5fc3<\/title>\n<meta name=\"robots\" content=\"index, follow, max-snippet:-1, max-image-preview:large, max-video-preview:-1\" \/>\n<link rel=\"canonical\" href=\"https:\/\/www.wsisp.com\/helps\/35908.html\" \/>\n<meta property=\"og:locale\" content=\"zh_CN\" \/>\n<meta property=\"og:type\" content=\"article\" \/>\n<meta property=\"og:title\" content=\"\u4e00\u6587\u8be6\u89e3LLaMa\u7cfb\u5217\u6a21\u578b\uff1a\u539f\u7406\u4ecb\u7ecd\u3001\u4ee3\u7801\u89e3\u8bfb - \u7f51\u7855\u4e92\u8054\u5e2e\u52a9\u4e2d\u5fc3\" \/>\n<meta property=\"og:description\" content=\"\u6587\u7ae0\u6d4f\u89c8\u9605\u8bfb897\u6b21\uff0c\u70b9\u8d5e18\u6b21\uff0c\u6536\u85cf21\u6b21\u3002\u4e0b\u56fe\u5f88\u76f4\u89c2\u7684\u5c55\u793a\u4e86\u65cb\u8f6c\u53d8\u6362\u7684\u8fc7\u7a0b\uff1a\u65cb\u8f6c\u7f16\u7801 RoPE \u53ef\u4ee5\u6709\u6548\u5730\u4fdd\u6301\u4f4d\u7f6e\u4fe1\u606f\u7684\u76f8\u5bf9\u5173\u7cfb\uff0c_swigu layer\" \/>\n<meta property=\"og:url\" content=\"https:\/\/www.wsisp.com\/helps\/35908.html\" \/>\n<meta property=\"og:site_name\" content=\"\u7f51\u7855\u4e92\u8054\u5e2e\u52a9\u4e2d\u5fc3\" \/>\n<meta property=\"article:published_time\" content=\"2025-05-07T01:09:34+00:00\" \/>\n<meta property=\"og:image\" content=\"https:\/\/www.wsisp.com\/helps\/wp-content\/uploads\/2025\/05\/20250507010932-681ab2ccd23ff.png\" \/>\n<meta name=\"author\" content=\"admin\" \/>\n<meta name=\"twitter:card\" content=\"summary_large_image\" \/>\n<meta name=\"twitter:label1\" content=\"\u4f5c\u8005\" \/>\n\t<meta name=\"twitter:data1\" content=\"admin\" \/>\n\t<meta name=\"twitter:label2\" content=\"\u9884\u8ba1\u9605\u8bfb\u65f6\u95f4\" \/>\n\t<meta name=\"twitter:data2\" content=\"11 \u5206\" \/>\n<script type=\"application\/ld+json\" class=\"yoast-schema-graph\">{\"@context\":\"https:\/\/schema.org\",\"@graph\":[{\"@type\":\"WebPage\",\"@id\":\"https:\/\/www.wsisp.com\/helps\/35908.html\",\"url\":\"https:\/\/www.wsisp.com\/helps\/35908.html\",\"name\":\"\u4e00\u6587\u8be6\u89e3LLaMa\u7cfb\u5217\u6a21\u578b\uff1a\u539f\u7406\u4ecb\u7ecd\u3001\u4ee3\u7801\u89e3\u8bfb - \u7f51\u7855\u4e92\u8054\u5e2e\u52a9\u4e2d\u5fc3\",\"isPartOf\":{\"@id\":\"https:\/\/www.wsisp.com\/helps\/#website\"},\"datePublished\":\"2025-05-07T01:09:34+00:00\",\"dateModified\":\"2025-05-07T01:09:34+00:00\",\"author\":{\"@id\":\"https:\/\/www.wsisp.com\/helps\/#\/schema\/person\/358e386c577a3ab51c4493330a20ad41\"},\"breadcrumb\":{\"@id\":\"https:\/\/www.wsisp.com\/helps\/35908.html#breadcrumb\"},\"inLanguage\":\"zh-Hans\",\"potentialAction\":[{\"@type\":\"ReadAction\",\"target\":[\"https:\/\/www.wsisp.com\/helps\/35908.html\"]}]},{\"@type\":\"BreadcrumbList\",\"@id\":\"https:\/\/www.wsisp.com\/helps\/35908.html#breadcrumb\",\"itemListElement\":[{\"@type\":\"ListItem\",\"position\":1,\"name\":\"\u9996\u9875\",\"item\":\"https:\/\/www.wsisp.com\/helps\"},{\"@type\":\"ListItem\",\"position\":2,\"name\":\"\u4e00\u6587\u8be6\u89e3LLaMa\u7cfb\u5217\u6a21\u578b\uff1a\u539f\u7406\u4ecb\u7ecd\u3001\u4ee3\u7801\u89e3\u8bfb\"}]},{\"@type\":\"WebSite\",\"@id\":\"https:\/\/www.wsisp.com\/helps\/#website\",\"url\":\"https:\/\/www.wsisp.com\/helps\/\",\"name\":\"\u7f51\u7855\u4e92\u8054\u5e2e\u52a9\u4e2d\u5fc3\",\"description\":\"\u9999\u6e2f\u670d\u52a1\u5668_\u9999\u6e2f\u4e91\u670d\u52a1\u5668\u8d44\u8baf_\u670d\u52a1\u5668\u5e2e\u52a9\u6587\u6863_\u670d\u52a1\u5668\u6559\u7a0b\",\"potentialAction\":[{\"@type\":\"SearchAction\",\"target\":{\"@type\":\"EntryPoint\",\"urlTemplate\":\"https:\/\/www.wsisp.com\/helps\/?s={search_term_string}\"},\"query-input\":\"required name=search_term_string\"}],\"inLanguage\":\"zh-Hans\"},{\"@type\":\"Person\",\"@id\":\"https:\/\/www.wsisp.com\/helps\/#\/schema\/person\/358e386c577a3ab51c4493330a20ad41\",\"name\":\"admin\",\"image\":{\"@type\":\"ImageObject\",\"inLanguage\":\"zh-Hans\",\"@id\":\"https:\/\/www.wsisp.com\/helps\/#\/schema\/person\/image\/\",\"url\":\"https:\/\/gravatar.wp-china-yes.net\/avatar\/?s=96&d=mystery\",\"contentUrl\":\"https:\/\/gravatar.wp-china-yes.net\/avatar\/?s=96&d=mystery\",\"caption\":\"admin\"},\"sameAs\":[\"http:\/\/wp.wsisp.com\"],\"url\":\"https:\/\/www.wsisp.com\/helps\/author\/admin\"}]}<\/script>\n<!-- \/ Yoast SEO plugin. -->","yoast_head_json":{"title":"\u4e00\u6587\u8be6\u89e3LLaMa\u7cfb\u5217\u6a21\u578b\uff1a\u539f\u7406\u4ecb\u7ecd\u3001\u4ee3\u7801\u89e3\u8bfb - \u7f51\u7855\u4e92\u8054\u5e2e\u52a9\u4e2d\u5fc3","robots":{"index":"index","follow":"follow","max-snippet":"max-snippet:-1","max-image-preview":"max-image-preview:large","max-video-preview":"max-video-preview:-1"},"canonical":"https:\/\/www.wsisp.com\/helps\/35908.html","og_locale":"zh_CN","og_type":"article","og_title":"\u4e00\u6587\u8be6\u89e3LLaMa\u7cfb\u5217\u6a21\u578b\uff1a\u539f\u7406\u4ecb\u7ecd\u3001\u4ee3\u7801\u89e3\u8bfb - \u7f51\u7855\u4e92\u8054\u5e2e\u52a9\u4e2d\u5fc3","og_description":"\u6587\u7ae0\u6d4f\u89c8\u9605\u8bfb897\u6b21\uff0c\u70b9\u8d5e18\u6b21\uff0c\u6536\u85cf21\u6b21\u3002\u4e0b\u56fe\u5f88\u76f4\u89c2\u7684\u5c55\u793a\u4e86\u65cb\u8f6c\u53d8\u6362\u7684\u8fc7\u7a0b\uff1a\u65cb\u8f6c\u7f16\u7801 RoPE \u53ef\u4ee5\u6709\u6548\u5730\u4fdd\u6301\u4f4d\u7f6e\u4fe1\u606f\u7684\u76f8\u5bf9\u5173\u7cfb\uff0c_swigu layer","og_url":"https:\/\/www.wsisp.com\/helps\/35908.html","og_site_name":"\u7f51\u7855\u4e92\u8054\u5e2e\u52a9\u4e2d\u5fc3","article_published_time":"2025-05-07T01:09:34+00:00","og_image":[{"url":"https:\/\/www.wsisp.com\/helps\/wp-content\/uploads\/2025\/05\/20250507010932-681ab2ccd23ff.png"}],"author":"admin","twitter_card":"summary_large_image","twitter_misc":{"\u4f5c\u8005":"admin","\u9884\u8ba1\u9605\u8bfb\u65f6\u95f4":"11 \u5206"},"schema":{"@context":"https:\/\/schema.org","@graph":[{"@type":"WebPage","@id":"https:\/\/www.wsisp.com\/helps\/35908.html","url":"https:\/\/www.wsisp.com\/helps\/35908.html","name":"\u4e00\u6587\u8be6\u89e3LLaMa\u7cfb\u5217\u6a21\u578b\uff1a\u539f\u7406\u4ecb\u7ecd\u3001\u4ee3\u7801\u89e3\u8bfb - \u7f51\u7855\u4e92\u8054\u5e2e\u52a9\u4e2d\u5fc3","isPartOf":{"@id":"https:\/\/www.wsisp.com\/helps\/#website"},"datePublished":"2025-05-07T01:09:34+00:00","dateModified":"2025-05-07T01:09:34+00:00","author":{"@id":"https:\/\/www.wsisp.com\/helps\/#\/schema\/person\/358e386c577a3ab51c4493330a20ad41"},"breadcrumb":{"@id":"https:\/\/www.wsisp.com\/helps\/35908.html#breadcrumb"},"inLanguage":"zh-Hans","potentialAction":[{"@type":"ReadAction","target":["https:\/\/www.wsisp.com\/helps\/35908.html"]}]},{"@type":"BreadcrumbList","@id":"https:\/\/www.wsisp.com\/helps\/35908.html#breadcrumb","itemListElement":[{"@type":"ListItem","position":1,"name":"\u9996\u9875","item":"https:\/\/www.wsisp.com\/helps"},{"@type":"ListItem","position":2,"name":"\u4e00\u6587\u8be6\u89e3LLaMa\u7cfb\u5217\u6a21\u578b\uff1a\u539f\u7406\u4ecb\u7ecd\u3001\u4ee3\u7801\u89e3\u8bfb"}]},{"@type":"WebSite","@id":"https:\/\/www.wsisp.com\/helps\/#website","url":"https:\/\/www.wsisp.com\/helps\/","name":"\u7f51\u7855\u4e92\u8054\u5e2e\u52a9\u4e2d\u5fc3","description":"\u9999\u6e2f\u670d\u52a1\u5668_\u9999\u6e2f\u4e91\u670d\u52a1\u5668\u8d44\u8baf_\u670d\u52a1\u5668\u5e2e\u52a9\u6587\u6863_\u670d\u52a1\u5668\u6559\u7a0b","potentialAction":[{"@type":"SearchAction","target":{"@type":"EntryPoint","urlTemplate":"https:\/\/www.wsisp.com\/helps\/?s={search_term_string}"},"query-input":"required name=search_term_string"}],"inLanguage":"zh-Hans"},{"@type":"Person","@id":"https:\/\/www.wsisp.com\/helps\/#\/schema\/person\/358e386c577a3ab51c4493330a20ad41","name":"admin","image":{"@type":"ImageObject","inLanguage":"zh-Hans","@id":"https:\/\/www.wsisp.com\/helps\/#\/schema\/person\/image\/","url":"https:\/\/gravatar.wp-china-yes.net\/avatar\/?s=96&d=mystery","contentUrl":"https:\/\/gravatar.wp-china-yes.net\/avatar\/?s=96&d=mystery","caption":"admin"},"sameAs":["http:\/\/wp.wsisp.com"],"url":"https:\/\/www.wsisp.com\/helps\/author\/admin"}]}},"_links":{"self":[{"href":"https:\/\/www.wsisp.com\/helps\/wp-json\/wp\/v2\/posts\/35908","targetHints":{"allow":["GET"]}}],"collection":[{"href":"https:\/\/www.wsisp.com\/helps\/wp-json\/wp\/v2\/posts"}],"about":[{"href":"https:\/\/www.wsisp.com\/helps\/wp-json\/wp\/v2\/types\/post"}],"author":[{"embeddable":true,"href":"https:\/\/www.wsisp.com\/helps\/wp-json\/wp\/v2\/users\/2"}],"replies":[{"embeddable":true,"href":"https:\/\/www.wsisp.com\/helps\/wp-json\/wp\/v2\/comments?post=35908"}],"version-history":[{"count":0,"href":"https:\/\/www.wsisp.com\/helps\/wp-json\/wp\/v2\/posts\/35908\/revisions"}],"wp:featuredmedia":[{"embeddable":true,"href":"https:\/\/www.wsisp.com\/helps\/wp-json\/wp\/v2\/media\/35906"}],"wp:attachment":[{"href":"https:\/\/www.wsisp.com\/helps\/wp-json\/wp\/v2\/media?parent=35908"}],"wp:term":[{"taxonomy":"category","embeddable":true,"href":"https:\/\/www.wsisp.com\/helps\/wp-json\/wp\/v2\/categories?post=35908"},{"taxonomy":"post_tag","embeddable":true,"href":"https:\/\/www.wsisp.com\/helps\/wp-json\/wp\/v2\/tags?post=35908"},{"taxonomy":"topic","embeddable":true,"href":"https:\/\/www.wsisp.com\/helps\/wp-json\/wp\/v2\/topic?post=35908"}],"curies":[{"name":"wp","href":"https:\/\/api.w.org\/{rel}","templated":true}]}}
<\/p>\n","protected":false},"excerpt":{"rendered":"