📚 AI 工程与智能体教程站

English

简体中文

English

简体中文

Chapter 3 · Inside the LLM ​

Chapter goals: Dive into the internal structure of a generative LLM's Transformer architecture — load the model, inspect its input/output tensors, understand the decoding process of selecting tokens from probability distributions, and the generation speedup brought by KV caching.

This chapter corresponds to the official notebook for Chapter 3 of Hands-On Large Language Models (by Jay Alammar & Maarten Grootendorst), providing an in-depth internal inspection of the Transformer architecture for generative LLMs.

3.1 [Optional] Installing Dependencies on Colab ​

If you are viewing this notebook on Google Colab (or any other cloud environment), you need to uncomment and run the code block below to install the dependencies for this chapter:

💡 NOTE: Running the examples in this chapter requires a GPU. In Google Colab, go to Runtime > Change runtime type > Hardware accelerator > GPU > GPU type > T4.

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

3.2 Loading an 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,
)

Here we load microsoft/Phi-3-mini-4k-instruct — a small but capable instruct model — and wrap it into a text generator via pipeline. return_full_text=False means only the newly generated portion is returned; do_sample=False means deterministic decoding is used.

3.3 Post-Training Transformer LLM: Inputs and Outputs ​

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) prints the model's complete structure. You will see a typical generative LLM composed of two major parts:

  • model (backbone): a stack of dozens of Transformer blocks forming the backbone network, responsible for progressively transforming the input token sequence into hidden states;
  • lm_head (output head): a linear layer that projects the final hidden state back to vocabulary size, producing logits (unnormalized probability scores) for each token.

3.4 Selecting a Single Token from a Probability Distribution (Sampling / Decoding) ​

We can manually split these two parts apart and observe the model's full process of "selecting the next word":

python
prompt = "The capital of France is"

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

# Move to CUDA
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)

Taking argmax over the logits at the last position decodes to " Paris" — this is the most intuitive appearance of greedy decoding.

python
model_output[0].shape
python
lm_head_output.shape

Comparing the shapes of the two tensors makes the data flow clear: model_output[0] has shape [batch_size, sequence_length, hidden_dimension] (hidden states), while lm_head_output has shape [batch_size, sequence_length, vocabulary_size] — the last dimension changes from hidden dimension to vocabulary size, and every position corresponds to scores for all tokens in the vocabulary.

3.5 Accelerating Generation with KV Caching ​

When generating long text, without caching, every new token generated would require recomputing attention keys and values for the entire sequence across all layers. The caching mechanism stores and reuses these intermediate results:

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
)

Generate 100 tokens with both use_cache=True and use_cache=False and compare the timings; you will find that generation is much faster with caching enabled — this is exactly why all mainstream inference frameworks enable KV caching by default.

3.6 Chapter Summary ​

  • A generative LLM = Transformer backbone network (model) + language model output head (lm_head), where the latter projects hidden states into logits over the vocabulary;
  • model.model(input_ids) yields hidden states before the lm_head, with shape [batch, seq_len, hidden_dim]; after passing through lm_head it becomes [batch, seq_len, vocab_size];
  • Taking argmax over the logits at the last position gives greedy decoding — the answer to "The capital of France is" is " Paris";
  • KV caching (use_cache=True) reuses attention keys and values from previous tokens, avoiding full recomputation at every step, significantly accelerating long-text generation.

🧪 随堂测验

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

1. In a generative LLM, what is the role of lm_head?

2. What is the shape of the tensor returned by `model.model(input_ids)`?

3. `use_cache=True` accelerates generation because?

4. Setting `do_sample=False` when loading the pipeline in this chapter means?

🛠️ Hands-on Practice ​

  1. Replace the prompt in Section 3.4 with "The capital of China is" and "The largest planet in our solar system is"; decode the last token using argmax for each, verify the answers the model gives, and print lm_head_output[0,-1].topk(5) to observe the top 5 candidate tokens by score.
  2. Use the structure output by print(model) to count the number of Transformer blocks in Phi-3 (count how many Phi3DecoderLayers there are), and print model_output[0].shape from Section 3.3; write a note on the relationship between hidden dimension and number of layers.
  3. Reproduce the comparison experiment from Section 3.5: change max_new_tokens to 50, 100, and 200 respectively; record the wall-clock time for both use_cache=True and False groups; draw a simple comparison table summarizing how the caching benefit varies with generation length.