llama网络结构及源码

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_len