算法进阶 | PyTorch从零构建Llama 3（建议收藏！）
本文泉源公众号“算法进阶”，仅用于学术分享，侵权删，干货满满。原文链接：PyTorch从零构建Llama 3
大家好，本文将详细指导怎样从零开始构建完备的Llama 3模子架构，并在自界说数据集上执行训练和推理。
https://i-blog.csdnimg.cn/direct/999c6d4d741e4779a1d2e32718257d24.png
[图1]：Llama 3架构展示训练和推理流程。由于官方Llama 3论文中未提供相关图表。以是此图为大概架构图，阅读本文后你应能绘制出更为准确的架构图。
本文目标
通过本文。你可以了解到:
[*] 深入理解Llama 3模子各组件的底层工作原理。
[*] 编写代码构建Llama 3的每个组件，并将它们组装成一个功能完备的Llama 3模子。
[*] 编写代码使用新的自界说数据集训练模子。
[*] 编写代码执行推理，使Llama 3模子可以或许根据输入提示生成新文本。
1、输入模块
如图1所示，输入模块包含三个组件：文本/提示、分词器和嵌入。
输入模块内部工作流程
让我们通过下图了解输入模块内的工作流程。
https://i-blog.csdnimg.cn/direct/644e6377b33949cfaa6d595fc56795d8.png
[图2]：输入模块流程图，展示提示、分词器和嵌入流程。
首先，单个或批量文本/提示被输入模子。例如：图中的"Hello World"。
输入模子的必须是数字格式，由于模子无法直接处置惩罚文本。分词器将这些文本/提示转换为标记ID（词汇表中标记的索引号表现）。我们将使用Tiny Shakespeare数据集构建词汇表并训练模子。Llama 3模子使用TikToken作为分词器，这是一种子词分词器。但是我们这个实现将使用字符级分词器。这样做的重要原因是让我们可以或许自行构建词汇表和分词器，包括编码和解码函数，这样可以深入理解底层工作原理并完全掌控代码。
每个标记ID将被转换为128维的嵌入向量（原始Llama 3 8B中为4096维）。然后这些嵌入将被转达到下一个解码器模块。
输入模块代码实现：
# 导入必要的库
import torch
from torch import nn
from torch.nn import functional as F
import math
import numpy as np
import time
from dataclasses import dataclass
from typing import Optional, Tuple, List
import pandas as pd
from matplotlib import pyplot as plt
### 步骤1: 输入模块 ###
# 使用Tiny Shakespeare数据集实现字符级分词器。部分字符级分词器代码参考自Andrej Karpathy的GitHub仓库
# (https://github.com/karpathy/nanoGPT/blob/master/data/shakespeare_char/prepare.py)
# 加载tiny_shakespeare数据文件 (https://github.com/tamangmilan/llama3/blob/main/tiny_shakespeare.txt)
device: str = 'cuda' if torch.cuda.is_available() else 'cpu' # 根据可用性分配设备为cuda或cpu
# 加载tiny_shakespeare数据文件
with open('tiny_shakespeare.txt', 'r') as f:
data = f.read()
# 通过提取tiny_shakespeare数据中的所有唯一字符准备词汇表
vocab = sorted(list(set(data)))
# 训练Llama 3模型需要额外的标记，如<|begin_of_text|>、<|end_of_text|>和<|pad_id|>，将它们添加到词汇表中
vocab.extend(['<|begin_of_text|>','<|end_of_text|>','<|pad_id|>'])
vocab_size = len(vocab)
# 创建字符与词汇表中对应整数索引之间的映射。
# 这对于构建分词器的编码和解码函数至关重要。
itos = {i:ch for i, ch in enumerate(vocab)}
stoi = {ch:i for i, ch in enumerate(vocab)}
# 分词器编码函数：输入字符串，输出整数列表
def encode(s):
return for ch in s]
# 分词器解码函数：输入整数列表，输出字符串
def decode(l):
return ''.join(itos for i in l)
# 定义稍后在模型训练中使用的张量标记变量
token_bos = torch.tensor(], dtype=torch.int, device=device)
token_eos = torch.tensor(], dtype=torch.int, device=device)
token_pad = torch.tensor(], dtype=torch.int, device=device)
prompts = "Hello World"
encoded_tokens = encode(prompts)
decoded_text = decode(encoded_tokens)
### 输入模块代码测试 ###
# 取消下面的三重引号来执行测试
"""
print(f"Shakespeare文本字符长度: {len(data)}")
print(f"词汇表内容: {''.join(vocab)}\n")
print(f"词汇表大小: {vocab_size}")
print(f"编码后的标记: {encoded_tokens}")
print(f"解码后的文本: {decoded_text}")
"""
### 测试结果: ###
"""
Shakespeare文本字符长度: 1115394
词汇表内容:
!$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz<|begin_of_text|><|end_of_text|><|pad_id|>
词汇表大小: 68
编码后的标记:
解码后的文本: Hello World
""" 2、解码器模块
参照图1的架构图，解码器模块包含以下子组件：
[*] RMS归一化
[*] 旋转位置编码
[*] KV缓存
[*] 分组查询留意力
[*] 前馈网络
[*] 解码器块
RMS归一化（Root Mean Square Normalization）
RMSNorm的必要性
从图1可以看出，输入模块的输出（嵌入向量）颠末RMSNorm模块。这是由于嵌入向量具有多个维度（Llama3-8b中为4096维），大概出现不同范围的值。这会导致模子梯度爆炸或消失，从而导致收敛缓慢乃至发散。而RMSNorm将这些值归一化到肯定范围，有助于稳定和加速训练过程。这使得梯度具有更一致的幅度，从而加快模子收敛。
RMSNorm的工作原理
https://i-blog.csdnimg.cn/direct/9634bf9f3033475ca58f632b9b34d874.png
[图3]：对形状为的输入嵌入应用RMSNorm
雷同于层归一化，RMSNorm沿嵌入特性或维度应用。上图中的嵌入形状为，意味着每个标记有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代码实现：
# 步骤2: 解码器模块
# 注：由于Llama 3模型由Meta开发，为了与他们的代码库保持一致并考虑未来兼容性，
# 我将使用Meta GitHub上的大部分代码，并进行必要的修改以实现我们的目标。
# 定义参数数据类：我们将在模型构建、训练和推理过程中使用这些参数。
# 注：为了更快地看到训练和推理结果，而不是专注于高准确性，我们对大多数参数采用较低的值，
# 这些值在Llama 3模型中设置得更高。
@dataclass
class ModelArgs:
dim: int = 512 # 嵌入维度
n_layers: int = 8 # 模型解码器块的数量
n_heads: int = 8 # 查询嵌入的头数
n_kv_heads: int = 4 # 键和值嵌入的头数
vocab_size: int = len(vocab) # 词汇表长度
multiple_of: int = 256 # 用于计算前馈网络维度
ffn_dim_multiplier: Optional = None # 用于计算前馈网络维度
norm_eps: float = 1e-5 # RMSNorm计算的默认Epsilon值
rope_theta: float = 10000.0 # RePE计算的默认theta值
max_batch_size: int = 10 # 最大批量大小
max_seq_len: int = 256 # 最大序列长度
epochs: int = 2500 # 总训练迭代次数
log_interval: int = 10 # 打印日志和损失值的间隔数
device: str = 'cuda' if torch.cuda.is_available() else 'cpu' # 根据可用性分配设备为cuda或cpu
## 步骤2a: RMSNorm
class RMSNorm(nn.Module):
def __init__(self, dim: int, eps: float = 1e-6):
super().__init__()
device = ModelArgs.device
self.eps = eps
# 缩放参数gamma，初始化为1，参数数量等于d
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。
页:
[1]