为什么要关注数据预处理？

很多从零训练LLM的教程把重点放在模型架构上，却忽略了数据准备——这恰恰是决定模型能否收敛的关键。我在研究生阶段复现GPT-2时，前两次都因为tokenizer和数据集处理不当导致loss不降，后来花了三天排查才发现问题。今天用FareedKhan-dev/train-llm-from-scratch项目的方法，但我会把那些“隐藏的坑”补上，让你一次跑通。

核心原理：tokenizer + 数据集迭代器

Tokenizer的本质是建立字符/子词到整数id的映射。对于从零训练，你需要根据你的语料构建tokenizer，而不是直接套用现成的。原因很简单：语料分布不同，vocab大小影响计算效率。

上图展示了vocab大小对模型loss的影响，太小会导致序列过长，太大则稀疏浪费。实验表明，对于1B token以内的中文语料，vocab在8000-16000之间最优。我一般选8192，因为2的幂次方便GPU内存对齐。

实现步骤：代码片段与解释

1. 构建tokenizer

使用HuggingFace的tokenizers库训练一个BPE tokenizer：

from tokenizers import Tokenizer, models, trainers, pre_tokenizers

tokenizer = Tokenizer(models.BPE())
tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)

trainer = trainers.BpeTrainer(
    vocab_size=8192,
    special_tokens=["<|endoftext|>", "<|pad|>"],
    min_frequency=2
)

# 假设你的语料文件列表：list_of_files
tokenizer.train(files=list_of_files, trainer=trainer)
tokenizer.save("tokenizer.json")

为什么用ByteLevel？ 因为可以处理任何unicode字符，且不需要预定义单词边界。我对比过WordLevel和BPE，BPE在中文和代码混合语料上压缩率更好。min_frequency=2能过滤掉只出现一次的无意义token。

2. 高效数据集迭代器

直接加载所有tokenized数据到内存会爆显存。用memory-mapped方式：

import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader

class LLMDataset(Dataset):
    def __init__(self, bin_path, block_size=512):
        self.data = np.memmap(bin_path, dtype=np.uint16, mode='r')
        self.block_size = block_size
        
    def __len__(self):
        return len(self.data) // self.block_size
    
    def __getitem__(self, idx):
        start = idx * self.block_size
        chunk = self.data[start:start+self.block_size+1]
        x = torch.tensor(chunk[:-1], dtype=torch.long)
        y = torch.tensor(chunk[1:], dtype=torch.long)
        return x, y

这个设计的关键：一次只映射文件的某一部分，不占用额外内存。block_size要匹配模型的上下文长度。我通常设为512或1024，太小会丢失长程依赖，太大会降低训练吞吐。

3. 训练超参数配置

以下是经过验证的配置（以10M参数模型为例）：

# config.yaml
model:
  vocab_size: 8192
  hidden_size: 256
  num_layers: 6
  num_heads: 8
  max_position_embeddings: 512

training:
  batch_size: 64  # 每个GPU
  learning_rate: 3e-4
  warmup_steps: 1000
  max_steps: 100000
  weight_decay: 0.1
  beta1: 0.9
  beta2: 0.95

为什么用3e-4？ 对于小模型（<100M参数），3e-4通常是最佳起点。我实验过1e-4到1e-3的范围，3e-4在10000步后loss最低。warmup_steps设为1000是为了让模型在初期稳定，避免梯度爆炸。

实验结果对比

我在8M token的中文维基百科语料上训练了一个10M参数模型（vocab=8192, block_size=512, 训练10万步），与使用默认gpt2 tokenizer（vocab=50257）对比例：

自定义tokenizer不仅更快，而且loss更低——因为vocab更匹配语料，序列长度更短。

常见问题和避坑指南

坑1：tokenizer训练用的语料与实际训练语料不一致
我见过有人用英文wiki训练tokenizer，然后用于中文语料训练。后果：分词碎片化，很多token只出现一次，模型学不到语义。解决：始终用20%的训练语料单独训练tokenizer。

坑2：DataLoader的collate_fn忽略padding
如果batch内序列长度不等，需要padding到相同长度。但使用memmap固定block_size的方式已经对齐，不需要额外padding。注意：如果使用block_size截断，对长文本会损失信息，所以训练前最好按长度均匀切割语料。

坑3：学习率调度器设置不当
很多教程直接copy cosine schedule，但小模型用constant+linear decay即可。我对比过：对于100k步训练，cosine在最后30k步学习率过低导致loss停滞，而线性decay更鲁棒。

from transformers import get_linear_schedule_with_warmup
scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=1000, num_training_steps=100000)

坑4：显存溢出却不知道为什么
检查batch_size和block_size的乘积是否超过GPU内存上限。一个简单的公式：所需显存 ≈ batch_size * block_size * vocab_size * 4 (bytes) * 模型参数倍数。对于10M模型，batch_size=64, block_size=512, 约需 64512819242 ≈ 2.1GB，A100 40G绰绰有余。但如果随意增大batch_size到256，就会溢出。

我的实操心得

从零训练LLM最难的不是模型代码，而是数据流水线。建议第一步先写一个数据预处理脚本，把原始文本转成uint16二进制文件，并单独保存a tokenizer。第二步用一个小词表（5120）和短序列（128）快速验证整个流程是否跑通。如果10分钟能见到loss下降，再切换到正式配置。这样能节省大量调试时间。