📚 编程进阶教程站

第 3 章 · 深入 LLM 内部 ​

本章目标：深入 generative LLM 的 transformer 架构内部——加载模型、观察它的输入输出张量、理解从概率分布中选择 token 的解码过程，以及 KV 缓存带来的生成加速。

本章对应《Hands-On Large Language Models》（Jay Alammar & Maarten Grootendorst 著）第 3 章的官方笔记本，对 generative LLM 的 transformer 架构做一次深入的内部考察。

3.1 [可选] 在 Colab 上安装依赖 ​

如果你在 Google Colab（或任何其他云环境）中查看本笔记本，需要取消注释并运行下面的代码块来安装本章的依赖：

💡 NOTE：运行本章示例需要 GPU。在 Google Colab 中，进入 Runtime > Change runtime type > Hardware accelerator > GPU > GPU type > T4。

python
# %%capture
# !pip install transformers>=4.41.2 accelerate>=0.31.0

3.2 加载 LLM ​

python
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-4k-instruct")

model = AutoModelForCausalLM.from_pretrained(
    "microsoft/Phi-3-mini-4k-instruct",
    device_map="cuda",
    torch_dtype="auto",
    trust_remote_code=False,
)

# Create a pipeline
generator = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    return_full_text=False,
    max_new_tokens=50,
    do_sample=False,
)

这里加载的是 microsoft/Phi-3-mini-4k-instruct——一个小型但能力不错的 instruct 模型，并用 pipeline 封装成文本生成器。return_full_text=False 表示只返回新生成的部分，do_sample=False 表示使用确定性解码。

3.3 训练后的 Transformer LLM：输入与输出 ​

python
prompt = "Write an email apologizing to Sarah for the tragic gardening mishap. Explain how it happened."

output = generator(prompt)

print(output[0]['generated_text'])
python
print(model)

print(model) 会打印出模型的完整结构。你会看到一个典型的 generative LLM 由两大部分组成：

  • model（主体）：数十层 transformer block 堆叠而成的骨干网络，负责把输入 token 序列逐层加工成隐藏状态；
  • lm_head（输出头）：一个线性层，把最后的隐藏状态投影回词表大小，得到每个 token 的 logits（未归一化的概率分数）。

3.4 从概率分布中选择单个 token（采样 / 解码） ​

我们可以手动拆开这两个部分，观察模型「选择下一个词」的全过程：

python
prompt = "The capital of France is"

# Tokenize the input prompt
input_ids = tokenizer(prompt, return_tensors="pt").input_ids

# Tokenize the input prompt
input_ids = input_ids.to("cuda")

# Get the output of the model before the lm_head
model_output = model.model(input_ids)

# Get the output of the lm_head
lm_head_output = model.lm_head(model_output[0])
python
token_id = lm_head_output[0,-1].argmax(-1)
tokenizer.decode(token_id)

对最后一个位置的 logits 取 argmax，解码出来的正是 " Paris"——这就是贪婪解码（greedy decoding）最直观的样子。

python
model_output[0].shape
python
lm_head_output.shape

对比两个张量的形状可以看清数据流：model_output[0] 的形状是 [batch_size, sequence_length, hidden_dimension]（隐藏状态），而 lm_head_output 的形状是 [batch_size, sequence_length, vocabulary_size]——最后一维从隐藏维度变成了词表大小，每个位置都对应词表中所有 token 的分数。

3.5 利用 KV 缓存（Keys and Values Caching）加速生成 ​

生成很长的文本时，如果没有缓存，每生成一个新 token 都要对整段序列重新计算所有层的注意力键和值。缓存机制把这些中间结果存下来复用：

python
prompt = "Write a very long email apologizing to Sarah for the tragic gardening mishap. Explain how it happened."

# Tokenize the input prompt
input_ids = tokenizer(prompt, return_tensors="pt").input_ids
input_ids = input_ids.to("cuda")
python
%%timeit -n 1
# Generate the text
generation_output = model.generate(
  input_ids=input_ids,
  max_new_tokens=100,
  use_cache=True
)
python
%%timeit -n 1
# Generate the text
generation_output = model.generate(
  input_ids=input_ids,
  max_new_tokens=100,
  use_cache=False
)

分别用 use_cache=True 和 use_cache=False 生成 100 个 token 并计时对比，你会发现开启缓存后生成速度快得多——这正是所有主流推理框架都默认启用 KV 缓存的原因。

3.6 本章小结 ​

  • generative LLM = transformer 骨干网络（model）+ 语言模型输出头（lm_head），后者把隐藏状态投影为词表上的 logits；
  • model.model(input_ids) 给出 lm_head 之前的隐藏状态，形状为 [batch, seq_len, hidden_dim]；经 lm_head 后变为 [batch, seq_len, vocab_size]；
  • 对最后位置的 logits 做 argmax 即贪婪解码——"The capital of France is" 的答案是 " Paris"；
  • KV 缓存（use_cache=True）复用先前 token 的注意力键值，避免每步全量重算，显著加速长文本生成。

🧪 随堂测验

点击你认为正确的选项。答错时会展示正确答案与原因解析。

1. 在 generative LLM 中，lm_head 的作用是什么？

2. `model.model(input_ids)` 返回的张量形状是？

3. `use_cache=True` 之所以能加速生成，是因为？

4. 本章加载 pipeline 时设置 `do_sample=False`，这意味着？

🛠️ 动手实践 ​

  1. 把 3.4 节的 prompt 换成 "The capital of China is" 和 "The largest planet in our solar system is"，分别用 argmax 解码最后一个 token，验证模型给出的答案，并打印 lm_head_output[0,-1].topk(5) 观察得分最高的 5 个候选 token。
  2. 用 print(model) 输出的结构统计 Phi-3 的 transformer block 层数（数一数有多少个 Phi3DecoderLayer），并把 3.3 节的 model_output[0].shape 打印出来，对照 hidden dimension 与层数的关系写一段笔记。
  3. 复现 3.5 节的对比实验：把 max_new_tokens 分别改成 50、100、200，记录 use_cache=True 与 False 两组耗时，画一张简单的对比表总结缓存收益随生成长度的变化。