从零开始构建2亿参数的llm：使用python和the Pile数据集的实用指南
快速说明 — 我们将使用The Pile数据集从零开始训练一个2亿参数LLM。因此，我们得到一个在响应中输出完美语法和标点的LLM，较短的上下文是有意义的，但整个响应并不如此。
我正在寻找AI领域的博士职位。您可以查看我的简历或GitHub个人资料。
之前，我在Medium上写了一篇关于使用Tiny Shakespeare数据集创建一个230万参数LLM的文章，但输出没有意义。
ZELBETH:
Sey solmenter! tis tonguerered if
Vurint as steolated have loven OID the queend refore
Are been, good plmp:
Proforne, wiftes swleen, was no blunderesd a a quain beath!
Tybell is my gateer stalk smend as be matious dazest
我有一个想法，如果我将transformer架构缩小并减少复杂性，同时使训练数据更加多样化，会怎样？那么，一个人使用他们几乎耗尽的GPU，能够创建多大参数的模型，可以讲正确的语法并生成一些有意义的文本？
这是我们训练模型的输出，遵循此博客：
我发现1300万参数足以在正确的语法和标点方面开始有意义，这是一个积极的点。这意味着我们可以使用一个非常特定的数据集来进一步微调我们之前训练的模型，以便于更窄的任务。我们可能最终会得到一个不到10亿参数的模型，甚至大约5亿参数，特别适合我们的特定用例，尤其是在安全地运行私有数据时。
我建议您首先使用我GitHub仓库中提供的脚本训练一个1300万参数的模型。您将在一天内获得结果，而不是等待更长时间，或者如果您的本地GPU可能不足以训练一个十亿参数的模型。
GitHub代码概述
所有代码都可以在我的GitHub个人资料中找到：
代码库的组织结构如下：
train-llm-from-scratch/
├── src/
│ ├── models/
│ │ ├── mlp.py
│ │ ├── attention.py
│ │ ├── transformer_block.py
│ │ ├── transformer.py
├── config/
│ └── config.py
├── data_loader/
│ └── data_loader.py
├── scripts/
│ ├── train_transformer.py
│ ├── data_download.py
│ ├── data_preprocess.py
│ ├── generate_text.py
├── data/
│ ├── train/
│ └── val/
├── models/
scripts/目录包括用于数据集下载、数据预处理、模型训练和使用训练好的模型进行文本生成等任务的脚本。src/models/目录包含关键组件的实现，包括transformer模型、多层感知器(MLP)、注意力机制和transformer块。config/目录包含指定项目默认参数的配置文件。data_loader/目录提供创建数据加载器和迭代器的函数。
目录
- 前提条件和训练时间
- 安装模块
- 导入库
- 准备训练数据
- Transformer 概述
- 多层感知器 (MLP)
- 单头注意力
- 多头注意力
- Transformer 块
- 最终模型
- 批处理
- 训练参数
- 训练模型
- 保存训练好的模型
- 训练损失
- 生成文本
- 接下来是什么
先决条件和训练时间
确保您对面向对象编程（OOP）和神经网络（NN）有基本的理解。熟悉PyTorch也将有助于编码。
您需要一个GPU来训练您的模型。Colab或Kaggle T4适用于训练一个超过1300万参数的模型，但对于十亿参数的训练将会失败。请查看比较：
| GPU类型 | 支持的参数 |
|---|---|
| Colab T4 | 最多1300万 |
| Kaggle T4 | 最多1300万 |
| 其他GPU | 各有不同 |
GPU类型比较
- Colab T4
- Kaggle T4
- 其他GPU
注意：确保检查您使用的GPU的规格以获得最佳性能。
安装模块
确保您的环境中已安装Git。您需要先克隆该仓库：
git clone https://github.com/FareedKhan-dev/train-llm-from-scratch.git
cd train-llm-from-scratch
然后您可以安装所需的依赖项：
pip install -r requirements.txt
导入库
让我们导入在本博客中将使用的所需库：
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import h5py
import os
import argparse
import requests
from tqdm import tqdm
import json
import zstandard as zstd
import tiktoken
import math
准备训练数据
我们的训练数据集需要多样化，包含来自不同领域的信息，而 The Pile 数据集是合适的选择。尽管它的大小为 825 GB，但我们只会使用其中的一小部分，即 5%–10%。让我们首先下载数据集，看看它是如何工作的。我将下载在 HuggingFace 上可用的版本。
!wget https://huggingface.co/datasets/monology/pile-uncopyrighted/resolve/main/val.jsonl.zst
!wget https://huggingface.co/datasets/monology/pile-uncopyrighted/resolve/main/train/00.jsonl.zst
!wget https://huggingface.co/datasets/monology/pile-uncopyrighted/resolve/main/train/01.jsonl.zst
!wget https://huggingface.co/datasets/monology/pile-uncopyrighted/resolve/main/train/02.jsonl.zst
下载需要一些时间，但您也可以将训练数据集限制为仅一个文件 00.jsonl.zst，而不是三个。它已经被拆分为训练/验证/测试。一旦完成，请确保将文件正确放置在各自的目录中。
import os
import shutil
import glob
train_dir = "data/train"
val_dir = "data/val"
os.makedirs(train_dir, exist_ok=True)
os.makedirs(val_dir, exist_ok=True)
train_files = glob.glob("*.jsonl.zst")
for file in train_files:
if file.startswith("val"):
dest = os.path.join(val_dir, file)
else:
dest = os.path.join(train_dir, file)
shutil.move(file, dest)
我们的数据集采用 .jsonl.zst 格式，这是一种常用于存储大型数据集的压缩文件格式。它结合了 JSON Lines (.jsonl)，其中每一行代表一个有效的 JSON 对象，以及 Zstandard (.zst) 压缩。让我们读取下载文件的一个示例，看看它的样子。
in_file = "data/val/val.jsonl.zst"
with zstd.open(in_file, 'r') as in_f:
for i, line in tqdm(enumerate(in_f)):
data = json.loads(line)
print(f"Line {i}: {data}")
if i == 2:
break
行：0
{
"text": "Effect of sleep quality ... epilepsy.",
"meta": {
"pile_set_name": "PubMed Abstracts"
}
}
行：1
{
"text": "LLMops a new GitHub Repository ...",
"meta": {
"pile_set_name": "Github"
}
}
现在我们需要对数据集进行编码（分词）。我们的目标是拥有一个至少能够输出正确单词的 LLM。为此，我们需要使用一个已经可用的分词器。我们将使用 OpenAI 提供的 tiktoken 开源分词器。我们将使用 r50k_base 分词器，该分词器用于 ChatGPT（GPT-3）模型，对我们的数据集进行分词。
我们需要为此创建一个函数，以避免重复，因为我们将对训练和验证数据集进行分词。
def process_files(input_dir, output_file):
"""
处理指定输入目录中的所有 .zst 文件，并将编码的标记保存到 HDF5 文件中。
参数：
input_dir (str): 包含输入 .zst 文件的目录。
output_file (str): 输出 HDF5 文件的路径。
"""
with h5py.File(output_file, 'w') as out_f:
dataset = out_f.create_dataset('tokens', (0,), maxshape=(None,), dtype='i')
start_index = 0
for filename in sorted(os.listdir(input_dir)):
if filename.endswith(".jsonl.zst"):
in_file = os.path.join(input_dir, filename)
print(f"Processing: {in_file}")
with zstd.open(in_file, 'r') as in_f:
for line in tqdm(in_f, desc=f"Processing {filename}"):
data = json.loads(line)
text = data['text'] + "<|endoftext|>"
encoded = enc.encode(text, allowed_special={'<|endoftext|>'})
encoded_len = len(encoded)
end_index = start_index + encoded_len
dataset.resize(dataset.shape[0] + encoded_len, axis=0)
dataset[start_index:end_index] = encoded
start_index = end_index
关于这个函数有两个重要的点：
- 我们将标记化的数据存储在 HDF5 文件中，这使我们在训练模型时能够更快地访问数据。
- 附加
<|endoftext|>标记标志着每个文本序列的结束，向模型发出它已经到达有意义上下文的结束的信号，这有助于生成连贯的输出。
现在我们可以简单地使用以下代码对训练和验证数据集进行编码：
out_train_file = "data/train/pile_train.h5"
out_val_file = "data/val/pile_dev.h5"
enc = tiktoken.get_encoding('r50k_base')
process_files(train_dir, out_train_file)
process_files(val_dir, out_val_file)
让我们看看我们标记化数据的示例：
with h5py.File(out_val_file, 'r') as file:
tokens_dataset = file['tokens']
print(f"Dtype of 'tokens' dataset: {tokens_dataset.dtype}")
print("First few elements of the 'tokens' dataset:")
print(tokens_dataset[:10])
‘tokens’ 数据集的数据类型：int32
‘tokens’ 数据集的前几个元素：
[ 2725 6557 83 23105 157 119 229 77 5846 2429]
我们已经为训练准备好了数据集。现在我们将编写 transformer 架构，并相应地研究其理论。
Transformer概述
让我们快速看一下transformer架构是如何用于处理和理解文本的。它通过将文本拆分成称为标记的小块，并预测序列中的下一个标记来工作。一个transformer有许多层，称为transformer块，彼此堆叠，在最后有一层用于进行预测。
每个transformer块有两个主要组成部分：
- 自注意力头：这些确定输入的哪些部分对模型最重要。例如，在处理一个句子时，注意力头可以突出单词之间的关系，例如代词与其所指名词之间的关系。
- 多层感知器（MLP）：这是一个简单的前馈神经网络。它接收注意力头强调的信息并进一步处理。MLP有一个输入层，从注意力头接收数据，一个隐藏层为处理增加复杂性，以及一个输出层将结果传递给下一个transformer块。
总的来说，注意力头充当“思考什么”的部分，而MLP是“如何思考”的部分。堆叠多个transformer块使模型能够理解文本中的复杂模式和关系，但这并不总是得到保证。
与其查看原始论文图表，不如让我们可视化一个更简单、更易于编码的架构图。
由 Fareed Khan 提供的Transformer架构
让我们阅读一下我们将要编码的架构流程：
- 输入标记被转换为嵌入并与位置信息结合。
- 模型有64个相同的transformer块，按顺序处理数据。
- 每个块首先运行多头注意力，以查看标记之间的关系。
- 每个块然后通过一个MLP处理数据，该MLP扩展然后压缩数据。
- 每个步骤使用残差连接（快捷方式）来帮助信息流动。
- 在整个过程中使用层归一化以稳定训练。
- 注意力机制计算哪些标记应该相互关注。
- MLP将数据扩展到4倍大小，应用ReLU，然后再压缩回去。
- 模型使用16个注意力头来捕捉不同类型的关系。
- 最后一层将处理后的数据转换为词汇量大小的预测。
- 模型通过反复预测下一个最可能的标记来生成文本。
多层感知器 (MLP)
MLP是transformer前馈网络中的一个基本构建块。它的作用是引入非线性并学习嵌入表示中的复杂关系。在定义MLP模块时，一个重要的参数是n_embed，它定义了输入嵌入的维度。
MLP通常由一个隐藏线性层组成，该层将输入维度扩大一个因子（通常是4，我们将使用这个因子），后面跟着一个非线性激活函数，通常是ReLU。这个结构使我们的网络能够学习更复杂的特征。最后，一个投影线性层将扩展的表示映射回原始嵌入维度。这一系列变换使得MLP能够优化由注意力机制学习到的表示。
MLP by Fareed Khan
class MLP(nn.Module):
"""
A simple Multi-Layer Perceptron with one hidden layer.
This module is used within the Transformer block for feed-forward processing.
It expands the input embedding size, applies a ReLU activation, and then projects it back
to the original embedding size.
"""
def __init__(self, n_embed):
super().__init__()
self.hidden = nn.Linear(n_embed, 4 * n_embed)
self.relu = nn.ReLU()
self.proj = nn.Linear(4 * n_embed, n_embed)
def forward(self, x):
"""
Forward pass through the MLP.
Args:
x (torch.Tensor): Input tensor of shape (B, T, C), where B is batch size,
T is sequence length, and C is embedding size.
Returns:
torch.Tensor: Output tensor of the same shape as the input.
"""
x = self.forward_embedding(x)
x = self.project_embedding(x)
return x
def forward_embedding(self, x):
"""
Applies the hidden linear layer followed by ReLU activation.
Args:
x (torch.Tensor): Input tensor.
Returns:
torch.Tensor: Output after the hidden layer and ReLU.
"""
x = self.relu(self.hidden(x))
return x
def project_embedding(self, x):
"""
Applies the projection linear layer.
Args:
x (torch.Tensor): Input tensor.
Returns:
torch.Tensor: Output after the projection layer.
"""
x = self.proj(x)
return x
我们刚刚编写了MLP部分，其中__init__方法初始化了一个隐藏线性层，该层扩大输入嵌入大小（n_embed）和一个将其缩小的投影层。ReLU激活在隐藏层之后应用。forward方法定义了数据在这些层中的流动，通过forward_embedding应用隐藏层和ReLU，通过project_embedding应用投影层。
单头注意力
注意力头是我们模型的核心部分。它的目的是关注输入序列的相关部分。在定义 Head 模块时，一些重要的参数是 head_size、n_embed 和 context_length。head_size 参数决定了键、查询和值投影的维度，影响注意力机制的表示能力。
输入嵌入维度 n_embed 定义了这些投影层的输入大小。context_length 用于创建因果掩码，确保模型仅关注前面的标记。
在 Head 内部，键、查询和值的线性层 (nn.Linear) 被初始化为不带偏置。一个大小为 context_length x context_length 的下三角矩阵 (tril) 被注册为缓冲区，以实现因果掩码，防止注意力机制关注未来的标记。
单头注意力由 Fareed Khan
class Head(nn.Module):
"""
一个单一的注意力头。
该模块计算注意力得分并将其应用于值。
它包括键、查询和值的投影，并使用因果掩码
防止关注未来的标记。
"""
def __init__(self, head_size, n_embed, context_length):
super().__init__()
self.key = nn.Linear(n_embed, head_size, bias=False)
self.query = nn.Linear(n_embed, head_size, bias=False)
self.value = nn.Linear(n_embed, head_size, bias=False)
self.register_buffer('tril', torch.tril(torch.ones(context_length, context_length)))
def forward(self, x):
"""
通过注意力头的前向传递。
参数:
x (torch.Tensor): 形状为 (B, T, C) 的输入张量。
返回:
torch.Tensor: 应用注意力后的输出张量。
"""
B, T, C = x.shape
k = self.key(x)
q = self.query(x)
scale_factor = 1 / math.sqrt(C)
attn_weights = q @ k.transpose(-2, -1) * scale_factor
attn_weights = attn_weights.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
attn_weights = F.softmax(attn_weights, dim=-1)
v = self.value(x)
out = attn_weights @ v
return out
我们的注意力头类 __init__ 方法初始化了键、查询和值投影的线性层，每个层将输入嵌入 (n_embed) 投影到 head_size。基于 context_length 的下三角矩阵用于因果掩码。forward 方法通过缩放查询和键的点积来计算注意力权重，应用因果掩码，使用 softmax 归一化权重，并计算值的加权和以产生注意力输出。
多头注意力
为了捕捉输入序列中的多样关系，我们将使用多头注意力的概念。MultiHeadAttention模块管理多个独立的注意力头并行操作。
这里的关键参数是n_head，它决定了并行注意力头的数量。输入嵌入维度（n_embed）和context_length也是实例化各个注意力头所必需的。每个头独立处理输入，将其投影到大小为n_embed // n_head的低维子空间。通过拥有多个头，模型可以同时关注输入的不同方面。
多头注意力由 Fareed Khan
class MultiHeadAttention(nn.Module):
"""
多头注意力模块。
该模块并行组合多个注意力头。每个头的输出被连接以形成最终输出。
"""
def __init__(self, n_head, n_embed, context_length):
super().__init__()
self.heads = nn.ModuleList([Head(n_embed // n_head, n_embed, context_length) for _ in range(n_head)])
def forward(self, x):
"""
多头注意力的前向传播。
参数:
x (torch.Tensor): 形状为 (B, T, C) 的输入张量。
返回:
torch.Tensor: 在连接所有头的输出之后的输出张量。
"""
x = torch.cat([h(x) for h in self.heads], dim=-1)
return x
现在我们已经定义了MultiHeadAttention类，它组合了多个注意力头，__init__方法初始化了一系列Head实例（总共n_head个），每个实例的head_size为n_embed // n_head。forward方法将每个注意力头应用于输入x，并沿最后一个维度连接它们的输出，合并每个头学习到的信息。
Transformer Block
为了创建一个亿参数模型，我们绝对需要一个深度架构。为此，我们需要编写一个 transformer 块并将其堆叠起来。一个块的关键参数是 n_head、n_embed 和 context_length。每个块包含一个多头注意力层和一个前馈网络 (MLP)，在每个层之前应用层归一化，并在每个层之后添加残差连接。
层归一化由嵌入维度 n_embed 参数化，有助于稳定训练。多头注意力机制，如前所述，采用 n_head、n_embed 和 context_length。MLP 还利用嵌入维度 n_embed。这些组件共同工作以处理输入并学习复杂模式。
Transformer Block by Fareed Khan
class Block(nn.Module):
"""
A single Transformer block.
This block consists of a multi-head attention layer followed by an MLP,
with layer normalization and residual connections.
"""
def __init__(self, n_head, n_embed, context_length):
super().__init__()
self.ln1 = nn.LayerNorm(n_embed)
self.attn = MultiHeadAttention(n_head, n_embed, context_length)
self.ln2 = nn.LayerNorm(n_embed)
self.mlp = MLP(n_embed)
def forward(self, x):
"""
Forward pass through the Transformer block.
Args:
x (torch.Tensor): Input tensor.
Returns:
torch.Tensor: Output tensor after the block.
"""
x = x + self.attn(self.ln1(x))
x = x + self.mlp(self.ln2(x))
return x
def forward_embedding(self, x):
"""
Forward pass focusing on the embedding and attention parts.
Args:
x (torch.Tensor): Input tensor.
Returns:
tuple: A tuple containing the output after MLP embedding and the residual.
"""
res = x + self.attn(self.ln1(x))
x = self.mlp.forward_embedding(self.ln2(res))
return x, res
我们的 Block 类表示一个单独的 transformer 块。__init__ 方法初始化层归一化层 (ln1、ln2)、一个 MultiHeadAttention 模块和一个 MLP 模块，所有这些都由 n_head、n_embed 和 context_length 参数化。
forward 方法实现了块的前向传递，应用层归一化和多头注意力，带有残差连接，随后是另一个层归一化和 MLP，再次带有残差连接。forward_embedding 方法提供了一个替代的前向传递，专注于注意力和初始 MLP 嵌入阶段。
最终模型
到目前为止，我们已经编码了transformer模型的小组件。接下来，我们将标记和位置嵌入与一系列transformer块集成，以执行序列到序列任务。为此，我们需要编码几个关键参数：n_head、n_embed、context_length、vocab_size和N_BLOCKS。
vocab_size决定了标记嵌入层的大小，将每个标记映射到大小为n_embed的稠密向量。context_length参数对于位置嵌入层很重要，它编码输入序列中每个标记的位置，维度也是n_embed。注意力头的数量（n_head）和块的数量（N_BLOCKS）决定了网络的深度和复杂性。
这些参数共同定义了transformer模型的架构和容量，因此让我们来编码它。
由 Fareed Khan 提供的Transformer类
class Transformer(nn.Module):
"""
主要的Transformer模型。
该类将标记和位置嵌入与一系列Transformer块和最终的线性层结合用于语言建模。
"""
def __init__(self, n_head, n_embed, context_length, vocab_size, N_BLOCKS):
super().__init__()
self.context_length = context_length
self.N_BLOCKS = N_BLOCKS
self.token_embed = nn.Embedding(vocab_size, n_embed)
self.position_embed = nn.Embedding(context_length, n_embed)
self.attn_blocks = nn.ModuleList([Block(n_head, n_embed, context_length) for _ in range(N_BLOCKS)])
self.layer_norm = nn.LayerNorm(n_embed)
self.lm_head = nn.Linear(n_embed, vocab_size)
self.register_buffer('pos_idxs', torch.arange(context_length))
def _pre_attn_pass(self, idx):
"""
结合标记和位置嵌入。
参数：
idx (torch.Tensor): 输入标记索引。
返回：
torch.Tensor: 标记和位置嵌入的总和。
"""
B, T = idx.shape
tok_embedding = self.token_embed(idx)
pos_embedding = self.position_embed(self.pos_idxs[:T])
return tok_embedding + pos_embedding
def forward(self, idx, targets=None):
"""
通过Transformer的前向传播。
参数：
idx (torch.Tensor): 输入标记索引。
targets (torch.Tensor, optional): 用于损失计算的目标标记索引。默认为None。
返回：
tuple: logits和损失（如果提供了targets）。
"""
x = self._pre_attn_pass(idx)
for block in self.attn_blocks:
x = block(x)
x = self.layer_norm(x)
logits = self.lm_head(x)
loss = None
if targets is not None:
B, T, C = logits.shape
flat_logits = logits.view(B * T, C)
targets = targets.view(B * T).long()
loss = F.cross_entropy(flat_logits, targets)
return logits, loss
def forward_embedding(self, idx):
"""
关注嵌入和注意力块的前向传播。
参数：
idx (torch.Tensor): 输入标记索引。
返回：
tuple: 注意力块后的输出和残差。
"""
x = self._pre_attn_pass(idx)
residual = x
for block in self.attn_blocks:
x, residual = block.forward_embedding(x)
return x, residual
def generate(self, idx, max_new_tokens):
"""
给定起始序列生成新标记。
参数：
idx (torch.Tensor): 初始标记索引序列。
max_new_tokens (int): 要生成的标记数量。
返回：
torch.Tensor: 扩展后的标记序列。
"""
for _ in range(max_new_tokens):
idx_cond = idx[:, -self.context_length:]
logits, _ = self(idx_cond)
logits = logits[:, -1, :]
probs = F.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)
idx = torch.cat((idx, idx_next), dim=1)
return idx
我们的Transformer类的__init__方法初始化了标记和位置嵌入层（token_embed、position_embed）、一系列Block模块（attn_blocks）、一个最终的层归一化层（layer_norm）和一个用于语言建模的线性层（lm_head）。
_pre_attn_pass方法结合了标记和位置嵌入。forward方法通过嵌入层和一系列transformer块处理输入序列，应用最终的层归一化，并生成logits。如果提供了targets，它还计算损失。forward_embedding方法提供了一个中间前向传播，直到注意力块的输出，而generate方法实现了标记生成。
批量处理
当我们在大数据上训练深度学习模型时，由于GPU的可用性，我们以批量的方式处理数据。因此，让我们创建一个 get_batch_iterator 函数，接受指向HDF5文件的 data_path、所需的 batch_size、每个序列的 context_length 以及要加载数据的 device。
batch_size 决定了在训练期间并行处理的序列数量，而 context_length 指定了每个输入序列的长度。data_path 指向训练数据的位置。
def get_batch_iterator(data_path, batch_size, context_length, device="gpu"):
"""
Creates an iterator for generating batches of data from an HDF5 file.
Args:
data_path (str): Path to the HDF5 file containing tokenized data.
batch_size (int): Number of sequences in each batch.
context_length (int): Length of each sequence.
device (str, optional): Device to load the data onto ('cpu' or 'cuda'). Defaults to "cpu".
Yields:
tuple: A tuple containing input sequences (xb) and target sequences (yb).
"""
with h5py.File(data_path, 'r') as hdf5_file:
dataset = hdf5_file['tokens']
dataset_size = dataset.shape[0]
n_examples = (dataset_size - 1) // context_length
example_idxs = np.arange(n_examples)
np.random.shuffle(example_idxs)
epochs = 0
counter = 0
while True:
if counter + batch_size > n_examples:
np.random.shuffle(example_idxs)
counter = 0
print(f"Finished epoch {epochs}")
epochs += 1
random_indices = example_idxs[counter:counter+batch_size] * context_length
random_samples = torch.tensor(np.array([dataset[idx:idx+context_length+1] for idx in random_indices]))
xb = random_samples[:, :context_length].to(device)
yb = random_samples[:, 1:context_length+1].to(device)
counter += batch_size
yield xb, yb
我们的 get_batch_iterator 函数处理训练数据的加载和批量处理。它接受 data_path、batch_size、context_length 和 device 作为输入。该函数打开HDF5文件，打乱数据，然后进入无限循环以生成批量。在每次迭代中，它选择数据的随机子集以形成输入序列 (xb) 和相应的目标序列 (yb) 的批量。
训练参数
现在我们已经编码了我们的模型，我们需要定义训练参数，例如头数、块数等，以及数据路径。
VOCAB_SIZE = 50304
CONTEXT_LENGTH = 512
N_EMBED = 2048
N_HEAD = 16
N_BLOCKS = 64
TRAIN_PATH = "data/train/pile_val.h5"
DEV_PATH = "data/val/pile_val.h5"
T_BATCH_SIZE = 32
T_CONTEXT_LENGTH = 16
T_TRAIN_STEPS = 200000
T_EVAL_STEPS = 1000
T_EVAL_ITERS = 250
T_LR_DECAY_STEP = 50000
T_LR = 5e-4
T_LR_DECAYED = 5e-5
T_OUT_PATH = "models/transformer_B.pt"
DEVICE = 'cuda'
default_config = {
'vocab_size': VOCAB_SIZE,
'context_length': CONTEXT_LENGTH,
'n_embed': N_EMBED,
'n_head': N_HEAD,
'n_blocks': N_BLOCKS,
'train_path': TRAIN_PATH,
'dev_path': DEV_PATH,
't_batch_size': T_BATCH_SIZE,
't_context_length': T_CONTEXT_LENGTH,
't_train_steps': T_TRAIN_STEPS,
't_eval_steps': T_EVAL_STEPS,
't_eval_iters': T_EVAL_ITERS,
't_lr_decay_step': T_LR_DECAY_STEP,
't_lr': T_LR,
't_lr_decayed': T_LR_DECAYED,
't_out_path': T_OUT_PATH,
'device': DEVICE,
}
对于大多数参数，我使用了最常见的值，并将它们存储在字典中以便于访问。在这里，这些参数是针对一个亿参数模型的。如果您想训练一个具有百万参数的模型，可以减少主要参数，包括 CONTEXT_LENGTH、N_EMBED、N_HEAD 和 N_BLOCKS。但是，您也可以在我的 GitHub 个人资料中运行百万参数模型脚本。
训练模型
让我们初始化我们的transformer模型并检查其总参数数量。
model = Transformer(
n_head=config['n_head'],
n_embed=config['n_embed'],
context_length=config['context_length'],
vocab_size=config['vocab_size'],
N_BLOCKS=config['n_blocks']
).to(config['device'])
total_params = sum(p.numel() for p in model.parameters())
print(f"模型中的总参数数量: {total_params:,}")
2,141,346,251
现在我们有了2亿参数的模型，我们需要定义我们的Adam优化器和损失跟踪函数，这将帮助我们在训练过程中跟踪模型的进展。
optimizer = torch.optim.AdamW(model.parameters(), lr=config['t_lr'])
losses = []
AVG_WINDOW = 64
@torch.no_grad()
def estimate_loss(steps):
"""
在训练和开发数据集上评估模型并计算平均损失。
参数:
steps (int): 要评估的步骤数。
返回:
dict: 包含'train'和'dev'拆分的平均损失的字典。
"""
out = {}
model.eval()
for split in ['train', 'dev']:
data_path = config['train_path'] if split == 'train' else config['dev_path']
batch_iterator_eval = get_batch_iterator(
data_path, config['t_batch_size'], config['t_context_length'], device=config['device']
)
losses_eval = torch.zeros(steps)
for k in range(steps):
try:
xb, yb = next(batch_iterator_eval)
_, loss = model(xb, yb)
losses_eval[k] = loss.item()
except StopIteration:
print(f"警告: {split}的迭代器提前结束。")
break
out[split] = losses_eval[:k + 1].mean()
model.train()
return out
我们现在将初始化我们的批处理函数和训练循环，这将开始我们的训练。
batch_iterator = get_batch_iterator(
config['train_path'],
config['t_batch_size'],
config['t_context_length'],
device=config['device']
)
pbar = tqdm(range(config['t_train_steps']))
for step in pbar:
try:
xb, yb = next(batch_iterator)
_, loss = model(xb, yb)
losses.append(loss.item())
pbar.set_description(f"训练损失: {np.mean(losses[-AVG_WINDOW:]):.4f}")
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
if step % config['t_eval_steps'] == 0:
train_loss, dev_loss = estimate_loss(config['t_eval_iters']).values()
print(f"步骤: {step}, 训练损失: {train_loss:.4f}, 开发损失: {dev_loss:.4f}")
if step == config['t_lr_decay_step']:
print('衰减学习率')
for g in optimizer.param_groups:
g['lr'] = config['t_lr_decayed']
except StopIteration:
print("训练数据迭代器提前结束。")
break
保存训练好的模型
由于我们的训练循环具备处理错误的能力，因此如果循环抛出任何错误，它将保存我们部分训练的模型以避免损失。一旦训练完成，我们可以保存训练好的模型，以便后续进行推断。
train_loss, dev_loss = estimate_loss(200).values()
modified_model_out_path = config['t_out_path']
save_tries = 0
while os.path.exists(modified_model_out_path):
save_tries += 1
model_out_name = os.path.splitext(config['t_out_path'])[0]
modified_model_out_path = model_out_name + f"_{save_tries}" + ".pt"
torch.save(
{
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'losses': losses,
'train_loss': train_loss,
'dev_loss': dev_loss,
'steps': len(losses),
},
modified_model_out_path
)
print(f"Saved model to {modified_model_out_path}")
print(f"Finished training. Train loss: {train_loss:.4f}, Dev loss: {dev_loss:.4f}")
亿参数模型的最终训练损失为0.2314，开发损失为0.643。
训练损失
当我绘制百万参数和亿参数模型的损失时，它们看起来非常不同。
训练损失比较
亿参数模型起始时损失更高，开始时波动很大。它最开始迅速下降，但随后又摇摆不定，之后变得平滑。这表明更大的模型在开始时更难找到正确的学习方式。它可能需要更多的数据和更仔细的设置。当学习率降低时（红线），损失更稳定地下降，表明这有助于其微调。
百万参数模型的损失从一开始就更容易下降。它的波动没有更大模型那么明显。当学习率降低时，它对曲线的影响不大。这可能是因为较小的模型更容易训练，并且更快地找到好的解决方案。这个巨大的差异显示了训练非常大模型的难度。它们需要不同的方法，并可能需要更多的时间来良好学习。
我们现在有了保存的模型。我们终于可以用它进行推理，看看它如何生成文本。
生成文本
让我们创建一个函数，从我们保存的模型生成文本，该函数接受保存的模型路径和编码器作为输入，并返回生成的文本。
def generate_text(model_path, input_text, max_length=512, device="gpu"):
"""
使用基于给定输入文本的预训练模型生成文本。
Args:
- model_path (str): 模型检查点的路径。
- device (torch.device): 加载模型的设备（例如，'cpu'或'cuda'）。
- input_text (str): 用于生成的输入文本。
- max_length (int, optional): 生成文本的最大长度。默认为512。
Returns:
- str: 生成的文本。
"""
checkpoint = torch.load(model_path)
model = Transformer().to(device)
model.load_state_dict(checkpoint['model_state_dict'])
enc = tiktoken.get_encoding('r50k_base')
input_ids = torch.tensor(
enc.encode(input_text, allowed_special={'<|endoftext|>'}),
dtype=torch.long
)[None, :].to(device)
with torch.no_grad():
generated_output = model.generate(input_ids, max_length)
generated_text = enc.decode(generated_output[0].tolist())
return generated_text
我们之前定义的transformer需要在这里被调用以加载架构，然后我们将保存的模型作为该架构的状态加载。
让我们首先观察百万参数和十亿参数模型在没有提供任何输入的情况下生成的内容，并看看它们随机生成的内容。
Billion_model_path = 'models/transformer_B.pt'
Million_model_path = 'models/transformer_M.pt'
input_text = "<|endoftext|>"
B_output = generate_text(Billion_model_path, input_text)
M_output = generate_text(Million_model_path, input_text)
print(B_output)
print(M_output)
当上下文简短且简单时，两个LLM能够生成清晰准确的词语。例如，在百万参数的输出中，短语**“这些村庄与中国的城市直接相连”**是有意义的，并传达了一个清晰的想法。它容易理解，并且逻辑上将村庄与城市联系在一起。
然而，当上下文变得更长且更复杂时，清晰度开始减弱。在十亿参数的输出中，句子如**“从1037到73百万难民（假设）有两英里东海岸”和“铁匠、音乐家和精品酒店激励了交付给加拿大人的压力”**变得更难以理解。这些想法似乎不连贯，句子结构不自然流畅。虽然使用的词可能仍然是正确的，但整体意义变得混乱和不清晰。
积极的一点是，超过1300万参数的LLM也开始生成一些有意义的内容，且拼写正确。例如，当我使用主题输入文本时，它开始为我生成一封电子邮件。尽管显然更广泛的文本没有提供有意义的结果，但看看输出：
input_text = "Subject: "
m_output = generate_text(Million_model_path, input_text)
print(m_output)
我们的百万参数模型给了我们动力，让我们知道可以在1B以下拥有一个非常狭窄、目标导向的LLM，而我们的1B训练模型则向我们展示了架构需要深入编码并加以适当考虑。否则，与百万参数模型相比，它不会改善训练或性能。除非你为十亿规模的模型拥有一个深度架构，否则它只会过拟合数据。
接下来做什么
我建议您创建一个超过1300万参数的模型，然后通过增加下一个100个参数来开始扩展它，提高其处理较短上下文的能力。您可以根据特定任务决定要训练多少更多的参数。然后，对于剩余的不足1亿参数的部分，尝试在特定领域的数据上微调模型，例如撰写电子邮件或论文，并观察它如何生成文本。
