llama网络结构及源码

模型初始化

config

lm_head

transformer

wte

rms_1/rms_2

attn

c_attn

c_proj

线性层mlp

ln_f

rope_cache

mask_cache

kv_caches

tokenizer

tokenizer初始化

tokennizer.encoder

位置编码和mask

确定最大文本长度

建立rope_cache

建立mask_cache

确定RoPE和mask

模型前向传播

生成词嵌入

隐藏层计算

确定kv_cache

进入transformer的隐藏层

attention模块计算

RMSNorm前向

计算q,k,v

词嵌入+位置编码

更新kv_cache

计算注意力

再次经过RMSNorn归一化

经过MLP层

编辑

RMSNorm归一化

网络输出

生成下一分词的循环过程

temperature

选取前topk

torch.topk(input, num)

torch.where(condition, x, y)

选择当前时刻网络生成的分词

torch.nn.functional.softmax(input, dim)

torch.multinomial(input, num_samples)

将新生成的索引加入到输入文本中

更新input_pos

idx.index_copy(dim, index, source)

第二次循环

更新输入

更新旋转位置编码

更新mask

更新词嵌入

进入transfomer隐藏层

对新生成的分词进行RMSNorm

生成新的q,k,v

更新k,v及kv_cache

生成当前位置的注意力分数

生成第二次循环的隐藏层输出

一、模型初始化

代码地址，首先模型初始化，确定模型属性

class LLaMA(nn.Module):
    def __init__(self, config: LLaMAConfig) -> None:
        super().__init__()
        assert config.padded_vocab_size is not None
        self.config = config

        self.lm_head = nn.Linear(config.n_embd, config.padded_vocab_size, bias=False)
        self.transformer = nn.ModuleDict(
            dict(
                wte=nn.Embedding(config.padded_vocab_size, config.n_embd),
                h=nn.ModuleList(Block(config) for _ in range(config.n_layer)),
                ln_f=RMSNorm(config.n_embd),
            )
        )

        self.rope_cache: Optional[RoPECache] = None
        self.mask_cache: Optional[MaskCache] = None
        self.kv_caches: List[KVCache] = []

config

确定模型参数，config使用默认的LLaMAConfig类

class LLaMAConfig:
    block_size: int = 2048
    vocab_size: int = 32000
    padded_vocab_size: Optional[int] = None
    n_layer: int = 32
    n_head: int = 32
    n_embd: int = 4096

    def __post_init__(self):
        if self.padded_vocab_size is None:
            self.padded_vocab_size = find_multiple(self.vocab_size, 64)

    @classmethod
    def from_name(cls, name: str) -> Self:
        return cls(**llama_configs[name])


llama_configs = {
    "7B": dict(n_layer=32, n_head=32, n_embd=4096),
    "13B": dict(n_layer=40, n_head=40, n_embd=5120),
    "30B": dict(n_layer=60, n_head=52, n_embd=6656),
    "65B": dict(n_layer=80, n_head=64, n_embd=8192),
}

对于所有类型的网络，不变的超参数为：
block_size: 模型处理的最大文本块的大小为2048。
vocab_size: 词汇表的大小，即模型能识别的词汇总数为32000。

padded_vocab_size经过find_multiple函数确定，用于确保词汇表大小是指定数值(64)的倍数
vocab_size=32000/64=500，词汇表大小是指定数值(64)的倍数
NLP 面试八股：“Transformers / LLM 的词表应该选多大?“ 学姐这么告诉我答案_训练时词表大小多少合适-CSDN博客
padded_vocab_size为32000

def find_multiple(n: int, k: int) -> int:
    if n % k == 0:
        return n
    return n + k - (n % k)

根据网络的参数量不同，模型层数、维度和自主意力头数不同：选择7B模型的情况下，
n_layer: 模型的层数。32
n_head: 自注意力机制中的头数。32
n_embd: 词嵌入的维度或隐藏层的维度。4096

lm_head

线性层：输入维度为4096，输出维度为32000，没有偏置

self.lm_head = nn.Linear(config.n_embd, config.padded_vocab_size, bias=False)

transformer

nn.ModuleDict()：和python字典一样存在键值对，可根据key选取网络

包括有嵌入层wte，隐藏层h和归一化层ln_f

self.transformer = nn.ModuleDict(
            dict(
                wte=nn.Embedding(config.padded_vocab_size, config.n_embd),
                h=nn.ModuleList(Block(config) for _ in range(config.n_layer)),
                ln_f=RMSNorm(config.n_embd),
            )
        )

wte

嵌入层生成词向量，输入维度为填充后词典的大小padded_vocab_size(32000)，输出维度为词嵌入维度n_embd(4096)

h

nn.ModuleList()：可以通过迭代的方式创建网络，和list的用法一致
有n_layer(32)层网络块Block
其中网络块Block，包括前后两层归一化层RMSNorm，一层自主意力层CausalSelfAttention和一层线性层MLP

class Block(nn.Module):
    def __init__(self, config: LLaMAConfig) -> None:
        super().__init__()
        self.rms_1 = RMSNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.rms_2 = RMSNorm(config.n_embd)
        self.mlp = MLP(config)

rms_1/rms_2

RMSNorm 详解三种常用标准化 Batch Norm & Layer Norm & RMSNorm_layernorm rmsnorm-CSDN博客
Llama改进之——均方根层归一化RMSNorm-CSDN博客

BatchNorm是对一个 batch 单个特征的所有样本做归一化
LayerNorm是对单个样本的所有特征做归一化

class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization.

    Derived from https://github.com/bzhangGo/rmsnorm/blob/master/rmsnorm_torch.py. BSD 3-Clause License:
    https://github.com/bzhangGo/rmsnorm/blob/master/LICENSE.
    """

    def __init__(self, size: int, dim: int = -1, eps: float = 1e-5) -> None:
        super().__init__()
        self.scale = nn.Parameter(torch.ones(size))
        self.eps = eps
        self.dim = dim

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # NOTE: the original RMSNorm paper implementation is not equivalent
        # norm_x = x.norm(2, dim=self.dim, keepdim=True)
        # rms_x = norm_x * d_x ** (-1. / 2)
        # x_normed = x / (rms_x + self.eps)
        norm_x = torch.mean(x * x, dim=self.dim, keepdim=True)
        x_normed = x * torch.rsqrt(norm_x + self.eps)
        return self.scale * x_normed

attn

CausalSelfAttention
首先声明嵌入层/隐藏层可以被注意力头数整除
包括c_attn层、c_proj层
注意力头数n_head(32),隐藏层维度n_embd(4096) ,最大文本块的大小block_size(2048)

c_attn

Q,K,V对应的线性层，输入维度为隐藏层维度n_embd(4096)，输出维度为3倍的隐藏层维度n_embd(4096)*3分别对应Q,K,V，没有偏置

c_proj

当前模块Block的输出映射，输入维度为隐藏层维度n_embd(4096)，输出维度为隐藏层维度n_embd(4096)，没有偏置

线性层mlp

hidden_dim(4*4096=16384) ，n_hidden（int(2 * hidden_dim / 3)=10922)
判断是否能被256整除，对n_hidden进行修正，(n + k - (n % k))结果为11008
两个全连接层输入维度为4096，输出维度为11008
映射层输入维度为11008，输出维度为4096

class MLP(nn.Module):
    def __init__(self, config: LLaMAConfig) -> None:
        super().__init__()
        hidden_dim = 4 * config.n_embd
        n_hidden = int(2 * hidden_dim / 3)
        n_hidden = find_multiple(n_hidden, 256)

        self.c_fc1 = nn.Linear(config.n_embd, n_hidden, bias=False)
        self.c_fc2 = nn.Linear(config.n_embd, n_hidden, bias=False)
        self.c_proj = nn.Linear(n_hidden, config.n_embd, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = F.silu(self.c_fc1(x)) * self.c_fc2(x)
        x = self.c_proj(x)
        return x

ln_f

为RMSNorm，归一化维度为n_embd(4096)

rope_cache

存储或缓存与RoPE（旋转位置编码）相关的数据。Optional[RoPECache] 表示它可以是 RoPECache 类型的对象，也可以是 None

self.rope_cache: Optional[RoPECache] = None

mask_cache

用于缓存与掩码相关的数据。Optional[MaskCache] 表示它可以是 MaskCache 类型的对象，也可以是 None。

self.mask_cache: Optional[MaskCache] = None

kv_caches

用于存储多个 KVCache 类型的缓存对象，初始化为空列表

self.kv_caches: List[KVCache] = []

二、tokenizer

tokenizer初始化

class Tokenizer:
    """Tokenizer for LLaMA."""

    def __init__(self, model_path: Path) -> None:
        self.processor = SentencePieceProcessor(model_file=str(model_path))
        self.bos_id = self.processor.bos_id()
        self.eos_id = self.processor.eos_id()
        self.pad_id = self.processor.pad_id()

tokennizer.encoder

tokenizer.encode(prompt, bos=True, eos=False, device=fabric.device)

    def encode(
        self,
        string: str,
        bos: bool = True,
        eos: bool = False,
        max_length: int = -1,
        pad: bool = False,
        device: Optional[torch.device] = None
    ) -> torch.Tensor:
        tokens = self.processor.encode(string)
        if bos:
            tokens = [self.bos_id] + tokens
        if eos:
            tokens = tokens + [self.eos_id]
        if max_length > 0:
            tokens = tokens[:max_length]
        if pad and len(tokens) < max_length:
            tokens += [self.pad_id] * (max_length - len(tokens))

        return torch.tensor(tokens, dtype=torch.int, device=device)

三、位置编码和mask

logits = model(x, max_seq_length, input_pos)

输入参数包括输入文本对应词库的idx，最大序列长度，现在输入文本的位置

 def forward(
        self, idx: torch.Tensor, max_seq_length: Optional[int] = None, input_pos: Optional[torch.Tensor] = None
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[KVCache]]]:
        B, T = idx.size()

        block_size = self.config.block_size
        if max_seq_length is None:
            max_seq_length = block_size
        assert T <= max_seq_length, f"Cannot forward sequence of length {T}, max seq length is only {max_seq_length}"
        assert max_seq_length <= block_size, f"Cannot attend to {max_seq_length}, block size is only {block_size}"
        assert T <= block_size, f"Cannot forward sequence of length {T}, block size is only {block_size}"

        if self.rope_cache is None:
            self.rope_cache = self.build_rope_cache(idx)
        if self.mask_cache is None:
            self.mask_cache = self.build_mask_cache(idx)

        if input_pos is not None:
            rope = self.rope_cache.index_select(0, input_pos)
            mask = self.mask_cache.index_select(2, input_pos)
            mask = mask[:, :, :, :max_seq_length]
        else:
            rope = self.rope_cache[:T]
            mask = self.mask_cache[:, :, :T, :T]

        # forward the model itself
        x = self.transformer.wte(idx)  # token embeddings of shape (b, t, n_embd)

        if input_pos is None:  # proxy for use_cache=False
            for block in self.transformer.h:
                x, _ = block(x, rope, mask, max_seq_length)
        else:
            if not self.kv_caches:
                head_size = self.config.n_embd // self.config.n_head
                cache_shape = (B, self.config.n_head, max_seq_length, head_size)
                self.kv_caches = [
                    (torch.zeros(cache_shape, device=x.device, dtype=x.dtype), torch.zeros(cache_shape, device=x.device, dtype=x.dtype))
                    for _ in range(self.config.n_layer)
                ]
            for i, block in enumerate(self.transformer.h):
                x, self.kv_caches[i] = block(x, rope, mask, max_seq_length, input_pos, self.kv_caches[i])

        x = self.transformer.ln_f(x)

        logits = self.lm_head(x)  # (b, t, vocab_size)

        return logits

确定最大文本长度

将输入文本长度，本次文本的最大文本长度，模型能够处理的最长文本块大小进行对比，确定本次的最大文本长度

建立rope_cache

建立位置缓存

def build_rope_cache(
    seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device, base: int = 10000
) -> RoPECache:
    """Enhanced Transformer with Rotary Position Embedding.

    Derived from: https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/
    transformers/rope/__init__.py. MIT License:
    https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/license.
    """
    # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$
    theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, dtype=dtype, device=device) / n_elem))

    # Create position indexes `[0, 1, ..., seq_len - 1]`
    seq_idx = torch.arange(seq_len, dtype=dtype, device=device)

    # Calculate the product of position index and $\theta_i$
    idx_theta = torch.outer(seq_idx, theta).float()

    cache = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=-1)

    # this is to mimic the behaviour of complex32, else we will get different results
    if dtype in (torch.float16, torch.bfloat16, torch.int8):
        cache = cache.half()
    return cache

输入的参数：模型能够处理的最大文本块大小，embeddiing维度/注意力头数

步骤：
1.计算

2.计算和最大文本块各个位置的乘积

3.计算sin(idx_theta)和cos(idx_theta)

建立mask_cache

    def build_mask_cache(self, idx: torch.Tensor) -> MaskCache:
        ones = torch.ones((self.config.block_size, self.config.block_size), device=idx.device, dtype=torch.bool)
        return torch.tril(ones).unsqueeze(0).unsqueeze(0)

步骤：
1.建立最大文本块大小*最大文本块大小的全1(True,布尔型)矩阵
2.取对角线，在增加0，1维度，mask的size为(1，1，模型处理的最大文本块大小，模型处理的最大文本块大小)

确定RoPE和mask

根据输入文本位置确定旋转位置编码和mask

 if input_pos is not None:
            rope = self.rope_cache.index_select(0, input_pos)
            mask = self.mask_cache.index_select(2, input_pos)
            mask = mask[:, :, :, :max_seq_length]
        else:
            rope = self.rope_cache[:T]
            mask = self.mask_cache[:, :, :T, :T]

根据输入文本位置确定旋转位置编码

根据输入文本位置和本次的最大文本长度确定mask，size为(1，1，输入文本长度，本次输入的最大文本长度)

四、模型前向传播

生成词嵌入

 x = self.transformer.wte(idx)

embedding层生成词嵌入

隐藏层计算

 if input_pos is None:  # proxy for use_cache=False
            for block in self.transformer.h:
                x, _ = block(x, rope, mask, max_seq_length)
        else:
            if not self.kv_caches:
                head_size = self.config.n_embd // self.config.n_head
                cache_shape = (B, self.config.n_head, max_seq_length, head_size)
                self.kv_caches = [
                    (torch.zeros(cache_shape, device=x.device, dtype=x.dtype), torch.zeros(cache_shape, device=x.device, dtype=x.dtype))
                    for _ in range(self.config.n_layer)
                ]
            for i, block in enumerate(self.transformer.h):
                x, self.kv_caches[i] = block(x, rope, mask, max_seq_length, input_pos, self.kv_caches[i])

确定kv_cache

步骤：
1.确定注意力头的维度
embedding维度/注意力头数

2.k,v的维度为(batch_size，注意力头数，本次的最大文本长度，各注意力头的维度)
3.对n_layer(32)层自注意力层，生成全0的k,v

进入transformer的隐藏层

for i, block in enumerate(self.transformer.h):
    x, self.kv_caches[i] = block(x, rope, mask, max_seq_length, input_pos, self.kv_caches[i])

每一层循环的输入包括：
词嵌入，位置嵌入，mask，本次的最大文本长度，输入位置，本次循环对应的kv_cache

    def forward(
        self,
        x: torch.Tensor,
        rope: RoPECache,
        mask: MaskCache,
        max_seq_length: int,
        input_pos: Optional[torch.Tensor] = None,
        kv_cache: Optional[KVCache] = None,
    ) -> Tuple[torch.Tensor, Optional[KVCache]]:
        h, new_kv_cache = self.attn(self.rms_1(x), rope, mask, max_seq_length, input_pos, kv_cache)
        x = x + h
        x = x + self.mlp(self.rms_2(x))
        return x, new_kv_cache

attention模块计算

更新kv_cache，计算自注意力

输入包括：
归一化层，位置编码，mask，本次的最大文本长度，输入位置，kv_cache

RMSNorm前向

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # NOTE: the original RMSNorm paper implementation is not equivalent
        # norm_x = x.norm(2, dim=self.dim, keepdim=True)
        # rms_x = norm_x * d_x ** (-1. / 2)
        # x_normed = x / (rms_x + self.eps)
        norm_x = torch.mean(x * x, dim=self.dim, keepdim=True)
        x_normed = x * torch.rsqrt(norm_x + self.eps)
        return self.scale * x_normed

计算q,k,v

    def forward(
        self,
        x: torch.Tensor,
        rope: RoPECache,
        mask: MaskCache,
        max_seq_length: int,
        input_pos: Optional[torch.Tensor] = None,
        kv_cache: Optional[KVCache] = None,
    ) -> Tuple[torch.Tensor, Optional[KVCache]]:
        B, T, C = x.size()  # batch size, sequence length, embedding dimensionality (n_embd)

        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)

        head_size = C // self.n_head
        k = k.view(B, T, self.n_head, head_size)
        q = q.view(B, T, self.n_head, head_size)
        v = v.view(B, T, self.n_head, head_size)

        q = apply_rope(q, rope)
        k = apply_rope(k, rope)

        k = k.transpose(1, 2)  # (B, nh, T, hs)
        q = q.transpose(1, 2)  # (B, nh, T, hs)
        v = v.transpose(1, 2)  # (B, nh, T, hs)

        if kv_cache is not None:
            cache_k, cache_v = kv_cache
            # check if reached token limit
            if input_pos[-1] >= max_seq_length:
                input_pos = torch.tensor(max_seq_length - 1, device=input_pos.device)
                # shift 1 position to the left
                cache_k = torch.roll(cache_k, -1, dims=2)
                cache_v = torch.roll(cache_v, -1, dims=2)
            k = cache_k.index_copy(2, input_pos, k)
            v = cache_v.index_copy(2, input_pos, v)
            kv_cache = k, v

        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
        #  att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        #  att = att.masked_fill(mask[:,:,:T,:T] == 0, float('-inf'))
        #  att = F.softmax(att, dim=-1)
        #  y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)

        # efficient attention using Flash Attention CUDA kernels
        y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0)

        y = y.transpose(1, 2).contiguous().view(B, T, C)  # re-assemble all head outputs side by side

        # output projection
        y = self.c_proj(y)

        return y, kv_cache

经过线性层生成q,k,v，并分配到不同的注意力头
q,k,v的size均为（B，输入文本长度，注意力头数，各注意力头的维度）

词嵌入+位置编码

针对q,k，将词嵌入和位置编码结合到一起
q = apply_rope(q, rope)
k = apply_rope(k, rope)
两个步骤一样，只看一个步骤就行

def apply_rope(x: torch.Tensor, rope_cache: RoPECache) -> torch.Tensor:
    # truncate to support variable sizes
    T = x.size(1)
    rope_cache = rope_cache[:T]

    # cast because the reference does
    xshaped = x.float().reshape(*x.shape[:-1], -1, 2)
    rope_cache = rope_cache.view(1, xshaped.size(1), 1, xshaped.size(3), 2)
    x_out2 = torch.stack(
        [
            xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1],
            xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1],
        ],
        -1,
    )

    x_out2 = x_out2.flatten(3)
    return x_out2.type_as(x)

输入：词嵌入，位置编码
词嵌入维度：(B，输入文本长度，注意力头数，各注意力头的维度)

位置编码维度：（输入文本长度，theta长度，sin/cos）

词嵌入和位置编码分别review
词嵌入：(B，输入文本长度，注意力头数，各注意力头的维度)-->
(B，输入文本长度，注意力头数，各注意力头的维度/2(theta长度)，2(sin/cos))

位置编码：（输入文本长度，theta长度，sin/cos）-->
(1(不是batch_size,单单强制为1)，输入文本长度，1(强制为1)，各注意力头的维度/2(theta长度)，2(sin/cos))

提取正弦和余弦：
xshaped[..., 0] 和 xshaped[..., 1] 分别表示 xshaped 中正弦和余弦的值。
rope_cache[..., 0] 和 rope_cache[..., 1] 同样表示 rope_cache 中的正弦和余弦值。

计算旋转：通过旋转公式实现了在 xshaped 的每个元素上应用 rope_cache 中存储的旋转信息。
实部 = xshaped[...,0] ∗ ropecache[...,0] − xshaped[...,1] ∗ ropecache[...,1]
虚部 = xshaped[...,1] ∗ ropecache[...,0] + xshaped[...,0] ∗ ropecache[...,1]

Stack 操作：最终的 torch.stack 操作将计算得到的实部和虚部按最后一个维度（-1）组合成一个新的张量。这样，输出张量 x_out2 的大小将是 (B，输入文本长度，注意力头数，各注意力头的维度/2(theta长度)，2(sin/cos))

xshaped = x.float().reshape(*x.shape[:-1], -1, 2)
rope_cache = rope_cache.view(1, xshaped.size(1), 1, xshaped.size(3), 2)
x_out2 = torch.stack(
        [
            xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1],
            xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1],
        ],
        -1,)

最后，将x_out2 resize为(B，输入文本长度，注意力头数，各注意力头的维度)，和输入向量x的形状一致

更新kv_cache

k = k.transpose(1, 2)  # (B, nh, T, hs)
q = q.transpose(1, 2)  # (B, nh, T, hs)
v = v.transpose(1, 2)  # (B, nh, T, hs)

if kv_cache is not None:
    cache_k, cache_v = kv_cache
    # check if reached token limit
    if input_pos[-1] >= max_seq_length:
        input_pos = torch.tensor(max_seq_length - 1, device=input_pos.device)
        # shift 1 position to the left
        cache_k = torch.roll(cache_k, -1, dims=2)
        cache_v = torch.roll(cache_v, -1, dims=2)
    k = cache_k.index_copy(2, input_pos, k)
    v = cache_v.index_copy(2, input_pos, v)
    kv_cache = k, v

1.将k,v resize为和kv_cache相匹配的形状（B，注意力头数，输入文本长度，各注意力头的维度）

2.将 k或v 的值放入 cache_k 中 input_pos 指定的索引位置，更新kv_cache （B，注意力头数，本次的最大文本长度，各注意力头的维度）

计算注意力

y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0)

# causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
#  att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
#  att = att.masked_fill(mask[:,:,:T,:T] == 0, float('-inf'))
#  att = F.softmax(att, dim=-1)
#  y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)

计算缩放点积注意力（Scaled Dot-Product Attention），利用了高效的 Flash Attention CUDA 内核以加速计算

q的size为(B，注意力头数，输入文本长度，各注意力头的维度)
k,v的size为(B，注意力头数，本次的最大文本长度，各注意力头的维度)
mask的size为(1，1，输入文本长度，本次输入的最大文本长度)

y = y.transpose(1, 2).contiguous().view(B, T, C)

转换维度，将各注意力头的维度连接到一起后，输出维度为(B，输入文本长度，embedding维度)

经过线性层维度不变，返回本次循环Block的自注意力计算结果y和kv_cache

y = self.c_proj(y)
return y, kv_cache

残差连接

x = x + h
x = x + self.mlp(self.rms_2(x))

x是进入隐藏层之前的向量，即由文本对应索引经过embedding后的词嵌入
h是经过隐藏层计算注意力分数后，由x更新的隐藏层输出

由词嵌入和隐藏层输出残差连接后，再和反馈层进行残差连接

经过MLP层

经过RMSNorn归一化，作为MLP层的输入

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # NOTE: the original RMSNorm paper implementation is not equivalent
        # norm_x = x.norm(2, dim=self.dim, keepdim=True)
        # rms_x = norm_x * d_x ** (-1. / 2)
        # x_normed = x / (rms_x + self.eps)
        norm_x = torch.mean(x * x, dim=self.dim, keepdim=True)
        x_normed = x * torch.rsqrt(norm_x + self.eps)
        return self.scale * x_normed

输出zise为 (B，输入文本长度，embedding维度)

x = F.silu(self.c_fc1(x)) * self.c_fc2(x)

分为F.silu(self.c_fc1(x))和self.c_fc2(x)两个部分相乘

维度大小的选取查看前面模型初始化的mlp

激活函数ReLU和SiLU的区别-CSDN博客
激活函数 Relu,Gelu,Mish,SiLU,Swish,Tanh,Sigmoid_gelu和silu-CSDN博客
大模型基础｜激活函数｜从ReLU 到SwiGLU
torch.nn.functional.silu — PyTorch 2.4 documentation

，其中

*: 这里的乘法是元素级的，即逐元素相乘。这意味着 F.silu(self.c_fc1(x)) 和 self.c_fc2(x) 的每个对应元素相乘，产生一个新的张量。

最后经过线性层后将11008维，还原为4096维(embedding维度)

本次循环的transfomer Block结束，进入下一次循环

RMSNorm归一化

x = self.transformer.ln_f(x)

经过32层transfomer Block后，再经过RMSNorm层进行归一化

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # NOTE: the original RMSNorm paper implementation is not equivalent
        # norm_x = x.norm(2, dim=self.dim, keepdim=True)
        # rms_x = norm_x * d_x ** (-1. / 2)
        # x_normed = x / (rms_x + self.eps)
        norm_x = torch.mean(x * x, dim=self.dim, keepdim=True)
        x_normed = x * torch.rsqrt(norm_x + self.eps)
        return self.scale * x_normed

网络输出

logits = self.lm_head(x)

输出维度为词库的大小，将输出的logits经过归一化计算选择词库中各个词的概率

五、生成下一分词的循环过程

temperature

温度越低，结果的差距越大，会使概率分布更加尖锐，从而使得模型更倾向于选择最高概率的类别。

选取前topk

        if top_k is not None:
            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
            logits = torch.where(logits < v[[-1]], -float("Inf"), logits)

torch.topk(input, num)

返回input中指定的前num个值和对应的索引

torch.where(condition, x, y)

根据条件张量选择输出。对于每个位置，condition 为 True 时取 x，为 False 时取 y。
v[[-1]] 选择 v 中的最后一个值，这个值是第 k 大的值。使用双重方括号是为了保持 v[[-1]] 的维度与 logits 相匹配。

选择当前时刻网络生成的分词

probs = torch.nn.functional.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1).to(dtype=dtype)

torch.nn.functional.softmax(input, dim)

对 input 张量的指定维度 dim 应用 softmax 操作
将 logits 转化为一个概率分布，每个值在 [0, 1] 之间，总和为 1。这样，模型的输出可以解释为各类别的预测概率

torch.multinomial(input, num_samples)

这个函数从 input 张量中根据给定的概率分布进行随机采样。
input: 在这里是 probs，是一个概率分布张量，通常表示每个类别的预测概率。
num_samples: 设定要从概率分布中采样的数量。num_samples=1 表示只需要采样一个类别

将新生成的索引加入到输入文本中

input_pos = input_pos[-1:] + 1

if idx.device.type == "xla":
    xm.mark_step()

# concatenate the new generation
idx = idx.index_copy(0, input_pos, idx_next)
# if <eos> token is triggered, return the output (stop generation)
if idx_next == eos_id:
    return idx[:input_pos]  # include the EOS token

更新input_pos

idx.index_copy(dim, index, source)

dim: 这个参数指定了要更新的维度。在这里，dim=0 表示更新张量的第一维。
index: 这是一个包含索引的张量，指示 source 张量的值应该被复制到 idx 张量的哪些位置。source: 这是一个包含要插入值的张量。在这里是 idx_next，它是一个表示新值的张量

第二次循环

更新输入

根据更新的input_pos更新输入

更新旋转位置编码

rope_cache已经保存了可处理最大数据块大小对应的各个位置的旋转位置编码
根据位置直接选取，对应位置的旋转位置编码即可

更新mask

mask_cache已经保存各个位置对应的mask
根据位置直接选取对应的mask，然后根据本次的最大文本长度对mask进行截取即可

更新词嵌入

根据文本对应索引经过embedding，生成新分词对应的词嵌入