基于 LLAMA 的大语言模型 LLM 推理过程快速入门

大模型笔记：以 LLAMA 为例，快速入门 LLM 的推理过程

本文借助 LLAMA 模型快速入门 LLM（大语言模型）的推理过程，技术细节很多都是通用的，也适合其他 LLM。LLM 的任务就是在阅读前 n 个单词后预测句子中下一个单词，输出取决于过去和现在输入，与未来无关。

什么是 LLM

LLM 模型主要由两个块组成：

1. 编码器（Encoder）：接收输入并构建其表示形式（特征）。模型被优化为从输入中获取理解（比如判断文本情感）。
2. 解码器（Decoder）：使用编码器的表示形式以及其他输入来生成目标序列。模型被优化用于生成输出。

LLAMA 属于Decoder-only models，只有 decoder 层。每次输入模型会带上上一次输出的结果，这与 CV 模型不同，CV 模型输入只需要一次即可得到结果。

这些部分都可以根据任务独立使用：

• Encoder-only models：适用于需要理解输入的任务，例如句子分类和命名实体识别。
• Decoder-only models：适用于生成性任务，如文本生成。
• Encoder-decoder models or sequence-to-sequence models：适用于需要输入的生成性任务，例如翻译或摘要。

LLAMA 相关的知识点

LLAMA 的 decoder 部分的结构取自 Transformer。对于 LLAMA 来说，只用了 decoder 部分，重点关注以下几个概念：

• Tokenization（分词器）
• Embedding（嵌入层）
• Positional Encoding（位置编码）
• Self-attention（自注意力机制）
• Multi-head attention（多头注意力与采用掩码机制的多头注意力）
• Batch Norm & Layer Norm（批标准化/层标准化，LLAMA 用的是 RMSNorm）
• ResNet（残差网络）

拿到 LLAMA 模型后，主要关注：

1. 模型的结构，包含哪些算子哪些 op，模型复杂度。
2. 模型的前后处理，前后处理实现细节，模型的执行方式。
3. 模型各种参数配置以及其他一些细节。

分词器、Token、Embedding

主要是分词、编码、Tokenizer（tokenization）、Embed（embedding）的过程。

什么是分词？

分词器可将原始文本转换为由 token 组成的文本的初始数值表征。分词器之所以是模型的重要构成部分之一，是因为模型可借此妥善应对人类语言的复杂性。例如，分词器可将凝集性语言中的词分解为更易管理的组成部分、处理原始语料库中不存在的新词或外来词/特殊字符，并确保模型生成紧凑的文本表征。

大部分基于 Transformer 的架构均使用经过训练的分词器，WordPiece（应用于 BERT）、SentencePiece（应用于 T5 或 RoBerta）等分词器同样具有多个变体。

代码示例

首先看 tokenizer，运行 LLAMA 的时候我们会调用 tokenizer = AutoTokenizer.from_pretrained(args.model, use_fast=False)。

如果我们模型传入的是 LLAMA 的某个模型（llama-7b），那么返回的就是 LlamaTokenizer：

class LlamaTokenizer(PreTrainedTokenizer):
    """
    Construct a Llama tokenizer. Based on byte-level Byte-Pair-Encoding.
    ...

这个类是 LLAMA 模型的分词器（tokenizer）的实现，基于字节级的字节对编码（Byte-Pair Encoding）。这个分词器的主要功能是将文本字符串转换为模型可以理解的数字序列，反之亦然。

具体我们看干了些啥，创建好 tokenizer 之后我们执行：input_ids = tokenizer.encode(args.text, return_tensors="pt").to(dev)，这里又分两步：

1. 第一步 Converts a string in a sequence of tokens (string)，这里调用 self.sp_model.encode(text, out_type=str)，sp_model 就是 sentencepiece 中的一个函数，执行完出来变为 ['▁"', 'this', '▁is', '▁a', '▁python', '▁code', ':"']。
2. 第二步将 token string 转变为 token id -> Converts a token string (or a sequence of tokens) in a single integer id (or a sequence of ids)，具体就是个 for 循环，对之前分好的 tokens 一个一个转。

最后我们得到：

input_ids    
tensor([[   0,  376, 1366,  338,  263, 3017,  775, 6160]], device='cuda:0')
input_ids.shape
torch.Size([1, 8])

至于如何转换为 embedding，之后会调用：inputs_embeds = self.embed_tokens(input_ids)，其中 embeds 的 shape 是 torch.Size([1, 8, 4096])。

在自然语言处理（NLP）中，嵌入（Embedding）是一种将离散变量（如单词、短语、或者文档）转换为连续向量的方法。这种转换的目的是让计算机能更好地理解和处理自然语言数据。embedding 矩阵的本质是一个查找表，每个单词会定位这个表中的某一行，而这一行就是这个单词学习到的在嵌入空间的语义。

自注意力 Self-Attention

Transformer 模型的一个关键特点是使用了称为注意力层的特殊层。'Attention Is All You Need'。

这一层会告诉模型，在处理每个单词的表示时，要对你传递给它的句子中某些单词特别关注（并且忽略其他单词）。

Self-attention 是 Transformer 的核心，其允许模型考虑到序列中的其他标记，以便更好地理解每个标记的上下文。每个标记的新表示形式是由它自己和其他标记的交互得到的。

位置编码

由于 Transformer 的结构没有考虑到标记的顺序，所以我们需要加入位置编码来给模型提供词元在序列中的位置信息。这些编码会被添加到词嵌入向量中。

多头注意力 (Multi-head Attention)

多头注意力是对自注意力机制的扩展。它将自注意力分解为多个'头'，每个头在不同的表示空间中学习和应用自注意力。这允许模型同时捕捉到各种不同类型的信息。在有掩码的多头注意力中，掩码被用于阻止模型查看某些不应该看到的信息，例如在生成新的标记时阻止查看未来的信息。现在基本都使用 MHA，一般不用单头。

批标准化 (Batch Norm) & 层标准化 (Layer Norm)

这些都是用于正规化激活的技术，可以加速学习，提高模型的性能。

批标准化是在整个批次的数据上进行标准化，而层标准化则是在单个数据样本上进行标准化。RMSNorm 是一种新的归一化方法，是对 LayerNorm 的一个改进，没有做 re-center 操作（移除了其中的均值项），可以看作 LayerNorm 在均值为 0 时的一个特例。

残差网络 (ResNet)

老熟人了。通过在网络中添加跳跃连接（或称为'skip'连接），可以使得模型更容易地学习到恒等映射，从而避免了训练深度网络时常见的梯度消失问题。在 Transformer 中，每个子层（如自注意力层和前馈神经网络层）都有一个对应的残差连接，并且每个子层的输出都会进行层标准化。

LLAMA 的模型结构

我们可以很轻易的通过 huggingface 代码库中看到 llama 的模型结构。

以 hugging 库中的 7B 模型为例，运行 model = LlamaForCausalLM.from_pretrained(model, torch_dtype='auto') 后，可以通过 print 看模型结构：

LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(32000, 4096, padding_idx=31999)
    (layers): ModuleList(
      (0-31): 32 x LlamaDecoderLayer(
        (self_attn): LlamaAttention(
          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (rotary_emb): LlamaRotaryEmbedding()
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (down_proj): Linear(in_features=11008, out_features=4096, bias=False)
          (up_proj): Linear(in_features=4096, out_features=11008, bias=False)
          (act_fn): SiLUActivation()
        )
        (input_layernorm): LlamaRMSNorm()
        (post_attention_layernorm): LlamaRMSNorm()
      )
    )
    (norm): LlamaRMSNorm()
  )
  (lm_head): Linear(in_features=4096, out_features=32000, bias=False)
)

7B 有 32 个 LlamaDecoderLayer，每个 Decoder 包含一个 LlamaAttention 和 LlamaMLP，然后是 LlamaRMSNorm 和 head 部分，核心的结构是 transformer。

看下 7B 模型的 config，可以看到模型类型为 float16，use_cache 设置为 true：

{
    "architectures": ["LLaMAForCausalLM"],
    "bos_token_id": 0,
    "eos_token_id": 1,
    "hidden_act": "silu",
    "hidden_size": 4096,
    "intermediate_size": 11008,
    "initializer_range": 0.02,
    "max_sequence_length": 2048,
    "model_type": "llama",
    "num_attention_heads": 32,
    "num_hidden_layers": 32,
    "pad_token_id": 0,
    "rms_norm_eps": 1e-06,
    "torch_dtype": "float16",
    "transformers_version": "4.27.0.dev0",
    "use_cache": true,
    "vocab_size": 32000
}

运行 Pipeline

LLAMA 的运行流程和大多数的 LLM 一样，流程如下：

1. 分词 encode，相当于预处理。
2. 输入 input_ids 后模型开始运行，这里会 for 循环运行好多次。
3. 运行完后得到 logits 进行后处理预测下一个 token。
4. 循环往复直到所有要求数量的 token 都输出或者输出遇到了 end_id。

再回顾下 LLAMA 的模型组成：

class LlamaModel(LlamaPreTrainedModel):
    def __init__(self, config: LlamaConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.layers = nn.ModuleList([LlamaDecoderLayer(config) for _ in range(config.num_hidden_layers)])
        self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.gradient_checkpointing = False
        self.post_init()

主要部分就三个：

• embed_tokens 嵌入层
• layers num_hidden_layers 个解码器
• norm RMSNorm 归一化函数

第一步 分词

调用 tokenizer.encode(args.text, return_tensors="pt").to(DEV)，具体流程如下：

输入 prompt -> '"this is a python code:"' -> ['▁"', 'this', '▁is', '▁a', '▁python', '▁code', ':"'] -> [376, 1366, 338, 263, 3017, 775, 6160] -> {'input_ids': tensor([[ 0, 376, 1366, 338, 263, 3017, 775, 6160]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1]])} -> 然后开始 generation，开始第二步

第二步 配置

主要就是配置一些信息和做预处理。

设置 min_length、max_length、top_p、temperature 进入 model.generate。 处理、配置、验证 generation_config 设置 generation parameters。 设置模型的输入：

inputs_tensor, model_input_name, model_kwargs = self._prepare_model_inputs(
    inputs, generation_config.bos_token_id, model_kwargs
)

bos_token_id 是 0，inputs 就是刚才传过来的 input_ids。

设置模型的参数：

model_kwargs["attention_mask"] = self._prepare_attention_mask_for_generation(
    inputs_tensor, generation_config.pad_token_id, generation_config.eos_token_id
)

其中 pad_token_id 为 -1，eos_token_id 为 1，一般从 config 中获取。

决定 generation mode，prepare distribution pre_processing samplers，prepare stopping criteria，设置好 sample_gen_mode、prepare logits warper。

第三步 Sample

继续准备模型的输入，调用 self.prepare_inputs_for_generation（如果提供了 past_key_values，那么 input_ids = input_ids[:, -1:]，同时根据 attention_mask 和是否提供 past_key_values 计算出 position_ids，也就是说提供了 past_key_values 的话，input_ids 可以少计算很多）返回 model_inputs。

开始进入 auto-regressive generation 的循环，是个 while True。

进入 LlamaForCausalLM 函数，输入刚才的 model_inputs，主要包含 input_ids、attention_mask、position_ids、past_key_values。

根据是否存在 past_key_values 更新 past_key_values_length 和 seq_length_with_past。

判断 inputs_embeds 是否存在判断是否需要调用 self.embed_tokens，也就是说如果自己提供了 embeds 就不需要在这里单独 embed 了。

这里调用 inputs_embeds = self.embed_tokens(input_ids)，embeds 的 shape 是 torch.Size([1, 8, 4096])，8 代表输入 input_ids 的长度。

调用 _prepare_decoder_attention_mask 函数，调用后的 attention_mask 维度为 torch.Size([1, 1, 8, 8])。

进入一个 for 循环，因为 llama 有很多 self.layers = nn.ModuleList([LlamaDecoderLayer(config) for _ in range(config.num_hidden_layers)])，都是一模一样的 LlamaDecoderLayer。

LlamaDecoderLayer

函数的传入的参数：hidden_states [1,8,4096]、attention_mask [1,1,8,8]、position_ids [1,8]、past_key_value [[]] or None。

执行的程序也很简单，顺序执行以下步骤：

1. 首先归一化 self.input_layernorm(hidden_states)
2. 然后经过 self_attn
3. 残差 hidden_states = residual + hidden_states
4. 全连接 self.post_attention_layernorm(hidden_states) -> self.mlp(hidden_states)
5. 继续残差 hidden_states = residual + hidden_states

代码如下：

residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)
# Self Attention
hidden_states, self_attn_weights, present_key_value = self.self_attn(
    hidden_states=hidden_states,
    attention_mask=attention_mask,
    position_ids=position_ids,
    past_key_value=past_key_value,
    output_attentions=output_attentions,
    use_cache=use_cache,
)
hidden_states = residual + hidden_states
# Fully Connected
residual = hidden_states
hidden_states = self.post_attention_layernorm(hidden_states)
hidden_states = self.mlp(hidden_states)
hidden_states = residual + hidden_states
outputs = (hidden_states,)
if output_attentions:
    outputs += (self_attn_weights,)
if use_cache:
    outputs += (present_key_value,)
return outputs

LlamaAttention

这个就是 Multi-headed attention from 'Attention Is All You Need' paper。这个类的成员变量如下：

class LlamaAttention(nn.Module):
    def __init__(self, config: LlamaConfig):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads # head 的数量 这里是 32
        self.head_dim = self.hidden_size // self.num_heads  # head 的大小 这里是 128
        self.max_position_embeddings = config.max_position_embeddings
        if (self.head_dim * self.num_heads) != self.hidden_size:
            raise ValueError(...)
        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
        self.rotary_emb = LlamaRotaryEmbedding(self.head_dim, max_position_embeddings=self.max_position_embeddings)

其中，多头机制的自注意力：

• self.num_heads 定义了 attention head 的数量
• self.head_dim 定义了每个 head 的大小，是 hidden_size 除以 num_heads
• 线性层 self.q_proj, self.k_proj, self.v_proj 将输入 hidden_states 映射到 num_heads * head_dim 的维度，以分别获得查询、键、值 tensor。

操作代码如下：

def forward(
    self,
    hidden_states: torch.Tensor,
    attention_mask: Optional[torch.Tensor] = None,
    position_ids: Optional[torch.LongTensor] = None,
    past_key_value: Optional[Tuple[torch.Tensor]] = None,
    output_attentions: bool = False,
    use_cache: bool = False,
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
    bsz, q_len, _ = hidden_states.size()
    query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
    key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
    value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
    kv_seq_len = key_states.shape[-2]
    if past_key_value is not None:
        kv_seq_len += past_key_value[0].shape[-2]
    cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
    if past_key_value is not None:
        key_states = torch.cat([past_key_value[0], key_states], dim=2)
        value_states = torch.cat([past_key_value[1], value_states], dim=2)
    past_key_value = (key_states, value_states) if use_cache else None
    attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
    if attention_mask is not None:
        attn_weights = attn_weights + attention_mask
        attn_weights = torch.max(attn_weights, torch.tensor(torch.finfo(attn_weights.dtype).min))
    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
    attn_output = torch.matmul(attn_weights, value_states)
    attn_output = attn_output.transpose(1, 2)
    attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
    attn_output = self.o_proj(attn_output)
    if not output_attentions:
        attn_weights = None
    return attn_output, attn_weights, past_key_value

其中两点需要注意：

1. query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) 就是多头计算，得到的结果 query_states、key_states、value_states 的维度是 torch.Size([1, 32, 8, 128])，32 代表头的数量、8 是输入的 input_ids 长度，128 代表头的大小。
2. 如果提供了 past_key_value，则利用 cache 的机制，直接 torch.cat([past_key_value[0], key_states], dim=2) 即可，每次传入的 input_ids 只有最新的一个 id。

回到刚才的 pipeline： 接着 for 循环完 32 个 decoder 层之后，需要进行 norm 操作：hidden_states = self.norm(hidden_states)。

输出后得到 outputs.logits 维度为 torch.Size([1, 8, 32000])，接下来对这个 logits 进行操作。

主要的过程：

• 将 logits 传递给 logits_processor 和 logits_warper，在这两个方法中进行一些预处理过程，例如添加惩罚项或对概率分布进行修改，使得生成的结果更符合期望。
• TopKLogitsWarper 类是一个用于处理模型输出分数（scores）的工具，主要用于进行所谓的'Top-K 截断'。
• TopPLogitsWarper 类实现了被称为"Top-p（或 nucleus）抽样"的策略。
• 最后，使用 softmax 函数将经过预处理的 logits 转换为概率分布，并利用 multinomial 方法从中采样得到下一个 token。

得到 next_tokens 之后，执行以下代码：

input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
if streamer is not None:
    streamer.put(next_tokens.cpu())
model_kwargs = self._update_model_kwargs_for_generation(
    outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder
)
if eos_token_id_tensor is not None:
    unfinished_sequences = unfinished_sequences.mul(
        next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0)
    )
if unfinished_sequences.max() == 0 or stopping_criteria(input_ids, scores):
    break

最后通过分词器进行 decode 即可得到所有结果，这个是支持 batch 的：

print(tokenizer.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0])

落地相关

部署的时候，除了模型要搞好，确保模型可以输出正常结果的设置参数也需要整明白，需要暴露出来以便上游去调节。

对于 llama 来说，可调节的参数就有点多了。

举个例子，随便拿出一个 gradio 展示的 LLM 模型，可调节的参数如下：

• seed：seed 如果固定了，我们的输出也就固定了，我们使用 chatgpt 的时候，seed 每次应该不一样。
• gen mode sample：一般来说就是用这个 do_sample；greed 貌似这个最快；beam search 这个效果最好。
• end_id：模型训练的时候设置的结束 id，模型在预测的时候，如果预测下一个 token 是 end_id 的话就应该终止预测了，就是这段话说完了。
• start_id：一般来说是在 tokenizer 中设置，在将输入文本 encode 的时候，设置的第一个 token id。

Padding

LLM 和 CV 模型一样，组 batch 的时候需求输入大小一致，比如 [4,3,128,128]，而 NLP 中输入的是 input_ids，padding 的方法如下：

input_sentences = [
    "DeepSpeed is a machine learning framework",
    "He is working on",
    "He has a",
    "He got all",
    "Everyone is happy and I can",
    "The new movie that got Oscar this year",
    "In the far far distance from our galaxy,",
    "Peace is the only way",
]
tokenizer.pad_token = 0
input_tokens = tokenizer.batch_encode_plus(input_sentences, return_tensors="pt", padding=True)

Bad Words 和 Stop Words

LLM 在生成 token 的时候需要避免一些不想生成的 token，遇到就停止的 token。

KV-Cache

KV-Cache 是 LLM 推理过程的一个优化点，可以减少计算，不过需要额外显存去存这些 cache。

总结

可以参考 HuggingFace 搭建服务的规则和一些细节，一般要支持：

• 对于 stream 模式，要支持多个同时的请求。
• 对于非 stream 模式，需要支持多 batch 输入，也可以支持组 batch。
• 对于 stream 模式，如果服务端认为生成 token 到结尾了，则可以主动断开连接并且返回终止符给客户端。

一些概念

Unconditional Generation

在自然语言处理（NLP）中，"unconditional generation"是指模型在没有特定输入或提示的情况下生成文本。对比之下，"conditional generation"是指模型在给定某些输入或提示（例如，开头的一段文本或特定的任务描述）的情况下生成文本。

"start token"和"end token"是用来标识生成任务的开始和结束的特殊标记。

Context Len

在自然语言处理（NLP）中，"context window"（上下文窗口）是一种常见的概念，它指的是在处理某个词或词组时考虑的前后相关词汇的范围。这个范围可以是固定的，也可以是动态的，取决于具体的模型和任务。

较大的上下文窗口可能会捕获更多的长距离依赖关系，但也可能增加模型的计算复杂度。相反，较小的上下文窗口可能会减少计算复杂度，但可能无法捕获一些重要的长距离依赖关系。因此，选择合适的上下文窗口大小和处理策略通常需要根据具体的任务和数据进行调整。