LLaMA代码笔记 --基于lit-llama

代码来自：lit-llama
modelscope模型下载：llama-7b
下载后的模型需要转换为lit-llama使用的格式，详见 howto 文件夹下的 download_weights.md

文中代码为了方便说明，删减了一些内容，详细代码请查看源码。

generate

输入参数：

idx: 输入的prompt经过 tokenizer encode之后输出的序列tensor.使用了默认的输入，token长度为6。
max_new_tokens: 每次新生成的最大token数
max_seq_length: 输入的序列最大长度.
temperature: 温度越高，结果越多样性；温度越低，确定性越高。
top_k: 默认为200。topk越大，结果越多样性；topk越小，结果确定性越高。

@torch.no_grad()
def generate(
    model: LLaMA,
    idx: torch.Tensor,
    max_new_tokens: int,
    *,
    max_seq_length: Optional[int] = None,
    temperature: float = 1.0,
    top_k: Optional[int] = None,
    eos_id: Optional[int] = None,
) -> torch.Tensor:

    # create an empty tensor of the expected final shape and fill in the current tokens
    T = idx.size(0)
    T_new = T + max_new_tokens
    if max_seq_length is None:
        max_seq_length = min(T_new, model.config.block_size)

    device, dtype = idx.device, idx.dtype
    # 创建了一个空的tensor，包括输入的idx，加上允许生成的最大tokens 数，定义了最终结果变量
    empty = torch.empty(T_new, dtype=dtype, device=device)
    empty[:T] = idx
    idx = empty
    input_pos = torch.arange(0, T, device=device) #指明输入的数据在idx中的pos。

    # generate max_new_tokens tokens
    for _ in range(max_new_tokens):
        x = idx.index_select(0, input_pos).view(1, -1) #在结果变量idx中使用input_pos 取出当前输入。

        # forward
        logits = model(x, max_seq_length, input_pos)   #(1,seq,32000)
        logits = logits[0, -1] / temperature

        # optionally crop the logits to only the top k options
        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)

        probs = torch.nn.functional.softmax(logits, dim=-1)
        idx_next = torch.multinomial(probs, num_samples=1).to(dtype=dtype) #多项式采样

        # advance
        input_pos = input_pos[-1:] + 1，#下一个输入的pos

        # concatenate the new generation
        idx = idx.index_copy(0, input_pos, idx_next) #把生成结果copy到结果变量中

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

    return idx

LLAMA Model

使用默认配置，7B模型。block_size 定义了rope和mask的大小，自然也限制了最大输入长度，超过了block_size的输入，无法取得位置编码和mask。

@dataclass
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),
}

LLaMA模型主要有多层attention模块构成。
预测第一个token的时候，需要创建 build_rope_cache 和 build_mask_cache，以及 kv_caches。
然后从rope_cache 和mask_cache 根据 input_pos 取出对应位置的值。
kv_caches 即缓存模型中所有层的kv值，7B有32层，则 kv_caches 的长度为32.
kv的shape为(B, self.config.n_head, max_seq_length, head_size)，使用torch.zeros 初始化。
逐层运行，并将每层的kv值保存在kv_caches 中。
这里每次输入的长度肯定是小于max_seq_length，也就是只更新相应index的kv_caches中的值。
最后经过RMSNorm后，经过线性层，输出每个vocab的概率.

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] = []

    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)
        #从rope_cache 和 mask_cache 取出对应位置的rope和mask
        if input_pos is not None:
            rope = self.rope_cache.index_select(0, input_pos) #(6,64,2),(1,64,2)
            mask = self.mask_cache.index_select(2, input_pos) #(1,1,6,2048),(1,1,1,2048)
            mask = mask[:, :, :, :max_seq_length] #1,1,6,56),(1,1,1,56)
        else:#未给出input_pos，则根据输入长度
            rope = self.rope_cache[:T]
            mask = self.mask_cache[:, :, :T, :T]

        # embeddings
        x = self.transformer.wte(idx)  # token embeddings of shape (b, t, n_embd) #(1,1,4096)

        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: #创建kv_caches
                head_size = self.config.n_embd // self.config.n_head  #128
                cache_shape = (B, self.config.n_head, max_seq_length, head_size) #(1,32,56,128)
                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])
        #RMSNorm
        x = self.transformer.ln_f(x) #(1,6,4096)

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

        return logits
        
    def build_rope_cache(self, idx: torch.Tensor) -> RoPECache:
        return build_rope_cache(
            seq_len=self.config.block_size,
            n_elem=self.config.n_embd // self.config.n_head,
            dtype=idx.dtype,
            device=idx.device,
        )
       # mask_cache 的shape为 （block_size，block_size），右上角为False。
    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)

build_rope_cache

rope 按照下面的计算方法计算，有很多的shape转换，可以使用较小的维度对照公式逐步查看。
参考：一文通透位置编码：从标准位置编码、旋转位置编码RoPE到ALiBi、LLaMA 2 Long(含NTK-aware简介)
rope_cache 的shape为 (2048,64,2)，2048是模型定义的block_size，64 为 attention中每个head 的dim 再除以2。

旋转角度计算公式
在这里插入图片描述

def build_rope_cache( #2048,128=4096/32
    seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device, base: int = 10000
) -> RoPECache:

    # 上面的角度计算公式，2(i-1),i从1到d/2，就等于torch.arange(0, d, 2）
    theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, dtype=dtype, device=device) / n_elem))  #(64)

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

    # Calculate the product of position index and $\theta_i$
    #每个角度都乘以index
    idx_theta = torch.outer(seq_idx, theta).float()  #(2048,64)

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

    # 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

在这里插入图片描述
在attention中，只有q和k需要添加位置信息，按照上图来计算。

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

    # cast because the reference does
    xshaped = x.float().reshape(*x.shape[:-1], -1, 2) #(1,6,32,64,2)
    rope_cache = rope_cache.view(1, xshaped.size(1), 1, xshaped.size(3), 2)  #(1,6,1,64,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,
    )  #(1,6,32,64,2)

    x_out2 = x_out2.flatten(3) #(1,6,32,128)
    return x_out2.type_as(x)

![在这里插入图片描述](https://i-blog.csdnimg.cn/direct/0eff924a1239455e906050c003365315.png
图片是llama2的结构图，llama1并没有使用GQA，其他结构是一样的。

先RMSNorm，再attention，残差相加，然后再RMSNorm，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)

    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部分中，q和k添加了rope。
使用了kv_cache，kv_cache的初始值都是 0，这里就需要把计算出的k 和 v copy到 kv_cache中对应的index位置。
在generate的for循环中，第一次输入全部的prompt，假设长度为 6；第二次只输入生成的token，长度为1，也就是说第二次以后，每次的x size都是（1，1，4096），因为之前的kv值都已经存在kv_cache中了。
如果input_pos >= max_seq_length，cache_k 和cache_v 就要左移，丢弃最早的kv值。

class CausalSelfAttention(nn.Module):
    def __init__(self, config: LLaMAConfig) -> None:
        super().__init__()
        assert config.n_embd % config.n_head == 0

        # key, query, value projections for all heads, but in a batch
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=False)
        # output projection
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)

        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.block_size = config.block_size

    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) #(1,6,4096)

        # 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) #(1,6,4096)

        head_size = C // self.n_head #128
        k = k.view(B, T, self.n_head, head_size) #(1,6,32,128)
        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) (1,32,6,128)
        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 #(1,32,56,128),(1,32,56,128)
            # check if reached token limit
            if input_pos[-1] >= max_seq_length:
                input_pos = torch.tensor(max_seq_length - 1, device=input_pos.device)
                # 左移，丢弃最早的kv值
                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) #(1,32,56,128)
            v = cache_v.index_copy(2, input_pos, v) #(1,32,56,128)
            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)
        
        #这里使用了kv_cache后，kv的shape和q不再相同，
        #(1,32,6,128),(1,32,56,128),(1,32,56,128) => (1,32,6,128)
        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 #(1,6,4096)
        # output projection
        y = self.c_proj(y)
        return y, kv_cache

MLP使用了SwiGLU函数，n_hidden也是个奇怪的数。
Llama改进之——SwiGLU激活函数

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