训练流程与超参数
本文为nanoGPT源码教学节选,不是本项目训练指令,也未在本机复现实测。节选省略了取下一微批、清梯度和scaler更新等步骤,不能直接拼接执行。
train.py 全流程
1. 初始化配置(命令行参数覆盖默认值)
2. 设置随机种子、DDP(分布式训练)
3. 加载数据集(OpenWebText)
4. 初始化模型(或从 checkpoint 恢复)
5. 配置优化器(AdamW)
6. 训练循环:
├── 获取批次数据
├── 前向传播(计算损失)
├── 反向传播(计算梯度)
├── 梯度裁剪
├── 参数更新
└── 评估 & 检查点保存
---
核心超参数
# ===== 模型架构 =====
n_layer = 12 # Transformer 层数
n_head = 12 # 注意力头数
n_embd = 768 # 嵌入维度
block_size = 1024 # 上下文窗口长度(最大序列长度)
bias = False # 不使用偏置(更现代的选择)
dropout = 0.0 # 预训练不用 dropout
# ===== 训练规模 =====
batch_size = 12 # 每步微批次大小
gradient_accumulation_steps = 40 # 梯度累积步数
# 全局tokens/step = 每rank累积步数 × world_size × batch_size × block_size
# 源码DDP初始化会按world_size缩小每rank累积步数;默认全局约491520 tokens
# ===== 优化器 =====
learning_rate = 6e-4 # 峰值学习率
max_iters = 600000 # 总训练步数
weight_decay = 1e-1 # 权重衰减
beta1 = 0.9 # Adam 一阶矩
beta2 = 0.95 # Adam 二阶矩(GPT-3 选择,而非默认0.999)
grad_clip = 1.0 # 梯度裁剪阈值
# ===== 学习率调度 =====
decay_lr = True
warmup_iters = 2000 # 预热步数
lr_decay_iters = 600000 # 衰减终止步数
min_lr = 6e-5 # 最小学习率(= 最大的1/10)
---
学习率调度详解
def get_lr(it):
# 1. 线性预热阶段
if it < warmup_iters:
return learning_rate * it / warmup_iters
# 2. 超过衰减终点:使用最小学习率
if it > lr_decay_iters:
return min_lr
# 3. 余弦衰减
decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
return min_lr + coeff * (learning_rate - min_lr)
LR
↑ /‾‾‾\
| / \
| / \___________
|/ \____
+------------------------→ 步数
预热 | 余弦衰减 | 最小LR
---
混合精度训练
# 自动选择精度
dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16'
# bfloat16: A100/H100 推荐(数值范围大,训练更稳定)
# float16: 旧显卡(需要 loss scaling 防止下溢)
ctx = torch.amp.autocast(device_type=device_type, dtype=ptdtype)
scaler = torch.cuda.amp.GradScaler(enabled=(dtype == 'float16'))
---
分布式训练(DDP)
# 检测是否在 DDP 环境
ddp = int(os.environ.get('RANK', -1)) != -1
if ddp:
init_process_group(backend=backend)
ddp_rank = int(os.environ['RANK'])
ddp_local_rank = int(os.environ['LOCAL_RANK'])
ddp_world_size = int(os.environ['WORLD_SIZE'])
master_process = ddp_rank == 0 # 只有主进程保存检查点
---
梯度累积
为什么需要梯度累积?GPU 显存不够一次放入大批次。
# 将一个大批次拆成多个小批次,累积梯度后再更新
for micro_step in range(gradient_accumulation_steps):
with ctx:
logits, loss = model(X, Y)
loss = loss / gradient_accumulation_steps # 归一化
scaler.scale(loss).backward() # 累积梯度
# 所有微批次处理完后,统一更新参数
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
scaler.step(optimizer)
---
旧笔记训练结果(来源细节待复查,非本项目实测)
| 配置 | 硬件 | 时间 | 验证损失 |
|------|------|------|----------|
| GPT-2 (124M) | 8×A100 40GB | ~4天 | ~2.85 |
---
相关页面
来源
本轮核查来源与限制(2026-09-17)
只修本轮明确列出的陈述,不将旧页整篇升级为已验证。themis_only为本地保存和禁止外发的兼容标记,不代表本地模型推理。