婴儿如何学语言

一个婴儿从零开始学会说话，经历三步：

训练三阶段

大模型训练分三个阶段，就像培养一个人才：

Part 1: 核心思想

理解预训练的本质：Next Token Prediction。这是所有自回归语言模型的统一训练目标。

预测下一个字

模型的任务：给定前面的字，预测下一个字

定义：Next Token Prediction——给定一个 token 序列，模型对每个位置输出一个概率分布，预测下一个 token 是什么。

白话：就像手机输入法的"联想输入"——你打了"今天天"，它猜下一个字是"气"。但模型要同时在整个序列的每个位置做这件事。

滑动预测全景

一次前向传播，所有位置都在练习

关键洞察：序列长度为 N 时，一次前向传播产生 N-1 个训练信号。这就是 Transformer 高效的原因。

这里用"婴儿学说话"的比喻只是帮你建立直观。在计算机内部，它实际上是矩阵乘法 + softmax，对词表中每个 token 计算概率。

输入与目标对齐

logits 取前 N-1 个，labels 取后 N-1 个，逐位比较

输入:  [你,  好,  世,  界]    (input_ids)
预测:  [好', 世', 界', ?' ]   (logits)
       取 [:-1] ↑  ↑  ↑
标签:  [好,  世,  界,  ! ]    (labels)
              取 [1:] ↑  ↑  ↑

交叉熵 = 比较每个位置的预测与真实值

白话：模型看了"你"，要预测"好"；看了"你好"，要预测"世"。关键操作是——

对齐代码溯源

回顾模型 forward 中的 shift 操作

# model_minimind.py 中的 forward 方法（简化）
def forward(self, input_ids, labels=None):
    # 1. 词嵌入 + 位置编码
    x = self.model(input_ids)
    # 2. Transformer 层
    x = self.lm_head(x)           # (batch, seq_len, vocab_size)

    if labels is not None:
        # 3. 关键：移位对齐
        loss = F.cross_entropy(
            x.view(-1, x.size(-1))[..., :-1, :],
            labels.view(-1)[..., 1:]
        )
        return CausalLMOutput(loss=loss, logits=x)

三步走：

1. 定义：shift 操作让预测和标签错位对齐。
2. 白话：就像批改作业——第 i 个位置的学生答案（logits[i]）对第 i+1 个位置的标准答案（labels[i+1]）。
3. 代码：logits[..., :-1, :] 与 labels[..., 1:] 逐位计算交叉熵。

Part 2: 数据处理

把原始文本变成模型能吃的数字。理解 PretrainDataset 的完整流水线。

数据文件格式

JSONL——每行一个 JSON 对象，简单直接

{"text": "今天天气真好，我们去公园玩吧。"}
{"text": "人工智能是计算机科学的一个分支。"}
{"text": "深度学习是机器学习的子领域，它使用多层神经网络……"}

数据集类源码

PretrainDataset：把原始文本变成训练样本

class PretrainDataset(Dataset):
    def __init__(self, data_path, tokenizer, max_length=512):
        super().__init__()
        self.tokenizer = tokenizer
        self.max_length = max_length
        # 用 HuggingFace load_dataset 加载 JSONL 文件
        self.samples = load_dataset('json', data_files=data_path, split='train')

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, index):
        sample = self.samples[index]
        # 第1步：分词，留2个位置给 BOS 和 EOS
        tokens = self.tokenizer(str(sample['text']),
            add_special_tokens=False,
            max_length=self.max_length - 2,
            truncation=True).input_ids
        # 第2步：加 BOS（开头）和 EOS（结尾）
        tokens = [self.tokenizer.bos_token_id] + tokens + [self.tokenizer.eos_token_id]
        # 第3步：PAD 填充到固定长度
        input_ids = tokens + [self.tokenizer.pad_token_id] * (self.max_length - len(tokens))
        input_ids = torch.tensor(input_ids, dtype=torch.long)
        # 第4步：labels = input_ids 的副本，PAD 位置设为 -100
        labels = input_ids.clone()
        labels[input_ids == self.tokenizer.pad_token_id] = -100
        return input_ids, labels

四步处理流水线

从原始文本到训练样本的完整路径

每一步的目的：

1. 分词：文字→数字，模型只认数字。
2. 加 BOS/EOS：告诉模型"从这开始"和"到这结束"。
3. PAD 填充：凑齐固定长度，方便组成 batch。
4. labels 中 PAD→-100：让损失函数忽略填充部分。

完整数字例子

假设数据是 "你好"，max_length = 8，BOS=1，EOS=2，PAD=0

为什么是 -100？

深入理解 ignore_index 机制

# PyTorch 源码中 cross_entropy 的核心逻辑（简化）
def cross_entropy(input, target, ignore_index=-100):
    loss = 0
    count = 0
    for i in range(len(target)):
        if target[i] == ignore_index:
            continue    # 直接跳过，不计入损失
        loss += -log(softmax(input[i])[target[i]])
        count += 1
    return loss / count  # 只对有效位置求平均

max_length 取舍

max_length 影响显存和训练效果

MiniMind 默认 max_seq_len=340。中文约 1 token ≈ 1.5~1.7 个字符，340 token 约覆盖 500+ 字的段落。

Part 3: 训练循环

精讲 train_epoch 函数的每一行。理解余弦调度、梯度累积、混合精度、梯度裁剪。

训练主循环源码

train_pretrain.py 中的 train_epoch 函数

def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
    start_time = time.time()
    last_step = start_step
    for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
        input_ids = input_ids.to(args.device)
        labels = labels.to(args.device)
        last_step = step

        # ① 计算当前学习率（余弦调度）
        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

        # ② 前向传播（混合精度）
        with autocast_ctx:
            res = model(input_ids, labels=labels)
            loss = res.loss + res.aux_loss       # 主损失 + MoE 辅助损失
            loss = loss / args.accumulation_steps # 梯度累积：除以累积步数

        # ③ 反向传播
        scaler.scale(loss).backward()

        # ④ 每隔 accumulation_steps 步才真正更新参数
        if step % 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 or step == iters:
            spend_time = time.time() - start_time
            current_loss = loss.item() * args.accumulation_steps
            Logger(f'Epoch:[{epoch+1}/{args.epochs}]({step}/{iters}), '
                   f'loss:{current_loss:.4f}, lr:{lr:.8f}')

        # ⑥ 定期保存检查点
        if (step % args.save_interval == 0 or step == iters) and is_main_process():
            model.eval()
            # ... 保存逻辑 ...
            model.train()

循环流程图

六个阶段：取数据→调学习率→前向→反向→更新→保存

余弦学习率

来自 trainer_utils.py 的 get_lr 函数

定义：学习率按余弦曲线从初始值平滑衰减到终值。

def get_lr(current_step, total_steps, lr):
    """余弦学习率调度"""
    return lr * (0.1 + 0.45 * (1 + math.cos(math.pi * current_step / total_steps)))

白话：学习率就像开车——起步时速度适中，中间平稳行驶，快到终点时慢慢减速。最低不会降到 0（保留 10%），让模型始终有微调能力。

这个比喻只是为了帮你建立直观感受。在计算机内部，它实际上是余弦学习率调度，使用 math.cos 函数根据当前步数和总步数计算衰减系数。

余弦调度曲线

学习率随训练步数的变化

梯度累积

小显存也能享受大批量的好处

想用 batch_size=256 训练，但显卡只够放 batch_size=32？

梯度累积：每次只放 32 条数据，但累积 8 次梯度（32 × 8 = 256），等效于 batch_size=256。

# 关键代码（train_pretrain.py）
loss = loss / args.accumulation_steps  # ① loss 除以累积步数
scaler.scale(loss).backward()          # ② 反向传播（梯度自动累加）

if step % args.accumulation_steps == 0:  # ③ 每 8 步才更新一次
    scaler.step(optimizer)
    optimizer.zero_grad(set_to_none=True)

累积机制详解

梯度在内存中是如何一步步叠加的

PyTorch 的 backward() 默认行为是累加梯度，不是覆盖。所以只有调用 zero_grad() 时才清零。

混合精度训练

用更少的内存跑更大的模型

定义：混合精度 = 用低精度（16bit）做前向和反向，用高精度（32bit）维护参数主副本。

# train_pretrain.py 中的设置
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" \
               else torch.cuda.amp.autocast(dtype=dtype)
# GradScaler 仅 float16 需要（bfloat16 范围够大，不需要）
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))

三种精度对比

bfloat16 为什么是最佳选择？

白话：float16 的"量程"太小（±6.5万），稍大的梯度就溢出了，需要 GradScaler 帮忙。bfloat16 和 float32 一样大（±3.4e38），但只要一半内存。MiniMind 默认用 bfloat16。

GradScaler

float16 的"保镖"——防止梯度下溢

# GradScaler 的工作流程
# 1. 放大 loss，防止小梯度在 float16 下变成 0
scaler.scale(loss).backward()

# 2. 缩放回来，检查有没有 inf/nan
scaler.unscale_(optimizer)

# 3. 如果正常，更新参数；如果有 inf/nan，跳过这步
scaler.step(optimizer)

# 4. 调整缩放因子
scaler.update()

bfloat16 因为范围够大，不需要这个保镖。所以 MiniMind 的代码中 enabled=(args.dtype == 'float16')——只在 float16 时启用。

梯度裁剪

防止梯度爆炸的"限速器"

# train_pretrain.py 第 43 行
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
# grad_clip 默认 1.0
# 如果所有参数梯度范数 > 1.0，就等比例缩小到 1.0

三步走：

1. 定义：计算所有参数梯度的 L2 范数，超过阈值则等比缩放。
2. 白话：给梯度设上限，就像汽车限速——不管多着急，最快 120。遇到"特别难"的数据时，梯度可能突然爆炸，裁剪防住它。
3. 代码：clip_grad_norm_(params, max_norm=1.0)，超过 1.0 就缩到 1.0。

这个比喻只是为了帮你建立直观感受。在计算机内部，梯度裁剪实际上是计算所有参数梯度向量的 L2 范数，除以 max_norm 得到缩放系数。

尾部残余梯度

train_epoch 末尾的兜底处理

# train_epoch 函数的末尾（第 74-79 行）
if last_step > start_step and last_step % args.accumulation_steps != 0:
    # 如果 epoch 结束时还有未更新的残余梯度
    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)

假设 accumulation_steps=8，epoch 在 step=45 结束。45 % 8 = 5，还剩 5 步累积的梯度没有更新。

Part 4: 工程实战

训练脚本 9 步总览、工程基础设施、端到端数据流。

训练脚本 9 步

train_pretrain.py 的完整结构

从初始化到训练到清理，一共 9 步。每一步都不能跳过，否则训练可能静默出错。

步骤 1-3 初始化

分布式环境、模型配置、混合精度

# ① 初始化分布式环境和随机种子
local_rank = init_distributed_mode()
if dist.is_initialized():
    args.device = f"cuda:{local_rank}"
setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))

# ② 配置模型参数
lm_config = MiniMindConfig(
    hidden_size=args.hidden_size,        # 默认 768
    num_hidden_layers=args.num_hidden_layers,  # 默认 8
    use_moe=bool(args.use_moe)           # 默认 False
)

# ③ 设置混合精度
device_type = "cuda" if "cuda" in args.device else "cpu"
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" \
               else torch.cuda.amp.autocast(dtype=dtype)

步骤 4-5 定义组件

日志、模型、数据集、优化器

# ④ 配 SwanLab 日志（可选）
wandb = None
if args.use_wandb and is_main_process():
    import swanlab as wandb
    wandb_id = ckp_data.get('wandb_id') if ckp_data else None
    wandb.init(project=args.wandb_project, id=wandb_id,
               resume='must' if wandb_id else None)

# ⑤ 定义模型、数据、优化器
model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
train_ds = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)

步骤 6-7 恢复与包装

检查点恢复、torch.compile、DDP 包装

# ⑥ 从 checkpoint 恢复状态
start_epoch, start_step = 0, 0
if ckp_data:
    model.load_state_dict(ckp_data['model'])
    optimizer.load_state_dict(ckp_data['optimizer'])
    scaler.load_state_dict(ckp_data['scaler'])
    start_epoch = ckp_data['epoch']
    start_step = ckp_data.get('step', 0)

# ⑦ 编译和分布式包装
if args.use_compile == 1:
    model = torch.compile(model)  # PyTorch 2.0 加速
if dist.is_initialized():
    model = DistributedDataParallel(model, device_ids=[local_rank])

步骤 8-9 训练与清理

训练循环与 DataLoader 构建

# ⑧ 开始训练循环
for epoch in range(start_epoch, args.epochs):
    train_sampler and train_sampler.set_epoch(epoch)
    setup_seed(42 + epoch)
    indices = torch.randperm(len(train_ds)).tolist()

    # 关键：断点续训时跳过已训练的 batch
    skip = start_step if (epoch == start_epoch and start_step > 0) else 0
    batch_sampler = SkipBatchSampler(
        train_sampler or indices, args.batch_size, skip
    )
    loader = DataLoader(train_ds, batch_sampler=batch_sampler,
                        num_workers=args.num_workers, pin_memory=True)

    if skip > 0:
        Logger(f'跳过前{skip}个step，从step {skip+1}开始')
    train_epoch(epoch, loader, len(loader) + skip, skip, wandb)

# ⑨ 清理分布进程
if dist.is_initialized():
    dist.destroy_process_group()

工程基础设施

检查点续训与多卡训练——这里先建立直觉，深度留给阶段 11

上面 9 步里出现了两个工程必备件：lm_checkpoint（检查点）和 init_distributed_mode（分布式）。它们的具体实现都在 trainer_utils.py 里，是所有训练器共享的基础设施——预训练、SFT、DPO、RL 用的是同一套代码。这里先记住它们解决什么问题，细节留到阶段 11 系统讲。

端到端数据流

从文本到模型更新的完整路径

这就是训练的完整闭环。每一步你都已经在前面学过了，现在把它们串起来。

Part 5: 排错与练习

常见错误诊断和动手练习。训练出问题时，先来这里查。

OOM 与显存

显存 ≈ 模型参数 × 4 (float32)
     或 模型参数 × 2 (bfloat16)
     + batch_size × seq_len × hidden × 4

Loss 异常诊断

NaN 与恢复异常

排查流程图

遇到训练问题时，按此流程排查

动手练习

进阶练习

完整训练命令

从最简单到自定义到多卡

# 最简单的单卡训练（所有参数用默认值）
$ python trainer/train_pretrain.py

# 自定义参数
$ python trainer/train_pretrain.py \
    --epochs 2 \
    --batch_size 32 \
    --learning_rate 5e-4 \
    --max_seq_len 340 \
    --accumulation_steps 8 \
    --hidden_size 768 \
    --num_hidden_layers 8

# 4卡训练 + SwanLab 日志
$ torchrun --nproc_per_node=4 trainer/train_pretrain.py \
    --use_wandb --epochs 4 --batch_size 64

# 断点续训
$ python trainer/train_pretrain.py --from_resume 1

# 从已有权重继续预训练
$ python trainer/train_pretrain.py --from_weight pretrain

参数速查表

train_pretrain.py 所有参数一览

知识全景图