Skip to content

7. 预训练:训练 MiniQwen 0.8B ​

第 1-6 章介绍了 MiniQwen 的架构。本章讲如何实际训练一个 MiniQwen 0.8B 模型。

本章内容仅作教学 demo。我们的数据集极小(0.8B 模型只用 0.3B 数据)、训练平台简单(3 张 4090D)、训练框架不成熟,训练后的 MiniQwen 只能"不乱码输出",远达不到可用水平。真实预训练请参考 Qwen 技术报告与 Megatron/DeepSpeed 等工业框架。

7.1 概述 ​

什么是预训练? ​

预训练是大语言模型训练的第一阶段,目标是让模型学习语言知识——给定前 n-1 个 token,预测第 n 个 token。完整流程通常分三步:

阶段目标数据训练方式
预训练(本章)学习语言知识大规模无标注文本下一个 token 预测
SFT(第 9-10 章)学习指令跟随指令-回复对监督微调
RLHF/DPO/GRPO对齐人类偏好人类反馈强化学习

预训练任务:下一个 token 预测 ​

预训练的核心是自回归语言建模:给定 token 序列 x1,…,xt,最大化下一个 token 的概率:

LLM=−1T∑t=1Tlog⁡P(xt∣x1,…,xt−1)

这在代码里就是 model(input_ids=input_ids, labels=input_ids).loss(labels 复用 input_ids,模型内部自动右移一位)。7.5 节会详细讲 loss 曲线和它的下降过程。

数据集概览 ​

数据集文件大小样本数说明
pretrain_hqpretrain_hq.jsonl1.6 GB~140 万条高质量中文预训练数据
pretrain_hq_seq1024pretrain_hq_seq1024.pt2.3 GB~29.7 万条预 tokenize,seq_len=1024
pretrain_hq_seq2048pretrain_hq_seq2048.pt2.3 GB~11.5 万条预 tokenize,seq_len=2048
  • pretrain_hq_seq*:用 Qwen 官方 tokenizer 预处理出来的
  • 我们实际使用:pretrain_hq_seq1024.pt(预 tokenize,加载快)

数据格式与打包 ​

预训练数据是 JSONL:每行一个 JSON 对象。

jsonl
{"text": "鉴别一组中文文章的风格和特点...</think>好的,现在帮我查一下今天的天气怎么样?..."}

每条数据是多轮对话拼接,用 </think> 分隔(这里不区分角色,仅用特殊 token 切分)。

打包(Packing):把多条短文本拼成固定长度 seq_len 的序列,避免 padding 浪费。model.forward 内部自动做右移——logits[:,:-1] 预测 labels[:,1:]:

原始 tokens: [t1, t2, t3, t4, t5, t6, t7, t8, ...]
打包为 seq_len=4:
  seq_1 input/labels: [t1, t2, t3, t4]   # input/labels 相同
  seq_2 input/labels: [t5, t6, t7, t8]

右移后:
  shift_logits 预测: [t1, t2, t3]
  shift_labels 目标: [t2, t3, t4]

代码见 pretrain/data.py,支持 .pt 直读(快速)和 jsonl 在线 tokenize+pack 两条路径。文档分隔符是 <EOT>(id 248044),不是 tokenizer.eos_token——预训练用 <EOT> 分隔文档,chat 才用 <|im_end|>。

接下来 7.2-7.7 按"由浅入深"展开:先看最小可跑的训练循环(理解整体),再看配置(优化器/LR/累积/clip/调度),然后数据 → 分布式 → 断点续训 → 实际启动与监控。

7.2 最小可跑的训练循环 ​

我们先不看 6 个文件的完整框架,只看一个最小训练循环——核心就是 4 步:前向 → 反向 → 梯度累积更新 → 日志。

python
for epoch in range(epochs):
    for step, (input_ids, labels) in enumerate(train_loader):
        # 1. 前向:model 内部把 logits 右移一位, 和 labels 算交叉熵
        loss = model(input_ids=input_ids, labels=labels)["loss"] / grad_accum_steps
        # 2. 反向
        loss.backward()
        # 3. 每 grad_accum_steps 步才更新一次参数
        if (step + 1) % grad_accum_steps == 0:
            clip_grad_norm_(model.parameters(), max_norm=grad_clip)   # 防爆
            optimizer.step()
            optimizer.zero_grad(set_to_none=True)
        # 4. 日志
        if global_step % log_interval == 0:
            print(f"step {global_step} loss {loss.item():.4f}")

4 个关键点:

  • model(...)["loss"]:model.forward 接 labels 时内部做右移交叉熵(见 7.5)
  • / grad_accum_steps:累积 N 步再求平均才等于"大 batch 一次"的梯度
  • clip_grad_norm_:防止梯度爆炸(LLM 训练必备)
  • zero_grad(set_to_none=True):比 set_to_zero 省内存(把梯度指针设为 None)

我们的 train_pretrain.py 在此基础上加了 DDP 初始化、断点续训、定期保存。先理解上面这 4 步,下面的优化器/学习率/累积等就是给这 4 步"加料"。

7.3 训练配置 ​

我们使用 AdamW(pretrain/train_pretrain.py):

python
optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=3e-4,                # 学习率
    weight_decay=0.1,       # 权重衰减(L2 正则)
    betas=(0.9, 0.95),     # 一阶/二阶矩衰减率
)
参数值说明
lr3e-4LLM 通常 1e-4 ~ 5e-4
weight_decay0.1所有参数施加 L2 正则
beta1/20.9/0.95一阶/二阶矩衰减率(0.95 而非默认 0.999,LLM 常用)

学习率调度:Warmup + 余弦衰减 ​

LLM 训练不能直接用大学习率(前期梯度大,损失震荡),需要 warmup(线性升到最大值)然后余弦衰减:

python
def get_lr(step, total_steps, max_lr, min_lr=0):
    warmup_steps = min(total_steps // 20, 500)   # 最多前 5% 步 warmup, 上限 500 步
    if step < warmup_steps:
        return max_lr * step / warmup_steps       # 线性升
    progress = (step - warmup_steps) / (total_steps - warmup_steps)
    return min_lr + 0.5 * (max_lr - min_lr) * (1 + cos(pi * progress))  # 余弦降到 min_lr
lr
│        ╱‾‾‾‾‾‾‾╲
max_lr │       ╱     ╲
      │      ╱       ╲___
      │     ╱            ╲___
min_lr │____╱                  ‾‾‾‾──→ step
      0    warmup       cosine decay

为什么要 warmup:训练初期参数随机,梯度方向不稳定,小学习率先"热热身",避免前期震荡。为什么要余弦衰减:训练后期损失接近最小,大学习率会在最优点附近震荡,余弦平滑降到 0 便于收敛。

梯度累积:模拟大 batch ​

GPU 显存有限,batch_size 通常只能开 1-2,但 LLM 训练需要大 batch(1M-4M tokens)才稳定。梯度累积用多次小 batch 模拟大 batch:

python
# loss 要除以 grad_accum_steps, 保证 N 步累积的梯度平均 = 大 batch 一次
loss = outputs["loss"] / grad_accum_steps
loss.backward()
if (step + 1) % grad_accum_steps == 0:
    optimizer.step()
    optimizer.zero_grad()

有效 batch size = batch_size × num_gpus × grad_accum_steps。我们 3 卡 × bs=1 × accum=16 = 48。

梯度裁剪 ​

python
clip_grad_norm_(model.parameters(), max_norm=1.0)

为什么:LLM 训练后期,某个罕见大梯度会让所有参数步子迈太大,模型"飞出"局部最小。clip_grad_norm_ 把整个梯度向量的 L2 范数裁到 ≤ 1.0(按整体范数裁,不是按每个参数),保证步长有上界。

7.4 数据准备(数据集与打包) ​

详见 7.1 数据集概览,这里展开 PretrainDataset 怎么读数据。

代码文件

python
class PretrainDataset(torch.utils.data.Dataset):
    def __init__(self, tokenizer, data_path, seq_len=1024, max_samples=None):
        self.seq_len = seq_len
        if data_path.endswith(".pt"):
            self._load_pt(data_path, seq_len, max_samples)   # 快
            return
        self._tokenize_and_pack(tokenizer, data_path, seq_len, max_samples)

    def __len__(self):    return len(self.data)
    def __getitem__(self, idx):
        input_ids = self.data[idx]
        labels = input_ids.clone()                  # 自回归, labels=input_ids
        return input_ids, labels

两条路径:

  • .pt 直读:_load_pt 用 torch.load(weights_only=True) 读 2D 张量 (num_seq, seq_len),启动快(无 tokenize 开销)
  • jsonl 在线 tokenize+pack:_tokenize_and_pack 逐条 tokenizer.encode(text, add_special_tokens=False),尾部追加 <EOT>(id 248044),拼成长流再按 seq_len 切块

.pt vs jsonl 的 seq_len 不一致怎么办:训练时按 .pt 实际宽度训练(--seq_len 仅作用于 jsonl 路径),会打印 WARNING。

7.5 训练监控:Loss 曲线 ​

Loss 的含义 ​

预训练 loss = 下一个 token 预测的交叉熵。对 token t,它是 model 在 x1:t−1 条件下预测 xt 的负对数似然(取平均):

loss=−1T∑t=1Tlog⁡Pθ(xt∣x<t)

理论下界 = 0(完美预测)。随机基线 = log(vocab_size) = log(248320) ≈ 12.42(均匀预测)。

实时监控 ​

bash
# 训练日志
[22:32:16] Epoch [1/1] Step [3200/99070] Loss: 7.9453 Avg: 9.5783 LR: 1.19e-04 GS: 200 ETA: 2215.1min
[时间戳] 当前epoch  当前/总步数 当前loss 平均loss 学习率 全局步数 预计剩余时间

# GPU 实时监控
watch -n 1 nvidia-smi

GPU 利用率通常 60-100%,数据/梯度同步时下降。Avg Loss 是从训练开始到当前的滑动平均(不是最近一步),用于观察长期趋势;当前 Loss 是最近一步,可能因 batch 差异抖动。

训练时间估算 ​

T训练=Nsamples×NepochsBeff×Tstep
  • Beff=batch_size×num_gpus×grad_accum_steps
  • 我们的实测(3×4090D + 0.3B 数据):
训练量数据量预估时间
1%~3000 条~0.4 小时
10%~3 万条~4 小时
50%~15 万条~19 小时
100%~30 万条~38 小时

7.6 断点续训 ​

LLM 训练动辄几十小时,中途可能 OOM、断电、节点掉线。断点续训让你能从上次中断处继续。

代码文件

Checkpoint 结构 ​

python
state = {
    "epoch": epoch,
    "micro_step": step,
    "global_step": global_step,
    "loss": loss,
    "model_state_dict": model.state_dict(),         # DDP 用 model.module.state_dict()
    "optimizer_state_dict": optimizer.state_dict(),  # 含动量/方差/LR scheduler
}
torch.save(state, path)

为什么连 optimizer_state_dict 都要保存:Adam 内部有动量(mt)和方差(vt),如果只重载模型、优化器从零开始,前几步的优化方向会从零重算——浪费已经学到的"梯度惯性"。

保存策略 ​

  • 每 save_interval 步保存一次(如 2000 步):checkpoint-{global_step}.pt
  • 每个 epoch 末尾保存:checkpoint-epoch{epoch+1}.pt
  • 训练最终保存:final.pt,并导出 miniqwen_final/(HF 格式,可直接用 inference.py 推理)

启动恢复 ​

bash
bash run.sh pretrain_resume
# 自动找最新的 checkpoint-epoch*.pt; 没有再退到 checkpoint-*.pt
# 找不到则 exit 1

load_checkpoint 返回 (start_epoch, micro_step, global_step, last_loss),训练循环从这些值继续。

注意:中断时正在写的 checkpoint 文件可能损坏——我们的策略是"先写完再原子 rename",避免读到半截文件。

7.7 DDP 分布式训练 ​

代码文件 + pretrain/utils.py

什么是 DDP? ​

DDP(DistributedDataParallel)每个 GPU 持有完整模型副本,但处理不同数据分片,每步结束时通过 AllReduce 同步梯度。

     AllReduce 同步梯度
        │
   ┌────┼────┐
   ▼    ▼    ▼
 GPU0  GPU1  GPU2
 model model model  (各自完整副本)
 data0 data1 data2  (不同数据分片)

工作流程:1) 前向 2) 反向(各 GPU 独立算梯度)3) AllReduce(所有 GPU 梯度取平均)4) 各 GPU 用相同平均梯度更新参数。

启动 DDP 训练 ​

用 torchrun 启动多卡:

bash
CUDA_VISIBLE_DEVICES=1,3 torchrun --nproc_per_node=2 --master_port=29501 \
    pretrain/train_pretrain.py --ddp

--nproc_per_node=2 表示单机 2 进程,torchrun 自动设置 WORLD_SIZE=2、RANK=0/1、LOCAL_RANK=0/1、MASTER_ADDR/PORT。

关键代码片段 ​

python
# 1. 初始化
local_rank = int(os.environ.get("LOCAL_RANK", 0))
torch.cuda.set_device(local_rank)
dist.init_process_group(backend="nccl")    # NCCL=GPU 通信
device = torch.device(f"cuda:{local_rank}")

# 2. 包装模型
model = model.to(device)
model = DDP(model, device_ids=[local_rank], output_device=local_rank)

# 3. 数据分片:DistributedSampler 确保每 GPU 看不同数据
train_sampler = DistributedSampler(train_dataset, shuffle=True)
train_loader = DataLoader(train_dataset, sampler=train_sampler, batch_size=bs, ...)

# 4. 每个 epoch 重设 sampler(保证 shuffle 不同)
for epoch in range(epochs):
    train_sampler.set_epoch(epoch)
    ...

NCCL vs Gloo:GPU 用 nccl(Nvidia Collective Communication Library,快),CPU 用 gloo。DDP 与 DP 的区别:DDP 是多进程(每 GPU 一个 Python 进程),DP(nn.DataParallel)是单进程多线程——DP 有 GIL 瓶颈且只支持单机,DDP 是当前标准。

其他并行方式(补充阅读) ​

DDP 是数据并行(DP)的工业标准。但 LLM 训练规模大到单卡放不下时,需要更多并行方式:

方式切什么解决什么代表框架
DDP(数据并行)数据 batch扩 batch sizePyTorch DDP
TP(张量并行)单个矩阵单层放不下Megatron-LM, 4D 并行
PP(流水线并行)层数单卡放不下完整模型PipelineParallel, GPipe
ZeRO(优化器状态并行)优化器状态 + 梯度 + 参数减少每卡显存DeepSpeed
FSDP(全分片数据并行)参数+梯度+优化器全分片综合 ZeRO-1/2/3PyTorch FSDP
3D / 4D 并行DP+TP+PP(+Seq)超大模型Megatron-DeepSpeed
MoE 并行专家分片MoE 训练Megatron-MoE, Tutel

推荐阅读:

7.8 实际启动 ​

全量预训练命令(3×4090D) ​

bash
CUDA_VISIBLE_DEVICES=0,1,3 nohup torchrun \
    --nproc_per_node=3 \
    --master_port=29501 \
    pretrain/train_pretrain.py \
    --data_path data/pretrain_hq.jsonl \
    --tokenizer_path tokenizer \
    --output_dir checkpoints/miniqwen_pretrain \
    --seq_len 1024 \
    --batch_size 1 \
    --grad_accum_steps 16 \
    --epochs 1 \
    --lr 3e-4 \
    --weight_decay 0.1 \
    --grad_clip 1.0 \
    --log_interval 50 \
    --save_interval 2000 \
    --ddp \
    > checkpoints/miniqwen_pretrain/pretrain.log 2>&1 &

run.sh 快捷命令 ​

bash
bash run.sh pretrain          # 单卡快速验证 (seq_len=512, max_samples=10000)
bash run.sh pretrain_2gpu     # 2 卡 DDP (有效 batch 8)
bash run.sh pretrain_resume   # 断点续训

run.sh 的所有路径从 paths.py 导出(PRETRAIN_DATA/PRETRAIN_CHECKPOINT 等),改路径只改 paths.py 一处。

7.9 完整流程总结 ​

MiniQwen 0.8B 预训练流程(pretrain/)
│
├── 1. 数据准备
│   └── pretrain_hq_seq1024.pt(~30 万条, 预 tokenize)
│
├── 2. 启动训练(3×4090D, ~38 小时)
│   └── torchrun --nproc_per_node=3 pretrain/train_pretrain.py --ddp
│
├── 3. 训练循环(99070 步)
│   ├── 前向:input_ids → logits → loss(右移交叉熵)
│   ├── 反向 + 梯度累积(16 步一次更新)
│   ├── LR:warmup 500 步 → cosine decay 到 0
│   ├── grad_clip 1.0 防爆
│   ├── Loss: 11.0 → 4.63(详见 7.5)
│   └── 每 2000 步 + 每 epoch + 最终各保存一次
│
├── 4. 输出
│   └── checkpoints/miniqwen_pretrain/miniqwen_final/
│       (HF 格式, 可直接 inference.py 推理)
│
└── 5. 后续
    ├── 第 8 章:权重适配(也可跳过, 改用此 checkpoint)
    └── 第 9-10 章:SFT/DPO/GRPO/... 后训练

关键点:

  • Loss 是下一个 token 预测的交叉熵,11.0→4.6 的下降对应"从猜到学常见搭配到学长距离依赖"
  • 预训练 38 小时,受限于 0.3B 数据 + 0.8B 模型容量;训练后模型"不乱码"但远不够好
  • 断点续训 + DDP 是工程必备,分别解决"长训练易中断"和"单卡放不下大 batch"

上一章:6. MiniQwen | 下一章:8. 权重适配