MiniMind 是一个开源的小型大语言模型项目，采用 Transformer 架构实现。该项目完整实现了从预训练到监督微调的训练流程，是学习和理解大语言模型训练的优秀案例。

本文详细解析 MiniMind 的训练代码，包括 train_pretrain.py（预训练）和 train_full_sft.py（全量监督微调）两个训练脚本，涵盖训练流程、关键技术、模型架构以及实践经验。文档采用循序渐进的方式，从整体概览到详细实现，适合不同水平的读者学习参考。

实际使用 AnyDev 上提供的 NVIDIA L20 的开发机按照文档指引，进行预训练和全量监督微调，都只需要2小时。

一、整体概览

MiniMind 的训练分为两个阶段：

1. 预训练（Pretrain）: 从零开始，让模型学习语言的统计规律

   • 技术任务：预测下一个 token
   • 学习目标：掌握语法、常识、推理能力，学习"如何理解语言"——通用语言基础

2. 全量监督微调（SFT）: 基于预训练模型，学习对话和指令遵循

   • 技术任务：预测下一个 token（只有 assistant 的回答部分）
   • 学习目标：掌握对话格式、指令遵循，学习"如何回答问题"——实用对话技能

两阶段的训练流程和关键技术相同：

• 模型: MiniMindForCausalLM - Transformer 架构的因果语言模型
• 训练技术: 支持单卡/多卡 DDP、混合精度训练（AMP）、梯度累积
• 优化器: AdamW + 余弦退火学习率调度
• 损失函数: CrossEntropyLoss（自回归任务）
• 监控工具: SwanLab（类似 Weights & Biases，用于追踪训练指标）

graph LR
    A[数据准备] --> B[模型初始化]
    B --> C[数据加载]
    C --> D[训练循环]
    D --> E[保存模型]
    
    A -.-> A1[Dataset]
    B -.-> B1[创建模型]
    C -.-> C1[DataLoader]
    D -.-> D1[前向/反向传播]
    E -.-> E1[.pth文件]
    
    style A fill:#e1f5ff
    style B fill:#e1f5ff
    style C fill:#e1f5ff
    style D fill:#e1f5ff
    style E fill:#e1f5ff

尽管两个阶段目标和训练过程大致相似，但是在数据集等一些细节方面有些差异：

二、训练流程详解

说明: 本节基于 train_pretrain.py 的训练流程介绍，SFT 的训练流程基本相同，主要差异已在第一节的对比表中列出

Step 1: 初始化配置（第118-177行）

解析命令行参数

args = parser.parse_args()

关键参数:

• --epochs: 训练轮数（默认1）
• --batch_size: 批次大小（默认32）
• --learning_rate: 学习率（默认5e-4）
• --hidden_size: 模型维度（默认512）
• --num_hidden_layers: Transformer 层数（默认8）
• --use_moe: 是否使用专家混合（可选）
• --ddp: 是否启用分布式训练

配置模型

lm_config = MiniMindConfig(
    hidden_size=args.hidden_size,
    num_hidden_layers=args.num_hidden_layers,
    use_moe=args.use_moe
)

初始化分布式环境（如果启用）

if ddp:
    init_distributed_mode()  # 第105-114行
    # 设置每个GPU不同的随机种子
    torch.manual_seed(base_seed + rank)

初始化监控工具

代码中 import swanlab as wandb（第172行），实际使用的是 SwanLab，它是类似 Weights & Biases 的机器学习实验追踪工具，用于记录和可视化训练指标（loss、学习率、训练时间等）

if args.use_wandb:
    import swanlab as wandb  # 实际使用的是 SwanLab（国产ML实验追踪工具）
    wandb.init(project=args.wandb_project, name=args.wandb_run_name)

Step 2: 初始化模型和数据（第178-189行）

创建模型

model, tokenizer = init_model(lm_config)  # 第98-102行

初始化过程:

1. 加载 tokenizer（第99行）
2. 创建模型并移动到 GPU（第100行）
3. 打印可训练参数量（第101行）

创建数据集

train_ds = PretrainDataset(
    args.data_path,      # '../dataset/pretrain_hq.jsonl'
    tokenizer, 
    max_length=args.max_seq_len  # 默认512
)

自回归预训练原理:

预训练的目标是让模型学习预测下一个 token。数据处理过程：

1. 读取文本: 从 JSONL 文件读取纯文本

2. Token 编码: 使用 tokenizer 编码为 token IDs

3. 构造训练样本:

   # 原始序列
   input_ids = [token1, token2, token3, token4, pad, pad]
   
   # 构造自回归任务
   X = input_ids[:-1]      # [token1, token2, token3, token4, pad]
   Y = input_ids[1:]        # [token2, token3, token4, pad, pad]
   

4. 生成 loss_mask: 排除 padding，只对真实 token 计算损失

示例:

文本: "我 爱 你"

X (输入): [我, 爱]
Y (目标): [爱, 你]

模型学习：给定"我 爱"，预测下一个 token 是"你"

创建数据加载器

train_sampler = DistributedSampler(train_ds) if ddp else None
train_loader = DataLoader(
    train_ds,
    batch_size=args.batch_size,
    shuffle=(train_sampler is None),
    num_workers=args.num_workers,
    sampler=train_sampler  # DDP时确保每个GPU看到不同数据
)

关键点:

• DistributedSampler: DDP 模式下分配数据到不同GPU
• shuffle=False（当使用 sampler 时）: 由 sampler 控制顺序

Step 3: 设置优化器（第191-196行）

什么是优化器？ 优化器用于更新模型参数，使损失函数逐渐减小。训练过程的核心是：

1. 计算梯度（反向传播）
2. 优化器根据梯度更新参数
3. 重复上述过程，模型性能逐渐提升

MiniMind 使用 AdamW 优化器，它是 Adam 的改进版，在自适应学习率和权重衰减方面表现优异，是训练大语言模型的常用选择。

混合精度训练

scaler = torch.cuda.amp.GradScaler(
    enabled=(args.dtype in ['float16', 'bfloat16'])
)

作用:

• 前向传播用 bfloat16（更快、省内存）
• 反向传播时保持 float32 精度
• GradScaler 防止梯度下溢

优化器

optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)

DDP 包装模型

if ddp:
    # 忽略位置编码的buffer（避免同步开销）
    model._ddp_params_and_buffers_to_ignore = {"freqs_cos", "freqs_sin"}
    model = DistributedDataParallel(model, device_ids=[ddp_local_rank])

Step 4: 训练循环（第198-201行）

iter_per_epoch = len(train_loader)
for epoch in range(args.epochs):
    train_sampler and train_sampler.set_epoch(epoch)  # 每个epoch打乱数据
    train_epoch(epoch, wandb)

三、训练循环详解 —— train_epoch() 函数

train_epoch() 函数是训练的核心，每个 epoch 调用一次，循环处理所有数据。单步训练过程如下：

Step 1: 准备数据（第36-39行）

X, Y, loss_mask = next(train_loader)
X = X.to(device)
Y = Y.to(device)
loss_mask = loss_mask.to(device)

Step 2: 计算学习率并前向传播（第41-45行）

# 余弦退火学习率调度
lr = get_lr(epoch * iter_per_epoch + step, total_steps, base_lr)
optimizer.param_groups[-1]['lr'] = lr

# 混合精度前向传播
with ctx:  # torch.cuda.amp.autocast()
    res = model(X)

Step 3: 计算损失（第47-52行）

# 对每个位置计算交叉熵损失
loss_fct = nn.CrossEntropyLoss(reduction='none')
loss = loss_fct(res.logits.view(-1, vocab_size), Y.view(-1))
loss = loss.view(Y.size())

# 应用 loss_mask 排除 padding，然后平均
loss = (loss * loss_mask).sum() / loss_mask.sum()
loss += res.aux_loss  # MOE辅助损失（平衡专家负载）

# 为梯度累积缩放损失
loss = loss / args.accumulation_steps

关键点:

• reduction='none': 返回每个位置的损失，允许手动应用 mask
• loss_mask: 确保只有真实token参与损失计算
• aux_loss: 确保MOE各专家负载均衡

Step 4: 反向传播和梯度更新（第55-64行）

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)

Step 5: 日志和保存（第66-95行）

if step % args.log_interval == 0:
    Logger(f'Epoch:[{epoch}/{epochs}]({step}/{iter}) loss:{loss} lr:{lr}')
    wandb.log({"loss": loss, "lr": lr})

if step % args.save_interval == 0:
    state_dict = {k: v.half() for k, v in model.state_dict().items()}
    torch.save(state_dict, f'{save_dir}/pretrain_{hidden_size}.pth')

四、关键训练技术

1. 混合精度训练（AMP）

什么是混合精度训练？

混合精度训练（Automatic Mixed Precision, AMP）是一种加速训练的技术，在训练过程中同时使用 fp32（float32）和 fp16/bfloat16（float16）两种精度。模型的计算使用 fp16，但梯度更新和关键操作使用 fp32 保持精度。

为什么用 bfloat16?

• 更快: 计算比 fp32 快 2-3 倍
• 省内存: 内存占用减半
• 精度: 数值范围与 fp32 相似，不易溢出

实现:

with torch.cuda.amp.autocast():
    output = model(input)

2. 梯度累积（Gradient Accumulation）

什么是梯度累积？

梯度累积是一种突破 GPU 内存限制的技术。当 GPU 内存不足无法使用足够大的 batch size 时，可以先计算多个小 batch 的梯度并累积，然后一次性更新参数，模拟更大的 batch size。

问题: GPU 内存限制，batch size 太小

解决: 累积分批梯度

loss = loss / accumulation_steps  # 缩放
loss.backward()  # 累积梯度

if (step + 1) % accumulation_steps == 0:
    optimizer.step()  # 等效 batch_size *= accumulation_steps

效果: batch_size=32, accumulation_steps=8 → 等效 batch_size=256

3. 梯度裁剪（Gradient Clipping）

什么是梯度裁剪？

梯度裁剪是一种防止梯度爆炸的技术。当梯度范数超过预设阈值时，按比例缩放梯度，使其保持在合理范围内，保证训练过程的稳定性。

实现:

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

作用:

• 计算所有梯度的范数
• 如果 > 1.0，按比例缩放
• 防止梯度爆炸，保持训练稳定

4. Loss Mask 机制

什么是 Loss Mask？

Loss Mask 是一种控制损失计算范围的机制。通过 mask 数组标记哪些位置需要计算损失，哪些位置应该被忽略（如 padding）。这确保模型只学习有意义的内容，而不会被 padding 等干扰信息影响。

实现:

loss = (loss * loss_mask).sum() / loss_mask.sum()

目的:

• 只有真实token参与损失计算
• padding token 被忽略
• 避免模型学习填充符号

5. 学习率调度（Cosine Annealing）

什么是学习率调度？

学习率调度是指在训练过程中动态调整学习率的策略。一个好的学习率调度可以让模型更快收敛并获得更好的性能。MiniMind 使用 Warmup + 余弦退火策略，训练初期用较小的学习率稳定训练，然后按照余弦曲线逐渐衰减到接近 0。

策略: 余弦退火 + Warmup

graph LR
    A[开始训练] --> B[Warmup阶段]
    B --> C[学习率: lr/10 → lr]
    C --> D[Cosine Decay]
    D --> E[学习率: lr → 0]
    
    style B fill:#ffeb3b
    style D fill:#4caf50

学习率变化曲线:

• Warmup阶段: 从 lr/10 逐渐上升到 lr（训练前期稳定启动）
• Cosine Decay阶段: 从 lr 按余弦曲线衰减到接近 0（训练后期平滑收敛）

优点:

• 训练初期稳定（warmup）
• 训练后期平滑收敛（cosine）

6. 分布式训练（DDP）

什么是分布式训练？

分布式数据并行（Distributed Data Parallel, DDP）是一种在多 GPU 或多机器上并行训练模型的技术。每个 GPU 都有模型的完整副本，独立处理不同的数据批次，在反向传播时同步梯度，从而实现高效的并行训练。

工作原理:

GPU 0:  batch [0,1,2,3]   → forward → 梯度 sync → update
GPU 1:  batch [4,5,6,7]   → forward → 梯度 sync → update
GPU 2:  batch [8,9,10,11] → forward → 梯度 sync → update
GPU 3:  batch [12,13,14,15] → forward → 梯度 sync → update

关键:

• 每个GPU独立前向传播
• 反向传播时同步梯度
• DistributedSampler 确保数据不重复

五、Pretrain vs SFT 关键差异详解

说明: 前面已经概述了 Pretrain 和 SFT 的核心差异，这里进一步详解最重要的技术差异。完整对比见表（参见第一节）。

1. loss_mask 机制差异（最重要！）

Pretrain (简单):

# dataset/lm_dataset.py 第46行
loss_mask = (input_ids != pad_token_id)  # 排除padding

SFT (智能):

# dataset/lm_dataset.py 第84-100行
def _generate_loss_mask(self, input_ids):
    loss_mask = [0] * len(input_ids)  # 先全为0
    
    # 找到 <|assistant|> 标记
    # 只对assistant的回答部分设为1
    # 忽略system和user的内容

可视化:

Pretrain: [1 1 1 1 0 0]        # 排除padding
SFT:      [0 0 0 0 1 1 1]      # 只计算assistant回答

效果: SFT让模型只学习如何回答，不学习如何提问

2. 模型初始化差异

Pretrain (第98-102行):

model = MiniMindForCausalLM(lm_config).to(device)  # 随机初始化

SFT (第96-106行):

model = MiniMindForCausalLM(lm_config)
# 加载预训练权重！
ckp = f'{save_dir}/pretrain_{hidden_size}.pth'
state_dict = torch.load(ckp, map_location=device)
model.load_state_dict(state_dict, strict=False)  # 加载权重
model = model.to(device)

3. 学习率差异的原因

执行命令

# 1. 先预训练
python trainer/train_pretrain.py \
    --learning_rate 5e-4 \
    --batch_size 32 \
    --data_path dataset/pretrain_hq.jsonl

# 2. 再SFT（自动加载预训练权重）
python trainer/train_full_sft.py \
    --learning_rate 5e-7 \
    --batch_size 16 \
    --data_path dataset/sft_mini_512.jsonl

为什么需要两步？

• Pretrain: 从零学习语言基础（需要大数据量）
• SFT: 基于预训练模型学习对话（小数据量，微调即可）
• 跳过pretrain直接用SFT → 效果很差（缺乏语言基础）

六、模型架构简析

这部分详细内容在 model/model_minimind.py

整体结构

Input (token_ids)
    ↓
Embedding          # 第388行: self.embed_tokens
    ↓
Transformer Blocks × N  # 第390行: self.layers
    ↓ 每个 Block:
    ├─ RMSNorm       # 第368行: self.input_layernorm
    ├─ Attention    # 第365行: self.self_attn
    ├─ Residual     # 第378行: hidden_states += residual
    ├─ RMSNorm     # 第369行: self.post_attention_layernorm
    ├─ FeedForward # 第370行: self.mlp (可选 MOE)
    └─ Residual
    ↓
RMSNorm             # 第391行: self.norm
    ↓
LM Head             # 第446行: self.lm_head
    ↓
Output (logits)

关键组件

1. Attention（第150-222行）

• Multi-head self-attention
• 支持 Flash Attention（加速）
• 使用 RoPE 位置编码
• GQA（Grouped Query Attention）支持

2. FeedForward（第225-238行）

• 标准前馈网络
• gate_proj → SwiLU 激活 → up_proj → down_proj
• 使用 SwiLU 激活函数

3. MOEFeedForward（第297-356行，可选）

专家混合架构:

• 多个专家网络（experts）
• 门控网络（gate）选择激活哪些专家
• 辅助损失（aux_loss）确保负载均衡

优势:

• 参数多，但只激活部分参数
• 可以用更少的计算训练更大的模型

4. MiniMindBlock（第359-380行）

单个 Transformer 层:

def forward(self, hidden_states, ...):
    # Attention + Residual
    residual = hidden_states
    hidden_states = self.input_layernorm(hidden_states)
    attn_output, _ = self.self_attn(hidden_states, ...)
    hidden_states = residual + attn_output
    
    # FeedForward + Residual
    residual = hidden_states
    hidden_states = self.post_attention_layernorm(hidden_states)
    ffn_output = self.mlp(hidden_states)
    hidden_states = residual + ffn_output
    
    return hidden_states

5. MiniMindModel（第383-436行）

完整的模型:

• 堆叠 N 个 MiniMindBlock
• 生成所有层的 hidden states 和 KV cache
• 计算总 aux_loss（如果有 MOE）

6. MiniMindForCausalLM（第439-470行）

最终的语言模型:

• 添加 lm_head（线性层）
• 输出 next token 的 logits
• 与 embed_tokens 共享权重

七、训练实践与常见问题

训练细节

1. 损失计算的特殊处理

loss_fct = nn.CrossEntropyLoss(reduction='none')

为什么用 reduction='none'?

• 返回每个位置的损失（不自动平均）
• 可以手动应用 loss_mask
• 更灵活地控制哪些位置参与计算

2. 半精度保存

state_dict = {k: v.half() for k, v in state_dict.items()}
torch.save(state_dict, ckp)

原因:

• 模型用 float32 训练，保存时转为 float16
• 减少 50% 存储空间
• 加载后可以转换为 float32 保持精度

3. 随机种子设置

base_seed = 1337
torch.manual_seed(base_seed)
torch.cuda.manual_seed(base_seed)

# DDP 时，每个GPU用不同的种子
if ddp:
    torch.manual_seed(base_seed + rank)
    torch.cuda.manual_seed(base_seed + rank)

目的: 确保实验结果可复现

4. model.train() 和 model.eval()

model.train()   # 训练: 启用 dropout
# ... 训练代码 ...
model.eval()    # 保存: 关闭 dropout

区别:

• train(): 启用 dropout、batch norm 的训练模式
• eval(): 关闭 dropout、使用固定的归一化统计

超参数建议

常见问题

Q1: 为什么用 DDP 而不是其他方式？

A: DDP 是当前推荐的多卡训练方式

• DDP2 已被弃用
• DP（DataParallel）效率低（单进程多线程）
• DDP 多进程，效率高

Q2: aux_loss 是什么？

A: MOE 模型的辅助损失

• 确保专家负载均衡
• 防止某些专家被选中频率过低
• aux_loss_alpha 控制权重

Q3: 为什么要设置 train_sampler.set_epoch(epoch)?

A: 每个 epoch 打乱数据顺序

• DDP 需要手动调用
• 确保每个 GPU 在不同 epoch 看到不同数据

Q4: Flash Attention 是什么？

A: 优化的 attention 实现

• 减少内存占用（O(N) 而不是 O(N²)）
• 加速计算
• 需要 PyTorch >= 2.0

八、总结

MiniMind 的训练框架采用了一系列现代深度学习的最佳实践。在训练效率方面，使用混合精度训练（bfloat16）显著提升速度和减少内存占用，配合 Flash Attention 优化注意力计算，并通过 DDP 分布式训练实现高效的多卡参数同步。

在训练稳定性方面，梯度裁剪防止梯度爆炸，余弦退火学习率调度确保模型稳定收敛，损失掩码机制确保只有真实 token 参与计算。工程层面，梯度累积突破 GPU 内存限制，半精度存储减少模型体积，严格的随机种子保证实验可复现。

整个框架支持 MOE（专家混合）架构，可在不显著增加计算量的情况下训练更大参数的模型。通过自回归预训练，模型学习预测下一个 token，逐步掌握语言统计规律、语法结构和语义关系。预训练为后续的 SFT 提供基础，SFT 阶段进一步学习对话格式和指令遵循，形成从 “预训练 → 语言能力 → 对话技能 → 实际应用” 的完整训练路径。