第五章-自己搭建大模型_第3部分-预训练一个小型LLM
总目录
目录
- 预训练一个小型LLM
- 3.1 数据下载与处理
- 3.2 构建 Dataset
- 3.3 预训练流程
- 3.4 SFT 微调
- 3.5 模型推理与生成
- 总结与资源](#4-总结与资源)
3. 预训练一个小型LLM
现在进入实战环节：训练一个约2亿参数的小型语言模型。我们将经历完整的数据准备、预训练、监督微调和推理流程。
3.1 数据下载与处理
数据集选择
预训练数据：出门问问序列猴子数据集
- 规模：约10B tokens
- 来源：网页、百科、代码、书籍等
- 特点：高质量中文语料
SFT数据：BelleGroup中文对话数据集
- 规模：350万条对话
- 格式：人机对话
- 特点：覆盖多种任务场景
下载脚本
import os
# 下载预训练数据
os.system(
"modelscope download --dataset ddzhu123/seq-monkey "
"mobvoi_seq_monkey_general_open_corpus.jsonl.tar.bz2 "
"--local_dir ./data"
)
# 解压
os.system("tar -xvf ./data/mobvoi_seq_monkey_general_open_corpus.jsonl.tar.bz2")
# 下载SFT数据
os.system(
"huggingface-cli download --repo-type dataset "
"BelleGroup/train_3.5M_CN --local-dir ./BelleGroup"
)
数据预处理
import json
from tqdm import tqdm
def split_text(text: str, chunk_size: int = 512) -> list:
"""将长文本切分为固定长度的块"""
return [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]
# 处理预训练数据
print("处理预训练数据...")
with open('seq_monkey_processed.jsonl', 'w', encoding='utf-8') as out_f:
with open('mobvoi_seq_monkey_general_open_corpus.jsonl', 'r', encoding='utf-8') as in_f:
for line in tqdm(in_f):
data = json.loads(line)
text = data['text']
# 切分长文本
chunks = split_text(text, chunk_size=512)
for chunk in chunks:
out_f.write(json.dumps({'text': chunk}, ensure_ascii=False) + '\n')
# 处理SFT数据
print("处理SFT数据...")
def convert_to_messages(conversations: list) -> list:
"""转换为标准消息格式"""
messages = [{"role": "system", "content": "你是一个AI助手"}]
for conv in conversations:
if conv['from'] == 'human':
messages.append({'role': 'user', 'content': conv['value']})
elif conv['from'] == 'assistant':
messages.append({'role': 'assistant', 'content': conv['value']})
return messages
with open('belle_processed.jsonl', 'w', encoding='utf-8') as out_f:
with open('BelleGroup/train_3.5M_CN.json', 'r', encoding='utf-8') as in_f:
for line in tqdm(in_f):
data = json.loads(line)
messages = convert_to_messages(data['conversations'])
out_f.write(json.dumps(messages, ensure_ascii=False) + '\n')
print("数据处理完成！")
3.2 构建 Dataset
PretrainDataset（预训练数据集）
import torch
import numpy as np
from torch.utils.data import Dataset
class PretrainDataset(Dataset):
"""
预训练数据集：自回归语言建模
目标：学习预测下一个token
"""
def __init__(self, data_path: str, tokenizer, max_length: int = 512):
super().__init__()
self.tokenizer = tokenizer
self.max_length = max_length
self.padding_id = 0
# 加载所有数据到内存
with open(data_path, 'r', encoding='utf-8') as f:
self.data = f.readlines()
def __len__(self):
return len(self.data)
def __getitem__(self, index: int):
# 解析数据
sample = json.loads(self.data[index])
text = f"{self.tokenizer.bos_token}{sample['text']}"
# Tokenize
input_ids = self.tokenizer(text, add_special_tokens=False)['input_ids']
input_ids = input_ids[:self.max_length]
# Padding
seq_len = len(input_ids)
padding_len = self.max_length - seq_len
input_ids = input_ids + [self.padding_id] * padding_len
# Loss mask：标记哪些位置需要计算损失
loss_mask = [1] * seq_len + [0] * padding_len
# 构造输入和目标
# X: [BOS, T1, T2, ..., Tn-1]
# Y: [T1, T2, T3, ..., Tn]
X = np.array(input_ids[:-1], dtype=np.int64)
Y = np.array(input_ids[1:], dtype=np.int64)
loss_mask = np.array(loss_mask[1:], dtype=np.int64)
return torch.from_numpy(X), torch.from_numpy(Y), torch.from_numpy(loss_mask)
数据流示意：
原始: [BOS, T1, T2, T3, T4, PAD, PAD, PAD]
X: [BOS, T1, T2, T3, T4, PAD, PAD]
Y: [T1, T2, T3, T4, PAD, PAD, PAD]
Mask: [1, 1, 1, 1, 0, 0, 0 ]
模型学习：给定X的前k个token，预测第k+1个token。
SFTDataset（监督微调数据集）
class SFTDataset(Dataset):
"""
SFT数据集：多轮对话
目标：只在assistant回复部分计算损失
"""
def __init__(self, data_path: str, tokenizer, max_length: int = 512):
super().__init__()
self.tokenizer = tokenizer
self.max_length = max_length
self.padding_id = 0
with open(data_path, 'r', encoding='utf-8') as f:
self.data = f.readlines()
def __len__(self):
return len(self.data)
def generate_loss_mask(self, input_ids: list) -> list:
"""
生成loss mask：只在assistant回复部分为1
识别模式：<|im_start|>assistant\n ... <|im_end|>
"""
mask = [0] * len(input_ids)
# <|im_start|>assistant\n 的token IDs（需根据实际tokenizer调整）
assistant_start = [3, 1074, 537, 500, 203] # 示例IDs
start_len = len(assistant_start)
i = 0
while i <= len(input_ids) - start_len:
# 检查是否匹配assistant开始标记
if input_ids[i:i+start_len] == assistant_start:
# 找到对应的结束标记 <|im_end|> (ID=4)
for j in range(i + start_len, len(input_ids)):
if input_ids[j] == 4: # <|im_end|>
# 标记assistant回复区间
for pos in range(i + start_len, j + 1):
mask[pos] = 1
break
i += start_len
else:
i += 1
return mask
def __getitem__(self, index: int):
# 解析对话
messages = json.loads(self.data[index])
# 应用聊天模板
text = self.tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=False
)
# Tokenize
input_ids = self.tokenizer(text, add_special_tokens=False)['input_ids']
input_ids = input_ids[:self.max_length]
# Padding
seq_len = len(input_ids)
padding_len = self.max_length - seq_len
input_ids = input_ids + [self.padding_id] * padding_len
# 生成loss mask
loss_mask = self.generate_loss_mask(input_ids)
# 构造输入输出
X = np.array(input_ids[:-1], dtype=np.int64)
Y = np.array(input_ids[1:], dtype=np.int64)
loss_mask = np.array(loss_mask[1:], dtype=np.int64)
return torch.from_numpy(X), torch.from_numpy(Y), torch.from_numpy(loss_mask)
SFT数据流示意：
对话:
<|im_start|>user\n你好<|im_end|>
<|im_start|>assistant\n你好！<|im_end|>
Mask:
[0, 0, ..., 0, 1, 1, 1, 1, 1, 0]
└─user─┘ └─assistant─┘
只有assistant的回复参与损失计算，这样模型只学习生成回复，不学习生成用户输入。
3.3 预训练流程
完整的预训练代码包含学习率调度、梯度累积、混合精度等技术。
import argparse
import time
import math
from torch import optim
from torch.utils.data import DataLoader
from contextlib import nullcontext
def get_lr(current_step: int, total_steps: int, warmup_steps: int, max_lr: float) -> float:
"""
余弦退火学习率调度
包含warmup阶段
"""
min_lr = max_lr / 10
# Warmup阶段：线性增长
if current_step < warmup_steps:
return max_lr * current_step / warmup_steps
# 训练结束后：保持最小学习率
if current_step > total_steps:
return min_lr
# 余弦退火阶段
progress = (current_step - warmup_steps) / (total_steps - warmup_steps)
cosine_decay = 0.5 * (1.0 + math.cos(math.pi * progress))
return min_lr + (max_lr - min_lr) * cosine_decay
def train_epoch(epoch: int, model, train_loader, optimizer, scaler, args):
"""训练一个epoch"""
model.train()
start_time = time.time()
for step, (X, Y, loss_mask) in enumerate(train_loader):
# 数据迁移到GPU
X = X.to(args.device)
Y = Y.to(args.device)
loss_mask = loss_mask.to(args.device)
# 动态学习率
current_step = epoch * len(train_loader) + step
total_steps = args.epochs * len(train_loader)
lr = get_lr(current_step, total_steps, args.warmup_steps, args.learning_rate)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
# 前向传播（混合精度）
with torch.cuda.amp.autocast(enabled=(args.dtype != 'float32')):
output = model(X, Y)
loss = output.last_loss / args.accumulation_steps
# 应用loss mask
loss_mask = loss_mask.view(-1)
loss = torch.sum(loss * loss_mask) / (loss_mask.sum() + 1e-8)
# 反向传播
scaler.scale(loss).backward()
# 梯度累积：每accumulation_steps步更新一次
if (step + 1) % args.accumulation_steps == 0:
# 梯度裁剪
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
# 更新参数
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True)
# 日志
if step % args.log_interval == 0:
elapsed = time.time() - start_time
eta = elapsed / (step + 1) * (len(train_loader) - step) / 60
print(f'Epoch {epoch+1}/{args.epochs} | '
f'Step {step}/{len(train_loader)} | '
f'Loss {loss.item() * args.accumulation_steps:.4f} | '
f'LR {lr:.2e} | '
f'ETA {eta:.1f}min')
# 定期保存
if (step + 1) % args.save_interval == 0:
save_checkpoint(model, args, step)
def save_checkpoint(model, args, step):
"""保存模型检查点"""
model.eval()
state_dict = model.module.state_dict() if hasattr(model, 'module') else model.state_dict()
checkpoint_path = f'{args.output_dir}/checkpoint_step{step}.pth'
torch.save(state_dict, checkpoint_path)
print(f'✓ 已保存检查点: {checkpoint_path}')
model.train()
# 主训练循环
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--data_path", type=str, required=True)
parser.add_argument("--output_dir", type=str, default="./output")
parser.add_argument("--epochs", type=int, default=1)
parser.add_argument("--batch_size", type=int, default=64)
parser.add_argument("--learning_rate", type=float, default=2e-4)
parser.add_argument("--warmup_steps", type=int, default=1000)
parser.add_argument("--accumulation_steps", type=int, default=8)
parser.add_argument("--grad_clip", type=float, default=1.0)
parser.add_argument("--dtype", type=str, default="bfloat16")
parser.add_argument("--num_workers", type=int, default=8)
parser.add_argument("--log_interval", type=int, default=100)
parser.add_argument("--save_interval", type=int, default=5000)
args = parser.parse_args()
# 设置设备
args.device = "cuda" if torch.cuda.is_available() else "cpu"
# 模型配置
config = ModelConfig(dim=1024, n_layers=18, vocab_size=6144)
# 初始化
torch.manual_seed(42)
model = Transformer(config).to(args.device)
tokenizer = AutoTokenizer.from_pretrained('./tokenizer_k/')
# 多GPU
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
print(f"使用 {torch.cuda.device_count()} 个GPU训练")
# 数据集
train_dataset = PretrainDataset(args.data_path, tokenizer)
train_loader = DataLoader(
train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True
)
# 优化器和混合精度
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype != 'float32'))
# 开始训练
print(f"\n{'='*50}")
print(f"开始预训练")
print(f"{'='*50}")
print(f"模型参数: {sum(p.numel() for p in model.parameters())/1e6:.1f}M")
print(f"训练样本: {len(train_dataset)}")
print(f"Batch size: {args.batch_size}")
print(f"Effective batch size: {args.batch_size * args.accumulation_steps}")
print(f"{'='*50}\n")
for epoch in range(args.epochs):
train_epoch(epoch, model, train_loader, optimizer, scaler, args)
运行命令：
python train_pretrain.py \
--data_path ./seq_monkey_processed.jsonl \
--output_dir ./pretrain_output \
--batch_size 32 \
--accumulation_steps 8 \
--epochs 1 \
--dtype bfloat16
训练技巧：
-
梯度累积：模拟更大的batch size
- 有效batch = batch_size × accumulation_steps
- 例：32 × 8 = 256
-
混合精度：使用bfloat16加速训练
- 节省约50%显存
- 提速约2-3倍
-
学习率调度：warmup + 余弦退火
- Warmup避免训练初期梯度过大
- 余弦退火帮助收敛
-
梯度裁剪：防止梯度爆炸
- 将梯度范数限制在1.0以内
资源估算：
| 配置 | 显存 | 预计时间 |
|---|---|---|
| 单卡(batch=4) | ~7GB | 500h+ |
| 4卡(batch=32) | ~24GB | 120h |
| 8卡(batch=64) | ~48GB | 50h |
3.4 SFT 微调
SFT代码与预训练类似，主要区别：
- 使用SFTDataset
- 加载预训练权重
- 学习率通常更小（1e-5到5e-5）
# 加载预训练权重
def load_pretrained(model, checkpoint_path):
"""加载预训练权重"""
state_dict = torch.load(checkpoint_path, map_location='cpu')
# 清理可能的前缀
unwanted_prefix = '_orig_mod.'
for k in list(state_dict.keys()):
if k.startswith(unwanted_prefix):
state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
model.load_state_dict(state_dict, strict=False)
print(f"✓ 已加载预训练权重: {checkpoint_path}")
# SFT训练主循环
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--pretrain_checkpoint", type=str, required=True)
parser.add_argument("--data_path", type=str, required=True)
parser.add_argument("--output_dir", type=str, default="./sft_output")
parser.add_argument("--learning_rate", type=float, default=2e-5)
# ... 其他参数同预训练 ...
args = parser.parse_args()
# 初始化模型
config = ModelConfig(dim=1024, n_layers=18, vocab_size=6144)
model = Transformer(config).to(args.device)
# 加载预训练权重
load_pretrained(model, args.pretrain_checkpoint)
# 使用SFT数据集
tokenizer = AutoTokenizer.from_pretrained('./tokenizer_k/')
train_dataset = SFTDataset(args.data_path, tokenizer)
train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
# 优化器（注意学习率更小）
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
# 训练
for epoch in range(args.epochs):
train_epoch(epoch, model, train_loader, optimizer, scaler, args)
运行命令：
python train_sft.py \
--pretrain_checkpoint ./pretrain_output/checkpoint_step50000.pth \
--data_path ./belle_processed.jsonl \
--output_dir ./sft_output \
--learning_rate 2e-5 \
--batch_size 32 \
--epochs 1
SFT注意事项：
- 学习率：SFT阶段学习率应比预训练小5-10倍，避免破坏预训练知识
- Epoch数：通常1-3个epoch即可，过多会过拟合
- 数据质量：SFT数据质量比数量更重要
3.5 模型推理与生成
训练完成后，我们可以使用模型进行文本生成。
import torch
from transformers import AutoTokenizer
from k_model import Transformer, ModelConfig
class TextGenerator:
"""文本生成器"""
def __init__(
self,
checkpoint_path: str,
tokenizer_path: str = './tokenizer_k/',
device: str = None
):
self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
# 加载tokenizer
self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
# 加载模型
config = ModelConfig(dim=1024, n_layers=18, vocab_size=6144)
self.model = Transformer(config)
state_dict = torch.load(checkpoint_path, map_location=self.device)
unwanted_prefix = '_orig_mod.'
for k in list(state_dict.keys()):
if k.startswith(unwanted_prefix):
state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
self.model.load_state_dict(state_dict, strict=False)
self.model.eval()
self.model.to(self.device)
# 统计参数
num_params = sum(p.numel() for p in self.model.parameters())
print(f"模型参数量: {num_params/1e6:.1f}M")
def generate(
self,
prompt: str,
max_new_tokens: int = 256,
temperature: float = 0.7,
top_k: int = 50,
use_chat_template: bool = True
) -> str:
"""
生成文本
Args:
prompt: 输入提示
max_new_tokens: 最多生成token数
temperature: 采样温度（0=贪婪，1=完全随机）
top_k: Top-K采样
use_chat_template: 是否使用聊天模板
Returns:
生成的文本
"""
# 应用聊天模板
if use_chat_template:
messages = [
{"role": "system", "content": "你是一个AI助手"},
{"role": "user", "content": prompt}
]
text = self.tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True
)
else:
text = prompt
# 编码
input_ids = self.tokenizer(text, return_tensors='pt')['input_ids']
input_ids = input_ids.to(self.device)
# 生成
with torch.no_grad():
output_ids = self.model.generate(
input_ids,
stop_id=self.tokenizer.eos_token_id,
max_new_tokens=max_new_tokens,
temperature=temperature,
top_k=top_k
)
# 解码
generated_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
return generated_text
# 使用示例
if __name__ == "__main__":
print("="*60)
print("LLaMA2 文本生成演示")
print("="*60)
# 初始化生成器
generator = TextGenerator(
checkpoint_path='./sft_output/checkpoint_final.pth'
)
# 测试问题
test_prompts = [
"你好，请介绍一下你自己",
"什么是深度学习？",
"中国的首都是哪里？",
"请用Python写一个冒泡排序",
"1+1等于几？"
]
for i, prompt in enumerate(test_prompts, 1):
print(f"\n{'='*60}")
print(f"问题 {i}: {prompt}")
print(f"{'-'*60}")
response = generator.generate(
prompt,
max_new_tokens=200,
temperature=0.7,
top_k=50
)
print(f"回答: {response}")
预期输出：
============================================================
LLaMA2 文本生成演示
============================================================
模型参数量: 370.5M
============================================================
问题 1: 你好,请介绍一下你自己
------------------------------------------------------------
回答: <|im_start|>system
你是一个AI助手<|im_end|>
<|im_start|>user
你好,请介绍一下你自己<|im_end|>
<|im_start|>assistant
你好!我是一个AI助手,很高兴认识你。我可以帮助你回答问题、提供信息、进行对话交流等。我的目标是为用户提供有用、准确的帮助。有什么我可以帮助你的吗?<|im_end|>
============================================================
问题 2: 什么是深度学习?
------------------------------------------------------------
回答: <|im_start|>system
你是一个AI助手<|im_end|>
<|im_start|>user
什么是深度学习?<|im_end|>
<|im_start|>assistant
深度学习是机器学习的一个分支,它使用多层神经网络来学习数据的表示。深度学习模型可以自动从原始数据中提取特征,而不需要人工特征工程。常见的深度学习模型包括卷积神经网络(CNN)、循环神经网络(RNN)和Transformer等。<|im_end|>
============================================================
问题 3: 中国的首都是哪里?
------------------------------------------------------------
回答: <|im_start|>system
你是一个AI助手<|im_end|>
<|im_start|>user
中国的首都是哪里?<|im_end|>
<|im_start|>assistant
中国的首都是北京。<|im_end|>
============================================================
问题 4: 请用Python写一个冒泡排序
------------------------------------------------------------
回答: <|im_start|>system
你是一个AI助手<|im_end|>
<|im_start|>user
请用Python写一个冒泡排序<|im_end|>
<|im_start|>assistant
好的,这是Python实现的冒泡排序:
```python
def bubble_sort(arr):
n = len(arr)
for i in range(n):
for j in range(0, n-i-1):
if arr[j] > arr[j+1]:
arr[j], arr[j+1] = arr[j+1], arr[j]
return arr
# 测试
arr = [64, 34, 25, 12, 22, 11, 90]
print(bubble_sort(arr))
```<|im_end|>
============================================================
问题 5: 1+1等于几?
------------------------------------------------------------
回答: <|im_start|>system
你是一个AI助手<|im_end|>
<|im_start|>user
1+1等于几?<|im_end|>
<|im_start|>assistant
1+1等于2。<|im_end|>
采样策略对比：
| 策略 | Temperature | Top-K | 特点 |
|---|---|---|---|
| 贪婪解码 | 0.0 | - | 确定性，可能重复 |
| 标准采样 | 1.0 | - | 多样性高，可能不连贯 |
| 温度采样 | 0.7 | - | 平衡多样性和连贯性 |
| Top-K | 0.7 | 50 | 避免低概率词，更流畅 |
4. 总结与资源
恭喜你完成了从零实现LLaMA2大模型的完整旅程！让我们回顾一下核心要点。
4.1 核心技术总结
模型架构创新：
- RMSNorm：去除均值中心化，提升7-10%计算效率
- RoPE：旋转位置编码，更好的长度外推能力
- GQA：分组查询注意力，显存节省40%，推理提速2倍
- SwiGLU：改进的激活函数，提升模型表达能力
训练技巧：
- 混合精度训练：节省50%显存，提速2-3倍
- 梯度累积：模拟大batch训练
- 学习率调度：Warmup + 余弦退火
- 梯度裁剪：稳定训练过程
数据处理：
- BPE Tokenizer：平衡词表大小和分词粒度
- 预训练vs SFT：自回归建模 vs 对话对齐
- Loss Masking：精确控制学习目标
4.2 性能指标
我们训练的模型达到：
| 指标 | 数值 |
|---|---|
| 参数量 | 215M |
| 词表大小 | 6,144 |
| 最大序列长度 | 512 |
| 训练数据 | 10B tokens (预训练) + 3.5M对话 (SFT) |
| 训练时长 | 46h (预训练) + 24h (SFT) |
| 硬件 | 8×4090 |
4.3 进阶方向
模型优化：
- 量化：INT8/INT4量化降低部署成本
- 剪枝：移除冗余参数
- 知识蒸馏：从大模型学习到小模型
- LoRA微调：低秩适配器高效微调
能力增强：
- 多模态：扩展到视觉、音频
- 长文本：扩展上下文长度到32k+
- 工具调用：集成外部工具和API
- 思维链：增强推理能力
部署优化：
- ONNX导出：跨平台部署
- TensorRT加速：GPU推理优化
- 量化部署：移动端部署
- 分布式推理：模型并行
4.4 常见问题
Q: 显存不足怎么办？
A:
- 减小batch_size（最低可到2）
- 增大accumulation_steps（保持有效batch不变）
- 使用梯度检查点（gradient checkpointing）
- 使用更小的模型配置
Q: 训练loss不下降？
A:
- 检查数据格式和loss_mask
- 降低学习率
- 增加warmup步数
- 检查梯度范数（是否exploding/vanishing）
Q: 生成效果不好？
A:
- 确保SFT数据质量
- 调整采样参数（temperature, top_k）
- 增加训练数据量
- 尝试多轮对话微调
Q: 如何评估模型？
A:
- 困惑度（Perplexity）：衡量语言建模能力
- BLEU/ROUGE：衡量生成质量
- 人工评估：最终标准
4.5 资源链接
预训练模型下载：
- ModelScope: https://modelscope.cn/
- HuggingFace: https://huggingface.co/
学习资源：
- Andrej Karpathy的视频教程
- HuggingFace NLP课程
- Stanford CS224N课程
开源项目：
- llama2.c: https://github.com/karpathy/llama2.c
- minimind: https://github.com/jingyaogong/minimind
- transformers: https://github.com/huggingface/transformers
4.6 致谢
本教程基于以下优秀开源项目：
- Meta的LLaMA2模型架构
- Andrej Karpathy的llama2.c项目
- HuggingFace的transformers库
- BelleGroup的中文对话数据集
感谢开源社区的无私贡献！
附录
A. 完整代码结构
project/
├── k_model.py # 模型定义
├── dataset.py # 数据集定义
├── train_tokenizer.py # Tokenizer训练
├── train_pretrain.py # 预训练脚本
├── train_sft.py # SFT训练脚本
├── model_sample.py # 推理脚本
├── tokenizer_k/ # Tokenizer文件
│ ├── tokenizer.json
│ ├── tokenizer_config.json
│ └── special_tokens_map.json
├── data/ # 数据目录
│ ├── seq_monkey_processed.jsonl
│ └── belle_processed.jsonl
├── pretrain_output/ # 预训练输出
│ └── checkpoint_*.pth
└── sft_output/ # SFT输出
└── checkpoint_*.pth
B. 环境配置
# 创建虚拟环境
conda create -n llm python=3.10
conda activate llm
# 安装PyTorch（根据CUDA版本选择）
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
# 安装其他依赖
pip install transformers tokenizers datasets
pip install numpy tqdm
# 可选：安装实验跟踪工具
pip install swanlab wandb
C. 调试技巧
# 1. 检查数据
dataset = PretrainDataset('data.jsonl', tokenizer)
X, Y, mask = dataset[0]
print(f"X shape: {X.shape}")
print(f"Y shape: {Y.shape}")
print(f"Mask sum: {mask.sum()}") # 应该 > 0
# 2. 检查模型输出
model.eval()
with torch.no_grad():
output = model(X.unsqueeze(0), Y.unsqueeze(0))
print(f"Logits shape: {output.logits.shape}")
print(f"Loss shape: {output.last_loss.shape if output.last_loss is not None else None}")
# 3. 检查梯度
for name, param in model.named_parameters():
if param.grad is not None:
print(f"{name}: grad_norm={param.grad.norm():.4f}")
# 4. 可视化注意力
# 在Attention.forward中添加：
# self.attention_weights = scores # 保存注意力权重
# 然后可视化：
import matplotlib.pyplot as plt
plt.imshow(model.layers[0].attention.attention_weights[0, 0].cpu())
plt.colorbar()
plt.show()
最后的话
从零实现一个大模型是一段充满挑战但极具收获的旅程。你不仅学会了代码实现，更重要的是理解了大模型背后的原理和工程实践。
共勉：
- 理论是基础，实践是关键
- 遇到问题不要气馁，debug调试是学习的一部分
- 保持好奇心，持续关注领域进展
- 分享你的经验，帮助他人成长
祝你在AI的道路上越走越远！
参考文献
[1] Touvron, H., et al. (2023). LLaMA: Open and Efficient Foundation Language Models.
[2] Touvron, H., et al. (2023). LLaMA 2: Open Foundation and Fine-Tuned Chat Models.
[3] Vaswani, A., et al. (2017). Attention is All You Need.
[4] Su, J., et al. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding.
[5] Ainslie, J., et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.
[6] Shazeer, N. (2020). GLU Variants Improve Transformer.
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
