训练流程与超参数

本文为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 |

---

相关页面

  • 02-技术/nanoGPT代码解析 — model.py 解析
  • 02-技术/Flash-Attention优化 — 速度优化
  • 01-概念/预训练与微调 — 训练概念
  • 来源

  • [nanoGPT train.py](https://github.com/karpathy/nanoGPT/blob/master/train.py)
  • 本轮核查来源与限制(2026-09-17)

  • https://raw.githubusercontent.com/karpathy/nanoGPT/master/train.py
  • 只修本轮明确列出的陈述,不将旧页整篇升级为已验证。themis_only为本地保存和禁止外发的兼容标记,不代表本地模型推理。