手把手构建语言模型训练循环：从数据切分到梯度裁剪与检查点恢复

1. 先别急着调参，想清楚课程要你用训练循环证明什么

CS336这门课的最初设计思路，是让你不借助TorchTrainer、HF Trainer这一层现成封装，把语言模型的训练链路亲手搭一遍。拿到Assignment 1时，我先入为主地以为难点全在模型结构：多头注意力、LayerNorm、残差连接、位置编码，这些才是大头。真正动手做Training Loop的时候才发现，训练循环是最容易“看起来全对、实际埋雷”的地方。模型结构写错，基本一步forward就会报错或者loss非常离谱；训练循环写错，往往不是立刻崩，而是几千步之后loss突然变成nan，或者你换了数据集后发现怎么调都复现不出别人那份loss曲线。

所以这条笔记我会把完整的思路写下来，从数据组织、forward/loss的构造、optimizer与scheduler，到checkpoint和日志体系，按“我认为一个合格Training Loop必须具备的检查点”来推进。这不是CS336官方solution，而是我自己动手做复现过程中的一条可验证路径。

先说结论：训练循环真正要“证明”的不是你会调库，而是你能控制几个关键信号。第一，初始loss落在合理范围，不要出现“一看就离谱但我选择继续训练”的情况；第二，能在单个小batch上过拟合，说明forward/backward/update这条链路本身没有断裂；第三，训练中期loss下降趋势平缓且可复现，换随机种子结果不会像过山车。把一个训练循环做到这三点，比堆一堆花哨的监控面板更重要。

我用来做冒烟测试的数据集是TinyStories的英文子集。原因很直接：单条样本短、领域单一、语料干净。如果你手头已有其他1B级别左右的中英文混合文本，也可以继续用，但第一轮调训练循环时尽量避开又多又杂的多语言混合语料，否则问题会被复杂数据源掩盖。模型刚开始不用追求GPT-2规模，建议用d_model=128、num_heads=4、num_layers=4、vocab_size=50257，seq_len设为512。这个配置在一块24GB消费级显卡上能跑，并且能快速暴露训练循环的问题。

2. 数据管线的“形状”先对了，训练的坑能少一半

2.1 我们需要准备哪几个tensor，每个是什么含义

很多人一上来就想着怎么按batch把文本切好，然后丢进模型。其实Training Loop真正要的是四类张量：

• input_ids：模型实际读取的token序列，形状是[batch_size, seq_len]；
• labels：用于计算交叉熵的目标token序列，和input_ids尺寸一致；
• loss_mask或ignore_index：有些位置不参与loss计算，比如padding位置、或者跨文档时不想让模型负担的边界片段；
• position_ids：如果你用的是类似GPT的绝对位置编码，需要显式传入。如果模型内部根据input_ids自动生成position_ids，可以省。

这里最容易忽略的是labels的构造方式。语言模型的训练目标是根据前t个token预测第t+1个token，所以label序列等于把input_ids整体右移一位后的序列。不要单独创建一个“右移之后”的tensor再拼回去，更省事且不容易错的做法是：labels = input_ids.clone()，然后用shift逻辑计算logits与labels的交叉熵。很多实现里是直接把logits[:, :-1, :]和labels[:, 1:]对齐，这样每一段的最后一个位置没有预测目标，但通常无伤大雅，因为下一段会补上。CS336这类课程通常会在最后一个token位置继续预测下一个sequence的起始token，也就是用一个连续token流切块来训练，所以不存在“需要忽略末尾token”这种麻烦。

2.2 文本拼成长流再切块，比单独sample更合理

如果你还记得GPT-2原文的做法，它其实不是把每条文本单独padding到一个固定长度再送进去，而是把一堆文档拼成一个超长token流，再按seq_len切块。TinyStories里的每篇是小故事，完全可以直接用tokenizer自带的<|endoftext|>（在GPT-2词表中是token 50256）拼在一起，形成一条长达数百万token的id数组，然后按照seq_len依次切成样本。

我第一次做时是按story单独处理，每个样本都自己截断或padding到512，这样做的后果是：大量文本被截断，而且每个batch里padding位置太多，计算浪费明显。改成拼成长流后，不仅输入密度高，labels天然就是右移后的下一个token，几乎不用额外维护复杂的对齐关系。

具体流程大致是：先加载全部语料，给每条样本末尾补一个eos token后用tokenizer批量编码，最后把所有id拼接成一个大数组，缓存成uint16的numpy数组。为什么要缓存成numpy而不是每次实时tokenize？因为实时编码会让数据管线成为训练瓶颈，而且调试时重复跑同一个tokenize过程很浪费时间。

生成样本的伪代码如下：

python复制import numpy as np

# ids: 已拼接好的完整token流，dtype=np.uint16
seq_len = 512
num_tokens = len(ids)
num_samples = num_tokens // seq_len

# 切成长度为seq_len的连续块
ids = ids[: num_tokens - (num_tokens % seq_len)]
data = ids.reshape(num_samples, seq_len)

这会把data按行顺序变成训练样本。如果你希望模型不是总从固定故事开头开始学，可以在训练前对整条token流做一次随机偏移，比如跳过开头的一个随机长度，然后再切块，这样样本切分不是永远对齐文档边界。随机偏移这件事最好只做一次并记录offset，保证后续实验可复现。

2.3 不要直接给一个dataloader开大shuffle，除非你确定labels不会错位

训练循环里最常见的“隐蔽错误”来自shuffle。如果你的数据是连续切块得到的二维数组，shuffle时只能按行随机打乱，不能把每个样本内部的token顺序打乱。用PyTorch的DataLoader时，可以在dataset里重排index，但务必保证同一个index读到的input_ids和labels来自同一行。

我的习惯是绕开这个坑：直接把数据做成一个简单的torch.utils.data.Dataset，返回第i个切块对应的input tensor和label tensor。labels不需要额外创建，只需要在loss函数里做一个shift。

python复制class TokenChunkDataset(torch.utils.data.Dataset):
    def __init__(self, data):
        self.data = data  # shape: [num_samples, seq_len]

    def __len__(self):
        return self.data.shape[0]

    def __getitem__(self, idx):
        x = self.data[idx].astype(np.int64)
        y = self.data[idx].astype(np.int64)
        return torch.from_numpy(x), torch.from_numpy(y)

batch时再加一行：inputs, labels = batch，不需要在criterion里做任何padding mask，因为整条流没有padding。如果你在早期使用带padding的样本，那还得额外传一个attention_mask给模型，让注意力不去看padding位置。能用纯文本流解决的问题，就不要让padding增加排查负担。

3. train_step的每一行，都要能回答“为什么这么写”

3.1 先看最朴素的更新流程

模型forward之后计算loss，再backward，optimizer.step，最后清零梯度。这是训练循环的最小内核。但单独看这四步，很多细节会被忽略：

python复制def train_step(model, batch, optimizer, scheduler, device, grad_clip=1.0):
    model.train()
    inputs, labels = batch
    inputs, labels = inputs.to(device), labels.to(device)

    logits = model(inputs)  # [batch, seq_len, vocab_size]
    loss = cross_entropy(logits, labels)

    optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
    optimizer.step()
    scheduler.step()

    return loss.item()

如果使用AdamW，一般不需要在zero_grad之前再额外做一次decoupled weight decay，optimizer内部会处理。梯度裁剪要放在backward之后、step之前，这个顺序我知道大家都会背，但我在调试时真遇到过把clip放在optimizer.step之后的代码，效果是梯度没有被裁剪，loss照样发散。你最好在train_step里加一行断言，检查反向传播后是否存在非有限梯度，这在早期排查NaN时有奇效。

3.2 初始loss应该接近ln(vocab_size)，这是一个免费的“健康检查”

一个被很多人忽略的判断点：模型刚初始化、还没训练时，如果我们用的是随机初始化的输出层和交叉熵loss，期望loss会非常接近ln(vocab_size)。GPT-2词表大小是50257，所以ln(50257)约等于10.82。TinyStories上如果你用自己的tokenizer且vocab_size不一样，先把ln(vocab_size)算出来作为基准值。

我第一次跑通时打印的第一个loss是10.85左右，和10.82的偏差很小，这说明logits的数值范围正常，没有出现初始化过大或过小的问题；如果第一个loss是4或者15，我不会继续训练，而是先检查模型输出层和embedding层是否共享权重、初始化的标准差是否合理。这一个数字能挡住大量“模型结构有问题但训练循环假装在跑”的情况。

3.3 单batch过拟合：给训练循环做“短路测试”

训练循环是否有效，最直接的检验方法不是完整训练一个epoch，而是拿一个batch，反复训练几十到几百步，看loss能否降到接近0、accuracy能否接近100%。

如果单batch都过拟合不了，大概率是梯度流断了、优化器配置错误、或者模型结构里某个mask把事情搞坏了。我常用的是16个样本组成的一个小batch，seq_len=512，训练约100步，学习率固定在3e-4。对一个小模型来说，100步后loss通常能降到接近0.1以内。

不过单batch过拟合有个容易误判的地方：如果你的序列特别长，或者batch里存在大量相同token，模型可能在几轮内“记住”了位置而不是依赖语义。这个很难完全避免，但对训练循环本身的验证足够了。它只能证明训练链路是通的，不能证明泛化能力正常。

3.4 loss.backward()之后先看梯度范数，再决定是否clip

在完整实验里，梯度裁剪可以防止单batch的异常大梯度把参数推出正常区域。但如果你一上来就无脑clip，可能会把真正的问题掩盖掉。成熟的调试流程是：先不要clip，训练几步，把每个step的grad_norm打印出来，看看量级。

正常的grad_norm大约在1~10这个区间，取决于模型尺寸和loss scale。如果grad_norm达到了100以上，甚至几千，那先别急着用grad clip去“修”，而是回头检查是否数值不稳定、是不是embedding的梯度在累计误差。Clip的值可以设为1.0，但只是为了限制单步波动，不是根治。

4. 优化器、学习率与调度器的细节：训练循环质量的决定因素

4.1 AdamW参数匹配，不是抄个betas就完事

预训练语言模型基本默认用AdamW，b1=0.9、b2=0.95、eps=1e-8、weight_decay=0.1。这个组合本身不是万能的，但它和特定学习率配合时表现稳定。Assignment 1阶段我不建议你自由发挥修改这些超参数，先把标准组合跑通，再谈优化。

需要注意的一点是，不同代码库中AdamW的weight_decay实现细节并不完全一致，有的会对所有参数做decay，有的默认只对非bias和非LayerNorm参数做。CS336这类手写实现里，通常会把bias和LayerNorm参数从“需要weight decay的集合”里排除。这个细节会体现在优化器参数分组上：

python复制decay_params = []
no_decay_params = []
for name, param in model.named_parameters():
    if not param.requires_grad:
        continue
    if param.ndim <= 1 or name.endswith(".bias"):
        no_decay_params.append(param)
    else:
        decay_params.append(param)

optimizer = torch.optim.AdamW(
    [
        {"params": decay_params, "weight_decay": 0.1},
        {"params": no_decay_params, "weight_decay": 0.0},
    ],
    lr=3e-4,
    betas=(0.9, 0.95),
    eps=1e-8,
)

这里的判断逻辑是：LayerNorm里的gamma/beta以及所有bias通常不参与weight decay。判断条件用param.ndim <= 1比较常见，因为embedding矩阵是2维、权重矩阵一般也是2维以上，而所有bias和LayerNorm参数是1维。当然，如果你要更严谨，可以用参数名精确排除LayerNorm相关项，上面的条件对于GPT-2结构通常已经够用。

4.2 学习率调度：warmup解决的是早期稳定性，不是玄学

很多人问warmup为什么能提升训练稳定性。本质原因是，Adam在训练初期基于很少的梯度统计量去自适应调整步长，如果直接用很大的学习率，那些还未被充分估计的二阶矩（分母）会让更新步长被严重放大，导致loss突然冲出稳定区。warmup期间把学习率从0线性升到最大值，是在给优化器积累统计量的时间。

CS336的Assignment 1不一定强制要求实现cosine schedule，但做一个并不难。我常用的组合是：前200步从0线性涨到最大学习率，随后按cosine曲线衰减到峰值的1/10或者0。Pytorch里有现成的LambdaLR，但为了便于调试，我会手写一个极简schedule函数：

python复制def warmup_cosine_schedule(step, warmup_steps, total_steps, peak_lr, min_lr_ratio=0.1):
    if step < warmup_steps:
        return peak_lr * step / warmup_steps
    progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
    coeff = 0.5 * (1.0 + math.cos(math.pi * progress))
    return min_lr_ratio * peak_lr + (peak_lr - min_lr_ratio * peak_lr) * coeff

然后每个train_step调用一次，并把当前lr打印到日志里。请注意，scheduler.step()不能多调少调，尤其在你开启梯度累积时，可能一个“逻辑step”里执行了多个micro-step。此时应当只在累积梯度完成、真正更新参数时才推进scheduler，否则学习率会比预期下降得快很多。

4.3 梯度裁剪设1.0还是5.0，取决于你的loss量级

对语言模型来说，clip_grad_norm的默认max_norm=1.0是一个比较稳妥的下限；如果你发现训练后期梯度范数本身就在0.5以下，clip不会起什么作用。我个人在GPT-2级别模型上通常设1.0，因为大模型很容易出现个别参数梯度过大的情况。

训练循环里NaN问题的最常见来源是fp16混合精度下loss溢出，不是梯度爆炸。如果你没有启用AMP，出现NaN大概率是你学习率太高、或者是初始loss超出常规。先看是不是出现NaN前的几步loss快速升高；如果是，把学习率减半；如果是突然跳动，优先怀疑数据中是否有异常样本、或者某个位置上label出现非法值。

4.4 梯度累积模块里常见的hidden bug

在小显存环境里跑大模型，很自然会想到梯度累积。它的逻辑是：积累多个micro-batch的梯度后再更新一次参数。这部分最容易出的bug是：忘记把loss除以累积步数。

因为每个micro-batch的loss都会对当前计算图做backward，梯度会直接累加到param.grad上。如果不缩放，积累N步后的梯度会比单batch等效的大约N倍，这时直接用相同学习率，会严重不稳定。正确做法是loss = micro_loss / grad_accum_steps后再backward。另一个容易漏掉的点是只在真正更新参数时调用optimizer.step、scheduler.step和optimizer.zero_grad。

为了降低踩坑概率，早期调试可以先不开梯度累积。等单卡全batch训练稳定了，再把它加入，并且确保“单条样本平均看到的数据”不变。换句话说，如果你原来batch_size=32、grad_accum=1，现在改成batch_size=8、grad_accum=4，学习率可以保持不变，因为有效更新次数相同。

5. 训练循环的“运营系统”：日志、checkpoint和eval不能靠事后补救

5.1 为什么每次只保存model.state_dict()是在给自己挖坑

训练循环做的不只是更新参数。你还需要在实验意外中断后恢复训练。如果只保存model.state_dict()，那么optimizer里的动量、scheduler里的步数、数据顺序、随机数状态全部丢失，事后想复现曲线会非常困难。我习惯把checkpoint设计成一个最小化的状态快照，至少包含：

• model.state_dict()
• optimizer.state_dict()
• scheduler的状态（或直接保存当前step与lr）
• 当前global_step
• 训练数据读取器当前epoch/offset（如果可恢复）
• 混合精度训练时的GradScaler状态（如果用了AMP）

把整个字典用一个save_checkpoint函数包起来：

python复制checkpoint = {
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "step": global_step,
    "lr": current_lr,
    "config": config,
    "rng_state": torch.get_rng_state(),
    "scaler": scaler.state_dict() if scaler else None,
}
torch.save(checkpoint, path)

恢复训练时，先load checkpoint，再显式执行model.load_state_dict和optimizer.load_state_dict，继续从step开始跑。这里特别容易输错step的边界：如果保存step=1000，恢复后应该从1001继续，不要重复执行第1000步，否则调度器往前多走一步。

5.2 eval时的model.eval()与no_grad必须同时出现

训练循环里通常每个一定步数会跑一次验证集。PyTorch中model.eval()只是切换dropout和LayerNorm的运行模式，它不等于禁止梯度计算。如果忘记包with torch.no_grad():，验证集上照样会构建计算图，显存很快溢出；反过来，如果你只用了no_grad，但忘了model.eval()，模型里的dropout仍在工作，验证loss会有随机波动，且不同步数与不同随机种子之间很难对比。

验证时最好每次都用同一批固定样本，不要每次都从验证集随机抽样。固定一个验证子集可以让你更干净地比较不同训练step之间的loss变化，排除批次随机性干扰。实际操作时，我会从验证集里固定取32个样本，写到一个单独的buffer里，每次eval时反复用同一批。

5.3 日志公式：不要只看loss，要把grad norm和lr一起记录下来

多年训练经验告诉我，loss曲线是结果，grad_norm和lr是原因。单独看loss下降，经常无法定位一个训练为什么在5000步后突然变差。我在日志里会记录下面几项：

• step
• train_loss
• grad_norm
• current_lr
• tokens_per_second
• throughput(样本/秒)

记录grad_norm有个额外好处：loss还没有发散之前，grad_norm往往已经出现异常峰值。如果你在训练循环里把grad_norm打到日志，发现某一步grad_norm=1000，下一步也许loss还能维持，但它是一个重要预警。

一种轻量实用的实现是维护一个list，每N步把它以纯文本append到日志文件里，而不是用wandb等工具。第一轮调试时用CSV日志就够，等训练规模变大了再接入实验管理平台。避免为了记录日志引入额外代码依赖，这会拉长你的最小复现链路。

6. 放大到更长训练时，最容易出的三个问题

6.1 训练一百步没崩，不代表三千步不会崩

训练循环的初步验证常常只跑几百步。但很多不稳定性是隐藏的。我最常遇到的一种情况是：第100步到第500步loss稳定下降，第1000步附近突然出现一个大幅loss spike，然后恢复；再过一段时间开始周期性出现spike且程度加深，最后在某个时刻变成NaN。

针对这种“延迟爆炸”，不要只想着调低学习率。优先检查梯度的长期分布：把训练日志里的grad_norm最小值、中位数、max值拉出来看。如果最大值远大于中位数，说明存在少数异常step，最可能的原因是数据里出现了超长重复片段或某种异常模式。例如TinyStories虽然整体干净，但某些文本可能包含大量重复的标点、异常空格，导致模型在那些位置产生极大loss。此时可以给loss做一个异常值截断，或先检查并清洗掉这类样本。

6.2 改变batch size后，学习率是否需要同步调整

我踩过的第二个坑来自线性缩放学习率的误解。很多人说增大batch size时，需要同步调大学习率。这个结论在理想凸优化条件下成立，但在非凸深度网络里并不能直接照搬。更稳健的做法是：batch size翻倍时，学习率先保持不变，训练一段时间观察loss曲线——如果loss下降幅度明显变慢，再把学习率乘以1.5~2倍左右；如果loss波动变大甚至发散，则说明学习率增幅太大，应该回退。

在做训练循环Assignment时，固定batch size并保持学习率不变是更干净的控制变量方法。只有在需要把训练时间压缩到极致时，才去尝试“batch size和学习率联动”的调参策略。

6.3 可复现性：不要再忽略“随机种子固定”的小事

课程里做对比实验时，我们希望换一个初始种子仍然能得到相似趋势的验证loss。要实现这一点，不只是给PyTorch设torch.manual_seed(42)，还要设置Python的random、numpy、以及CUDA的随机种子。如果用了DataLoader的shuffle，还需要给DataLoader传一个generator=torch.Generator().manual_seed(0)。

很多人在训练循环里已经固定了种子，但仍发现验证loss在相同配置下波动。排查思路是先检查是否每次启动时，token流切块的random offset不一致；其次是模型是否使用了无法保证确定性的算子（例如某些fused attention kernel或者cudnn的benchmark模式）。如果是实验性质的小模型，可以把torch.backends.cudnn.deterministic设为True，虽然牺牲一点速度，但结果更稳定。

完整初始化种子代码：

python复制def set_seed(seed=42):
    random.seed(seed)
    numpy.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

7. 我最后实际采用的训练循环模板

这段把前面所有讨论浓缩成一个可以直接套用到Assignment 1的代码模板。模板里不包含模型定义，只假设模型接口是model(input_ids)返回logits；labels和input_ids等同。模板的规模偏小，单卡即可测试。你要在自己的数据集上跑更长的训练时，再把checkpoint、日志处理接进去。

python复制import datetime
import math
import torch
import torch.nn as nn
from torch.utils.data import DataLoader

def cross_entropy(logits, labels):
    vocab_size = logits.size(-1)
    logits = logits[:, :-1, :].contiguous()
    labels = labels[:, 1:].contiguous()
    loss_fct = nn.CrossEntropyLoss()
    return loss_fct(logits.view(-1, vocab_size), labels.view(-1))

def train_loop(model, train_dataset, val_sample, total_steps, args):
    model.train()
    device = args.device

    decay_params = [p for p in model.parameters() if p.ndim >= 2]
    no_decay_params = [p for p in model.parameters() if p.ndim < 2]
    optimizer = torch.optim.AdamW(
        [
            {"params": decay_params, "weight_decay": 0.1},
            {"params": no_decay_params, "weight_decay": 0.0},
        ],
        lr=args.peak_lr,
        betas=(0.9, 0.95),
        eps=1e-8,
    )

    train_loader = DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=2,
        drop_last=True,
    )

    global_step = 0
    log_interval = 20
    eval_interval = 200

    optimizer.zero_grad()
    for batch in train_loader:
        if global_step >= total_steps:
            break
        loss = train_step(model, batch, optimizer, args)
        grad_norm = compute_grad_norm(model)

        if global_step % log_interval == 0:
            lr = optimizer.param_groups[0]["lr"]
            print(f"step={global_step}, loss={loss:.3f}, grad_norm={grad_norm:.3f}, lr={lr:.2e}")

        if global_step % eval_interval == 0:
            val_loss = evaluate(model, val_sample, device)
            print(f"validation loss at step {global_step}: {val_loss:.3f}")

        global_step += 1

代码里省略了train_step和eval细节，但整体结构我建议按这个骨架展开。你可以在每一轮迭代开始时检查global_step超过total_steps就退出。数据集的遍历循环不要简单使用for epoch in range(num_epochs)，否则当total_steps不是epoch整数倍时，容易多跑或少跑一个循环。

最后分享一点个人感觉：CS336这一系列作业真正锻炼人的，不是你能否背出AdamW公式或手写attention，而是你能不能在复杂状态机里保持清醒——数据流在动、优化器状态在变、学习率在衰减、checkpoint在覆盖。Training Loop是这个复杂系统中唯一同时控制它们的地方。如果一开始就觉得“只是写个for循环而已”，后面很多隐蔽问题都会在这里爆发。建议做完Assignment 1后，把优化器状态、调度器步数、随机数恢复都完整测一遍，让训练循环具备“第二天关机重启仍然接着跑”的能力。别嫌麻烦，后面更大的模型和更长训练时间会证明，这些前期投入非常值得。