【LLM】从零开始实现 LLaMA3

分词器

在这里，我们不会实现一个 BPE 分词器（但 Andrej Karpathy 有一个非常简洁的实现）。

BPE（Byte Pair Encoding，字节对编码）是一种数据压缩算法，也被用于自然语言处理中的分词方法。它通过逐步将常见的字符或子词组合成更长的词元（tokens），从而有效地表示文本中的词汇。

在自然语言处理中的 BPE 分词器的工作原理如下：

1. 初始化：首先，将所有词汇表中的单词分解为单个字符或符号。例如，单词 “hello” 会被表示为 ["h", "e", "l", "l", "o"]。

2. 统计频率：接下来，统计所有字符对（相邻字符组合）的出现频率。例如，如果 “l” 和 “l” 出现在一起的频率最高，那么它们会被作为一个新的词元 “ll”。

3. 合并频率最高的字符对：将出现频率最高的字符对合并成一个新的词元。然后重复这个过程，直到达到预定义的词元数量或不能再合并为止。

4. 生成词汇表：最终生成的词汇表包含了从单个字符到更复杂的子词的所有词元，这些词元可以组合成原始的单词和短语。

优点：

• BPE 分词器的优点在于它能够有效处理词汇量巨大的语言，尤其是在处理罕见词汇时。这种方法能够通过将罕见词分解为更常见的子词来避免产生过多的不在词表的词（OOV, Out-of-Vocabulary）。
• BPE 也适用于多语言模型，因为它不依赖于特定语言的词汇表，能够有效地分解和编码来自不同语言的文本。

应用：

• 在大规模预训练语言模型（如 GPT、BERT、LLaMA）中，BPE 被广泛使用，作为一种分词方法，它能有效地将文本转换为模型可以处理的词元序列，同时保持文本的连贯性。

正则表达式

正则表达式的各部分通过逻辑或 (|) 操作符连接，这意味着它会依次尝试匹配每一个部分，直到找到一个匹配为止。以下是各部分的匹配优先顺序和使用方式的说明：

1. (?i:'s|'t|'re|'ve|'m|'ll|'d)：

   • 用途：首先检查是否有常见的英语缩写。这是优先级最高的，因为这些缩写通常会出现在词的末尾，直接匹配缩写可以避免误分割。
   • 示例：在处理 it's 时，匹配到 's，这个部分将被识别为一个独立的词元。

2. [^\r\n\p{L}\p{N}]?\p{L}+：

   • 用途：如果没有匹配到缩写，它会检查是否有以字母为主的词语，并允许在字母前有一个可选的非字母、非数字字符（如标点符号）。
   • 示例：对于 O'Connor，O 和 'Connor 会被分别识别。

3. \p{N}{1,3}：

   • 用途：用于匹配由 1 到 3 个数字组成的短数字序列。
   • 示例：数字 42 会被单独识别为一个词元。

4. ?[^\s\p{L}\p{N}]+[\r\n]*：

   • 用途：匹配非字母、非数字的符号或特殊字符，可能包括一个前导空格，并考虑到可能的换行符。
   • 示例：句号 . 或逗号 , 会被识别为单独的词元。

5. \s*[\r\n]+：

   • 用途：用于识别和分割由换行符组成的行间空白。
   • 示例：在换行符 \n 或回车符 \r\n 前后的空白会被识别并处理。

6. \s+(?!\S)：

   • 用途：匹配位于行尾的空白字符。这部分通过负向先行断言确保这些空白后面不会跟随非空白字符，主要用于捕获行尾的多余空格。
   • 示例：行尾的多个空格 " " 会被识别并处理。

7. \s+：

   • 用途：用于匹配普通的空白字符序列，确保文本中的空格或制表符被正确识别和分割。
   • 示例：词与词之间的空格 " " 会被识别为一个独立的词元。

应用示例

假设输入文本是 "It's 42!"：

8. (?i:'s|'t|'re|'ve|'m|'ll|'d) 匹配 's，因此 "It's" 会被分割为 "It" 和 " 's"。
9. [^\r\n\p{L}\p{N}]?\p{L}+ 匹配 "It"，将其作为一个单独的词元。
10. \p{N}{1,3} 匹配 "42"，将其作为一个独立的数字词元。
11. ?[^\s\p{L}\p{N}]+[\r\n]* 匹配 "!"，将其作为一个独立的符号词元。

最终结果是 "It's 42!" 会被分割为四个词元：["It", "'s", "42", "!"]。

这个正则表达式通过一系列有序的匹配规则来确保文本被合理地分割为不同类型的词元，以便进一步处理或分析。

代码详解

加载和使用一个预训练的 tokenizer（分词器），并对文本进行编码和解码

from pathlib import Path
import tiktoken
from tiktoken.load import load_tiktoken_bpe
import torch
import json
import matplotlib.pyplot as plt

# 定义 tokenizer 路径和特殊 tokens
tokenizer_path = "Meta-Llama-3-8B/tokenizer.model"
special_tokens = [
            "<|begin_of_text|>",
            "<|end_of_text|>",
            "<|reserved_special_token_0|>",
            "<|reserved_special_token_1|>",
            "<|reserved_special_token_2|>",
            "<|reserved_special_token_3|>",
            "<|start_header_id|>",
            "<|end_header_id|>",
            "<|reserved_special_token_4|>",
            "<|eot_id|>",  # end of turn
        ] + [f"<|reserved_special_token_{i}|>" for i in range(5, 256 - 5)]

# 返回一个包含 token 和其对应 rank 的字典
mergeable_ranks = load_tiktoken_bpe(tokenizer_path)

tokenizer = tiktoken.Encoding(# 创建一个 tokenizer 实例。
    name=Path(tokenizer_path).name, #  tokenizer 的名称，通常是模型文件的名称。
    pat_str=r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+", # 正则表达式模式，用于定义如何将文本分割成 tokens。
    mergeable_ranks = mergeable_ranks, # 之前加载的 BPE 模型。
    special_tokens = {token: len(mergeable_ranks) + i for i, token in enumerate(special_tokens)}, # 将特殊 tokens 映射到唯一的 ID，这些 ID 从 mergeable_ranks 的长度开始递增。
)

# tokenizer.encode(...)：将文本编码为 token IDs。
# tokenizer.decode(...): 将 token IDs 解码回文本。

print(tokenizer.decode(tokenizer.encode("hello world!")))
print(tokenizer.encode("hello world!"))

读取模型文件

通常，读取模型文件依赖于：模型类的编写方式 以及 类中变量的名称。

加载一个名为 "Meta-Llama-3-8B/consolidated.00.pth" 的 PyTorch 模型，并打印出该模型的前20个参数的名称。这有助于了解模型的结构及其包含的各个参数。

model = torch.load("Meta-Llama-3-8B/consolidated.00.pth") # 加载 PyTorch 模型权重文件
print(json.dumps(list(model.keys())[:20], indent=4))

• model: 加载后的模型通常是一个字典（dict）

• 其中键是模型参数的名称，值是对应的参数（如权重矩阵、偏置等）

• print(json.dumps(list(model.keys())[:20], indent=4))：

  • model.keys() 获取模型字典中所有参数的名称。每个参数（例如权重、偏置等）都会有一个对应的名称作为键（key）。

  • list(model.keys()) 将这些键转换成一个列表（model.keys() 返回的是一个字典的视图，需要用 list() 转换成列表）。

  • [:20] 表示只取这个列表中的前20个元素，限制输出的参数键数量。

  • json.dumps(..., indent=4) 将这个列表格式化为一个 JSON 字符串，并且使用 4 个空格进行缩进，方便阅读。

  • print(...) 将格式化后的 JSON 字符串打印出来。

• 输出：

  [
      "tok_embeddings.weight",
      "layers.0.attention.wq.weight",
      "layers.0.attention.wk.weight",
      "layers.0.attention.wv.weight",
      "layers.0.attention.wo.weight",
      "layers.0.feed_forward.w1.weight",
      "layers.0.feed_forward.w3.weight",
      "layers.0.feed_forward.w2.weight",
      "layers.0.attention_norm.weight",
      "layers.0.ffn_norm.weight",
      "layers.1.attention.wq.weight",
      "layers.1.attention.wk.weight",
      "layers.1.attention.wv.weight",
      "layers.1.attention.wo.weight",
      "layers.1.feed_forward.w1.weight",
      "layers.1.feed_forward.w3.weight",
      "layers.1.feed_forward.w2.weight",
      "layers.1.attention_norm.weight",
      "layers.1.ffn_norm.weight",
      "layers.2.attention.wq.weight"
  ]
  

看原生 8b 的参数设置

这段代码的作用是从 params.json 文件中读取配置，并将其转化为 Python 字典（config），方便在程序中使用。如果你想查看文件的具体内容，可以打印 config 变量。

with open("Meta-Llama-3-8B/params.json", "r") as f:
    config = json.load(f)
config

# 将 config 中的内容取出来
dim = config["dim"]
n_layers = config["n_layers"]
n_heads = config["n_heads"]
n_kv_heads = config["n_kv_heads"]
vocab_size = config["vocab_size"]
multiple_of = config["multiple_of"]
ffn_dim_multiplier = config["ffn_dim_multiplier"]
norm_eps = config["norm_eps"]
rope_theta = torch.tensor(config["rope_theta"])

这段代码的作用是读取并解析 Meta-Llama-3-8B/params.json 文件中的 JSON 数据，并将其存储到一个 Python 字典对象中。具体的操作过程如下：

1. with open("Meta-Llama-3-8B/params.json", "r") as f:
   • 使用 Python 的 open 函数以只读模式（"r"）打开位于 "Meta-Llama-3-8B/params.json" 路径下的文件。open 会返回一个文件对象，通常我们会用 with 语句来确保文件在操作完成后能够自动关闭。
   • f 是打开文件后返回的文件对象。
2. config = json.load(f)
   • json.load(f) 会读取文件 f 中的 JSON 数据，并将其解析为 Python 对象。通常，JSON 数据会转换为字典（如果是一个对象）或列表（如果是一个数组）。这里，config 就会存储解析后的 JSON 数据，通常是一个字典，包含 params.json 文件中的所有键值对。
3. config
   • 这行代码是直接访问变量 config，它包含了 JSON 文件中的内容，通常是一个 Python 字典，可以通过这个字典来访问和处理 JSON 文件中的数据。

• 输出：

  {'dim': 4096,
   'n_layers': 32,
   'n_heads': 32,
   'n_kv_heads': 8,
   'vocab_size': 128256,
   'multiple_of': 1024,
   'ffn_dim_multiplier': 1.3,
   'norm_eps': 1e-05,
   'rope_theta': 500000.0}
  

这是一个典型的 Transformer 模型（如 LLaMA 或其他类似模型）的配置文件，描述了模型的结构和超参数。以下是对每个参数的解释：

1. dim

• 含义：模型的隐藏层维度（hidden dimension），即每个 Transformer 层的输入和输出的向量维度。
• 解释：dim=4096 表示每个 token 的嵌入向量和隐藏状态的维度是 4096。
• 作用：决定了模型的容量和计算复杂度。维度越大，模型表达能力越强，但计算开销也越大。

2. n_layers

• 含义：Transformer 模型的层数（即堆叠的 Transformer 块的数量）。
• 解释：n_layers=32 表示模型有 32 层 Transformer 块。
• 作用：层数越多，模型可以学习更复杂的特征，但训练和推理的计算成本也会增加。

3. n_heads

• 含义：多头注意力机制（Multi-Head Attention）中的注意力头数。
• 解释：n_heads=32 表示模型有 32 个注意力头。
• 作用：每个注意力头可以关注输入序列的不同部分，多头机制增强了模型的表达能力。

4. n_kv_heads

• 含义：键值（Key-Value）注意力头的数量。
• 解释：n_kv_heads=8 表示在计算注意力时，键和值的注意力头数为 8。
• 作用：这是一种优化技术，可以减少计算量。通常，n_kv_heads 小于 n_heads，通过分组共享键值头来降低计算复杂度。

5. vocab_size

• 含义：词汇表的大小，即模型可以处理的不同 token 的数量。
• 解释：vocab_size=128256 表示词汇表中有 128,256 个不同的 token。
• 作用：词汇表大小决定了模型的输入和输出空间。较大的词汇表可以更好地表示语言，但也会增加模型的计算和存储开销。

6. multiple_of

• 含义：前馈神经网络（FFN）隐藏层维度的倍数约束。
• 解释：multiple_of=1024 表示 FFN 的隐藏层维度必须是 1024 的倍数。
• 作用：这种约束通常用于优化硬件计算效率（如 GPU 的矩阵计算），确保维度对齐。

7. ffn_dim_multiplier

• 含义：前馈神经网络（FFN）隐藏层维度的乘数。
• 解释：ffn_dim_multiplier=1.3 表示 FFN 的隐藏层维度是 dim 的 1.3 倍。
• 作用：FFN 的隐藏层维度通常比输入维度大，以增强模型的表达能力。

8. norm_eps

• 含义：层归一化（Layer Normalization）中的 epsilon 值，用于数值稳定性。
• 解释：norm_eps=1e-05 表示在层归一化中，分母加上一个很小的值（1e-5）以避免除零错误。
• 作用：确保归一化计算的稳定性，避免梯度爆炸或消失。

9. rope_theta

• 含义：旋转位置编码（Rotary Position Embedding, RoPE）中的基数参数。
• 解释：rope_theta=500000.0 是 RoPE 的一个超参数，用于控制位置编码的旋转频率。
• 作用：RoPE 是一种改进的位置编码方法，rope_theta 的值会影响位置编码的表示能力。

将文本转换为 tokens

将文本转换为 tokens 的过程使用了 tiktoken 库（它是 OpenAI 提供的一个工具库）。在自然语言处理中，tokenizer 是将文本转化为模型可以理解的 token 的工具，而 tiktoken 是一个专门为 OpenAI 的语言模型（如 GPT 系列）设计的 tokenization 库。

1. 什么是 Token（标记）？
   • 在自然语言处理中，token 是文本的基本单位，可以是一个单词、子词或符号。语言模型在处理文本时，不是直接处理字面上的字符或单词，而是将文本切分成一个个 token，然后通过这些 token 来理解和生成语言。
   • 例如，句子 “I love coding” 可能会被分解成以下 tokens：["I", "love", "coding"] 或 ["I", " love", " coding"]，具体取决于 tokenizer 的设计。
2. tiktoken 是什么？
   • tiktoken 是 OpenAI 提供的一个用于将文本转化为 token 的库，专门为 OpenAI 的 GPT 模型（例如 GPT-3、GPT-4 等）设计。
   • 这个库使用了一种非常高效的 tokenization 方法，能够快速地将大量文本分解成适合模型处理的 token 格式。
   • tiktoken 库能够根据不同的 GPT 模型（如 GPT-2、GPT-3）使用相应的 tokenization 规则，确保生成的 token 与模型的训练方式一致。
3. tokenizer 的作用
   • Tokenizer（分词器） 是一个将原始文本转换为 tokens 的工具。文本必须先经过 tokenization，才能被输入到语言模型中进行处理。
   • tiktoken 就是这个分词器，它的任务是把原始文本转化成模型能够理解的 token。

这段代码的主要功能是将一个字符串 prompt 转换为 tokens，并在每个 token 之间加上一个特殊的 token 128000，然后通过解码过程将这些 tokens 还原成字符串。最后，代码将打印出这些 tokens 和它们解码后的文本。

prompt = "the answer to the ultimate question of life, the universe, and everything is "
tokens = [128000] + tokenizer.encode(prompt)
print(tokens)

tokens = torch.tensor(tokens)
prompt_split_as_tokens = [tokenizer.decode([token.item()]) for token in tokens]

print(prompt_split_as_tokens)

• 输出：

[128000, 1820, 4320, 311, 279, 17139, 3488, 315, 2324, 11, 279, 15861, 11, 323, 4395, 374, 220]
['<|begin_of_text|>', 'the', ' answer', ' to', ' the', ' ultimate', 'uestion', ' of', ' life', ',', ' the', ' universe', ',', ' and', 'everything', ' is', ' ']

将 token 转换为它们的嵌入向量

embedding_layer = torch.nn.Embedding(vocab_size, dim) # 128256, 4096
embedding_layer.weight.data.copy_(model["tok_embeddings.weight"])
token_embeddings_unnormalized = embedding_layer(tokens).to(torch.bfloat16) # accelerate
token_embeddings_unnormalized.shape

1. 初始化嵌入层

embedding_layer = torch.nn.Embedding(vocab_size, dim) # 初始时是随机的，但它是可学习的。

• 这行代码创建了一个 嵌入层（Embedding Layer）。
• vocab_size：表示词汇表的大小，也就是有多少个不同的 token（词语、子词等）。在这里，它是 128,256，说明有 128,256 个不同的 token。
• dim：表示每个 token 的嵌入维度，也就是说每个 token 会被表示为一个 dim 维的向量。在这段代码中，dim 是 4096，所以每个 token 会被映射成一个 4096 维的向量。

2. 复制预训练权重

embedding_layer.weight.data.copy_(model["tok_embeddings.weight"])

• 这行代码的作用是将预训练模型中的嵌入权重加载到这个嵌入层中。
• model["tok_embeddings.weight"] 是从预训练模型中获取嵌入层的权重（即 token 嵌入矩阵）。
• .weight.data 访问嵌入层权重的原始数据。
• .copy_() 将预训练模型中的嵌入权重复制到当前的 embedding_layer 中。这样做是为了初始化嵌入层，使其使用预训练的权重。

3. 获取 token 嵌入

token_embeddings_unnormalized = embedding_layer(tokens).to(torch.bfloat16)

• embedding_layer(tokens)：这行代码执行了嵌入查找。tokens 是输入的 token ID（通常是一个 token ID 的张量），嵌入层会根据这些 ID 查找并返回对应的嵌入向量。
• .to(torch.bfloat16)：将生成的嵌入向量转换为 bfloat16 数据类型，这样可以节省内存并加速计算。bfloat16 是一种 16 位浮点数格式，通常用于加速训练，特别是在支持 bfloat16 的硬件（如 Google TPU 或某些支持的 GPU）上。

4. 查看 token 嵌入的形状

token_embeddings_unnormalized.shape

• 这一行返回 token_embeddings_unnormalized 张量的形状。

• 这个张量的形状通常是

  (batch_size, seq_len, dim)
  

  ，其中：

  • batch_size：批次大小，即一次处理的序列数量。
  • seq_len：每个序列的长度，也就是 token 数量。
  • dim：嵌入的维度，在这里是 4096，表示每个 token 对应的嵌入向量的大小。

• 输出：

  torch.Size([17, 4096])
  

然后我们使用 RMS 归一化对嵌入向量进行归一化

需要注意的是，在这一步之后，张量的形状不会发生改变，只有值被归一化。为了避免出现除以零的情况，我们需要使用一个 norm_eps（来自配置文件）。这是因为我们不希望意外地将 RMS 设置为 0，从而导致除以 0 的错误。以下是公式：

这段代码实现了一个 RMS Normalization（均方根归一化）的操作，通常用于深度学习中的模型训练，尤其是在神经网络的标准化层中。

def rms_norm(tensor, norm_weights):
    return (tensor * torch.rsqrt(tensor.pow(2).mean(-1, keepdim=True) + norm_eps)) * norm_weights

$$RMS(a)=\sqrt{\frac{{\sum }_{i=1}^n{a}_i^2}{n}}$$

构建 Transformer 的第一层

token_embeddings = rms_norm(token_embeddings_unnormalized, model["layers.0.attention_norm.weight"])
token_embeddings.shape

这行代码的作用是对 token_embeddings_unnormalized 张量进行 RMS 归一化操作，并使用 model["layers.0.attention_norm.weight"] 作为归一化的权重。最后，它返回归一化后张量的形状。

从零实现注意力机制

# GQA: 4q -> 1k 几个组 concat 之后得到 wo
print(
    model["layers.0.attention.wq.weight"].shape,
    model["layers.0.attention.wk.weight"].shape,
    model["layers.0.attention.wv.weight"].shape,
    model["layers.0.attention.wo.weight"].shape
)

• 输出：

  torch.Size([4096, 4096]) 
  torch.Size([1024, 4096]) 
  torch.Size([1024, 4096])
  torch.Size([4096, 4096])
  

多头 query

从多个注意力头中解包查询（query），解包后的结果形状为 [32x128x4096]。其中，32 代表 LLaMA 3 中注意力头的数量，128 是查询向量的大小，而 4096 是 token 嵌入向量的大小。

q_layer0 = model["layers.0.attention.wq.weight"]
head_dim = q_layer0.shape[0] // n_heads # 4096 / 32 = 128，一个头的维度是：128 维
q_layer0 = q_layer0.view(n_heads, head_dim, dim) # 32 * 128 == 4096
q_layer0.shape # [32x128x4096]

第一层的第一个注意力头

首先访问第一层第一个注意力头的查询权重矩阵，该矩阵的大小为 [128x4096]。

q_layer0_head0 = q_layer0[0]
q_layer0_head0.shape # torch.Size([128, 4096])

我们现在将查询权重 $W^q$ 与 token 嵌入向量相乘，以生成 token 的 $Q$

结果的形状是 [17x128]，这是因为我们有 17 个 token，每个 token 对应一个长度为 128 的查询向量。

q_per_token = torch.matmul(token_embeddings, q_layer0_head0.T) # [17, 4096] [4096, 128]
q_per_token.shape # torch.Size([17, 128])

位置编码（Positional Encoding）

我们现在已经为提示中的每个 token 生成了查询向量，但如果你仔细想想，单独的查询向量并不知道它在提示中的位置。

查询："the answer to the ultimate question of life, the universe, and everything is "

在我们的提示中，我们使用了三次 “the”，我们需要所有三个 “the” token 的查询向量（每个大小为 [1x128]）根据它们在查询中的位置而有所不同。我们使用 RoPE（旋转位置编码，Rotary Positional Embedding） 来实现这种旋转。

为什么需要位置编码？

• 问题：Transformer 模型本身是 位置无关的，即它无法区分输入序列中 token 的顺序。例如，句子 “A B C” 和 “C B A” 会被模型视为相同的输入。
• 解决方案：通过位置编码，为每个 token 添加位置信息，使模型能够理解 token 的顺序。

RoPE（旋转位置编码）

RoPE 是一种高效的位置编码方法，通过旋转向量的方式将位置信息注入到查询（query）和键（key）向量中。

1. 核心思想

   • 对查询向量和键向量进行 旋转，旋转的角度与 token 的位置相关。

   • 旋转后的向量既保留了原始信息，又包含了位置信息。

   • 横轴 $x$ 为实数，纵轴 $y$ 为虚数

2. 公式

   • $RoPE$ 的旋转公式如下：
     $$RoPE(x,m)=x\cdot \cos (m\theta )+rotate(x)\cdot \sin (m\theta )$$

     • $x$ 是查询向量或键向量
     • $m$ 是 $token$ 的位置（例如：第一个 token、第二个 token 等）
     • $\theta$ 是一个与维度有关的参数‘

3. 效果

   • 对于不同位置的相同 token（例如三个 “the”），RoPE 会生成不同的查询向量。
   • 旋转后的查询向量既包含了 token 的语义信息，也包含了位置信息。

实现步骤

1. 输入：
   • 查询向量：形状为 [17x128]（17 个 token，每个 token 的查询向量维度为 128）。
   • 位置信息：每个 token 的位置索引（例如 [0, 1, 2, ..., 16]）。
2. 旋转操作：
   • 对每个 token 的查询向量应用 RoPE 旋转。
   • 旋转后的查询向量形状仍为 [17x128]，但每个向量现在包含了位置信息。
3. 输出：
   • 旋转后的查询向量：形状为 [17x128]，每个 token 的查询向量根据其位置进行了调整。

假设原始的 q_per_token 张量的形状是 (batch_size, length, features)，其中 batch_size 是批次大小，length 是序列长度，features 是特征维度。通过 view(q_per_token.shape[0], -1, 2)，张量的形状将被重塑成 (batch_size, new_length, 2)，其中 new_length 是通过将原始的 length * features 拆分成 2 的倍数得到的。

q_per_token_split_into_pairs = q_per_token.float().view(q_per_token.shape[0], -1, 2)
q_per_token_split_into_pairs.shape # torch.Size([17, 64, 2])

这段代码的目的是将 q_per_token 张量中的数据按 2 个一组进行重新排列。

在上面的步骤中，我们将查询向量（query vectors）分成一对一对的，然后对每一对应用一个旋转角度偏移！

现在我们有一个形状为 [17x64x2] 的向量，这是将每个 token 的 128 维查询向量拆分为 64 对！每一对都会根据 token 的位置（m）旋转一个角度 m * theta，其中 m 是当前 token 的位置，theta 是预定义的旋转角度。

详细解释

1. 拆分查询向量：

   • 原始查询向量的形状是[17, 128]其中：
     • 17 是 prompt 中的 token 数量。
     • 128 是每个 token 的查询向量的长度。
   • 我们将每个 128 维的查询向量拆分为 64 对（128 / 2 = 64），因此形状变为 [17, 64, 2]。

2. 旋转角度偏移：

   • 对每一对查询向量（[2] 的形状），我们根据 token 的位置 m 旋转一个角度 m * theta。
   • m 是当前 token 的位置（例如，第一个 token 的 m = 0，第二个 token 的 m = 1，依此类推）。
   • theta 是一个预定义的角度值，通常与位置编码（positional encoding）相关。

3. 旋转的目的：

   • 旋转操作是为了将位置信息（positional information）注入到查询向量中。
   • 这样，模型可以区分不同位置的 token，即使它们的查询向量在内容上相同。

4. 旋转的实现：

   • 对于每一对 [x, y]，旋转角度 m * theta 后，新的值可以通过二维平面中向量的旋转公式计算：
     $$\begin{gathered} x^{{}^{\prime }}=x\cdot \cos (m\cdot \theta )-y\cdot \sin (m\cdot \theta ) \\ {y}^{{}^{\prime }}=x\cdot \sin (m\cdot \theta )+y\cdot \cos (m\cdot \theta ) \end{gathered}$$

示例

假设：

• 有一个 prompt，包含 3 个 token。
• 每个 token 的查询向量长度为 4（为了简化，实际中是 128）。
• 预定义的 theta = 0.1。

1. 原始查询向量：

   [
       [x1, y1, x2, y2],  # Token 1
       [x3, y3, x4, y4],  # Token 2
       [x5, y5, x6, y6]   # Token 3
   ]
   

2. 拆分为两两一对：

   [
       [[x1, y1], [x2, y2]],  # Token 1
       [[x3, y3], [x4, y4]],  # Token 2
       [[x5, y5], [x6, y6]]   # Token 3
   ]
   

3. 旋转操作：

   • 对于 Token 1（m = 0），旋转角度为 0 * theta = 0，因此查询向量不变。
   • 对于 Token 2（m = 1），旋转角度为 1 * theta = 0.1，应用旋转公式。
   • 对于 Token 3（m = 2），旋转角度为 2 * theta = 0.2，应用旋转公式。

4. 旋转后的结果：

   [
       [[x1, y1], [x2, y2]],  # Token 1（未旋转）
       [[x3', y3'], [x4', y4']],  # Token 2（旋转 0.1）
       [[x5', y5'], [x6', y6']]   # Token 3（旋转 0.2）
   ]
   

使用复数的点积来旋转向量

zero_to_one_split_into_64_parts 是一个张量，包含从 0 到 1 之间均匀分布的 64 个数值，其中每个数值相差 1/64。这个张量的范围是 [0, 1)，代表了 64 个等间距的分割点。

zero_to_one_split_into_64_parts = torch.tensor(range(64))/64