在自然语言处理(NLP)领域，大型语言模型(LLM)如GPT-3虽然强大，但对计算资源的需求也极其庞大。相对而言，小型语言模型(SLM) 为在标准硬件上进行学习和实验提供了一个更实用的选择。本文将深入探讨如何利用TinyStories数据集，从零开始构建一个高效且可训练的小型语言模型(SLM)。我们将详细解析数据加载、分词(Tokenization)、高效存储、批次生成以及使用Transformer架构进行模型训练的各个环节，为初学者和经验丰富的从业者提供全面的实践指导。
1. 环境搭建与依赖项安装
构建 小型语言模型(SLM) 的第一步是搭建合适的开发环境，并安装所需的依赖库。文章中使用了Hugging Face的datasets库来加载TinyStories数据集，以及tiktoken库进行Byte Pair Encoding (BPE分词(Tokenization))。
!pip install -U datasets
!pip install tiktoken
!pip install numpy matplotlib tqdm
!pip install accelerate wandb tensorboard
- datasets库: 简化了数据集的加载和预处理，内置对大型数据集和并行处理的支持。例如，使用
load_dataset("roneneldan/TinyStories")一行代码即可轻松加载TinyStories数据集，无需手动下载和解析。 - tiktoken库: 由于其速度和与GPT-2 分词(Tokenization) 的兼容性而被选中。GPT-2的分词方案非常适合英语文本，并产生紧凑的词汇表，这对于小型模型至关重要。
- 其他依赖: NumPy用于数值计算，matplotlib用于可视化，tqdm用于显示进度条，accelerate用于加速训练，wandb和tensorboard用于模型训练的可视化与监控。
这种方式的优点在于，Hugging Face的datasets库提供了一种统一的接口来访问各种数据集，极大地降低了数据处理的复杂性。tiktoken库则提供了高效的 BPE分词(Tokenization) 方案，在词汇表大小和文本表示的灵活性之间取得了平衡。
2. TinyStories数据集加载与探索
TinyStories数据集是一个由大约200万个训练故事和2万个验证故事组成的数据集，每个故事都是一段简短且连贯的文本。它被设计成能够测试语言模型是否真正理解了语言的结构和含义，而不是简单地记住训练数据。
from datasets import load_dataset
df = load_dataset("roneneldan/TinyStories")
print(df)
# DatasetDict({
# train: Dataset({
# features: ['text'],
# num_rows: 2119719
# })
# validation: Dataset({
# features: ['text'],
# num_rows: 21990
# })
# })
选择TinyStories作为 小型语言模型(SLM) 的训练数据是因为它具有以下优势：
- 故事简短： 降低了计算需求，使得在资源有限的硬件上进行训练成为可能。
- 内容简单： 便于模型学习基本的语言结构和语义关系。
- 结构连贯： 确保模型生成的文本也具有一定的可读性和逻辑性。
Hugging Face datasets库的集成简化了数据下载和缓存的复杂性，确保了对数据集的有效访问。 虽然datasets库支持流式处理，但该示例中未使用它。
3. Byte Pair Encoding (BPE) 分词(Tokenization)
分词(Tokenization) 是将原始文本转换为模型可以理解的数字表示的关键步骤。Byte Pair Encoding (BPE) 是一种常用的子词 分词(Tokenization) 算法，它通过将文本分解为频繁出现的词片段来平衡词汇表大小和表示灵活性。
import tiktoken
import numpy as np
from tqdm.auto import tqdm
encoding = tiktoken.get_encoding("gpt2")
def processing(sample_text):
ids = encoding.encode_ordinary(sample_text['text'])
out = {'ids': ids, 'len': len(ids)}
return out
tiktoken.get_encoding("gpt2")初始化了一个GPT-2 BPE分词(Tokenization) 器，它可以将原始文本转换为token ID。processing函数接收一个数据集样本（包含text键的字典），使用encode_ordinary将文本编码为token ID，然后返回一个包含token ID（ids）和长度（len）的字典。
与基于词的 分词(Tokenization) （产生大型词汇表）或基于字符的 分词(Tokenization) （产生长序列）不同，BPE分词(Tokenization) 提供了一种中间方案。它通过将罕见词分解为子词（例如，“unhappiness” → [“un”, “happi”, “ness”]）来处理它们，从而减小了词汇表大小，同时保留了含义。GPT-2 分词(Tokenization) 器针对英语进行了预训练和优化，具有大约50,257个token的词汇表，适合TinyStories数据集。
4. 将分词(Tokenization)后的数据存储到二进制文件中
将分词后的数据存储到二进制文件(.bin)中，可以优化内存使用和处理速度，这对于 小型语言模型(SLM) 训练至关重要。
import os
import numpy as np
from tqdm.auto import tqdm
if not os.path.exists("train.bin"):
tokenized = df.map(
processing,
remove_columns=['text'],
desc="tokenizing the splits",
num_proc=8,
)
for split, dset in tokenized.items():
arr_len = np.sum(dset['len'], dtype=np.uint64)
filename = f'{split}.bin'
dtype = np.uint16
arr = np.memmap(filename, dtype=dtype, mode='w+', shape=(arr_len,))
total_batches = 1024
idx = 0
for batch_idx in tqdm(range(total_batches), desc=f'writing {filename}'):
batch = dset.shard(num_shards=total_batches, index=batch_idx, contiguous=True).with_format('numpy')
arr_batch = np.concatenate(batch['ids'])
arr[idx : idx + len(arr_batch)] = arr_batch
idx += len(arr_batch)
arr.flush()
- 检查现有文件： 确保token化仅在二进制文件尚不存在时运行，从而避免了冗余处理。
df.map函数将处理函数应用于每个数据集样本，删除text列并添加ids和len。num_proc=8参数启用具有8个CPU内核的并行处理，以加快执行速度。- token的总数（
arr_len）是通过对分割中所有样本的len字段求和来计算的，使用np.uint64来处理大数。 - 使用
np.memmap创建内存映射的数组，用于存储token ID，因为GPT-2词汇表大小（约50,257）适合2¹⁶ (65,536)。 - 数据集分为1,024个分片（
total_batches = 1024），以处理可管理块中的数据，从而减少了内存开销。每个分片的token ID被连接并写入内存映射的数组。 arr.flush()确保所有数据都写入磁盘，从而完成.bin文件。
将token ID存储为二进制文件（np.uint16）比将文本或Python对象保留在RAM中更节省内存，这对于具有数百万个token的数据集至关重要。 np.memmap允许直接磁盘访问，防止在处理大型数据集期间RAM过载。
5. 为训练生成输入-输出对
为了训练 小型语言模型(SLM)，需要将数据转换为输入-输出对，其中输入是token序列，输出是序列中每个token的下一个token。
import torch
import numpy as np
def get_batch(split):
if split == 'train':
data = np.memmap('train.bin', dtype=np.uint16, mode='r')
else:
data = np.memmap('validation.bin', dtype=np.uint16, mode='r')
ix = torch.randint(len(data) - block_size, (batch_size,))
x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64)) for i in ix])
y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64)) for i in ix])
if device_type == 'cuda':
x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
else:
x, y = x.to(device), y.to(device)
return x, y
- 加载数据： 将
train.bin或validation.bin读取为内存映射的数组。 - 随机索引： 为
block_size的序列生成batch_size随机索引。 - 输入-输出对： 为下一个token预测创建输入
x(data[i:i+block_size]) 和输出y(data[i+1:i+1+block_size])。 - 张量转换： 转换为PyTorch张量 (
np.int64)。 - GPU优化： 使用
pin_memory()和non_blocking=True进行GPU传输。
这种方式确保了训练示例的多样性，高效处理大型数据集，并与自回归训练对齐。使用GPU优化可以减少数据传输开销。
6. Transformer 模型架构
小型语言模型(SLM) 的核心是Transformer架构，它通过自注意力机制捕捉文本中的复杂依赖关系。本文中使用的Transformer模型包括以下组件：
from dataclasses import dataclass
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
@dataclass
class GPTConfig:
block_size: int = 128
vocab_size: int = 50257
n_layer: int = 6
n_head: int = 6
n_embd: int = 384
dropout: float = 0.1
bias: bool = True
class GPT(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.transformer = nn.ModuleDict(dict(
wte=nn.Embedding(config.vocab_size, config.n_embd),
wpe=nn.Embedding(config.block_size, config.n_embd),
drop=nn.Dropout(config.dropout),
h=nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
ln_f=LayerNorm(config.n_embd, config.bias),
))
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
self.transformer.wte.weight = self.lm_head.weight # weight tying
self.apply(self._init_weights)
for pn, p in self.named_parameters():
if pn.endswith('c_proj.weight'):
nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * config.n_layer))
def _init_weights(self, module):
if isinstance(module, nn.Linear):
nn.init.normal_(module.weight, mean=0.0, std=0.02)
if module.bias is not None:
nn.init.zeros_(module.bias)
elif isinstance(module, nn.Embedding):
nn.init.normal_(module.weight, mean=0.0, std=0.02)
def forward(self, idx, targets=None):
device = idx.device
b, t = idx.size()
assert t <= self.config.block_size
pos = torch.arange(0, t, dtype=torch.long, device=device)
tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
x = self.transformer.drop(tok_emb + pos_emb)
for block in self.transformer.h:
x = block(x)
x = self.transformer.ln_f(x)
if targets is not None:
# training inference
logits = self.lm_head(x)
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
else:
# inference
# only consider the last token of shape (1,1,vocab_size)
logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
loss = None
return logits, loss
- GPTConfig: 定义了模型的超参数，包括
block_size(上下文窗口大小)、vocab_size(词汇表大小)、n_layer(Transformer块的数量)、n_head(注意力头的数量)、n_embd(嵌入维度)、dropout(dropout率)和bias(线性层中的偏置)。这些参数平衡了模型容量和计算效率。 - GPT: 是Transformer模型的核心，包含以下组件：
wte: token嵌入层，将token ID映射到384维向量。wpe: 位置嵌入层，编码token在上下文窗口中的位置。drop: dropout层，应用于token和位置嵌入的总和，以防止过拟合。h: 一系列Transformer块，每个块包含自注意力和前馈组件。ln_f: 最终的层归一化，在语言模型头之前对输出进行归一化。lm_head: 语言模型头，将嵌入投影到词汇大小的logits，用于下一个token预测。
- 权重绑定： token嵌入层 (
wte) 和输出头 (lm_head) 共享权重，减少了模型参数并提高了效率。 - 权重初始化： 线性层和嵌入层的权重使用正态分布初始化，以稳定训练。
- 前向传播： 接收一批token索引，嵌入token和位置，应用dropout，通过Transformer块处理，并进行归一化。如果提供了目标（训练模式），则计算logits和交叉熵损失以进行预测。否则（生成模式），仅输出最后一个token的logits。
自注意力机制能够捕获文本中的复杂依赖关系，因此Transformer架构非常适合语言建模。权重绑定通过将嵌入层减少大约50%的参数，在不牺牲性能的情况下提高了内存效率。像n_layer=6和n_embd=384这样的参数确保模型可以在标准硬件上运行，同时保持合理的性能。
7. Transformer 块组件
Transformer模型中的每个块都包含自注意力和前馈网络，它们共同作用以学习文本中的模式。
class Block(nn.Module):
def __init__(self, config):
super().__init__()
self.ln1 = LayerNorm(config.n_embd, config.bias)
self.attn = CausalSelfAttention(config)
self.ln2 = LayerNorm(config.n_embd, config.bias)
self.mlp = MLP(config)
def forward(self, x):
x = x + self.attn(self.ln1(x))
x = x + self.mlp(self.ln2(x))
return x
class CausalSelfAttention(nn.Module):
def __init__(self, config):
super().__init__()
assert config.n_embd % config.n_head == 0
self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
self.attn_dropout = nn.Dropout(config.dropout)
self.resid_dropout = nn.Dropout(config.dropout)
self.n_head = config.n_head
self.n_embd = config.n_embd
self.flash = hasattr(F, 'scaled_dot_product_attention')
if not self.flash:
self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
.view(1, 1, config.block_size, config.block_size))
def forward(self, x):
B, T, C = x.size() # batch size, time steps, embedding dimension (32, 128, 384)
# calculate query, key, values for all heads in batch and move head forward to be the batch dim
q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
# causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
if self.flash:
# efficient attention using Flash Attention CUDA kernels
y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.attn_dropout.p if self.training else 0.0, is_causal=True)
else:
# manual implementation of attention
att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
att = F.softmax(att, dim=-1)
att = self.attn_dropout(att)
y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
# output projection
y = self.resid_dropout(self.c_proj(y))
return y
class MLP(nn.Module):
def __init__(self, config):
super().__init__()
self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
self.gelu = nn.GELU()
self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
self.dropout = nn.Dropout(config.dropout)
def forward(self, x):
x = self.c_fc(x)
x = self.gelu(x)
x = self.c_proj(x)
x = self.dropout(x)
return x
- Block: 组合
CausalSelfAttention和MLP，每个都以层归一化(ln1、ln2)开头。使用残差连接（x + ...）将注意力和MLP输出添加到输入，从而保持信息流并防止梯度消失。输出形状保持为(batch_size, block_size, n_embd)（例如，32×128×384）。 - CausalSelfAttention:
- 初始化： 创建一个线性层(
c_attn)来计算查询、键和值(3×n_embd)，以及一个投影层(c_proj)。确保n_embd可被n_head整除（384/6=每个头的64维）。 - 前向传播： 将输入拆分为查询、键和值，重塑用于多头注意力（6个头），并应用缩放的点积注意力。
- Flash Attention： 如果可用，使用PyTorch的
scaled_dot_product_attention来提高效率；否则，使用因果掩码（三角形矩阵）手动计算注意力，以防止关注未来的token。 - 输出： 使用dropout投影回
n_embd维度。
- 初始化： 创建一个线性层(
- MLP: 将嵌入扩展到
4×n_embd(1536)，应用GELU激活以实现非线性，投影回n_embd，并应用dropout。
8. 模型训练
训练过程至关重要，它利用学习算法调整模型的参数，使其能够生成连贯且有意义的文本。以下代码段概述了模型训练过程，涉及优化器设置、学习速率调度和损失计算等各个方面：
import torch
from contextlib import nullcontext
from torch.optim.lr_scheduler import LinearLR, SequentialLR, CosineAnnealingLR
# --- 初始化、配置、设置 ---
learning_rate = 1e-4
max_iters = 10000
warmup_steps = 1000
min_lr = 5e-4
eval_iters = 500
batch_size = 32
block_size = 128
gradient_accumulation_steps = 32
device = "cuda" if torch.cuda.is_available() else "cpu"
device_type = 'cuda' if 'cuda' in device else 'cpu'
dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16' # 支持bfloat16自动设置为bfloat16
ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
ctx = nullcontext() if device_type == 'cpu' else torch.amp.autocast(device_type=device_type, dtype=ptdtype)
torch.set_default_device(device)
torch.manual_seed(42) # 设置随机种子保证实验可复现
# --- 模型实例化、优化器和学习率调整器定义 ---
config = GPTConfig(vocab_size=50257, block_size=128, n_layer=6, n_head=6, n_embd=384, dropout=0.1, bias=True)
model = GPT(config).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, betas=(0.9, 0.95), weight_decay=0.1, eps=1e-9)
scheduler_warmup = LinearLR(optimizer, total_iters=warmup_steps)
scheduler_decay = CosineAnnealingLR(optimizer, T_max=max_iters - warmup_steps, eta_min=min_lr)
scheduler = SequentialLR(optimizer, schedulers=[scheduler_warmup, scheduler_decay], milestones=[warmup_steps])
scaler = torch.cuda.amp.GradScaler(enabled=(dtype == 'float16')) # 梯度缩放
# --- 损失评估函数 ---
def estimate_loss(model):
out = {}
model.eval() # 设置为评估模式
with torch.inference_mode():
for split in ['train', 'val']:
losses = torch.zeros(eval_iters)
for k in range(eval_iters):
X, Y = get_batch(split)
with ctx:
logits, loss = model(X, Y)
losses[k] = loss.item()
out[split] = losses.mean()
model.train() # 恢复训练模式
return out
# --- 训练循环 ---
best_val_loss = float('inf') # 初始化最佳验证损失为无穷大
best_model_params_path = "best_model_params.pt"
train_loss_list, validation_loss_list = [], [] # 记录训练和验证损失
for epoch in tqdm(range(max_iters)):
# --- 定期评估与模型保存 ---
if epoch % eval_iters == 0 and epoch != 0:
losses = estimate_loss(model)
print(f"Epoch {epoch}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")
print(f"The current learning rate: {optimizer.param_groups[0]['lr']:.5f}")
train_loss_list += [losses['train']]
validation_loss_list += [losses['val']]
if losses['val'] < best_val_loss:
best_val_loss = losses['val']
torch.save(model.state_dict(), best_model_params_path)
# --- 前向传播、损失计算和反向传播 ---
X, y = get_batch("train") # 获取训练批次
X, y = X.to(device), y.to(device) # 将数据移至设备（GPU 或 CPU）
with ctx:
logits, loss = model(X, y) # 前向传播
loss = loss / gradient_accumulation_steps # 梯度累积：归一化损失
scaler.scale(loss).backward() # 缩放损失，进行反向传播
# --- 梯度累积与优化步骤 ---
if ((epoch + 1) % gradient_accumulation_steps == 0) or (epoch + 1 == max_iters):
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5) # 梯度裁剪
scaler.step(optimizer) # 更新优化器中的参数
scaler.update() # 更新梯度缩放器
optimizer.zero_grad(set_to_none=True) # 清空梯度
# --- 学习率更新 ---
scheduler.step()
- 超参数： 包括
learning_rate(学习率),max_iters(最大迭代次数),warmup_steps(预热步数),min_lr(最小学习率),eval_iters(评估迭代次数),batch_size(批次大小),block_size(块大小), 和gradient_accumulation_steps(梯度累积步数)。 - 设备和精度： 如果可用，使用 CUDA，并通过自动混合精度 (AMP) 使用 bfloat16 或 float16 以减少内存使用并加快训练速度。
- 优化器： 使用 AdamW，其中 betas=(0.9, 0.95)，weight_decay=0.1 用于正则化，eps=1e-9 用于数值稳定性。
- 学习率调度器：
结合了 LinearLR（用于 1,000 步的预热，逐渐增加学习率）和 CosineAnnealingLR（用于在剩余迭代中衰减到 minlr）。SequentialLR 在 warmupsteps 处从预热切换到衰减。 - 梯度缩放： GradScaler 可防止在混合精度训练期间 float16 中出现下溢。
9. 样本生成
训练完成后，可以使用模型生成新的文本。以下代码段演示了如何使用训练好的 小型语言模型(SLM) 生成文本：
@torch.no_grad()
def generate(model, idx, max_new_tokens, temperature=1.0, top_k=None):
"""
使用给定的模型从提供的索引开始生成文本。
参数：
model (torch.nn.Module): 经过训练的文本生成模型。
idx (torch.Tensor): 输入索引序列，用作生成过程的起始点。
max_new_tokens (int): 要生成的最大token数。
temperature (float, 可选): 用于控制输出随机性的温度。值越高，输出越随机。默认为1.0。
top_k (int, 可选): 如果指定，则只从概率最高的 top_k 个 token 中进行采样。默认为 None。
返回：
torch.Tensor: 包含输入索引和生成的 token 的完整索引序列。
"""
for _ in range(max_new_tokens):
# 使用模型生成下一个token
# (1) 上下文截断
idx_cond = idx if idx.size(1) <= config.block_size else idx[:, -config.block_size:]
# (2) 模型预测
logits, _ = model(idx_cond)
# (3) 调整温度
logits = logits[:, -1, :] / temperature
# (4) Top-K 过滤
if top_k is not None:
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
logits[logits < v[:, [-1]]] = -float('Inf')
# (5) 计算概率
probs = F.softmax(logits, dim=-1)
# (6) 采样
idx_next = torch.multinomial(probs, num_samples=1)
# (7) 拼接
idx = torch.cat((idx, idx_next), dim=1)
return idx
generate函数以自回归方式生成文本，它接收一个张量 idx （提示的token ID）并生成 max_new_tokens 个token。该函数使用最后 block_size=128 个token来符合模型的上下文限制。logits针对最后一个token计算，并按 temperature 缩放以控制随机性。如果指定，则限制采样到最高的k个logits，将其他的设置为负无穷，以实现集中输出。应用softmax以获得概率，并使用多项式采样下一个token。
10. 结论
本文提供了一个从头开始构建 小型语言模型(SLM) 的实践指南，涵盖了数据预处理、分词(Tokenization)、存储、批处理、Transformer架构、训练和生成等核心技术。利用二进制存储、混合精度训练和梯度累积等技术展示了效率的最佳实践，而 TinyStories 数据集和 PyTorch 确保了学习者的可访问性。
无论是探索语言建模的初学者，还是试验Transformer的经验丰富的开发人员，本文都希望成为有价值的资源。通过调整超参数或生成自己的数据集来查看模型运行情况！