26 KiB
26 KiB
5.3
在前面的章节中，我们熟悉了各种大模型的模型结构，以及如如何训练Tokenizer。在本节中，我们将动手训练一个小型的LLM。
5.3.1 训练Tokenizer
首先，我们需要为文本处理训练一个Tokenizer。Tokenizer的作用是将文本转换为数字序列，以便模型能够理解和处理。我们使用的数据集是 TinyStory ，它是一个由GPT-3.5和GPT-4生成的小型故事数据集，包含简短的故事，且词汇量有限。在这个任务中，我们采用字符级Tokenizer，将文本中的每个字符映射为对应的数字。通过以下命令可以下载数据集并训练Tokenizer。
python train_vocab.py --download True --vocab_size 4096
LLaMA2 的词表大小为 32,000，但由于 TinyStory 数据集较小，词汇量有限，我们将词表大小设置为 4,096。训练完成后，我们得到的 Tokenizer 能够将文本转换为数字序列，也可以将数字序列还原为文本。
def download_file(url: str, fname: str, chunk_size=1024):
"""发送HTTP GET请求以流式方式获取文件"""
···
def download():
"""执行 download_file 下载数据集"""
···
def train_vocab(vocab_size: int=32000, num_shards: int=20):
"""
vocab_size: int, 词汇表的大小，决定分词器的词汇量。
num_shards: int, 用于加快词汇表训练的效率，指定要处理的分片数量。
"""
# 确保词汇表大小为正数
assert vocab_size > 0, "Vocab size must be positive"
# SentencePiece 模型的前缀路径，将用于保存分词器
prefix = os.path.join(DATA_CACHE_DIR, f"tok{vocab_size}")
# 1) 将多个分片中的文本导出为单个文本文件 tiny.txt
tiny_file = os.path.join(DATA_CACHE_DIR, "tiny.txt")
data_dir = os.path.join(DATA_CACHE_DIR, "TinyStories_all_data")
shard_filenames = sorted(glob.glob(os.path.join(data_dir, "*.json")))
# 创建 tiny.txt 文件并写入指定数量的分片中的文本
print(f"Writing temporary file {tiny_file} with {num_shards} shards...")
with open(tiny_file, "w", encoding="utf-8") as of:
# 遍历前 num_shards 个分片
for shard in tqdm(shard_filenames[:num_shards]):
with open(shard, "r") as f:
data = json.load(f) # 读取分片中的JSON数据
# 遍历每个例子，将其中的故事文本写入 tiny.txt 文件
for example in data:
text = example["story"]
text = text.strip() # 去除文本首尾的空白字符
of.write(text + "\n") # 每个文本写入一行
# 输出生成的 tiny.txt 文件的大小
print(f"Size is: {os.path.getsize(tiny_file) / 1024 / 1024:.2f} MB")
# 2) 使用 SentencePiece 训练分词器
print("Will now train the vocab...")
spm.SentencePieceTrainer.train(
input=tiny_file, # 输入文件为之前生成的 tiny.txt
model_prefix=prefix, # 模型前缀路径
model_type="bpe", # 使用 Byte-Pair Encoding (BPE) 训练分词器
vocab_size=vocab_size, # 词汇表大小
self_test_sample_size=0, # 自测样本大小设置为 0
input_format="text", # 输入文件格式为纯文本
character_coverage=1.0, # 覆盖所有字符（包括非常见字符）
num_threads=os.cpu_count(), # 使用 CPU 的线程数
split_digits=True, # 拆分数字
allow_whitespace_only_pieces=True, # 允许仅由空格组成的词元
byte_fallback=True, # 启用字节级回退
unk_surface=r" \342\201\207 ", # UNK token 表示未知字符的方式
normalization_rule_name="identity" # 使用“identity”归一化规则
)
# 3) 可选的清理操作，询问用户是否删除临时文件 tiny.txt
dec = input(f"Delete the temporary file {tiny_file}? [y/N] ")
if dec.lower() == "y":
os.remove(tiny_file) # 删除临时文件
print(f"Deleted {tiny_file}")
# 输出模型保存的路径
print(f"Trained tokenizer is in {prefix}.model")
print("Done.")
在本部分中，我们使用了 SentencePiece 库来训练自定义的 Tokenizer。首先，我们需要从 TinyStory 数据集中提取文本内容，作为训练的输入数据。SentencePiece 是一种基于子词单元的分词算法，能够有效处理不同语言中的词汇碎片化问题。
训练 Tokenizer 时，SentencePiece 会自动生成两个文件：tok4096.model 和 tok4096.vocab，其中 tok4096.model 是我们训练好的模型文件，位于 data 目录下。这个文件可以用于将文本数据转换为 Token 序列，也可以将 Token 序列还原为文本。
为了更便捷地使用这个 Tokenizer，我们还在 tokenizer.py 文件中定义了一个 Tokenizer 类。这个类封装了 Tokenizer 的常用操作，例如文本编码和解码功能，并支持加载我们训练好的模型文件。通过这个类，我们可以轻松地将文本转换为模型可接受的数字序列，或将预测结果转化为可读的文本。
具体的代码实现和细节可以在 tokenizer.py 文件中找到，接下来我们将进一步展示如何使用该类来处理 TinyStory 数据集中的故事文本。
class Tokenizer:
def __init__(self, tokenizer_model=None):
"""
初始化分词器。加载预训练的SentencePiece模型，并设置一些特殊的token ID。
参数:
tokenizer_model: str, 可选，分词器模型的路径，如果不指定则使用默认路径 TOKENIZER_MODEL。
"""
# 如果提供了分词器模型路径，使用该路径；否则使用默认模型路径
model_path = tokenizer_model if tokenizer_model else TOKENIZER_MODEL
# 确保模型文件存在
assert os.path.isfile(model_path), model_path
# 加载 SentencePiece 模型
self.sp_model = SentencePieceProcessor(model_file=model_path)
self.model_path = model_path
# 获取分词器的特殊token和词汇表大小
self.n_words: int = self.sp_model.vocab_size() # 词汇表大小
self.bos_id: int = self.sp_model.bos_id() # 句子开头 (BOS) 的ID
self.eos_id: int = self.sp_model.eos_id() # 句子结尾 (EOS) 的ID
self.pad_id: int = self.sp_model.pad_id() # 填充 (PAD) 的ID
# 验证分词器词汇表大小是否正确
assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()
def encode(self, s: str, bos: bool, eos: bool) -> List[int]:
"""
将字符串编码为词元ID列表。可以选择是否添加句子开头 (BOS) 和句子结尾 (EOS) 标记。
参数:
s: str, 要编码的字符串。
bos: bool, 是否在编码的词元列表前添加 BOS 标记。
eos: bool, 是否在编码的词元列表末尾添加 EOS 标记。
返回:
List[int]: 编码后的词元ID列表。
"""
# 确保输入是字符串类型
assert type(s) is str
# 使用SentencePiece将字符串编码为词元ID
t = self.sp_model.encode(s)
# 如果需要BOS标记，将其添加到词元列表开头
if bos:
t = [self.bos_id] + t
# 如果需要EOS标记，将其添加到词元列表末尾
if eos:
t = t + [self.eos_id]
return t
def decode(self, t: List[int]) -> str:
"""
将词元ID列表解码为字符串。
参数:
t: List[int], 词元ID列表。
返回:
str: 解码后的字符串。s
"""
return self.sp_model.decode(t)
在这个 Tokenizer 类中，我们首先初始化了一些特殊的 token ID，这些特殊 tokens 在自然语言处理任务中有着重要作用，分别用于填充、处理未识别的词汇、表示句子的开头和结尾等。在模型训练和推理过程中，正确处理这些特殊 tokens 对于提升模型性能至关重要。
接着，我们定义了两个关键方法：
-
encode 方法：该方法负责将输入文本转换为 token ID 序列。通过加载预训练的 Tokenizer 模型，我们可以对文本进行分词，将其拆解为词或子词，并将其映射为相应的数字表示。这个数字序列可以被模型接受用于训练和推理。
-
decode 方法：与 encode 方法相反，decode 方法用于将 token ID 序列还原为可读的文本。它将数字序列转换回对应的 tokens，并拼接成完整的文本，从而可以对模型的输出进行解释和展示。
这些方法的定义使得我们在使用过程中，可以非常方便地在文本与数字序列之间进行转换，为模型的输入与输出提供接口。大家可以使用以下代码测试 Tokenizer 的功能，验证其是否能够正确地将文本转换为数字序列，或者将数字序列还原为文本。
# 测试 Tokenizer
enc = Tokenizer('./data/tok4096.model') # 加载分词器
tetx = 'Hello, world!' # 测试文本
print(enc.encode(text, bos=True, eos=True)) # 编码文本
print(enc.decode(enc.encode(text, bos=True, eos=True))) # 解码文本
OUTPUT:
[1, 346, 2233, 4010, 1475, 4021, 2]
Hello, world!
5.3.2 数据预处理
在训练模型之前，首先需要对数据进行预处理。这一步的核心任务是将文本数据转换为模型能够理解的数字序列。具体来说，文本中的每个字符、单词或子词都需要被映射为一个唯一的数字 ID，这样模型才能处理这些数据。
# 定义分片处理函数
def process_shard(args, vocab_size, tokenizer_model_path):
""" 处理数据分片，将其中的文本进行分词并保存为二进制文件 """
···
# 定义预处理函数，用于对多个数据分片进行批量处理
def pretokenize(vocab_size):
""" 预处理所有的数据分片，并将分词后的数据保存为二进制文件 """
···
class PretokDataset(torch.utils.data.IterableDataset):
"""从磁盘加载已预处理的分词数据，并将其以 PyTorch 张量的形式返回。"""
def __init__(self, split, max_seq_len, vocab_size, vocab_source):
"""
初始化数据集。
参数:
split: str, 数据集的分割方式（'train' 或 'test'）。
max_seq_len: int, 最大序列长度，用于生成输入输出序列。
vocab_size: int, 词汇表的大小。
vocab_source: str, 词汇表的来源（'llama2' 或 'custom'）。
"""
super().__init__()
self.split = split # 数据集划分（训练集或测试集）
self.max_seq_len = max_seq_len # 最大序列长度
self.vocab_size = vocab_size # 词汇表大小
self.vocab_source = vocab_source # 词汇表来源
def __iter__(self):
"""
返回迭代器，按批次加载数据并生成模型输入/输出。
"""
# 获取DataLoader的worker信息（用于并行数据加载）
worker_info = torch.utils.data.get_worker_info()
worker_id = worker_info.id if worker_info else 0 # worker ID
# 获取分布式训练的rank信息（用于多GPU训练）
rank = dist.get_rank() if dist.is_initialized() else 0
# 基于worker_id和rank生成唯一的随机数种子，确保数据在每个worker和rank之间是唯一的
seed = 42 + worker_id + 1337 * rank
rng = random.Random(seed)
print(f"Created a PretokDataset with rng seed {seed}")
# 根据词汇表来源决定数据路径
if self.vocab_source == "llama2":
# 如果使用 Llama 2 词汇表，.bin 文件和 .json 文件在同一目录下
bin_dir = os.path.join(DATA_CACHE_DIR, "TinyStories_all_data")
shard_filenames = sorted(glob.glob(os.path.join(bin_dir, "*.bin")))
elif self.vocab_source == "custom":
# 如果使用自定义词汇表，.bin 文件在 tok{N} 目录下
bin_dir = os.path.join(DATA_CACHE_DIR, f"tok{self.vocab_size}")
shard_filenames = sorted(glob.glob(os.path.join(bin_dir, "*.bin")))
# 根据数据集划分使用不同的分片文件
# 训练集使用所有分片文件，测试集只使用第一个分片
shard_filenames = shard_filenames[1:] if self.split == "train" else shard_filenames[:1]
assert len(shard_filenames) > 0, f"在 {bin_dir} 中未找到任何 .bin 文件"
while True:
# 随机打乱分片文件
rng.shuffle(shard_filenames)
for shard in shard_filenames:
# 使用 memmap 读取文件，使得数据留在磁盘上，减少内存占用
m = np.memmap(shard, dtype=np.uint16, mode="r")
# 计算该分片中的批次数量
num_batches = len(m) // self.max_seq_len
num_batches -= 1 # 去掉最后一个不完整的批次
assert num_batches > 0, "这个分片文件太小了？请检查。"
# 随机打乱批次索引
ixs = list(range(num_batches))
rng.shuffle(ixs)
# 对每个批次生成输入 x 和目标输出 y
for ix in ixs:
start = ix * self.max_seq_len # 批次起始索引
end = start + self.max_seq_len + 1 # 批次结束索引
# 将数据转换为 NumPy 数组并拷贝到 RAM 中
chunk = torch.from_numpy((m[start:end]).astype(np.int64))
# 模型输入 x 是当前批次的前 max_seq_len 个词元
x = chunk[:-1]
# 模型输出 y 是下一个词元
y = chunk[1:]
# 生成 x, y 对
yield x, y
class Task:
@staticmethod
def iter_batches(batch_size, device, num_workers=0, **dataset_kwargs):
ds = PretokDataset(**dataset_kwargs)
dl = torch.utils.data.DataLoader(
ds, batch_size=batch_size, pin_memory=True, num_workers=num_workers
)
for x, y in dl:
x = x.to(device, non_blocking=True)
y = y.to(device, non_blocking=True)
yield x, y
在这部分中，首先定义了 process_shard 函数，用于处理数据分片。该函数的主要功能是将文本数据分词后，转换为更高效的二进制文件格式，以便后续更快速地加载和处理数据。
接下来，我们定义了 pretokenize 函数，用于批量处理多个数据分片。通过这一函数，所有数据可以并行处理，进一步加快预处理的速度。
然后，我们设计了一个 PretokDataset 类，用于加载已预处理好的数据集。我们继承了 torch.utils.data.IterableDataset 来定义该数据集，这使得我们可以更灵活、高效地处理数据。在这个类中，核心是 __iter__ 方法，它负责生成用于训练的数据批次。
最后，我们还定义了一个 Task 类，专门用于迭代数据集，并生成模型所需的输入和目标输出。这一部分的设计确保了数据流的顺畅对接，为模型训练提供了标准化的数据输入。可以通过以下代码来测试预处理后的数据集。
5.3.3 训练模型
在数据预处理完成后，我们就可以开始训练模型了。我们使用的模型是一个和LLama2结构一样的 Decoder only Transformer模型，使用Pytorch实现。相关代码在model.py文件中。此处不再赘述，源码中有详细的中文注释，且我们在之前的文章中也有详细的介绍。
在模型这一部分可以重点看一下生成式模型是如何实现生成token的，可以查看model.py文件中的Transforerm类中的generate方法。
在完成数据预处理后，我们就可以开始训练模型了。我们使用的模型是一个与 LLaMA2 结构相同的 Decoder-only Transformer 模型，采用 PyTorch 实现。具体的实现细节已经包含在 model.py 文件中，在此不再赘述。该源码中包含详细的中文注释，此外我们在之前的文章中也对模型架构进行了深入介绍。
在模型部分，建议重点关注生成式模型如何生成 token 的过程。可以参考 model.py 文件中的 Transformer 类，尤其是 generate 方法的实现，它展示了模型如何基于已有的上下文生成后续 token 的机制。
@torch.inference_mode()
def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
"""
给定输入序列 idx（形状为 (bz,seq_len) 的长整型张量），通过多次生成新 token 来完成序列。
在 model.eval() 模式下运行。效率较低的采样版本，没有使用键k/v cache。
"""
for _ in range(max_new_tokens):
# 如果序列上下文过长，截断它到最大长度
idx_cond = idx if idx.size(1) <= self.args.max_seq_len else idx[:, -self.args.max_seq_len:]
# 前向传播获取序列中最后一个位置的 logits
logits = self(idx_cond)
logits = logits[:, -1, :] # 只保留最后一个时间步的输出
if temperature == 0.0:
# 选择最有可能的索引
_, idx_next = torch.topk(logits, k=1, dim=-1)
else:
# 缩放 logits 并应用 softmax
logits = logits / temperature
if top_k is not None:
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
logits[logits < v[:, [-1]]] = -float('Inf')
probs = F.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)
# 将采样的索引添加到序列中并继续
idx = torch.cat((idx, idx_next), dim=1)
return idx
在 generate 方法中，我们首先获取序列中最后一个位置的 logits，然后基于这些 logits 生成新的 token。接着，生成的新 token 会被添加到序列中，模型随后会继续生成下一个 token。通过这种迭代过程，我们能够生成完整的文本。接下来，您可以使用以下命令开始训练模型。
python train.py
在 train.py 中我们定义了很多超参数，包括但不限于模型的维度，层数，学习率等等。如下所示，更多的内容大家可以在源码中查看，源码加了很详细的中文注释，相信大家可以很容易看懂。
# -----------------------------------------------------------------------------
# I/O 配置，用于定义输出目录和训练时的日志记录与评估设置
out_dir = "output" # 模型输出保存路径
eval_interval = 2000 # 评估间隔步数
log_interval = 1 # 日志记录间隔步数
eval_iters = 100 # 每次评估时迭代的步数
eval_only = False # 如果为True，脚本在第一次评估后立即退出
always_save_checkpoint = False # 如果为True，在每次评估后总是保存检查点
init_from = "scratch" # 可以选择从头开始训练（'scratch'）或从已有的检查点恢复（'resume'）
# 数据配置
batch_size = 8 # 每个微批次的样本数量，如果使用梯度累积，实际批次大小将更大
max_seq_len = 256 # 最大序列长度
vocab_size = 4096 # 自定义词汇表大小
# 模型配置
dim = 288 # 模型的隐藏层维度
n_layers = 8 # Transformer的层数
n_heads = 8 # 注意力头的数量
n_kv_heads = 4 # 模型分组
multiple_of = 32 # 在某些层的维度必须是该数的倍数
dropout = 0.0 # Dropout概率
# AdamW优化器配置
gradient_accumulation_steps = 4 # 梯度累积步数，用于模拟更大的批次
learning_rate = 5e-4 # 最大学习率
max_iters = 100000 # 总的训练迭代次数
weight_decay = 1e-1 # 权重衰减系数
beta1 = 0.9 # AdamW优化器的β1参数
beta2 = 0.95 # AdamW优化器的β2参数
grad_clip = 1.0 # 梯度裁剪阈值，0表示不裁剪
# 学习率衰减配置
decay_lr = True # 是否启用学习率衰减
warmup_iters = 1000 # 学习率预热的步数
# 系统设置
device = "cuda:0" # 设备选择：'cpu'，'cuda'，'cuda:0'等
dtype = "bfloat16" # 数据类型：'float32'，'bfloat16'，'float16'
5.3.4 使用模型生成文本
在模型训练完成后，会在output目录下生成一个ckpt.pt文件，这个文件就是我们训练好的模型。我们可以使用以下命令生成文本。
python sample.py --prompt "One day, Lily met a Shoggoth"
我们来看下sample.py文件中的代码，这个文件中定义了一个TextGenerator类，用于生成文本。
class TextGenerator:
def __init__(self,
checkpoint='output/ckpt.pt', # 模型检查点路径
tokenizer_model_path='tok4096.model', # 分词器模型路径
seed=1337, # 随机种子，确保可重复性
device=None, # 设备，优先使用 CUDA，如果没有可用的 CUDA，则使用 CPU
dtype="float32"): # 数据类型，默认为 float32，可以选择 float16 或 bfloat16
"""
初始化 TextGenerator 类，加载模型、设置设备和分词器等。
"""
# 模型加载配置
self.checkpoint = checkpoint # 保存的模型检查点路径
self.tokenizer_model_path = tokenizer_model_path # 分词器模型文件路径
self.seed = seed # 随机数种子，用于生成的可重复性
self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu') # 根据硬件条件选择设备
self.dtype = dtype # 模型的浮点数类型
self.device_type = 'cuda' if 'cuda' in self.device else 'cpu' # 判断当前设备是否为 CUDA
# 设置随机种子，确保生成的可重复性
torch.manual_seed(seed) # 设置 CPU 随机种子
torch.cuda.manual_seed(seed) # 设置 CUDA 随机种子
torch.backends.cuda.matmul.allow_tf32 = True # 允许 CUDA 使用 TF32 精度进行矩阵乘法运算
torch.backends.cudnn.allow_tf32 = True # 允许 cuDNN 使用 TF32 精度加速
# 根据 dtype 选择适当的自动混合精度上下文
ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[self.dtype]
self.ctx = nullcontext() if self.device_type == 'cpu' else torch.amp.autocast(device_type=self.device_type, dtype=ptdtype)
# 加载模型检查点文件
checkpoint_dict = torch.load(self.checkpoint, map_location=self.device) # 加载模型参数
gptconf = ModelArgs(**checkpoint_dict['model_args']) # 初始化模型参数
self.model = Transformer(gptconf) # 实例化 Transformer 模型
state_dict = checkpoint_dict['model'] # 获取模型状态字典
# 去除状态字典中的不必要前缀
unwanted_prefix = '_orig_mod.' # 这个前缀在保存时可能被添加，现在要去除它
for k, v in list(state_dict.items()):
if k.startswith(unwanted_prefix):
state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k) # 去除不必要的前缀
# 加载模型参数到模型中
self.model.load_state_dict(state_dict, strict=False)
# 计算模型参数量
num_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
print(f"Model has {num_params} parameters.")
# 设置模型为评估模式（evaluation mode），防止训练模式下的 dropout 等操作影响结果
self.model.eval()
# 将模型放置到正确的设备上（GPU 或 CPU）
self.model.to(self.device)
# 初始化分词器
self.tokenizer = Tokenizer(tokenizer_model=self.tokenizer_model_path) # 根据指定的路径加载分词器
def sample(self,
start="Hello!", # 生成文本的起始提示词，可以是任意字符串
num_samples=3, # 生成样本的数量，默认生成 3 个样本
max_new_tokens=256, # 每个样本生成的最大 token 数，默认最多生成 256 个 token
temperature=1.0, # 控制生成的随机性，1.0 为标准，值越大越随机
top_k=300): # 保留概率最高的 top_k 个 token，限制生成时的选择范围
"""
根据给定的起始文本生成样本。
:param start: 生成文本的起始提示词
:param num_samples: 要生成的文本样本数
:param max_new_tokens: 每个样本生成的最大 token 数
:param temperature: 控制生成的随机性，值越小生成越确定，值越大生成越随机
:param top_k: 限制生成时选择的 token 范围
:return: 生成的文本样本列表
"""
# 如果 start 是以 'FILE:' 开头，表示从文件中读取起始文本
if start.startswith('FILE:'):
with open(start[5:], 'r', encoding='utf-8') as f:
start = f.read() # 读取文件内容作为起始文本
# 将起始文本编码为 token id 序列
start_ids = self.tokenizer.encode(start, bos=True, eos=False) # bos=True 表示加上句首标记，eos=False 表示不加句尾标记
x = (torch.tensor(start_ids, dtype=torch.long, device=self.device)[None, ...]) # 将编码后的 token id 转为 PyTorch 张量
generated_texts = [] # 用于保存生成的文本样本
with torch.no_grad(): # 禁用梯度计算，提升效率
with self.ctx: # 进入自动混合精度的上下文（如果是 GPU 并使用 float16 时）
for k in range(num_samples): # 循环生成指定数量的样本
y = self.model.generate(x, max_new_tokens, temperature=temperature, top_k=top_k) # 生成文本
generated_texts.append(self.tokenizer.decode(y[0].tolist())) # 解码生成的 token 序列为可读文本
return generated_texts # 返回生成的文本样本
参考文献