SmolGPT：从零训练小型语言模型的最小实现

这是一个极简的 PyTorch 实现，用于从零开始训练属于你自己的小型 LLM。项目以教学和简洁为设计目标，同时支持高效训练、Flash Attention 和现代采样技术。

核心特性

• 极简代码库：纯 PyTorch 实现，没有额外的抽象开销
• 现代架构：GPT 模型具备以下特性：
  • Flash Attention（环境支持时启用）
  • RMSNorm 和 SwiGLU
  • 高效的 top-k/p/min-p 采样
  • 旋转位置编码 RoPE（可选）
• Flash Attention（环境支持时启用）
• 高效的 top-k/p/min-p 采样
• 旋转位置编码 RoPE（可选）
• 训练特性：
  • 混合精度（bfloat16/float16）
  • 梯度累积
  • 带 warmup 的学习率衰减
  • 权重衰减和梯度裁剪
• 混合精度（bfloat16/float16）
• 梯度累积
• 带 warmup 的学习率衰减
• 权重衰减和梯度裁剪
• 数据集支持：内置 TinyStories 数据集处理功能
• 自定义 Tokenizer：集成 SentencePiece tokenizer 训练

安装

pip install -r requirements.txt

环境要求

• PyTorch 2.0+，支持 CUDA
• 现代 GPU（推荐）

选项一：完整训练流程

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

python train.py

训练损失和验证损失会记录在 out/logs/ 中。要使用 TensorBoard 进行可视化，请运行：

tensorboard --logdir=/out/logs

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

选项二：使用预训练模型

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

# Download pre-trained checkpoint
wget https://huggingface.co/OmAlve/TinyStories-SmolGPT/resolve/main/ckpt.pt -P out/

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

预训练模型详情 🔍

提供的 checkpoint 基于 TinyStories 数据集训练。

• 4096 个 token 的词表
• 512 维 embedding
• 使用约 40 亿个 token 训练，耗时约 18.5 小时
• 验证损失：1.0491

Prompt: One day, Lily met a unicorn

Output:
One day, Lily met a unicorn in the park. The unicorn had shiny fur and a pretty dress. Lily asked the unicorn, "Where did you come from?"
The unicorn replied, "I came from the forest and wanted to meet you, Lily. I am here to make sure you are safe."
Lily and the unicorn played together and had lots of fun. But then, the unicorn started to act funny. Lily didn't know what was happening. Suddenly, the unicorn turned into a little girl and said, "I was under a spell, and your kindness broke it. Thank you for breaking it."
Lily was surprised and happy. She gave the unicorn a big hug and they played together all day. From that day on, the unicorn would always come to play with Lily - her new friend.

Prompt: The dragon flew over the mountains

Output:
The dragon flew over the mountains, over the rivers and over the rivers. He was very brave and strong.
One day, the dragon saw something very strange. It was a big, shiny rock. He wanted to know what it was, so he flew down and touched it with his nose. Suddenly, the rock began to move!
The dragon was so surprised! He had never seen anything like it before. He looked around and saw that it was a little mouse! The mouse was very scared and started to run away.
The dragon was very sad. He wanted to help the mouse, so he decided to try and make friends. He flew around and around until he found the mouse. He said hello to the mouse and asked if he wanted to be friends.
The mouse was so happy! He said yes, and they played together all day long. From then on, the dragon and the mouse were the best of friends. They had lots of fun together and the dragon was never lonely again.

关键参数

可在 config.py 中修改：

GPTConfig(
    block_size=512,    # Context length
    n_layer=8,         # Number of transformer layers
    n_head=8,          # Number of attention heads
    n_embed=512,       # Embedding dimension
    dropout=0.2,       # Dropout rate
    bias=False,        # Use bias in layers
    use_rotary=False,  # Toggle rotary embeddings
)

TrainingConfig(
    batch_size=64,
    max_iters=30000,
    learning_rate=6e-4,
    weight_decay=0.1,
    grad_clip=1.0,
    warmup_iters=1000
)

项目结构

om-alve-smolgpt/
├── config.py       - Model & training configuration
├── dataset.py      - Data loading & preprocessing
├── model.py        - GPT model implementation
├── preprocess.py   - Dataset preparation scripts
├── sample.py       - Text generation script
├── tokenizer.py    - Tokenizer wrapper
└── train.py        - Main training loop

参与贡献

欢迎贡献！如果你希望改进以下内容，请提交 issue 或 PR：

• 性能改进

训练设备规格

训练设备通过 LightningAI 租用。

• GPU：NVIDIA L4 Tensor Core（针对 AI 工作负载优化）

注意：此实现受到现代 LLM 训练实践的启发，并针对教学用途进行了调整。如果用于生产环境，请考虑扩大模型规模和数据集规模。