TinyLLM学习日记

发表于 2025-05-22 更新于 2025-05-24 分类于 LLM实践， TinyLLM 本文字数： 5.6k 阅读时长 ≈ 10 分钟

我希望通过训练一个TinyLLM来打好大模型训练的基础，过程中会遇到很多问题，因此在这里记录学习日记。比较重要的是，要学习理解好一些封装好的接口的使用，避免知其然不知其所以然。

本部分的教程使用的是datawhalechina/tiny-universe: 《大模型白盒子构建指南》：一个全手搓的Tiny-Universe中的TinyLLM部分，不会再对原教程赘述，只记录相关探索的笔记。

Step 1: 训练Tokenizer

SentencePiece库的使用

预备知识

TinyLLM使用了 SentencePiece 库来训练自定义的 Tokenizer，我希望增进对其的理解，参考资料：大模型词表扩充必备工具SentencePiece - 知乎。

Tokenizer有三种粒度：word/character/subword
subword平衡了两种方法，常见的子词算法有Byte-Pair Encoding (BPE) / Byte-level BPE（BBPE）、Unigram LM、WordPiece、SentencePiece等。
- BPE，即字节对编码。其核心思想是从字母开始，不断找词频最高、且连续的两个token合并，直到达到目标词数。
- BBPE的核心思想是将BPE从字符级别扩展到子节（Byte）级别。BPE的一个问题是如果遇到了unicode编码，基本字符集可能会很大。BBPE就是以一个字节为一种“字符”，不管实际字符集用了几个字节来表示一个字符。这样的话，基础字符集的大小就锁定在了256（2^8）。采用BBPE的好处是可以跨语言共用词表，显著压缩词表的大小。而坏处就是，对于类似中文这样的语言，一段文字的序列长度会显著增长。因此，BBPE based模型可能比BPE based模型表现的更好。然而，BBPE sequence比起BPE来说略长，这也导致了更长的训练/推理时间。BBPE其实与BPE在实现上并无大的不同，只不过基础词表使用256的字节集。

SentencePiece 特性

固定最终词汇表大小
使用原始句子训练
空格被视为基本符号 “▁” ，因此可以无歧义地对文本进行detokenize

SentencePiece 实验

使用一个简单的示例进行测试：“aa bb cc aab abbd bb.”

测试代码如下：

import sentencepiece as spm
dataset_path = './demo.txt'
vocab_size=10
spm.SentencePieceTrainer.train(input=dataset_path,model_type="bpe", model_prefix='demo', vocab_size=vocab_size)

sp = spm.SentencePieceProcessor()
sp.load('demo.model')

text='aa bb cc aab abbd bb.'
print(sp.encode_as_pieces(text))
print(sp.encode_as_ids(text))

for i in range(vocab_size):
    print(i,sp.id_to_piece(i))

使用该代码可以看到分词器对这一简单句子的分词结果。

设置vocab_size小于9时，会报错

RuntimeError: Internal: src/trainer_interface.cc(582) [(static_cast<int>(required_chars_.size() + meta_pieces_.size())) <= (trainer_spec_.vocab_size())] Vocabulary size is smaller than required_chars. 8 vs 9. Increase vocab_size or decrease character_coverage with --character_coverage option.

这是因为SentencePieceTrainer会自动添加未知符：、BOS：<s>、EOS：</s>、▁，加上这个例子本来的5个字符，需要至少9个字符才能分词。

而设置的上限即分词算法能计算到最大标记总数，例如，考虑一个简单的例子“ab ac bc cd de”，其上限是：字符总数（5）+

”▁？“型（4）+”▁？？“型（5）+”？？“型（5）+自动添加（4）=23。

一般遇到设定词典大小过大的问题时，可能是数据不够丰富导致的，这时可以选择增加数据或者减少词典大小。

分词器训练与使用

训练：（参数文档参考sentencepiece/doc/options.md at master · google/sentencepiece）

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”归一化规则
    )

加载：

1	sp_model = SentencePieceProcessor(model_file=model_path)

编码：(s：str)

1	sp_model.encode(s)

解码：(t: List[int])

1	sp_model.decode(t)

Step 2: 数据预处理

functools.partial

functools.partial 是 Python 标准库中 functools 模块提供的一个高阶函数，主要用于部分应用函数参数。它允许固定函数的部分参数，生成一个新的简化版函数，从而减少后续调用时的参数传递量。

代码将process_shard(args, vocab_size, tokenizer_model_path)

封装为fun = partial(process_shard, vocab_size=vocab_size, tokenizer_model_path=TOKENIZER_MODEL)

则后续调用时形如fun((0,'path'))，传入一个元组

预处理

将文本数据使用Step1训练的分词器转换为数字序列，并编码为可训练的格式（为每一段文本添加BOS），最后以二进制形式保存

加载已预处理好的数据集

TinyLLM中设计了一个 PretokDataset 类

核心加载数据的代码如下：

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

可以看出这里使用的是步长与窗口大小相等的滑动窗口采样方法，之后可以尝试修改这部分数据加载机制以更大程度地利用数据。

Step 3: 训练模型

TInyLLM使用的模型是与 LLaMA2 结构相同的 Decoder-only Transformer 模型，此部分根据源码进行解读分析。

在最基本的大模型架构基础上，使用了以下策略：

对残差投影进行特殊的缩放初始化

1
2
3

for pn, p in self.named_parameters():
    if pn.endswith('w3.weight') or pn.endswith('wo.weight'):
        torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * args.n_layers))

这是对DecoderLayer内的MLP层的第三层线性变换和Attention层的输出权重矩阵进行放缩

旋转编码

参考十分钟读懂旋转编码（RoPE）

TinyLLM这里的实现与LLAMA里的一致，之后再专门研究学习一下编码。

学习率调整

包括线性预热、余弦退火和最小学习率限制。

自动混合精度训练

参考【Trick2】torch.cuda.amp自动混合精度训练 —— 节省显存并加快推理速度_torch.cuda.amp.gradscaler()-CSDN博客。

因为在某些上下文中torch.FloatTensor有优势，有的torch.HalfTensor有优势。动态估计的原理就是在不出现inf或者NaN梯度值的情况下尽可能增大scaler的值。在每次scaler.step(optimizer)中，都会检查是否有inf或NaN的梯度出现：

如果出现了inf或者NaN，scaler.step(optimizer)会忽略此次的权重更新（optimizer.step() )，并且将scaler的大小缩小（乘上backoff_factor）；
如果没有出现inf或者NaN，那么权重正常更新，并且当连续多次（growth_interval指定）没有出现inf或者NaN，则scaler.update()会将scaler的大小增加（乘上growth_factor）。

# 实例化一个GradScaler对象
scaler = amp.GradScaler(enabled=True)
# 将梯度放大 防止梯度消失
scaler.scale(loss).backward()
# 更新优化器和梯度缩放器
scaler.step(optimizer)
scaler.update()

Step 4: 使用模型生成文本

推理时，为提示词加上BOS，然后逐个字符生成，可以使用temperature、top_k来控制生成的随机性。