算法进阶 | PyTorch从零构建Llama 3（建议收藏！）

本文泉源公众号“算法进阶”，仅用于学术分享，侵权删，干货满满。
原文链接：PyTorch从零构建Llama 3
大家好，本文将详细指导怎样从零开始构建完备的Llama 3模子架构，并在自界说数据集上执行训练和推理。

[图1]：Llama 3架构展示训练和推理流程。由于官方Llama 3论文中未提供相关图表。以是此图为大概架构图，阅读本文后你应能绘制出更为准确的架构图。
本文目标

通过本文。你可以了解到:

• 深入理解Llama 3模子各组件的底层工作原理。
  
  
• 编写代码构建Llama 3的每个组件，并将它们组装成一个功能完备的Llama 3模子。
  
  
• 编写代码使用新的自界说数据集训练模子。
  
  
• 编写代码执行推理，使Llama 3模子可以或许根据输入提示生成新文本。
  

1、输入模块

如图1所示，输入模块包含三个组件：文本/提示、分词器和嵌入。
输入模块内部工作流程
让我们通过下图了解输入模块内的工作流程。

[图2]：输入模块流程图，展示提示、分词器和嵌入流程。
首先，单个或批量文本/提示被输入模子。例如：图中的"Hello World"。
输入模子的必须是数字格式，由于模子无法直接处置惩罚文本。分词器将这些文本/提示转换为标记ID（词汇表中标记的索引号表现）。我们将使用Tiny Shakespeare数据集构建词汇表并训练模子。Llama 3模子使用TikToken作为分词器，这是一种子词分词器。但是我们这个实现将使用字符级分词器。这样做的重要原因是让我们可以或许自行构建词汇表和分词器，包括编码和解码函数，这样可以深入理解底层工作原理并完全掌控代码。
每个标记ID将被转换为128维的嵌入向量（原始Llama 3 8B中为4096维）。然后这些嵌入将被转达到下一个解码器模块。
输入模块代码实现：

1.  # 导入必要的库  
2.  import torch  
3.  from torch import nn  
4.  from torch.nn import functional as F  
5.  
6.  import math  
7.  import numpy as np  
8.  import time  
9.  from dataclasses import dataclass  
10.  from typing import Optional, Tuple, List  
11.  import pandas as pd  
12.  from matplotlib import pyplot as plt
13.  
14.  ### 步骤1: 输入模块 ###  
15.  
16.  # 使用Tiny Shakespeare数据集实现字符级分词器。部分字符级分词器代码参考自Andrej Karpathy的GitHub仓库
17.  # (https://github.com/karpathy/nanoGPT/blob/master/data/shakespeare_char/prepare.py)
18.  # 加载tiny_shakespeare数据文件 (https://github.com/tamangmilan/llama3/blob/main/tiny_shakespeare.txt)  
19.  
20.  device: str = 'cuda' if torch.cuda.is_available() else 'cpu'   # 根据可用性分配设备为cuda或cpu  
21.  
22.  # 加载tiny_shakespeare数据文件  
23.  with open('tiny_shakespeare.txt', 'r') as f:  
24.    data = f.read()  
25.  
26.  # 通过提取tiny_shakespeare数据中的所有唯一字符准备词汇表  
27.  vocab = sorted(list(set(data)))  
28.  
29.  # 训练Llama 3模型需要额外的标记，如<|begin_of_text|>、<|end_of_text|>和<|pad_id|>，将它们添加到词汇表中  
30.  vocab.extend(['<|begin_of_text|>','<|end_of_text|>','<|pad_id|>'])  
31.  vocab_size = len(vocab)  
32.  
33.  # 创建字符与词汇表中对应整数索引之间的映射。
34.  # 这对于构建分词器的编码和解码函数至关重要。
35.  itos = {i:ch for i, ch in enumerate(vocab)}  
36.  stoi = {ch:i for i, ch in enumerate(vocab)}  
37.  
38.  # 分词器编码函数：输入字符串，输出整数列表  
39.  def encode(s):  
40.    return [stoi[ch] for ch in s]  
41.  
42.  # 分词器解码函数：输入整数列表，输出字符串  
43.  def decode(l):  
44.    return ''.join(itos[i] for i in l)  
45.  
46.  # 定义稍后在模型训练中使用的张量标记变量  
47.  token_bos = torch.tensor([stoi['<|begin_of_text|>']], dtype=torch.int, device=device)  
48.  token_eos = torch.tensor([stoi['<|end_of_text|>']], dtype=torch.int, device=device)  
49.  token_pad = torch.tensor([stoi['<|pad_id|>']], dtype=torch.int, device=device)  
50.  
51.  prompts = "Hello World"  
52.  encoded_tokens = encode(prompts)  
53.  decoded_text = decode(encoded_tokens)  
54.  
55.  ### 输入模块代码测试 ###  
56.  # 取消下面的三重引号来执行测试  
57.  """  
58.  print(f"Shakespeare文本字符长度: {len(data)}")  
59.  print(f"词汇表内容: {''.join(vocab)}\n")  
60.  print(f"词汇表大小: {vocab_size}")  
61.  print(f"编码后的标记: {encoded_tokens}")  
62.  print(f"解码后的文本: {decoded_text}")  
63.  """  
64.  ### 测试结果: ###  
65.  """  
66.  Shakespeare文本字符长度: 1115394  
67.  词汇表内容:  
68.   !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz<|begin_of_text|><|end_of_text|><|pad_id|>  
69.  
70.  词汇表大小: 68  
71.  编码后的标记: [20, 43, 50, 50, 53, 1, 35, 53, 56, 50, 42]  
72.  解码后的文本: Hello World  
73.  """

复制代码

2、解码器模块

参照图1的架构图，解码器模块包含以下子组件：

• RMS归一化
  
  
• 旋转位置编码
  
  
• KV缓存
  
  
• 分组查询留意力
  
  
• 前馈网络
  
  
• 解码器块
  

RMS归一化（Root Mean Square Normalization）

RMSNorm的必要性
从图1可以看出，输入模块的输出（嵌入向量）颠末RMSNorm模块。这是由于嵌入向量具有多个维度（Llama3-8b中为4096维），大概出现不同范围的值。这会导致模子梯度爆炸或消失，从而导致收敛缓慢乃至发散。而RMSNorm将这些值归一化到肯定范围，有助于稳定和加速训练过程。这使得梯度具有更一致的幅度，从而加快模子收敛。
RMSNorm的工作原理

[图3]：对形状为[3,3]的输入嵌入应用RMSNorm
雷同于层归一化，RMSNorm沿嵌入特性或维度应用。上图中的嵌入形状为[3,3]，意味着每个标记有3个维度。
示例：对第一个标记X1的嵌入应用RMSNorm：
X1标记在每个维度上的值（x11、x12和x13）分别除以所有这些值的均方根。公式如图3所示。
为制止除以零并包管数值稳定性，在均方根中参加一个小常数E（Epsilon）。乘以一个缩放参数Gamma (Y)。每个特性都有一个独特的Gamma参数（如图中d1维度的Y1、d2维度的Y2和d3维度的Y3），这是一个学习参数，可以向上或向下缩放以进一步稳定归一化。gamma参数初始化为1（如上面的盘算所示）。
如示例所示，嵌入值本来较大且分布范围宽。应用RMSNorm后，值变小且范围缩小。盘算使用实际的RMSNorm函数完成。
RMSNorm相比层归一化的上风
如上例所示没有盘算任何均值或方差，而这在层归一化中是必须的。以是RMSNorm通过制止盘算均值和方差减少了盘算开销。根据作者的研究，RMSNorm在不影响准确性的同时提供了性能上风。
RMSNorm代码实现：

1.  # 步骤2: 解码器模块  
2.  # 注：由于Llama 3模型由Meta开发，为了与他们的代码库保持一致并考虑未来兼容性，
3.  # 我将使用Meta GitHub上的大部分代码，并进行必要的修改以实现我们的目标。
4.  
5.  # 定义参数数据类：我们将在模型构建、训练和推理过程中使用这些参数。
6.  # 注：为了更快地看到训练和推理结果，而不是专注于高准确性，我们对大多数参数采用较低的值，
7.  # 这些值在Llama 3模型中设置得更高。
8.  
9.  @dataclass  
10.  class ModelArgs:  
11.      dim: int = 512              # 嵌入维度  
12.      n_layers: int = 8           # 模型解码器块的数量  
13.      n_heads: int = 8            # 查询嵌入的头数  
14.      n_kv_heads: int = 4         # 键和值嵌入的头数  
15.      vocab_size: int = len(vocab) # 词汇表长度  
16.      multiple_of: int = 256        # 用于计算前馈网络维度  
17.      ffn_dim_multiplier: Optional[float] = None  # 用于计算前馈网络维度  
18.      norm_eps: float = 1e-5                       # RMSNorm计算的默认Epsilon值  
19.      rope_theta: float = 10000.0   # RePE计算的默认theta值  
20.  
21.      max_batch_size: int = 10     # 最大批量大小  
22.      max_seq_len: int = 256         # 最大序列长度  
23.  
24.      epochs: int = 2500             # 总训练迭代次数  
25.      log_interval: int = 10        # 打印日志和损失值的间隔数    
26.      device: str = 'cuda' if torch.cuda.is_available() else 'cpu'   # 根据可用性分配设备为cuda或cpu
27.  
28.  ## 步骤2a: RMSNorm  
29.  
30.  class RMSNorm(nn.Module):  
31.    def __init__(self, dim: int, eps: float = 1e-6):  
32.      super().__init__()  
33.      device = ModelArgs.device  
34.      self.eps = eps  
35.      # 缩放参数gamma，初始化为1，参数数量等于d

复制代码

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。

继续阅读请点击广告