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

今天我们就结合一份具体数据完整走一遍从环境准备、数据处理、模型构建到模型训练的全过程训练一个参数量约0.1B也就是1亿参数的小型语言模型。需要说明的是小型语言模型和大语言模型的底层训练逻辑是一致的都是基于大量文本数据进行Token预测训练。区别主要在于大语言模型数据规模更大、模型参数更多、算力需求更高以及训练优化策略更复杂。这里使用一个真实公开数据集完整走一遍小型GPT模型的训练过程。核心流程包括环境准备下载真实数据集数据清洗与格式化构建Tokenizer构造训练样本搭建Transformer模型模型训练文本生成测试一、环境准备首先需要准备Python深度学习环境。最基础的环境包括conda create -n mini_llm python3.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_okTrue) 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(fdata/processed/{split_name}.txt, w, encodingutf-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( filespaths, vocab_size8000, min_frequency2, special_tokens[ pad, s, /s, unk, mask, ], ) os.makedirs(tokenizer, exist_okTrue) 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]也就是标签整体向右移动一位。在训练代码里可以这样构造batchimport 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, biasFalse) self.query nn.Linear(n_embd, head_size, biasFalse) self.value nn.Linear(n_embd, head_size, biasFalse) 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_embd256, n_layer6, num_heads8, dropout0.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, targetsNone): B, T idx.shape tok_emb self.token_embedding_table(idx) pos_emb self.position_embedding_table( torch.arange(T, deviceidx.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_samples1) idx torch.cat((idx, idx_next), dim1) return idx这个小模型已经包含了GPT的主要结构Token EmbeddingPosition EmbeddingMasked Multi-Head Self-AttentionFeed Forward NetworkResidual ConnectionLayerNormLinear输出层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, encodingutf-8) as f: train_text f.read() with open(data/processed/validation.txt, r, encodingutf-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, dtypetorch.long) val_data torch.tensor(val_ids, dtypetorch.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_sizevocab_size, block_sizeblock_size, n_embdn_embd, n_layern_layer, num_headsnum_heads, dropoutdropout ).to(device) optimizer torch.optim.AdamW(model.parameters(), lrlearning_rate) for step in tqdm(range(max_steps)): if step % eval_interval 0: losses estimate_loss() print( fstep {step}: ftrain loss {losses[train]:.4f}, fval loss {losses[val]:.4f} ) xb, yb get_batch(train) logits, loss model(xb, yb) optimizer.zero_grad(set_to_noneTrue) 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.5879Loss下降说明模型正在学习训练语料中的语言规律。这里的训练不是让模型背诵文本而是让它不断学习在当前上下文中下一个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_sizevocab_size, block_sizeblock_size, n_embdn_embd, n_layern_layer, num_headsnum_heads, dropoutdropout ).to(device) model.load_state_dict( torch.load(checkpoints/mini_gpt.pt, map_locationdevice) ) model.eval() prompt Artificial intelligence input_ids tokenizer.encode(prompt).ids idx torch.tensor([input_ids], dtypetorch.long).to(device) out model.generate(idx, max_new_tokens200) 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时代风口轻松解锁职业新可能希望大家都能把握机遇实现薪资/职业跃迁这份完整版的大模型 AI 学习资料已经上传CSDN朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费】