在PyTorch中调试小型GPT:用于隔离各类故障的测试方法。
一种用于在PyTorch中调试字符级GPT的逐步工作流程,涵盖从令牌ID和目标偏移到梯度、NaN损失及检查点的处理。
损失、困惑度以及重复率等评估指标能够表明小型GPT模型存在异常,但很少能说明具体原因。人们往往会选择调整学习率、增加一层结构后再试。本指南用一种可重复的工作流程取代了这种猜测方式,用于训练基于WikiText-2的紧凑型字符级GPT(Mini-GPT):将每种异常现象对应到可能的原因,对每个处理阶段进行针对性检查,并按问题出现的顺序加以解决。最终你会得到一组判定标准以及一个在每次耗时的训练任务之前运行的诊断脚本。
为何GPT处理流程可能出现错误却不会崩溃
语言模型的训练涉及众多连续的转换步骤,每一步都需要使用前一步的输出作为输入:
Raw dataset
→ cleaned text
→ tokenizer
→ token IDs
→ training batches
→ embeddings
→ Transformer blocks
→ vocabulary logits
→ cross-entropy loss
→ gradients
→ optimizer
→ checkpoints
→ generation
任何一处缺陷都会影响后续的所有处理。假设目标数据与输入数据之间没有相差一个位置的情况:
Input: The cat
Target: The cat
现在,网络因复制它已看到的令牌而非预测下一个令牌而获得奖励。由于复制操作很简单,因此不会出现任何异常,损失值也仍有可能下降。模型只是在优化错误的目标函数。这与调试Web服务有着本质区别,在Web服务中,错误的返回值通常会导致测试或页面出错:
即使代码能够完整运行,也仍可能训练出有问题的模型。
从数据流的起点逐步向终点排查
需按照数据流经的顺序逐一验证各个阶段:
1. Environment
2. Files
3. Tokenizer
4. Token IDs
5. Training batches
6. Model shapes
7. Initial loss
8. Gradients
9. Optimizer
10. Validation behavior
11. Checkpoints
12. Generation
在确认数据流程之前就判断生成质量是浪费时间的行为,因为不良样本可能源自前11个阶段中的任何一个。在每个阶段,只需提出一个具体问题,并通过通过或失败的检测来得到答案。
四种故障类型
这类模型出现的几乎所有问题都属于以下四类之一,了解所属类别就能缩小搜索范围。
- 正确性故障:代码存在逻辑错误。目标位置没有发生偏移,因果掩码导致位置能够看到后续的标记,损失函数使用了布局错误的张量,或者分词器的标识与模型训练时所用的词汇表不一致。
- 数值故障:数学运算变得不稳定。损失值变为
NaN,梯度数值激增,逻辑斯蒂值溢出到无穷大,或者软最大化函数接收到没有有效数值的行。 - 优化故障:实现方式正确,但由于学习率过高或过低、模型规模过小,或是训练时间过短,导致学习效果不佳。
从简单的调试配置开始
使用完整运行环境进行调试会让每一次假设验证都变成漫长的等待。应定义一个能在几秒内记住少量示例的小模型:
debug_config = MiniGPTConfig(
vocab_size=tokenizer.vocab_size,
block_size=32,
embedding_dim=64,
num_heads=4,
num_layers=2,
expansion_factor=4,
dropout=0.0,
)
再配合小批量数据使用:
debug_batch_size = 8
以及短暂的运行时间:
debug_steps = 200
故意关闭 Dropout 功能:
dropout = 0.0
Dropout 会将随机激活值置零,从而导致相同运行结果出现差异。取消该功能(并设置固定种子)可让每次测试结果均可复现。一旦流程通过,再恢复生产环境设置。
确认运行时环境
在接触模型之前,先打印出Python和PyTorch的版本信息,以及是否可以使用CUDA或苹果的Metal后端(MPS):
import platform
import torch
print("Python:", platform.python_version())
print("PyTorch:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
if hasattr(torch.backends, "mps"):
print(
"MPS available:",
torch.backends.mps.is_available(),
)
有一个辅助函数会选择最佳的硬件设备,优先考虑CUDA,其次是MPS,最后才是CPU。通过hasattr检查可以确保在没有MPS后端的旧版本上也能正常运行:
def get_device():
if torch.cuda.is_available():
return torch.device("cuda")
if (
hasattr(torch.backends, "mps")
and torch.backends.mps.is_available()
):
return torch.device("mps")
return torch.device("cpu")
调用该函数一次并将结果记录下来:
device = get_device()
print("Selected device:", device)
当训练速度异常缓慢时,往往是因为程序期望使用GPU,但由于驱动程序或安装问题而被迫转而使用CPU。
确保所有输入文件都存在
要确认分词器定义以及编码后的训练集、验证集和测试集数据都已存在,若缺失应尽早抛出明确的FileNotFoundError错误,而非在训练循环深处出现令人困惑的错误:
from pathlib import Path
required_paths = [
Path("tokenizer/char_tokenizer.json"),
Path("data/encoded/train_ids.pt"),
Path("data/encoded/val_ids.pt"),
Path("data/encoded/test_ids.pt"),
]
for path in required_paths:
if not path.exists():
raise FileNotFoundError(
f"Required file not found: {path}"
)
print("Found:", path)
同时也要打印出各数据集的大小:
for path in required_paths:
print(
path,
path.stat().st_size,
"bytes",
)
空文件或异常小的文件通常意味着预处理任务被中断,留下了不完整的输出结果。
单独测试分词器
加载字符分词器:
tokenizer = CharTokenizer.from_file(
"tokenizer/char_tokenizer.json"
)
检查词汇表的大小以及字符列表的两端,查看是否有缺失或损坏的内容:
print("Vocabulary size:", tokenizer.vocab_size)
print("First tokens:", tokenizer.chars[:20])
print("Last tokens:", tokenizer.chars[-20:])
往返特性
无损分词器在编码和解码后会返回完全相同的输入内容。使用repr函数打印时,会显示出诸如尾随空格之类的不可见字符:
sample = "The history of"
encoded = tokenizer.encode(sample)
decoded = tokenizer.decode(encoded)
print("Encoded:", encoded)
print("Decoded:", repr(decoded))
assert decoded == sample
正在检查的不变条件:
decode(encode(text)) = text
出现错误意味着至少有一个字符无法被准确表示,通常是词汇表中缺少该符号。完整脚本中的encode方法在遇到这种情况时会抛出KeyError异常,而非像预期的那样默默忽略该字符。
捕获分词器与检查点不匹配的问题
词汇表大小相同是必要条件,但并非充分条件。两个词汇表都可能包含100个字符,但仍会为它们分配不同的ID:
Tokenizer A: "a" → 10
Tokenizer B: "a" → 24
使用一种映射方式训练模型,却用另一种映射方式来调用它时,虽然所有张量的形状都看似正确,但输出的结果却毫无意义。应将词汇表与检查点一起保存,或者至少保存其指纹信息。对序列化的字符列表进行SHA-256哈希处理即可;正确设置separators和ensure_ascii参数能确保同一列表始终序列化为相同的字节:
import hashlib
import json
def tokenizer_fingerprint(chars):
payload = json.dumps(
chars,
ensure_ascii=False,
separators=(",", ":"),
).encode("utf-8")
return hashlib.sha256(
payload
).hexdigest()
为已加载的分词器计算该值:
fingerprint = tokenizer_fingerprint(
tokenizer.chars
)
print("Tokenizer fingerprint:", fingerprint)
保存时将其存入检查点字典中:
checkpoint[
"tokenizer_fingerprint"
] = fingerprint
加载时进行比较,若不匹配则拒绝继续处理。不过 is not None 条件仍允许那些没有指纹信息的旧检查点被加载:
saved_fingerprint = checkpoint.get(
"tokenizer_fingerprint"
)
if (
saved_fingerprint is not None
and saved_fingerprint != fingerprint
):
raise ValueError(
"Checkpoint and tokenizer do not match"
)
查看编码后的令牌ID
将训练数据以64位整数的形式加载到CPU上,因为类型嵌入和交叉熵函数都期望这种格式:
train_ids = torch.load(
"data/encoded/train_ids.pt",
map_location="cpu",
).long()
打印数据的形状、数据类型以及取值范围:
print("Shape:", train_ids.shape)
print("Dtype:", train_ids.dtype)
print("Minimum ID:", train_ids.min().item())
print("Maximum ID:", train_ids.max().item())
每个ID都必须属于词汇表中的内容:
0 ≤ token ID < vocabulary size
以断言的形式进行验证:
assert train_ids.min().item() >= 0
assert (
train_ids.max().item()
< tokenizer.vocab_size
)
超出范围的ID会破坏嵌入查找机制,在GPU上该错误可能会以与实际原因无关的、难以理解的设备端断言形式出现。常见原因包括错误的分词器文件、损坏的编码文件、编码后重新构建的词汇表,或是不一致的特殊标记处理方式。
将存储的数据读回为文本
即使数值在有效范围内,也仍可能编码出错误的文本,因此请解码几百个ID并读取其内容:
sample_ids = train_ids[:500]
sample_text = tokenizer.decode(
sample_ids.tolist()
)
print(sample_text)
你应该能看到格式正常、包含换行符、标题和标点符号的可读WikiText,且不会出现连续重复或损坏的字符。如果样本显示异常,请立即停止:任何模型修改都无法弥补有问题的分词器或数据集带来的问题。
验证输入与目标之间的单个标记偏移
手动构建一个示例,使目标窗口从稍后的位置开始:
block_size = 32
start = 100
inputs = train_ids[
start:
start + block_size
]
targets = train_ids[
start + 1:
start + block_size + 1
]
将两者解码后进行比较:
input_text = tokenizer.decode(
inputs.tolist()
)
target_text = tokenizer.decode(
targets.tolist()
)
print("Input: ", repr(input_text))
print("Target:", repr(target_text))
目标数据应与输入数据类似,只是去掉第一个字符并添加一个新字符。然后对张量之间的关系进行断言:
assert torch.equal(
inputs[1:],
targets[:-1],
)
项目中的一些检查能够发现更严重的错误。该不变性为:
inputs[1:] == targets[:-1]
每个目标位置都存储与对应输入位置之后的标记,这正是下一个标记预测所需要的。
相同切片错误
常见的错误是两个切片使用相同的边界:
inputs = data[
start:
start + block_size
]
targets = data[
start:
start + block_size
]
这样模型就会学习到恒等映射关系:
Current token → current token
将目标切片的起始位置向后移动一个标记即可解决该问题:
targets = data[
start + 1:
start + block_size + 1
]
从而恢复预期的功能:
Current context → next token
这种错误的明显迹象是损失值在初期会异常快速下降。
检查批次形状、数据类型和设备
抽取一个真实的批次:
inputs, targets = get_batch(
data=train_ids,
batch_size=8,
block_size=32,
device=device,
)
打印所有可能出问题的信息:
print("Input shape:", inputs.shape)
print("Target shape:", targets.shape)
print("Input dtype:", inputs.dtype)
print("Target dtype:", targets.dtype)
print("Input device:", inputs.device)
print("Target device:", targets.device)
对于批次大小为8、块大小为32的情况,预期结果如下:
Input shape: [8, 32]
Target shape: [8, 32]
Dtype: torch.int64
Device: same as model
将这些预期结果固定下来:
assert inputs.shape == targets.shape
assert inputs.dtype == torch.long
assert targets.dtype == torch.long
assert inputs.device == device
assert targets.device == device
解决设备不匹配问题
在PyTorch开发中,这种错误经常出现:
Expected all tensors to be on the same device
某个操作接收到了位于不同设备上的张量,比如CPU和GPU。打印出参数和批次所在的位置:
model_device = next(
model.parameters()
).device
print("Model device:", model_device)
print("Input device:", inputs.device)
明确地将模型及每个批次移动到指定设备:
model = model.to(device)
inputs = inputs.to(device)
targets = targets.to(device)
模型的内部存在一个更隐蔽的问题:torch.arange默认在CPU上运行,因此需要从传入的参数中获取设备信息:
positions = torch.arange(
sequence_length,
device=token_ids.device,
)
否则,在 CUDA 或 MPS 上将位置嵌入加到令牌嵌入上会失败。将设备与输入绑定也有助于保持模型的可移植性。
验证前向传播
运行一个带有目标值的批次,以便模型返回对数概率和损失值:
logits, loss = model(
inputs,
targets,
)
检查这些数值:
print("Logits shape:", logits.shape)
print("Loss shape:", loss.shape)
print("Loss value:", loss.item())
对于每个位置,对数概率需要为每个词汇表条目提供一个分数,而损失值必须是标量:
assert logits.shape == (
inputs.size(0),
inputs.size(1),
tokenizer.vocab_size,
)
assert loss.ndim == 0
形状为 [B, V, T] 的对数概率意味着转置或重塑操作使维度顺序错误。由于 F.cross_entropy 在第二维接受类别分数,顺序错误的张量有时会无误地传入损失函数并计算出无意义的数值。
将初始损失与随机基准值进行比较
一个权重较小的新初始化模型会预测出近乎均匀的分布,而针对V个类别的均匀分布计算的交叉熵值为log(V):
import math
expected_loss = math.log(
tokenizer.vocab_size
)
print("Expected loss:", expected_loss)
print("Actual loss:", loss.item())
轻微的偏差属于正常现象,而较大的偏差则是异常信号。初始损失远高于基准值可能表明逻辑斯蒂值极端、初始化不稳定、令牌ID无效、目标与词汇表不匹配,或是输出形状有误。初始损失远低于基准值则说明模型实际上并未学到任何知识,这可能是数据泄露、意外加载了训练好的权重、目标与输入相同、能够看到未来的令牌,或是错误地恢复了检查点。
证明因果掩码有效
GPT必须仅依据之前的标记来预测每个位置。构建两个具有相同前缀但不同后缀的序列;如果模型是因果型的,那么前缀的对数概率必须一致。在评估模式下会禁用丢弃机制,以避免随机性干扰:
model.eval()
prefix_length = 8
sequence_length = 16
sequence_a = torch.randint(
0,
tokenizer.vocab_size,
(1, sequence_length),
device=device,
)
sequence_b = sequence_a.clone()
sequence_b[
:,
prefix_length:
] = torch.randint(
0,
tokenizer.vocab_size,
(
1,
sequence_length - prefix_length,
),
device=device,
)
在不计算梯度的情况下运行这两个序列:
with torch.no_grad():
logits_a, _ = model(sequence_a)
logits_b, _ = model(sequence_b)
测量最大的前缀差异:
prefix_difference = (
logits_a[:, :prefix_length, :]
- logits_b[:, :prefix_length, :]
).abs().max().item()
print(
"Maximum prefix difference:",
prefix_difference,
)
在能够容纳浮点数噪声的较小容差范围内验证两者是否相等:
assert torch.allclose(
logits_a[:, :prefix_length, :],
logits_b[:, :prefix_length, :],
atol=1e-5,
)
出现故障意味着后续位置的信息会泄露到前面的位置,通常是由于掩码缺失、应用到了错误的维度上,或是基于错误的三角形结构构建的。端到端的测试比仅检查掩码张量更为有效。
对单个批次进行过拟合
如果从本指南中选择一种技术,那就选这一项:
具备足够容量的模型应该能够记住一小批数据。
它同时运用数据、模型、损失函数、反向传播以及优化器。先固定每一步都会重复使用的那批数据:
fixed_inputs, fixed_targets = get_batch(
data=train_ids,
batch_size=8,
block_size=32,
device=device,
)
构建一个没有 Dropout 的小型模型:
debug_config = MiniGPTConfig(
vocab_size=tokenizer.vocab_size,
block_size=32,
embedding_dim=64,
num_heads=4,
num_layers=2,
expansion_factor=4,
dropout=0.0,
)
debug_model = MiniGPT(
debug_config
).to(device)
反复使用该批数据进行训练,每 50 步记录一次情况:
optimizer = torch.optim.AdamW(
debug_model.parameters(),
lr=1e-3,
)
for step in range(500):
optimizer.zero_grad(
set_to_none=True
)
_, debug_loss = debug_model(
fixed_inputs,
fixed_targets,
)
debug_loss.backward()
optimizer.step()
if step % 50 == 0:
print(
step,
debug_loss.item(),
)
损失值应远低于基准值。如果并非如此,可能是损失函数有问题、梯度无法传递到某些参数、目标数据未做偏移、模型容量甚至对于当前任务来说都太小、学习率选择不当、因果掩码存在故障,或是优化器根本没有对任何参数进行更新。可将测试视为每次完整训练的验证环节。
确认梯度能够传递到所有参数
执行一次前向传播和一次反向传播:
optimizer.zero_grad(
set_to_none=True
)
_, loss = model(
inputs,
targets,
)
loss.backward()
报告所有 .grad 值仍为 None 的可训练参数,以及其余参数的范数:
for name, parameter in (
model.named_parameters()
):
if not parameter.requires_grad:
continue
if parameter.grad is None:
print(
"NO GRADIENT:",
name,
)
else:
print(
name,
parameter.grad.norm().item(),
)
梯度缺失通常意味着在 __init__ 中定义但在 forward 方法中未被使用的层、意外调用的 .detach() 方法、跳过了某些组件的前向传播路径、从断开的张量计算出的损失,或是设置了 requires_grad=False 的情况。
跟踪全局梯度范数
逐参数计算范数可以发现无效层;而一个汇总数值则可用于监测随时间变化的稳定性。此函数会将所有梯度 L2 范数合并起来,也就是用于数值裁剪的同一指标:
def calculate_gradient_norm(model):
squared_norm = 0.0
for parameter in model.parameters():
if parameter.grad is None:
continue
parameter_norm = (
parameter.grad
.detach()
.norm(2)
.item()
)
squared_norm += (
parameter_norm ** 2
)
return squared_norm ** 0.5
在每次反向传播后记录该数值:
gradient_norm = (
calculate_gradient_norm(model)
)
print("Gradient norm:", gradient_norm)
注意那些恰好为零、过大或出现剧烈波动的范数,以及 NaN 和 inf 值。
提前检测 NaN 和无穷大值
一旦张量中出现非有限值,辅助函数就会立即触发并给出该张量的名称:
def assert_finite_tensor(
tensor,
name,
):
if not torch.isfinite(
tensor
).all():
raise FloatingPointError(
f"{name} contains NaN or infinity"
)
可将其应用于logits和损失值:
assert_finite_tensor(
logits,
"logits",
)
assert_finite_tensor(
loss,
"loss",
)
以及backward()之后的所有梯度:
for name, parameter in (
model.named_parameters()
):
if parameter.grad is not None:
assert_finite_tensor(
parameter.grad,
f"gradient for {name}",
)
通过检查多个点可以找到无效数值首次出现的位置,这比在数百步之后才发现NaN损失值要有用得多。
为何损失值会变为NaN
常见原因包括:学习率过高、梯度爆炸、注意力矩阵中所有位置都被屏蔽、softmax输入无效、混合精度运算溢出、参数已损坏、除以零、对零或负数取对数,以及logits值为无穷大。出现这种情况时:
- 停止训练过程。
- 找到最后一个损失值为有限值的步骤。
- 降低学习率。
- 启用梯度裁剪功能。
一旦参数中出现NaN,就绝不能继续迭代;每次更新都会使错误扩散,因此应从上一个正常的检查点重新开始。
将梯度裁剪作为防护措施
当梯度的综合范数超过阈值时,裁剪功能会对其重新缩放,因此它应放在backward()和optimizer.step()之间执行:
loss.backward()
gradient_norm = (
torch.nn.utils.clip_grad_norm_(
model.parameters(),
max_norm=1.0,
)
)
optimizer.step()
clip_grad_norm_会返回裁剪前的范数值,这一数值也可用于监控:
print(
"Gradient norm before clipping:",
float(gradient_norm),
)
如果几乎每一步都会超过阈值,那就说明存在过高的学习率或数值不稳定等问题,裁剪功能只能起到临时防护作用,无法替代合理设置的学习率。
验证优化器是否改变了权重
在更新之前复制一个参数。.clone() 方法非常重要:如果没有使用它,before 变量会与原参数共享存储空间,从而导致其也会发生变化:
parameter_name, parameter = next(
model.named_parameters()
)
before = parameter.detach().clone()
执行一次训练步骤:
optimizer.zero_grad(
set_to_none=True
)
_, loss = model(
inputs,
targets,
)
loss.backward()
optimizer.step()
测量变化情况:
after = parameter.detach()
maximum_change = (
after - before
).abs().max().item()
print(
"Maximum parameter change:",
maximum_change,
)
并且需要满足以下条件:
assert maximum_change > 0
如果权重没有变化,可能是因为学习率过低、优化器是在未考虑模型参数的情况下构建的(例如在模型被替换之前)、缺少梯度、缺失 optimizer.step() 方法,或是参数被冻结了。
双向调整学习率
过高的学习率会导致损失值迅速上升或剧烈波动、梯度范数过大、出现 NaN 损失值,以及样本始终无法得到改善。初步的解决办法是降低学习率,例如:
max_learning_rate = 1e-4
而不是:
max_learning_rate = 1e-3
在损失值旁边记录学习率。在使用热身阶段时,不稳定性往往恰好出现在峰值处,而仅通过损失曲线图则无法发现这一点。
如果学习率过低,情况会有所不同:尽管存在梯度,损失值下降得非常缓慢,参数几乎不会移动,甚至单批次测试也需要很多步才能完成。此时应提高学习率,例如设为:
max_learning_rate = 3e-4
而不是:
max_learning_rate = 1e-5
并没有一种通用的正确值;它取决于模型和批量的大小、优化器以及数据集。可以进行一些仅改变学习率的短期实验。
防止损失值持续上升的检查清单
按顺序思考以下问题。数据方面:
Are targets shifted by one token?
Are token IDs within range?
Does decoded input look correct?
模型方面:
Are logits shaped [B, T, V]?
Is the causal mask valid?
Are positions on the correct device?
损失值方面:
Does cross-entropy receive raw logits?
Are logits and targets flattened correctly?
将 softmax 概率直接输入 cross_entropy 是一个常见错误,因为该函数本身会计算对数 softmax。梯度情况如下:
Do all important parameters receive gradients?
Are gradient norms finite and nonzero?
优化器:
Is the learning rate positive?
Does optimizer.step() run?
Do parameters change?
容量:
Can the model overfit one batch?
This order avoids random trial and error.
这种排序方式可以逐个排除某一类原因,而无需依赖试错法。
识别过拟合与欠拟合
过拟合的表现为曲线发散:
Training loss: continues decreasing
Validation loss: stops decreasing or increases
量化差距大小:
generalization_gap = (
validation_loss
- training_loss
)
解决方法包括保留最佳的验证集检查点、提高 dropout 或权重衰减系数、缩小模型规模、增加更多或更具多样性的数据,以及提前停止训练。应使用验证集来确定停止点;使用测试集会导致信息泄露并抬高最终得分。
欠拟合的表现为两条曲线始终保持在较高水平:
Training loss: remains high
Validation loss: remains similarly high
可能的原因包括容量不足、训练轮次过少、学习率过低、上下文窗口过短、数据对模型架构来说过于复杂,或是分词方式导致上下文丢失。解决方案包括增加训练步数、提升嵌入维度、添加更多Transformer层、扩大上下文长度、使用字节对编码而非字符编码,以及重新调整学习率。首先重新运行单批次测试:如果模型无法记住一个批次的数据,问题出在正确性或优化方面,而非容量问题。
诊断重复生成现象
重复生成的典型表现如下:
the the the the
或者,使用WikiText标题标记的话则是:
= = = = = = =
造成重复生成的原因包括贪婪解码、温度值过低、top-k阈值过小、模型训练不足或过拟合、数据中存在重复结构,以及上下文窗口过短。可以尝试采用更均衡的采样方式:
temperature = 0.8
top_k = 20
top_p = 0.9
如果惩罚力度适中,重复惩罚反而能起到帮助作用:
repetition_penalty = 1.05
在字符模型中,过重的惩罚会阻碍字母的重复使用,进而迅速破坏拼写正确性。如果每次解码过程仍需循环处理,问题出在模型本身而非采样器。关于这些设置之间的相互作用,可参阅我们关于温度、top-k和top-p的指南。
诊断混乱的生成结果
另一种故障表现则为出现异常符号、断词、过多标点、话题突然跳转以及难以理解的字符串。可能的原因是温度值过高、未启用top-k或top-p过滤、分词器不匹配、使用了错误的检查点、模型训练不足导致验证损失过高,或是权重从未成功加载。可以尝试采用更严格的采样策略:
temperature = 0.6
top_k = 10
top_p = 0.9
确认权重确实来自检查点:
model.load_state_dict(
checkpoint["model_state_dict"]
)
并且采样时dropout功能处于关闭状态:
model.eval()
输出忽略提示词时
提示词可能非常简短或与训练数据不符;模型可能规模较小、训练不足、处理长距离依赖的能力较弱,或受限于较短的上下文长度;此外字符标记会增加语义模式学习的难度。请测试更长的、类似WikiText风格的提示词,并对比最简短的提示词与更丰富的提示词:
"The "
以及更复杂的提示词:
"The history of the city began"
后者能为模型提供更多参考信息。如果续写结果仍出现偏差,需检查验证损失及注意力机制的实现情况。
无法加载的检查点
出现的错误很常见:
Missing key(s) in state_dict
Unexpected key(s) in state_dict
Size mismatch
这意味着你构建的模型与保存的模型不一致:其配置、层数量、嵌入维度、词汇表大小或权重绑定发生了变化,类别或属性被重命名,或者正在加载来自其他架构的优化器状态。请检查保存的配置:
print(
checkpoint["config"]
)
应基于该配置而非当前默认设置来构建模型:
config = MiniGPTConfig(
**checkpoint["config"]
)
model = MiniGPT(config)
然后加载状态字典。如果用新配置构建模型却期望旧的权重适用,就会导致大多数此类错误。
列出缺失和意外的键
仅用于诊断,以非严格模式加载并打印差异:
load_result = model.load_state_dict(
checkpoint["model_state_dict"],
strict=False,
)
print(
"Missing keys:",
load_result.missing_keys,
)
print(
"Unexpected keys:",
load_result.unexpected_keys,
)
这些列表通常会显示问题原因,比如子模块被重命名了。在推理或继续训练时需严格控制加载过程,这样不兼容的检查点会明确报错,而不会让各层保持随机初始值。
将优化器状态移至正确的设备
恢复优化器后,其内部的张量(如AdamW的矩估计值)可能位于与模型不同的设备上。此工具会将状态中的所有张量移至正确位置:
def move_optimizer_to_device(
optimizer,
device,
):
for state in optimizer.state.values():
for key, value in state.items():
if torch.is_tensor(value):
state[key] = value.to(
device
)
应在加载后立即调用它:
optimizer.load_state_dict(
checkpoint[
"optimizer_state_dict"
]
)
move_optimizer_to_device(
optimizer,
device,
)
当在一台机器上保存模型而在另一台机器上继续训练时,这一点尤为重要,例如从CUDA切换到MPS或CPU。
将关键检查整合到一个健康状态函数中
把最重要的断言集中到同一个函数中:令牌的数据类型与范围、目标偏移量、logits的形状、损失值是否有限,以及与随机基线的比较结果:
def run_model_health_checks(
model,
tokenizer,
train_ids,
device,
):
model.eval()
assert train_ids.dtype == torch.long
assert train_ids.min().item() >= 0
assert (
train_ids.max().item()
< tokenizer.vocab_size
)
batch_size = 4
block_size = min(
32,
model.config.block_size,
)
inputs, targets = get_batch(
data=train_ids,
batch_size=batch_size,
block_size=block_size,
device=device,
)
assert inputs.shape == targets.shape
assert torch.equal(
inputs[:, 1:],
targets[:, :-1],
)
with torch.no_grad():
logits, loss = model(
inputs,
targets,
)
assert logits.shape == (
batch_size,
block_size,
tokenizer.vocab_size,
)
assert torch.isfinite(loss)
expected_loss = math.log(
tokenizer.vocab_size
)
print("Current loss:", loss.item())
print(
"Random baseline:",
expected_loss,
)
print("Model health checks passed.")
对于已训练好的检查点,损失值应明显低于基准值;否则说明权重未能正确加载或分词器不匹配。
利用前向钩子定位数值问题
当出现未知位置的NaN值时,前向钩子会在训练过程中检查每个模块的输出。该钩子能够处理单个张量或元组,一旦检测到第一个非有限值,就会以该模块的类名作为错误信息抛出异常:
def finite_output_hook(
module,
inputs,
output,
):
tensors = []
if torch.is_tensor(output):
tensors = [output]
elif isinstance(output, tuple):
tensors = [
item
for item in output
if torch.is_tensor(item)
]
for tensor in tensors:
if not torch.isfinite(
tensor
).all():
raise FloatingPointError(
"Nonfinite output detected in "
f"{module.__class__.__name__}"
)
将其附加到每一个线性层、层归一化层和嵌入层上,并保存对应的处理对象:
hooks = []
for module in model.modules():
if isinstance(
module,
(
torch.nn.Linear,
torch.nn.LayerNorm,
torch.nn.Embedding,
),
):
hooks.append(
module.register_forward_hook(
finite_output_hook
)
)
执行一次前向传播;由于各层是按顺序执行的,首次出现的异常就会指示出出现问题的第一个层类型。之后再移除这些钩子:
for hook in hooks:
hook.remove()
钩子在每次前向调用时都会执行,从而降低模型运行速度,因此仅应在排查错误时使用它们。要获取准确的模块路径,请在注册时记录 named_modules() 返回的名称。
完整的诊断脚本
所有检查功能都整合在一个命令行工具中。将其保存为:
debug_mini_gpt.py
该脚本定义了一个最简版的CharTokenizer、设备选择机制、指纹识别功能、批量处理辅助工具以及各种数值计算函数。它根据检查点中的配置来构建模型,会排除词汇表大小不同的分词器,随后在固定种子的前提下运行八项测试:分词器的往返处理、标记范围检测、批量移位测试、前向传播测试、因果独立性测试、梯度计算测试、优化器更新测试;此外还可以选择对新的调试模型进行单批次过拟合测试。有两点值得注意:分词器测试会解码存储的ID后再重新编码,以此验证真实数据集;优化器测试则使用全新的AdamW实例,避免旧状态产生干扰。
import argparse
import hashlib
import json
import math
from pathlib import Path
import torch
from mini_gpt import MiniGPT
from mini_gpt import MiniGPTConfig
class CharTokenizer:
def __init__(self, chars):
self.chars = chars
self.vocab_size = len(chars)
self.stoi = {
char: index
for index, char in enumerate(chars)
}
self.itos = {
index: char
for index, char in enumerate(chars)
}
@classmethod
def from_file(cls, path):
with open(
path,
"r",
encoding="utf-8",
) as file:
data = json.load(file)
return cls(data["chars"])
def encode(self, text):
return [
self.stoi[char]
for char in text
]
def decode(self, token_ids):
return "".join(
self.itos[int(token_id)]
for token_id in token_ids
)
def get_device():
if torch.cuda.is_available():
return torch.device("cuda")
if (
hasattr(torch.backends, "mps")
and torch.backends.mps.is_available()
):
return torch.device("mps")
return torch.device("cpu")
def tokenizer_fingerprint(chars):
payload = json.dumps(
chars,
ensure_ascii=False,
separators=(",", ":"),
).encode("utf-8")
return hashlib.sha256(
payload
).hexdigest()
def get_batch(
data,
batch_size,
block_size,
device,
):
start_positions = torch.randint(
low=0,
high=len(data) - block_size,
size=(batch_size,),
)
inputs = torch.stack([
data[
position:
position + block_size
]
for position in start_positions
])
targets = torch.stack([
data[
position + 1:
position + block_size + 1
]
for position in start_positions
])
return (
inputs.to(device),
targets.to(device),
)
def assert_finite_tensor(
tensor,
name,
):
if not torch.isfinite(
tensor
).all():
raise FloatingPointError(
f"{name} contains NaN or infinity"
)
def calculate_gradient_norm(model):
squared_norm = 0.0
for parameter in model.parameters():
if parameter.grad is None:
continue
norm = (
parameter.grad
.detach()
.norm(2)
.item()
)
squared_norm += norm ** 2
return squared_norm ** 0.5
def load_model(
checkpoint_path,
device,
):
checkpoint = torch.load(
checkpoint_path,
map_location=device,
)
config = MiniGPTConfig(
**checkpoint["config"]
)
model = MiniGPT(config)
model.load_state_dict(
checkpoint["model_state_dict"]
)
model = model.to(device)
return model, checkpoint
def test_tokenizer(
tokenizer,
train_ids,
):
print("\n1. Testing tokenizer")
print(
"Vocabulary size:",
tokenizer.vocab_size,
)
print(
"Tokenizer fingerprint:",
tokenizer_fingerprint(
tokenizer.chars
),
)
sample_ids = train_ids[:300]
sample_text = tokenizer.decode(
sample_ids.tolist()
)
round_trip_ids = tokenizer.encode(
sample_text
)
assert round_trip_ids == (
sample_ids.tolist()
)
print("Decoded sample:")
print(repr(sample_text))
print(
"Tokenizer round-trip test passed."
)
def test_token_ids(
tokenizer,
train_ids,
):
print("\n2. Testing token IDs")
print("Shape:", train_ids.shape)
print("Dtype:", train_ids.dtype)
minimum_id = train_ids.min().item()
maximum_id = train_ids.max().item()
print("Minimum ID:", minimum_id)
print("Maximum ID:", maximum_id)
assert minimum_id >= 0
assert maximum_id < (
tokenizer.vocab_size
)
print("Token ID test passed.")
def test_batch(
tokenizer,
train_ids,
block_size,
device,
):
print("\n3. Testing batches")
inputs, targets = get_batch(
data=train_ids,
batch_size=4,
block_size=block_size,
device=device,
)
print("Input shape:", inputs.shape)
print("Target shape:", targets.shape)
assert inputs.shape == targets.shape
assert inputs.dtype == torch.long
assert targets.dtype == torch.long
assert torch.equal(
inputs[:, 1:],
targets[:, :-1],
)
input_text = tokenizer.decode(
inputs[0].cpu().tolist()
)
target_text = tokenizer.decode(
targets[0].cpu().tolist()
)
print("Input sample:")
print(repr(input_text))
print("Target sample:")
print(repr(target_text))
print("Batch shift test passed.")
return inputs, targets
def test_forward_pass(
model,
tokenizer,
inputs,
targets,
):
print("\n4. Testing forward pass")
model.eval()
with torch.no_grad():
logits, loss = model(
inputs,
targets,
)
print("Logits shape:", logits.shape)
print("Loss:", loss.item())
assert logits.shape == (
inputs.size(0),
inputs.size(1),
tokenizer.vocab_size,
)
assert_finite_tensor(
logits,
"logits",
)
assert_finite_tensor(
loss,
"loss",
)
print(
"Random baseline:",
math.log(tokenizer.vocab_size),
)
print("Forward-pass test passed.")
def test_future_independence(
model,
tokenizer,
device,
):
print(
"\n5. Testing causal independence"
)
model.eval()
sequence_length = min(
16,
model.config.block_size,
)
prefix_length = (
sequence_length // 2
)
sequence_a = torch.randint(
0,
tokenizer.vocab_size,
(1, sequence_length),
device=device,
)
sequence_b = sequence_a.clone()
sequence_b[
:,
prefix_length:
] = torch.randint(
0,
tokenizer.vocab_size,
(
1,
sequence_length
- prefix_length,
),
device=device,
)
with torch.no_grad():
logits_a, _ = model(sequence_a)
logits_b, _ = model(sequence_b)
difference = (
logits_a[
:,
:prefix_length,
:,
]
- logits_b[
:,
:prefix_length,
:,
]
).abs().max().item()
print(
"Maximum shared-prefix difference:",
difference,
)
assert torch.allclose(
logits_a[
:,
:prefix_length,
:,
],
logits_b[
:,
:prefix_length,
:,
],
atol=1e-5,
)
print(
"Causal independence test passed."
)
def test_gradients(
model,
inputs,
targets,
):
print("\n6. Testing gradients")
model.train()
model.zero_grad(
set_to_none=True
)
_, loss = model(
inputs,
targets,
)
loss.backward()
missing_gradients = []
nonfinite_gradients = []
for name, parameter in (
model.named_parameters()
):
if not parameter.requires_grad:
continue
if parameter.grad is None:
missing_gradients.append(name)
continue
if not torch.isfinite(
parameter.grad
).all():
nonfinite_gradients.append(
name
)
print(
"Gradient norm:",
calculate_gradient_norm(model),
)
if missing_gradients:
print(
"Missing gradients:",
missing_gradients,
)
if nonfinite_gradients:
print(
"Nonfinite gradients:",
nonfinite_gradients,
)
assert not missing_gradients
assert not nonfinite_gradients
print("Gradient test passed.")
def test_optimizer_update(
model,
inputs,
targets,
):
print("\n7. Testing optimizer update")
optimizer = torch.optim.AdamW(
model.parameters(),
lr=1e-3,
)
name, parameter = next(
model.named_parameters()
)
before = parameter.detach().clone()
optimizer.zero_grad(
set_to_none=True
)
_, loss = model(
inputs,
targets,
)
loss.backward()
torch.nn.utils.clip_grad_norm_(
model.parameters(),
max_norm=1.0,
)
optimizer.step()
maximum_change = (
parameter.detach() - before
).abs().max().item()
print("Tracked parameter:", name)
print(
"Maximum parameter change:",
maximum_change,
)
assert maximum_change > 0
print(
"Optimizer update test passed."
)
def run_single_batch_overfit(
tokenizer,
train_ids,
device,
steps,
):
print(
"\n8. Running single-batch "
"overfitting test"
)
config = MiniGPTConfig(
vocab_size=tokenizer.vocab_size,
block_size=32,
embedding_dim=64,
num_heads=4,
num_layers=2,
expansion_factor=4,
dropout=0.0,
)
model = MiniGPT(config).to(device)
inputs, targets = get_batch(
data=train_ids,
batch_size=8,
block_size=config.block_size,
device=device,
)
optimizer = torch.optim.AdamW(
model.parameters(),
lr=1e-3,
)
initial_loss = None
final_loss = None
for step in range(steps):
optimizer.zero_grad(
set_to_none=True
)
_, loss = model(
inputs,
targets,
)
if initial_loss is None:
initial_loss = loss.item()
assert_finite_tensor(
loss,
"single-batch loss",
)
loss.backward()
torch.nn.utils.clip_grad_norm_(
model.parameters(),
max_norm=1.0,
)
optimizer.step()
final_loss = loss.item()
if (
step % 50 == 0
or step == steps - 1
):
print(
f"Step {step:4d}: "
f"loss {final_loss:.4f}"
)
print(
"Initial loss:",
initial_loss,
)
print(
"Final loss:",
final_loss,
)
assert final_loss < initial_loss
print(
"Single-batch overfitting "
"test passed."
)
def parse_args():
parser = argparse.ArgumentParser(
description=(
"Run Mini-GPT diagnostic tests"
)
)
parser.add_argument(
"--checkpoint",
type=str,
default=(
"checkpoints/mini_gpt_best.pt"
),
)
parser.add_argument(
"--tokenizer",
type=str,
default=(
"tokenizer/char_tokenizer.json"
),
)
parser.add_argument(
"--train-data",
type=str,
default=(
"data/encoded/train_ids.pt"
),
)
parser.add_argument(
"--overfit-steps",
type=int,
default=300,
)
parser.add_argument(
"--skip-overfit",
action="store_true",
)
return parser.parse_args()
def main():
args = parse_args()
torch.manual_seed(42)
device = get_device()
print("Using device:", device)
tokenizer = CharTokenizer.from_file(
args.tokenizer
)
train_ids = torch.load(
args.train_data,
map_location="cpu",
).long()
checkpoint_path = Path(
args.checkpoint
)
if not checkpoint_path.exists():
raise FileNotFoundError(
f"Checkpoint not found: "
f"{checkpoint_path}"
)
model, checkpoint = load_model(
checkpoint_path=checkpoint_path,
device=device,
)
if (
tokenizer.vocab_size
!= model.config.vocab_size
):
raise ValueError(
"Tokenizer and model vocabulary "
"sizes do not match"
)
print(
"Checkpoint step:",
checkpoint.get("step"),
)
test_tokenizer(
tokenizer,
train_ids,
)
test_token_ids(
tokenizer,
train_ids,
)
test_block_size = min(
32,
model.config.block_size,
)
inputs, targets = test_batch(
tokenizer=tokenizer,
train_ids=train_ids,
block_size=test_block_size,
device=device,
)
test_forward_pass(
model=model,
tokenizer=tokenizer,
inputs=inputs,
targets=targets,
)
test_future_independence(
model=model,
tokenizer=tokenizer,
device=device,
)
test_gradients(
model=model,
inputs=inputs,
targets=targets,
)
test_optimizer_update(
model=model,
inputs=inputs,
targets=targets,
)
if not args.skip_overfit:
run_single_batch_overfit(
tokenizer=tokenizer,
train_ids=train_ids,
device=device,
steps=args.overfit_steps,
)
print(
"\nAll requested diagnostics passed."
)
if __name__ == "__main__":
main()
最新的 PyTorch 版本改变了 torch.load 的默认行为,使其仅用于加载权重。因此根据您的版本及检查点内容,可能需要显式设置 weights_only 参数;请查阅最新文档。
运行诊断测试
使用默认路径运行完整测试套件:
python debug_mini_gpt.py
如需快速检查,可跳过过拟合检测阶段。对应的参数为 --skip-overfit,前面有两个连字符:
python debug_mini_gpt.py - skip-overfit
使用另一个检查点:
python debug_mini_gpt.py \
--checkpoint checkpoints/mini_gpt_latest.pt
为过拟合检测增加更多步骤:
python debug_mini_gpt.py \
--overfit-steps 500
尾随的反斜杠用于在类 Unix 风格的shell中延续命令;如果您的shell不支持该功能,请将命令写在同一行。
实际调试步骤
当模型出现异常时,请按顺序逐步排查,不要跳过任何步骤。
步骤1至5:数据与形状
读取解码后的数据:
Does the tokenized dataset decode correctly?
检查目标偏移情况:
Does inputs[:, 1:] equal targets[:, :-1]?
检查令牌范围:
Are all IDs between 0 and vocab_size - 1?
检查张量形状:
Inputs: [B, T]
Targets: [B, T]
Logits: [B, T, V]
检查初始损失值:
Is it near log(vocab_size) for a new model?
步骤6至10:模型行为与训练
检查因果独立性:
Can changing the future affect prefix logits?
唯一可接受的答案是否。检查梯度值:
Are gradients present, finite, and nonzero?
检查参数更新情况:
Does optimizer.step() change weights?
记住一个批次的数据:
Can the model memorize a tiny fixed batch?
只有之后才启动完整训练:
Only after all earlier tests pass should you invest in a long training run.
各阶段的检查清单
训练前:
□ Dataset files exist
□ Tokenizer round-trip works
□ Token IDs are within vocabulary range
□ Decoded data looks correct
□ Inputs and targets are shifted by one
□ Batch tensors use torch.long
□ Model and batch use the same device
□ Logits have shape [B, T, V]
□ Initial loss is near log(V)
□ Future-independence test passes
□ All important parameters receive gradients
□ Optimizer changes parameters
□ Model can overfit one batch
训练过程中:
□ Loss remains finite
□ Gradient norms remain finite
□ Learning rate follows the intended schedule
□ Training loss decreases
□ Validation loss is evaluated in eval mode
□ Best checkpoint updates when validation improves
□ Samples become more structured
生成过程中:
□ Best checkpoint is loaded
□ Matching tokenizer is loaded
□ Model is in evaluation mode
□ Context is cropped to block size
□ Only final-position logits are sampled
□ Temperature is positive
□ Top-k does not exceed vocabulary size
□ Repetition is measured, not only observed
使调试更加困难的习惯
一次性更改过多设置
如果某次实验同时改变了所有这些因素,就无法将结果归因于其中任何一个:
Learning rate
Batch size
Dropout
Model size
Context length
每次实验只修改一个主要变量。
仅凭样本判断
质量不佳的文本可能源于训练不足、解码错误、使用了错误的检查点或分词器、过拟合或欠拟合,而样本无法区分这些原因。应先查看指标和流程测试结果。
忽略警告信息
涉及张量形状改变、切换到其他设备、出现NaN或无穷大值,以及检查点键不匹配的警告通常表明存在真实错误;在屏蔽这些警告之前需先理解其含义。
跳过简单测试
在进行长时间运行之前,先从小规模测试开始:
Tiny model
Tiny batch
Short context
Few training steps
然后再逐步扩大规模。
将截断处理视为解决方案
裁剪机制可以处理一次过大的更新,但频繁出现裁剪现象则表明有更深层的问题需要关注:学习率、初始化参数、损失缩放比例、数值精度或数据中的异常值。
练习:故意破坏流程
在看到测试失败后,人们会更信任该测试,因此每个练习都会设置一个已知的错误。
移除目标偏移
使输入和目标完全相同,确认批量测试失败后,再恢复目标偏移。
注入超出范围的标记
将某个标记的ID设置为:
tokenizer.vocab_size
范围检查必须失败,因为有效的最高ID为:
vocab_size - 1
禁用因果掩码
暂时移除掩码并运行与未来无关性相关的测试;此时仅更改后缀应能改变前缀的逻辑值。
使用极端的学习率
设置为:
learning_rate = 0.1
记录损失值、梯度范数、参数值以及非有限值的检测情况,并尽量缩短运行时间。
冻结模型
应用以下操作,观察梯度与优化器测试的结果:
for parameter in model.parameters():
parameter.requires_grad = False
加载到错误的架构中
将检查点加载到在某些方面存在差异的模型中,然后查看缺失键、意外键以及大小不匹配等错误信息:
Vocabulary size
Embedding dimension
Number of layers
比较dropout设置
使用两种不同的值进行单批次测试,对比两者记忆数据的速度:
dropout = 0.0
dropout = 0.2
编写调试报告
将这些结果以JSON格式保存,以便进行对比和复现故障:
Tokenizer fingerprint
Vocabulary size
Token range
Batch shape
Initial loss
Expected baseline
Gradient norm
Missing gradients
Parameter update size
Single-batch final loss
关键要点
- 即使任务无错误完成,模型仍可能学错内容。
[B, T, V],且新模型的初始值应接近log(vocab_size)。完整流程如下:
Environment
→ files
→ tokenizer
→ token IDs
→ batch shifting
→ shapes
→ initial loss
→ causal independence
→ gradients
→ optimizer updates
→ one-batch overfitting
→ full training
→ evaluation
→ generation
所有方法背后都遵循同一条规则:不要凭直觉进行调试;应编写能够针对某一假设进行验证的测试,确认之后再继续下一步。
接下来的自然步骤是采用更好的数据表示方式。字符标记会生成很长的序列,而词汇表规模也会变得极为庞大;字节对编码则能学习常见字符序列的合并方式,从而缩短序列长度,使相同的上下文窗口能够容纳更多文本。采用这种编码方式意味着需要训练合并规则、重新对WikiText-2进行编码,并调整模型的词汇表规模,而此处的所有验证步骤依然保持不变:
Character tokens
→ learned subword merges
→ shorter sequences
→ better use of the context window
相关阅读
- 下一个令牌循环:在接触智能体之前了解LLM的思维模型 — 通过简单的离线Python示例,学习令牌、上下文窗口、采样及生成循环的工作原理,同时解释为何会有RAG、ReAct和LangGraph这些技术。
- 从单节点聊天机器人到LangGraph中的MCP支持智能体 — 逐步构建LangGraph应用:状态与还原器、边、工具循环、检查点线程、三种流式模式,以及通过MCP提供的工具。