最近，我发现了一个特别有趣的项目——smolGPT，它是一个极简的 PyTorch 实现，能够让我们从零训练一个属于自己的小型 GPT 模型。该项目的核心目标是教育用途，代码精炼，没有多余的抽象封装，使人可以直接理解大语言模型（LLM）是如何训练的。今天，我就带大家一起看看这个项目是如何工作的，并分享一下如何用它训练自己的 GPT 模型。

为什么选择 smolGPT？

在探索 LLM 训练过程中，通常会遇到很多复杂的框架，比如 Hugging Face Transformers、DeepSpeed 等工具。虽然它们强大，但内部封装过于深入，初学者很难直接上手。而 smolGPT 采用了最小化代码实现，直接基于 PyTorch 编写，让我们更清晰地理解 GPT 模型的核心结构。

主要特点

• 现代 GPT 架构：支持 Flash Attention（如果可用）、RMSNorm、SwiGLU 以及高效的 top-k/p 采样。
• 训练优化：支持混合精度（bfloat16/float16）、梯度累积、学习率预热（warmup）、权重衰减和梯度裁剪等优化技术。
• 数据支持：内置 TinyStories 训练数据处理，也可以自定义数据集。
• 自定义 Tokenizer：集成了 SentencePiece，用于训练自己的分词器。

这些特性使得它成为希望深入了解 GPT 训练原理的开发者们的绝佳学习材料。

环境准备

想要运行 smolGPT，你需要一个现代 GPU（最好支持 CUDA），并安装以下环境：

pip install -r requirements.txt

基本要求：

• Python 3.8+
• PyTorch 2.0+（带 CUDA 支持）

训练一个自己的 GPT

我们可以用 smolGPT 来训练一个从零开始的 LLM，主要分为三个步骤：

1. 预处理数据

首先，我们需要准备训练数据。smolGPT 默认使用 TinyStories 数据集，这是一个适合小型 GPT 训练的短故事集合。也可以使用自己的数据集，格式为简单的文本文件，每行一个样本。

执行以下命令进行数据预处理，并训练一个 4096 词汇量的分词器：

python preprocess.py prepare-dataset --vocab-size 4096

2. 开始训练

数据准备好后，就可以启动训练了：

python train.py

smolGPT 采用了标准的 GPT 训练流程，包括：

• 采用 8 层 Transformer，每层有 8 个注意力头，嵌入维度为 512。
• 上下文窗口为 512，适用于小规模任务。
• 训练过程使用 warmup + 余弦学习率衰减，初始学习率为 6e-4。
• 批量大小 64，总共训练 30,000 轮。

如果想修改这些参数，可以直接编辑 config.py：

GPTConfig(
    block_size=512,    # 上下文长度
    n_layer=8,         # Transformer 层数
    n_head=8,          # 注意力头数
    n_embed=512,       # 词向量维度
    dropout=0.2,       # Dropout
    bias=False         # 是否使用偏置
)

3. 文本生成

训练完成后，我们可以用 smolGPT 生成文本：

python sample.py --prompt "Once upon a time" --num_samples 3 --temperature 0.7 --max_new_tokens 500

这里：

• --prompt 是输入的文本开头；
• --num_samples 指定生成多少个样本；
• --temperature 控制生成文本的随机性（值越高，文本越多样化）；
• --max_new_tokens 限制生成的最大 token 数。

4. 使用预训练模型

如果不想自己训练，也可以直接使用别人训练好的模型，比如官方提供的 TinyStories-SmolGPT：

# 下载 tokenizer
wget https://huggingface.co/OmAlve/TinyStories-SmolGPT/resolve/main/tok4096.model -P data/

# 下载已训练好的 checkpoint
wget https://huggingface.co/OmAlve/TinyStories-SmolGPT/resolve/main/ckpt-v1.pt -P out/

然后执行：

python sample.py \
    --prompt "Once upon a time" \
    --tokenizer_path data/tok4096.model \
    --ckpt_dir out/ \
    --num_samples 3 \
    --max_new_tokens 200 \
    --temperature 0.7

这样，我们就可以用别人训练好的模型生成文本。

模型训练细节

训练数据

官方提供的 TinyStories 数据集，训练过程中大约使用了 40 亿个 token，训练时间 18.5 小时，最终的验证损失（Validation Loss）约为 1.0491。

模型结构

smolGPT 的默认模型架构如下：

• 词汇量：4096
• Transformer 层数：8
• 注意力头数：8
• 嵌入维度：512
• 训练数据量：40 亿 token
• 训练时间：18.5 小时

在 config.py 中，可以调整这些参数以训练更大或更小的模型。

smolGPT VS 其他 LLM

可以简单对比一下 smolGPT 与其他 GPT 训练方案的特点：

如果你是初学者，想要了解 GPT 训练的基本原理，smolGPT 是一个非常好的入门选择。如果需要一个可以大规模训练的 LLM，nanoGPT 和 Hugging Face 可能更合适。

引用链接

[1] smolGPT: https://github.com/Om-Alve/smolGPT