了解Llama3.1如何工作——深入探讨模型流程

原文标题：Understand How Llama3.1 Works — A Deep Dive Into the Model Flow

原文链接：https://medium.com/@yuxiaojian/understand-how-llama3-1-works-a-deep-dive-into-the-model-flow-b149aba04bed

翻译：岁月月宝贝?

了解Llama3.1如何工作——深入探讨模型流程

像Llama 3.1这样的大型语言模型功能强大，但理解它们的内部工作机制可能很复杂，尤其是当理论脱离实际应用时。在这次深入探讨中，我们将采取一种独特的方法，从反向角度探索模型。通过逆向追溯工作流程，我们将揭示驱动Llama 3.1的复杂过程，提供对这个模型如何运作的深入、实际的理解。

一、模型架构⛪

让我们加载llama3.1模型并仔细查看它。为了减少内存使用，我们对所有线性变换使用了4位量化（Linear4bit）。量化不会影响我们对模型工作原理的理解。

4位量化（4-bit quantization）是一种模型压缩技术，它将模型中的权重和激活值从传统的32位或16位精度降低到4位。这种技术可以显著减少模型在内存中的占用空间，降低计算成本，同时尽可能保持模型的性能。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
# 定义基础模型的名称
base_model = "meta-llama/Meta-Llama-3.1-8B-Instruct"
# 加载与基础模型相对应的分词器
tokenizer = AutoTokenizer.from_pretrained(base_model)
# 创建量化配置对象
bnb_config = BitsAndBytesConfig(
 load_in_4bit=True, # 模型将被加载为4位量化版本
 bnb_4bit_quant_type="nf4", # 指定4位量化的类型为"nf4"，基于分位数的量化方法
 bnb_4bit_compute_dtype=torch.bfloat16 # 设置计算时使用的精度类型为bfloat16
)
# 加载并量化模型
base_model_bnb_4b = AutoModelForCausalLM.from_pretrained(
 base_model, # 基础模型名称
 quantization_config=bnb_config, # 传入量化配置
 device_map='auto' # 设备映射设置为'auto'，自动决定在哪个设备上运行模型
)
# 变量base_model_bnb_4b表示加载并量化后的模型

打印 base_model_bnb_4b 将使用模型结构。

LlamaForCausalLM(
 
 (model): LlamaModel( # Llama模型
 
 # 1.1词嵌入层，输入词汇量为128256，嵌入维度为4096 
 (embed_tokens): Embedding(128256, 4096) 
 
 # 多层解码器层 
 (layers): ModuleList( 
 (0-31): 32 x LlamaDecoderLayer( # 包含32层Llama解码器层 
 (self_attn): LlamaSdpaAttention( # 自注意力机制
 (q_proj): Linear4bit(in_features=4096, out_features=4096, bias=False) # 3.1查询投影，4位量化线性层
 (k_proj): Linear4bit(in_features=4096, out_features=1024, bias=False) # 3.2键投影，4位量化线性层
 (v_proj): Linear4bit(in_features=4096, out_features=1024, bias=False) # 3.3值投影，4位量化线性层
 (o_proj): Linear4bit(in_features=4096, out_features=4096, bias=False) # 5.输出投影，4位量化线性层
 (rotary_emb): LlamaRotaryEmbedding() # 4.旋转位置编码,然后应该是Attention(?) Mechanism
 )
 (mlp): LlamaMLP( # 前馈神经网络
 (gate_proj): Linear4bit(in_features=4096, out_features=14336, bias=False) # 7.门控投影，4位量化线性层
 (up_proj): Linear4bit(in_features=4096, out_features=14336, bias=False) # 8.上层投影，4位量化线性层
 (down_proj): Linear4bit(in_features=14336, out_features=4096, bias=False) # 10.下层投影，4位量化线性层
 (act_fn): SiLU() # 9.激活函数
 )
 (input_layernorm): LlamaRMSNorm((4096,), eps=1e-05) # 2.输入层归一化
 (post_attention_layernorm): LlamaRMSNorm((4096,), eps=1e-05) # 6.注意力后层归一化
 ) 
 )
 
 # 11.最终层归一化 
 (norm): LlamaRMSNorm((4096,), eps=1e-05) 
 
 # 1.2旋转位置编码
 (rotary_emb): LlamaRotaryEmbedding() 
 
 )
 
 (lm_head): Linear(in_features=4096, out_features=128256, bias=False) # 12.语言模型头部，输出词汇量为128256
 
)

如果我们将模型可视化，插图将如下所示。它展示了LlamaForCausalLM模型的整体结构：

输入的词元（tokens）首先通过嵌入层（embedding layer）进行处理，该层将词元ID转换为大小为4096的密集向量（dense vectors）。
对嵌入后的词元应用旋转嵌入（LlamaRotaryEmbedding）。
模型的核心由32个相同的Llama解码层（LlamaDecoderLayers）组成。每一层包括：
- 输入层归一化（Input layer normalization）
- 自注意力机制（Self-attention mechanism），具有用于查询（queries）、键（keys）、值（values）和输出（output）的独立投影
- 注意力后的层归一化（Post-attention layer normalization）
- 包含门控（gate）和上投影（up projections）、SiLU激活函数和下投影（down projection）的多层感知机（MLP）块
在32个解码层之后，有一个最终的层归一化（final layer normalization）。
模型主体的输出通过语言模型头部（LM head），将4096维的向量投影回词汇空间（128256维）。
最终输出表示序列中下一个词元的概率分布。

二、模型配置?

让我们检查一下Llama 3.1 8B Instruct模型的设置。

LlamaConfig {
 "_name_or_path": "meta-llama/Meta-Llama-3.1-8B-Instruct", # 模型的名称或路径
 
 "architectures": [ # 1.1模型的架构
 "LlamaForCausalLM" # 因果语言模型架构
 ],
 
 "attention_bias": false, # 是否使用注意力偏置
 "attention_dropout": 0.0, # 注意力机制中的dropout比例
 
 "bos_token_id": 128000, # 3.2序列开始（BOS）的词元ID
 "eos_token_id": [ # 3.3序列结束（EOS）的词元ID列表
 128001,
 128008,
 128009
 ],
 
 "hidden_act": "silu", # 5.1 隐藏层的激活函数
 "hidden_size": 4096, # 2.1 隐藏层的大小
 
 "initializer_range": 0.02, # 权重初始化的范围
 "intermediate_size": 14336, # 2.2前馈网络中间层的大小
 "max_position_embeddings": 131072, # 4.1最大位置嵌入数
 "mlp_bias": false, # MLP层是否使用偏置
 
 "model_type": "llama", # 1.2模型类型
 
 "num_attention_heads": 32, # 2.4注意力头的数量
 "num_hidden_layers": 32, # 2.3隐藏层的数量
 "num_key_value_heads": 8, # 2.5键值对头的数量
 
 "pretraining_tp": 1, # 预训练时的张量并行度
 
 "quantization_config": { # 量化配置
 "_load_in_4bit": true, # 加载4位量化模型
 "_load_in_8bit": false, # 不加载8位量化模型
 "bnb_4bit_compute_dtype": "bfloat16", # 6.2 4位量化计算时的数据类型
 "bnb_4bit_quant_storage": "uint8", # 4位量化存储时的数据类型
 "bnb_4bit_quant_type": "nf4", # 6.1 4位量化类型
 "bnb_4bit_use_double_quant": false, # 是否使用双重量化
 "llm_int8_enable_fp32_cpu_offload": false, # 是否启用INT8 CPU卸载
 "llm_int8_has_fp16_weight": false, # 是否有INT8 FP16权重
 "llm_int8_skip_modules": null, # 跳过的模块列表
 "llm_int8_threshold": 6.0, # INT8量化阈值
 "load_in_4bit": true, # 加载4位量化模型
 "load_in_8bit": false, # 不加载8位量化模型
 "quant_method": "bitsandbytes" # 量化方法
 },
 
 "rms_norm_eps": 1e-05, # 5.2 RMS归一化epsilon值
 
 "rope_scaling": { # 4.2 ROPE（旋转位置编码）缩放配置
 "factor": 8.0, # 缩放因子
 "high_freq_factor": 4.0, # 高频缩放因子
 "low_freq_factor": 1.0, # 低频缩放因子
 "original_max_position_embeddings": 8192, # 原始最大位置嵌入数
 "rope_type": "llama3" # ROPE类型
 },
 
 "rope_theta": 500000.0, # ROPE的theta值
 
 "tie_word_embeddings": false, # 7.3 是否绑定词嵌入
 "torch_dtype": "float16", # 7.1 PyTorch数据类型
 "transformers_version": "4.44.2", # transformers库的版本
 "use_cache": true, # 7.2 是否使用缓存
 
 "vocab_size": 128256 # 3.1词汇表大小
}

让我们分解关键组件:

模型架构:
"architectures": ["LlamaForCausalLM"]: 这是一个基于Llama架构的因果语言模型。
"model_type": "llama": 确认它是一个Llama模型.
模型大小和结构:
"hidden_size": 4096: 隐藏状态的维度.
"intermediate_size": 14336: 前馈层的维度.
"num_hidden_layers": 32:变换器层的数量.
"num_attention_heads": 32: 注意力头的数量.
"num_key_value_heads": 8: 分组查询注意力的键/值头的数量.
分词(器):
"vocab_size": 128256: 词汇表的大小.
"bos_token_id": 128000: 序列开始（BOS）标记的ID.
"eos_token_id": [128001, 128008, 128009]: 序列结束（EOS）标记的ID列表.
位置嵌入:
"max_position_embeddings": 131072: 模型能够处理的最大序列长度.
"rope_scaling": RoPE（旋转位置嵌入）缩放的配置，允许更长的上下文长度.
激活和归一化:
"hidden_act": "silu": 使用的激活函数（SiLU，也称为Swish）.
"rms_norm_eps": 1e-05: RMSNorm的epsilon值.
量化:
模型配置为使用bitsandbytes库进行4位量化.
"bnb_4bit_quant_type": "nf4": 使用NF4（归一化浮点4）量化.
"bnb_4bit_compute_dtype": "bfloat16": 计算以bfloat16精度进行.
其他值得注意的设置:
"torch_dtype": "float16": 模型权重以float16精度存储.
"use_cache": true: 启用过去键/值的使用，以实现更快的推理.
"tie_word_embeddings": false: 输入和输出嵌入没有绑定.

三、分词器（Tokenizer）?

分词器充当原始文本与模型可以处理的数值表示之间的字典。Llama 3.1 使用的分词器拥有128K个词元的词汇表。

# 获取词汇表大小
vocab_size = len(tokenizer.get_vocab())
# 打印词汇表大小
print(f"Vocabulary size: {vocab_size}")
# 输出：Vocabulary size: 128256

输入到模型中的提示（prompt）包含了一些特殊标记，这些标记用来指示每个部分的信息类型、开始和结束。下面是一个标记及其对应ID的映射示例。

# 定义一个函数，用于打印文本的分词和对应的ID
def print_tokens_with_ids(txt):
 # 使用tokenizer对文本进行分词，add_special_tokens=False表示不添加特殊标记
 tokens = tokenizer.tokenize(txt, add_special_tokens=False)
 # 将文本编码为对应的ID序列，add_special_tokens=False表示不添加特殊标记
 token_ids = tokenizer.encode(txt, add_special_tokens=False)
 # 打印分词和对应的ID序列
 print(list(zip(tokens, token_ids)))
# 定义一个包含特殊标记的示例文本
prompt = """<|begin_of_text|><|start_header_id|>system<|end_header_id|>
Based on the information provided, rewrite the sentence by changing its tense from past to future.<|eot_id|><|start_header_id|>user<|end_header_id|>
She played the piano beautifully for hours and then stopped as it was midnight.<|eot_id|><|start_header_id|>assistant<|end_header_id|>
"""
# 调用函数，打印示例文本的分词和对应的ID
print_tokens_with_ids(prompt)
# 输出：Token和Token ID的对应列表
# > [('<|begin_of_text|>', 128000), ... , ('<|end_header_id|>', 128007), ('ĊĊ', 271)]

模型的输入仅仅是词元的ID。

[128000, 128000, 128006, 9125, 128007, 271, 29815, 389, 279,
 2038, 3984, 11, 18622, 279, 11914, 555, 10223, 1202,
 43787, 505, 3347, 311, 3938, 13, 128009, 128006, 882,
 128007, 271, 8100, 6476, 279, 27374, 32719, 369, 4207,
 323, 1243, 10717, 439, 433, 574, 33433, 13, 128009,
 128006, 78191, 128007, 271]

使用模型的生成（generate）方法来“预测”输出词元。

# 使用模型的generate方法生成输出词元
outputs = base_model_bnb_4b.generate(
 input_ids=input_ids, # 输入的词元ID序列
 pad_token_id=tokenizer.eos_token_id, # 用于填充的词元ID，这里使用分词器的序列结束（EOS）词元ID
 max_new_tokens=200, # 生成的最大新词元数量，限制生成文本的长度
 do_sample=True, # 是否使用采样方法生成词元，而不是贪婪地选择概率最高的词元
 top_p=0.9, # 在采样时，从概率累积和达到0.9的词元集中选择，即只考虑累积概率最高的90%的词元
 temperature=0.1 # 控制生成文本的随机性，值越高，生成的文本越随机
)

输出的词元是下面的黑色部分。最后一个词元128009是<|eot_id|>，它表示输出的结束。

[128000, 128000, 128006, 9125, 128007, 271, 29815, 389, 279,
 2038, 3984, 11, 18622, 279, 11914, 555, 10223, 1202,
 43787, 505, 3347, 311, 3938, 13, 128009, 128006, 882,
 128007, 271, 8100, 6476, 279, 27374, 32719, 369, 4207,
 323, 1243, 10717, 439, 433, 574, 33433, 13, 128009,
 128006, 78191, 128007, 271, 8586, 374, 279, 11914, 59624,
 304, 279, 3938, 43787, 1473, 8100, 690, 1514, 279,
 27374, 32719, 369, 4207, 323, 1243, 3009, 439, 433,
 374, 33433, 13, 128009]

使用分词器将这些词元映射回文本。

# 使用分词器的batch_decode方法将生成的词元ID序列转换回文本
# outputs.detach().cpu().numpy()将生成的词元ID从计算图中分离，转移到CPU，并转换为NumPy数组
result = tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=False)[0]
# 打印转换后的文本结果
print(result)

输出的文本是

<|begin_of_text|><|begin_of_text|><|start_header_id|>system<|end_header_id|>
Based on the information provided, rewrite the sentence by changing its tense from past to future.<|eot_id|><|start_header_id|>user<|end_header_id|>
She played the piano beautifully for hours and then stopped as it was midnight.<|eot_id|><|start_header_id|>assistant<|end_header_id|>
Here is the sentence rewritten in the future tense:
She will play the piano beautifully for hours and then stop as it is midnight.<|eot_id|>

四、词元生成过程?‍?‍?‍?

LlamaForCausalLM（因果语言模型架构）模型，像其他基于变换器（transformer）的语言模型一样，使用一系列词元作为输入，并每次生成（或者更准确地说，预测）一个词元。在文本生成过程中，模型通常以自回归的方式使用：

从一个初始提示开始（可以是多个词元）。
模型预测下一个词元。
将这个预测的词元添加到输入序列中。
重复步骤2-3以生成更多词元，直到遇到停止词元或达到最大词元数量。

这里有一个简化的插图：

# 输入序列：[token1, token2, token3, ...]
# 一系列词元作为模型的输入
# 模型处理流程：
# |
# v
# +----------------+
# | |
# | Llama模型 |
# | |
# +----------------+
# 模型接收输入序列并进行处理
# |
# v
# 输出：[prob_dist1, prob_dist2, prob_dist3, ...]
#（每一步生成的概率分布都是基于之前所有生成的词元以及原始输入序列！！！）

这儿：

token1, token2, token3, ... 是输入的词元。
prob_dist1, prob_dist2, prob_dist3, ... 是每个位置下一个词元的概率分布。

对于生成过程：

步骤 1: [A, B, C] -> 模型 -> [_, _, _, D]
步骤 2: [A, B, C, D] -> 模型 -> [_, _, E]
步骤 3: [A, B, C, D, E] -> 模型 -> [_, _, _, _, F]
...

根据我们在分词器部分使用的例子，步骤将是：

Step 1:
Input: [128000, ..., 271]
Output: [128000, ..., 271, 8586]
 ^^^^
 New token
... (continuing for each new token) ...
Final Step:
Input: [128000, ..., 374, 33433, 13]
Output: [128000, ..., 374, 33433, 13, 128009]

上下文窗口（最大序列长度）决定了可以用作输入的词元数量。Llama3.1模型可以处理的最大序列长度是131072，即 "max_position_embeddings": 131072。

`Temperature and Top_p`

生成参数 top_p 和 temperature 在控制模型在文本生成过程中如何选择下一个词元方面至关重要。它们影响输出中创造性和确定性之间的平衡。让我们来分解一下这个问题：

1.温度参数（Temperature）：

温度参数是一个超参数，它通过在应用softmax函数之前对logits进行缩放来控制预测的随机性。

公式：$ P(x_i) = \frac{\exp(z_i / T)}{\sum \exp(z_j / T)} $其中 $z_i $ 是logits，$ T $ 是温度参数，$ P(x_i) $ 是词元i的概率。
范围：通常在0到2之间，1表示中性。
较低的温度（例如，0.1）会使分布“更尖锐”（更高峰）。这意味着输出更加确定性，倾向于产生更多重复和“安全”的文本。
较高的温度（例如，1.5）会使分布“更平坦”。这意味着输出更加随机，倾向于产生更多多样化和潜在创造性的文本。

视觉表现：

Low Temperature (0.1) High Temperature (1.5)
 | |
 |* |
 |* | * * *
 |* * * |* * * *
----+------------------ ----+------------------
 A B C D E A B C D E

在这张图中，星号（*）表示概率分布的形状。在低温度下，概率分布更集中，而在高温度下，概率分布更分散。这反映了模型在不同温度设置下生成文本的随机性和创造性。

2. Top-p（核采样）：

Top-p，也被称为核采样，设置了一个概率阈值，并且只考虑那些累积概率超过这个阈值的前几个词元。这个范围通常在0到1之间。

较低的p值（例如，0.1）考虑的选项较少。输出更加集中和确定性。
较高的p值（例如，0.9）考虑的选项更多。输出更加多样化。

3.温度参数和Top-p之间的相互作用：

首先应用温度参数，重塑概率分布。
然后，对重塑后的分布应用top-p过滤。

在示例中

do_sample=True, top_p=0.9, temperature=0.1

do_sample=True：启用采样（而不总是选择最有可能的词元）。
top_p=0.9：考虑广泛范围的词元（90%的概率质量）。
temperature=0.1：非常低的温度，使分布更加尖锐。

这种组合倾向于产生以下特性的输出：

相当确定性（由于低温度）
但由于高top-p，仍有一定的可变性
可能连贯且“安全”，但偶尔会有一些意想不到的选择。

五、逐步生成词元?

模型通过统一的生成接口将所有层封装起来。为了清晰地理解这个过程，我们可以使用已加载的模型 base_model_bnb_4b 来分解步骤。

1.创建词元ID和嵌入（embedding）

# 使用分词器将提示文本转换为词元ID，并且将结果转换为PyTorch张量，如果需要的话进行截断，以确保不超过模型的最大序列长度
input_ids = tokenizer(prompt, return_tensors="pt", truncation=True).input_ids.cuda()
#词元ID
# 创建位置ID，用于模型中的位置编码，这里创建了一个与输入序列长度相同的连续整数序列，并增加一个维度以匹配批处理维度
position_ids = torch.arange(0, input_ids.shape[1]).unsqueeze(0).cuda()
# 获取输入序列的长度
seq_length = input_ids.shape[1]
#嵌入
# 将词元ID转换为嵌入向量，这是模型的嵌入层，用于将词元ID转换为模型可以处理的数值表示
embeddings = base_model_bnb_4b.model.embed_tokens(input_ids)
# 打印嵌入向量的形状
print("Embeddings shape:", embeddings.shape)
# 输出：Embeddings shape: torch.Size([1, 49, 4096])
# 表示有1个序列，序列长度为49个词元，每个词元的嵌入维度为4096

2.迭代32层

# 逐层处理模型
for layer_idx, layer in enumerate(base_model_bnb_4b.model.layers):
 print(f"Processing layer {layer_idx}") # 打印当前处理的层索引
 
 # 1. 输入层归一化（Input LayerNorm）
 normalized_hidden_states = layer.input_layernorm(hidden_states)
 
 # 2. 自注意力机制（Self-attention mechanism）
 # 2.1 查询（Query）、键（Key）、值（Value）投影
 query_states = layer.self_attn.q_proj(normalized_hidden_states)
 key_states = layer.self_attn.k_proj(normalized_hidden_states)
 value_states = layer.self_attn.v_proj(normalized_hidden_states)
 
 # 2.2 重塑和转置Q、K、V
 # (batch_size, seq_length, num_heads, head_dim) -> (batch_size, num_heads, seq_length, head_dim)
 query_states = query_states.view(batch_size, seq_length, layer.self_attn.num_heads, layer.self_attn.head_dim).transpose(1, 2)
 key_states = key_states.view(batch_size, seq_length, layer.self_attn.num_key_value_heads, layer.self_attn.head_dim).transpose(1, 2)
 value_states = value_states.view(batch_size, seq_length, layer.self_attn.num_key_value_heads, layer.self_attn.head_dim).transpose(1, 2)
 
 # 2.3 应用旋转位置嵌入（Rotary positional embeddings）
 kv_seq_len = key_states.shape[-2]
 cos, sin = layer.self_attn.rotary_emb(value_states, position_ids)
 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
 
 # 2.4 处理分组查询注意力（Grouped-query attention）
 # 为每个查询组重复K和V
 key_states = repeat_kv(key_states, layer.self_attn.num_key_value_groups)
 value_states = repeat_kv(value_states, layer.self_attn.num_key_value_groups)
 
 # 2.5 计算注意力分数
 attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(layer.self_attn.head_dim)
 attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
 
 # 2.6 应用注意力到值
 attn_output = torch.matmul(attn_weights, value_states)
 
 # 2.7 重塑注意力输出
 # (batch_size, num_heads, seq_length, head_dim) -> (batch_size, seq_length, hidden_size)
 attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_length, hidden_size)
 
 # 2.8 输出投影
 attn_output = layer.self_attn.o_proj(attn_output)
 
 # 2.9 残差连接
 hidden_states = hidden_states + attn_output
 
 # 3. 注意力后的层归一化（Post-attention LayerNorm）
 normalized_hidden_states = layer.post_attention_layernorm(hidden_states)
 
 # 4. MLP（多层感知机，Feed-Forward Network）
 # 4.1 应用门控和上投影
 gate_output = layer.mlp.gate_proj(normalized_hidden_states)
 up_output = layer.mlp.up_proj(normalized_hidden_states)
 
 # 4.2 应用激活函数和逐元素乘法
 mlp_output = up_output * layer.mlp.act_fn(gate_output)
 
 # 4.3 下投影
 mlp_output = layer.mlp.down_proj(mlp_output)
 
 # 4.4 残差连接
 hidden_states = hidden_states + mlp_output
 
 print(f" Output shape: {hidden_states.shape}") # 打印当前层的输出形状

旋转位置嵌入需要这些辅助函数

import torch
import torch.nn.functional as F
import math
# 旋转位置嵌入（RoPE）和分组查询注意力的辅助函数
def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
 """
 应用旋转位置嵌入到查询和键张量。
 
 参数：
 q, k：查询和键张量
 cos, sin：旋转嵌入的余弦和正弦分量
 position_ids：位置ID张量
 
 返回：
 q_embed, k_embed：应用了旋转嵌入的查询和键张量
 """
 # 重塑余弦和正弦张量
 cos = cos.squeeze(1).squeeze(0) # [seq_len, dim]
 sin = sin.squeeze(1).squeeze(0) # [seq_len, dim]
 
 # 选择相关的位置嵌入
 cos = cos[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
 sin = sin[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
 
 # 应用旋转嵌入
 q_embed = (q * cos) + (rotate_half(q) * sin)
 k_embed = (k * cos) + (rotate_half(k) * sin)
 
 return q_embed, k_embed
def rotate_half(x):
 """
 旋转输入张量的一半隐藏维度。
 
 这个操作是旋转位置嵌入技术的一部分。
 
 参数：
 x：输入张量
 
 返回：
 张量的后一维度的一半被旋转
 """
 x1 = x[..., :x.shape[-1] // 2] # 隐藏维度的前半部分
 x2 = x[..., x.shape[-1] // 2:] # 隐藏维度的后半部分
 return torch.cat((-x2, x1), dim=-1) # 连接旋转后的两半
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
 """
 为分组查询注意力重复键和值状态。
 
 这个函数扩展键和值状态以匹配分组查询注意力机制中的查询头数。
 
 参数：
 hidden_states：形状为 (batch, num_key_value_heads, seqlen, head_dim) 的输入张量
 n_rep：重复次数（通常是 num_query_heads // num_key_value_heads）
 
 返回：
 形状为 (batch, num_attention_heads, seqlen, head_dim) 的张量
 """
 batch, num_key_value_heads, slen, head_dim = hidden_states.shape
 if n_rep == 1:
 return hidden_states # 如果 n_rep 为 1，则无需重复
 
 # 扩展并重塑以重复键/值状态
 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)

然后，输出将经过最后的过程:

# 最后一层归一化（Final LayerNorm）
# 对最后一层的隐藏状态进行归一化
hidden_states = base_model_bnb_4b.model.norm(hidden_states)
# 语言模型头部（Language Model Head）
# 将归一化的隐藏状态投影到词汇空间
lm_logits = base_model_bnb_4b.lm_head(hidden_states)
# 获取最后一个词元的logits
# 我们只对预测下一个词元感兴趣，所以我们取最后一个位置的logits
last_token_logits = lm_logits[:, -1, :]
# 注意：以下行被注释掉了，因为我们使用了更复杂的采样方法
# next_token = torch.argmax(last_token_logits, dim=-1)
# 第7步：应用温度参数（Apply temperature）
# 温度参数调整预测的随机性。较低的值使模型更有信心。
temperature = 0.1
scaled_logits = last_token_logits / temperature
# 第8步：应用Top-p（核）采样（Apply top-p (nucleus) sampling）
# 这种方法会截断最不可能的词元，其累积概率超过(1 - top_p)
top_p = 0.9
# 将logits按降序排序
sorted_logits, sorted_indices = torch.sort(scaled_logits, descending=True)
# 计算累积概率
cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
# 移除累积概率超过阈值的词元
sorted_indices_to_remove = cumulative_probs > top_p
# 调整索引以保留第一个超过阈值的词元
sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
sorted_indices_to_remove[..., 0] = 0
# 将排序后的张量散布回原始索引
indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
# 将移除索引的logits设置为负无穷
scaled_logits[indices_to_remove] = float('-inf')
# 第9步：从过滤后的分布中采样（Sample from the filtered distribution）
# 将logits转换为概率
probs = torch.softmax(scaled_logits, dim=-1)
# 根据计算出的概率随机采样一个词元
next_token = torch.multinomial(probs, num_samples=1)
# 打印输入词元和预测的下一个词元
print(f"\nInput tokens: {input_ids}")
print(f"Most likely next token ID: {next_token.item()}")

你可以继续这个过程，直到生成一个停止词元

# 将新生成的词元追加到input_ids
input_ids = torch.cat([input_ids, next_token], dim=-1)
# 检查是否生成了序列结束词元
if next_token.item() == tokenizer.eos_token_id:
 break

输出将会是

Input tokens: tensor([[128000, 128000, 128006, 9125, 128007, 271, 29815, 389, 279,
 2038, 3984, 11, 18622, 279, 11914, 555, 10223, 1202,
 43787, 505, 3347, 311, 3938, 13, 128009, 128006, 882,
 128007, 271, 8100, 6476, 279, 27374, 32719, 369, 4207,
 323, 1243, 10717, 439, 433, 574, 33433, 13, 128009,
 128006, 78191, 128007, 271]], device='cuda:0')
Most likely next token ID: 8586
Input tokens: tensor([[128000, 128000, 128006, 9125, 128007, 271, 29815, 389, 279,
 2038, 3984, 11, 18622, 279, 11914, 555, 10223, 1202,
 43787, 505, 3347, 311, 3938, 13, 128009, 128006, 882,
 128007, 271, 8100, 6476, 279, 27374, 32719, 369, 4207,
 323, 1243, 10717, 439, 433, 574, 33433, 13, 128009,
 128006, 78191, 128007, 271, 8586]], device='cuda:0')
Most likely next token ID: 374
...
Input tokens: tensor([[128000, 128000, 128006, 9125, 128007, 271, 29815, 389, 279,
 2038, 3984, 11, 18622, 279, 11914, 555, 10223, 1202,
 43787, 505, 3347, 311, 3938, 13, 128009, 128006, 882,
 128007, 271, 8100, 6476, 279, 27374, 32719, 369, 4207,
 323, 1243, 10717, 439, 433, 574, 33433, 13, 128009,
 128006, 78191, 128007, 271, 8586, 374, 279, 11914, 449,
 279, 43787, 5614, 311, 3938, 1473, 8100, 690, 1514,
 279, 27374, 32719, 369, 4207, 323, 1243, 3009, 439,
 433, 690, 387, 33433, 13]], device='cuda:0')
Most likely next token ID: 128009

这个过程会一直重复，直到模型预测出一个序列结束词元（如ID 128009），或者达到预设的序列长度限制。

你可以在GitHub仓库中找到Jupyter Notebook。

六、结论?

我们过去常常严重依赖数学概念来理解模型的工作原理，但这往往导致与实际代码的脱节。这次深入探讨采取了一种实践方法，从反向角度探索模型的工作流程。它提供了对模型在现实世界场景中如何运作的深入理解。

?请大家登录博客园，多多支持岁月月宝贝，本宝会继续分享高质量内容！?

Llama3.1 是怎么工作哒？原文翻译版