当前位置：首页 > news >正文

从零训练一个小型语言模型

news 2026/6/9 21:56:55

今天，我们就结合一份具体数据，完整走一遍从环境准备、数据处理、模型构建到模型训练的全过程，训练一个参数量约0.1B，也就是1亿参数的小型语言模型。

需要说明的是，小型语言模型和大语言模型的底层训练逻辑是一致的，都是基于大量文本数据进行Token预测训练。区别主要在于大语言模型数据规模更大、模型参数更多、算力需求更高，以及训练优化策略更复杂。

这里使用一个真实公开数据集，完整走一遍小型GPT模型的训练过程。核心流程包括：

环境准备
下载真实数据集
数据清洗与格式化
构建Tokenizer
构造训练样本
搭建Transformer模型
模型训练
文本生成测试

一、环境准备

首先需要准备Python深度学习环境。最基础的环境包括：

conda create -n mini_llm python=3.10 conda activate mini_llm pip install torch torchvision torchaudio pip install transformers datasets tokenizers tqdm numpy

如果有NVIDIA GPU，还需要确认PyTorch能识别CUDA：

# test.py import torch print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))

如果输出为：

True NVIDIA GPU名称(4090、H20-3e等）

说明GPU环境正常。如果没有GPU，也可以用CPU跑通流程，只是训练速度会慢很多。本次项目目录可以设计成：

mini_llm/ ├── data/ │ ├── raw/ │ ├── processed/ ├── tokenizer/ ├── checkpoints/ ├── prepare_data.py ├── train_tokenizer.py ├── model.py ├── train.py └── generate.py

二、选择一个真实数据集

为了演示从零训练语言模型，可以选择公开英文文本数据集，例如 WikiText-2。WikiText-2是一个常用的小规模语言模型训练数据集，来自维基百科文章，适合用来跑通小型语言模型训练流程。可以直接用Hugging Face datasets下载：

from datasets import load_dataset dataset = load_dataset("wikitext", "wikitext-2-raw-v1") print(dataset)

输出大致如下：

DatasetDict({ train: Dataset({ features: ['text'], num_rows: 36718 }) validation: Dataset({ features: ['text'], num_rows: 3760 }) test: Dataset({ features: ['text'], num_rows: 4358 }) })

三、数据处理

原始WikiText数据里有一些空行、标题、特殊符号，需要先做基础清洗。新建prepare_data.py：

# prepare_data.py from datasets import load_dataset import os dataset = load_dataset("wikitext", "wikitext-2-raw-v1") os.makedirs("data/processed", exist_ok=True) def save_split(split_name): texts = dataset[split_name]["text"] cleaned = [] for line in texts: line = line.strip() if len(line) == 0: continue cleaned.append(line) with open(f"data/processed/{split_name}.txt", "w", encoding="utf-8") as f: for line in cleaned: f.write(line + "\n") save_split("train") save_split("validation") save_split("test")

运行prepare_data.py，处理后得到：

data/processed/train.txt data/processed/validation.txt data/processed/test.txt

这一步的目标很简单：把原始数据集整理成模型可以读取的纯文本文件。

四、训练Tokenizer

大语言模型不能直接处理文本，它需要先把文本切成Token，再把Token 转换成数字ID。这里使用BPE Tokenizer。新建train_tokenizer.py：

# train_tokenizer.py from tokenizers import ByteLevelBPETokenizer import os paths = ["data/processed/train.txt"] tokenizer = ByteLevelBPETokenizer() tokenizer.train( files=paths, vocab_size=8000, min_frequency=2, special_tokens=[ "<pad>", "<s>", "</s>", "<unk>", "<mask>", ], ) os.makedirs("tokenizer", exist_ok=True) tokenizer.save_model("tokenizer")

运行train_tokenizer.py，生成：

tokenizer/vocab.json tokenizer/merges.txt

其中，vocab.json保存Token到 ID的映射；merges.txt保存BPE合并规则。例如一句话：

Artificial intelligence is changing science.

经过Tokenizer后会变成类似：

[523, 1782, 318, 2450, 3842, 17]

模型真正输入的就是这些数字ID。

五、构造训练样本

语言模型的训练目标是预测下一个Token。假设一段文本被编码成：

[10, 25, 31, 47, 88, 92, 103]

如果上下文长度设置为 block_size = 4，那么训练样本就是：

输入 x: [10, 25, 31, 47] 标签 y: [25, 31, 47, 88]

也就是标签整体向右移动一位。在训练代码里，可以这样构造batch：

import torch def get_batch(data, batch_size, block_size, device): ix = torch.randint(len(data) - block_size - 1, (batch_size,)) x = torch.stack([ data[i:i + block_size] for i in ix ]) y = torch.stack([ data[i + 1:i + block_size + 1] for i in ix ]) return x.to(device), y.to(device)

这一步是训练语言模型训练数据构造的核心，把训练输出处理成：输入前面的Token，让模型预测后面的Token。

六、构建Transformer语言模型

接下来搭建一个小型GPT模型。这里采用Decoder-only架构，因为GPT这一类大语言模型本质上是自回归语言模型，它的训练目标是根据前面的Token预测下一个Token。

新建model.py：

# model.py import torch import torch.nn as nn import torch.nn.functional as F class Head(nn.Module): def __init__(self, n_embd, head_size, block_size, dropout): super().__init__() self.key = nn.Linear(n_embd, head_size, bias=False) self.query = nn.Linear(n_embd, head_size, bias=False) self.value = nn.Linear(n_embd, head_size, bias=False) self.register_buffer( "tril", torch.tril(torch.ones(block_size, block_size)) ) self.dropout = nn.Dropout(dropout) def forward(self, x): B, T, C = x.shape k = self.key(x) q = self.query(x) wei = q @ k.transpose(-2, -1) * k.shape[-1] ** -0.5 wei = wei.masked_fill(self.tril[:T, :T] == 0, float("-inf")) wei = F.softmax(wei, dim=-1) wei = self.dropout(wei) v = self.value(x) out = wei @ v return out class MultiHeadAttention(nn.Module): def __init__(self, n_embd, num_heads, head_size, block_size, dropout): super().__init__() self.heads = nn.ModuleList([ Head(n_embd, head_size, block_size, dropout) for _ in range(num_heads) ]) self.proj = nn.Linear(n_embd, n_embd) self.dropout = nn.Dropout(dropout) def forward(self, x): out = torch.cat([h(x) for h in self.heads], dim=-1) out = self.dropout(self.proj(out)) return out class FeedForward(nn.Module): def __init__(self, n_embd, dropout): super().__init__() self.net = nn.Sequential( nn.Linear(n_embd, 4 * n_embd), nn.GELU(), nn.Linear(4 * n_embd, n_embd), nn.Dropout(dropout), ) def forward(self, x): return self.net(x) class Block(nn.Module): def __init__(self, n_embd, num_heads, block_size, dropout): super().__init__() head_size = n_embd // num_heads self.sa = MultiHeadAttention( n_embd, num_heads, head_size, block_size, dropout ) self.ffwd = FeedForward(n_embd, dropout) self.ln1 = nn.LayerNorm(n_embd) self.ln2 = nn.LayerNorm(n_embd) def forward(self, x): x = x + self.sa(self.ln1(x)) x = x + self.ffwd(self.ln2(x)) return x class GPTLanguageModel(nn.Module): def __init__( self, vocab_size, block_size, n_embd=256, n_layer=6, num_heads=8, dropout=0.1 ): super().__init__() self.block_size = block_size self.token_embedding_table = nn.Embedding(vocab_size, n_embd) self.position_embedding_table = nn.Embedding(block_size, n_embd) self.blocks = nn.Sequential(*[ Block(n_embd, num_heads, block_size, dropout) for _ in range(n_layer) ]) self.ln_f = nn.LayerNorm(n_embd) self.lm_head = nn.Linear(n_embd, vocab_size) def forward(self, idx, targets=None): B, T = idx.shape tok_emb = self.token_embedding_table(idx) pos_emb = self.position_embedding_table( torch.arange(T, device=idx.device) ) x = tok_emb + pos_emb x = self.blocks(x) x = self.ln_f(x) logits = self.lm_head(x) if targets is None: loss = None else: B, T, C = logits.shape logits = logits.view(B * T, C) targets = targets.view(B * T) loss = F.cross_entropy(logits, targets) return logits, loss @torch.no_grad() def generate(self, idx, max_new_tokens): for _ in range(max_new_tokens): idx_cond = idx[:, -self.block_size:] logits, loss = self(idx_cond) logits = logits[:, -1, :] probs = F.softmax(logits, dim=-1) idx_next = torch.multinomial(probs, num_samples=1) idx = torch.cat((idx, idx_next), dim=1) return idx

这个小模型已经包含了GPT的主要结构：

Token Embedding
Position Embedding
Masked Multi-Head Self-Attention
Feed Forward Network
Residual Connection
LayerNorm
Linear输出层

model.py是模型结构定义文件，也是后续最容易做模型创新的地方。如果只是跑通训练流程，直接使用标准Decoder-only Transformer就可以。但如果后面想做自己的模型改进，可以改写模型架构代码。例如改Attention机制、改位置编码方式、调整Transformer Block结构、增加新的模块，或者改变模型深度、宽度和上下文长度。

七、模型训练

新建train.py：

# train.py import torch from tokenizers import ByteLevelBPETokenizer from model import GPTLanguageModel from tqdm import tqdm batch_size = 32 block_size = 128 max_steps = 10000 eval_interval = 500 learning_rate = 3e-4 n_embd = 768 n_layer = 12 num_heads = 12 block_size = 512 dropout = 0.1 device = "cuda" if torch.cuda.is_available() else "cpu" tokenizer = ByteLevelBPETokenizer( "tokenizer/vocab.json", "tokenizer/merges.txt" ) vocab_size = tokenizer.get_vocab_size() with open("data/processed/train.txt", "r", encoding="utf-8") as f: train_text = f.read() with open("data/processed/validation.txt", "r", encoding="utf-8") as f: val_text = f.read() train_ids = tokenizer.encode(train_text).ids val_ids = tokenizer.encode(val_text).ids train_data = torch.tensor(train_ids, dtype=torch.long) val_data = torch.tensor(val_ids, dtype=torch.long) def get_batch(split): data = train_data if split == "train" else val_data ix = torch.randint(len(data) - block_size - 1, (batch_size,)) x = torch.stack([ data[i:i + block_size] for i in ix ]) y = torch.stack([ data[i + 1:i + block_size + 1] for i in ix ]) return x.to(device), y.to(device) @torch.no_grad() def estimate_loss(): model.eval() out = {} for split in ["train", "val"]: losses = torch.zeros(20) for k in range(20): x, y = get_batch(split) logits, loss = model(x, y) losses[k] = loss.item() out[split] = losses.mean() model.train() return out model = GPTLanguageModel( vocab_size=vocab_size, block_size=block_size, n_embd=n_embd, n_layer=n_layer, num_heads=num_heads, dropout=dropout ).to(device) optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate) for step in tqdm(range(max_steps)): if step % eval_interval == 0: losses = estimate_loss() print( f"step {step}: " f"train loss {losses['train']:.4f}, " f"val loss {losses['val']:.4f}" ) xb, yb = get_batch("train") logits, loss = model(xb, yb) optimizer.zero_grad(set_to_none=True) loss.backward() optimizer.step() torch.save(model.state_dict(), "checkpoints/mini_gpt.pt")

运行：

mkdir -p checkpoints python train.py

训练过程中会看到类似输出：

step 0: train loss 8.9872, val loss 8.9845 step 500: train loss 6.3128, val loss 6.4217 step 1000: train loss 5.4821, val loss 5.6534 step 5000: train loss 3.8924, val loss 4.1027 step 10000: train loss 3.2158, val loss 3.5879

Loss下降说明模型正在学习训练语料中的语言规律。这里的训练不是让模型背诵文本，而是让它不断学习在当前上下文中，下一个Token最可能是什么。

这里要重点强调一下，大语言模型训练和传统机器学习训练有一个明显区别，它通常不会像传统模型那样反复训练很多个epoch。

在传统机器学习任务中，数据集规模相对有限，所以我们经常会让模型在同一份训练集上反复学习很多轮。例如图像分类、表格预测或者普通神经网络训练，可能会设置50 个epoch、100个epoch，让模型多次看完整个数据集。

但大语言模型的预训练语料非常庞大，通常以Token总量来衡量训练规模。比如一个模型训练了300B Token、1T Token或10T Token，本质上指的是模型在训练过程中一共处理了多少个Token。

因此，大语言模型训练时更常关注的是：

训练了多少step
每个step处理多少Token
总共消耗了多少Token

而不是简单地说训练了多少个epoch。如果数据规模足够大，模型可能只完整看一遍语料，甚至还没有完整看完所有数据，训练就已经结束了。这和传统小数据集上反复训练多个epoch的方式完全不同。

八、生成文本测试

训练完成后，可以写一个generate.py测试模型效果。

# generate.py import torch from tokenizers import ByteLevelBPETokenizer from model import GPTLanguageModel n_embd = 768 n_layer = 12 num_heads = 12 block_size = 512 dropout = 0.1 device = "cuda" if torch.cuda.is_available() else "cpu" tokenizer = ByteLevelBPETokenizer( "tokenizer/vocab.json", "tokenizer/merges.txt" ) vocab_size = tokenizer.get_vocab_size() model = GPTLanguageModel( vocab_size=vocab_size, block_size=block_size, n_embd=n_embd, n_layer=n_layer, num_heads=num_heads, dropout=dropout ).to(device) model.load_state_dict( torch.load("checkpoints/mini_gpt.pt", map_location=device) ) model.eval() prompt = "Artificial intelligence" input_ids = tokenizer.encode(prompt).ids idx = torch.tensor([input_ids], dtype=torch.long).to(device) out = model.generate(idx, max_new_tokens=200) generated_ids = out[0].tolist() generated_text = tokenizer.decode(generated_ids) print(generated_text)

运行python generate.py，模型会基于输入开头继续生成文本。如果模型规模较小、训练数据较少，生成内容不会特别稳定，可能有语法错误或重复内容。

十、如果继续放大，需要改哪些地方？

如果要把这个小模型继续放大，需要重点改这些部分。

第一，数据规模要扩大

WikiText-2只适合演示流程，真正训练大模型需要更大规模、更高质量的数据，例如网页语料、论文语料、代码语料、专业领域语料等。

第二，Tokenizer要重新训练

如果是中文、英文、代码、分子 SMILES 或专业领域文本，最好基于自己的语料训练 Tokenizer，而不是直接套用通用词表。

第三，模型参数要增加

可以逐步增大：

n_embd n_layer num_heads block_size

例如：

n_embd = 1024 n_layer = 24 num_heads = 12 block_size = 1024

第四，训练效率要优化

模型变大后，普通训练方式会很慢，需要引入：

mixed precision gradient accumulation learning rate scheduler checkpoint resume distributed training FlashAttention DeepSpeed 或 FSDP

第五，要增加训练监控

真实训练不能只看Loss，还要记录：

train loss validation loss learning rate GPU memory tokens per second checkpoint sample generation

这些指标可以用TensorBoard、Weights & Biases或日志文件记录。

大模型训练的复杂性，不在于单个步骤多神秘，而在于每个步骤都被放大了很多倍。小模型用一个公开数据集就能跑通，真正的大模型需要海量数据、分布式训练、显存优化、训练监控和大量工程细节。

但只要理解了这个最小流程，再看千亿参数大模型训练，本质上就清楚了。所谓大模型训练，就是把数据处理、Tokenizer、Transformer、Loss和优化器这几件事，在更大的数据和算力上执行。

现在再单纯卷大模型参数，已经不是最现实的路线了。参数越大，训练和推理成本越高，部署门槛也越高；而且通用大模型虽然能力强，但未必真正适合具体行业和科研任务。

对多数团队来说，更有价值的方向是**垂域模型，**不盲目追求更大的参数，而是基于自己的领域数据，让模型更懂特定场景。

比如环境科学、化学分子、生物医药、材料设计、法律金融等领域，真正需要的往往不是一个什么都懂一点的通用模型，而是一个能理解专业术语、数据结构和任务需求的领域模型。

所以，这篇从0.1B小模型开始，并不是为了和千亿参数模型比能力，而是为了跑通完整训练流程。流程打通后，就可以把通用文本替换成自己的领域数据，进一步训练面向具体任务的垂域模型。

学AI大模型的正确顺序，千万不要搞错了

🤔2026年AI风口已来！各行各业的AI渗透肉眼可见，超多公司要么转型做AI相关产品，要么高薪挖AI技术人才，机遇直接摆在眼前！

有往AI方向发展，或者本身有后端编程基础的朋友，直接冲AI大模型应用开发转岗超合适！

就算暂时不打算转岗，了解大模型、RAG、Prompt、Agent这些热门概念，能上手做简单项目，也绝对是求职加分王🔋

📝给大家整理了超全最新的AI大模型应用开发学习清单和资料，手把手帮你快速入门！👇👇

学习路线:

✅大模型基础认知—大模型核心原理、发展历程、主流模型（GPT、文心一言等）特点解析
✅核心技术模块—RAG检索增强生成、Prompt工程实战、Agent智能体开发逻辑
✅开发基础能力—Python进阶、API接口调用、大模型开发框架（LangChain等）实操
✅应用场景开发—智能问答系统、企业知识库、AIGC内容生成工具、行业定制化大模型应用
✅项目落地流程—需求拆解、技术选型、模型调优、测试上线、运维迭代
✅面试求职冲刺—岗位JD解析、简历AI项目包装、高频面试题汇总、模拟面经

以上6大模块，看似清晰好上手，实则每个部分都有扎实的核心内容需要吃透！

我把大模型的学习全流程已经整理📚好了！抓住AI时代风口，轻松解锁职业新可能，希望大家都能把握机遇，实现薪资/职业跃迁～