JRY's digital garden

第5课 手写 Train.py 大模型代码逻辑

17 min read

内容总结与提炼

核心主题与目标

  • 核心主题: 详细讲解并手写大模型训练脚本 train.py 的完整逻辑。
  • 关键信息: 视频首先完成了 model.py 中用于推理的 generate 方法，然后逐步构建了 train.py，涵盖了数据加载、预处理、批次生成、模型实例化、优化器设置、训练循环以及损失评估等核心环节。
  • 学习目标: 观众应掌握从零开始编写一个完整的 PyTorch 大模型训练脚本的能力，理解训练过程中的关键步骤和代码实现细节。

关键知识点梳理

  • generate 方法实现:
    • 功能: 用于模型推理，根据输入的文本 (prompt) 循环生成新的 token。
    • 核心逻辑: 在一个循环中，每次都对当前序列进行处理，预测下一个 token，然后将新 token 拼接回原序列，再进行下一次预测。
    • 上下文截断: 每次送入模型的序列长度固定为 context_length，通过从后往前截取实现。
    • Token 采样: 讲解了如何从 logits 中通过 softmax 得到概率分布，并使用 torch.multinomial 进行采样，而不仅仅是取概率最高的 argmax。
  • train.py 脚本构建:
    • 数据加载与预处理: 使用 open 读取文本数据，利用 tiktoken 进行分词 (encode)，并将 token 索引转换为 PyTorch Tensor。
    • 数据集划分: 将全部数据按比例（如 90/10）切分为训练集 (train_data) 和验证集 (validation_data)。
    • get_batch 函数:
      • 功能: 从数据集中随机抽取批次数据 (batch)。
      • 实现: 随机生成起始索引，然后根据 context_length 切片，生成 x_batch 和 y_batch (y_batch 是 x_batch 向右平移一位的结果)。
      • 批处理: 使用 torch.stack 将多个样本堆叠成一个批次 Tensor。
  • 训练循环 (Training Loop):
    • 模型与优化器: 实例化 Model 类，并定义 AdamW 优化器，将其与模型参数绑定。
    • 核心步骤:
      1. 获取一个批次的数据 (get_batch)。
      2. 前向传播：将数据送入模型，计算 loss。
      3. 清零梯度 (optimizer.zero_grad())。
      4. 反向传播：计算梯度 (loss.backward())。
      5. 更新参数：根据梯度更新模型权重 (optimizer.step())。
  • 损失评估 (Loss Estimation):
    • 目的: 监控模型训练状态，防止过拟合。
    • estimate_loss 函数:
      • 功能: 在不更新梯度的情况下，计算训练集和验证集的平均损失。
      • 实现: 循环多次（eval_iters），对每个数据集分别取样计算 loss，最后求平均值。
    • 定期评估: 在训练循环中，每隔一定步数 (eval_interval) 调用此函数，并打印 train_loss 和 val_loss。

逻辑结构与关系

  • 承上启下: 视频首先完成了上一讲 model.py 中遗留的 generate 方法，然后以此为基础，进入本讲的核心——train.py 的编写。
  • 循序渐进: train.py 的编写遵循了标准的机器学习流程：数据准备 → 模型定义 → 训练循环 → 评估与保存。每个部分都进行了详细的代码讲解。
  • 原理-实践: 讲师在编写代码的同时，不断解释其背后的原理，如张量形状的变化、优化器的作用、损失评估的必要性等，将理论与实践紧密结合。

重要细节与论证

  • 张量形状 (Tensor Shape): 视频反复强调理解和操作 Tensor 形状是 PyTorch 编程的关键。详细解释了在截断、批处理、softmax 等操作中 Tensor 维度的变化。
  • 代码简化技巧: 演示了 Python 中将多行 if-else 语句改写为一行三元表达式的技巧，提高了代码的简洁性。
  • 参数配置: 讲解了将学习率 (learning_rate)、批次大小 (batch_size) 等超参数预先定义的好处，并提到未来会将它们统一放入配置文件中，便于管理。
  • torch.multinomial vs argmax: 解释了为什么在生成时使用带随机性的采样 (multinomial) 比确定性的 argmax 效果更好，可以增加生成文本的多样性。

总结与复习

核心要点回顾

  1. 推理 (generate) 逻辑: 推理是一个自回归（auto-regressive）的循环过程，不断预测并拼接下一个 token，同时保持上下文窗口长度固定。
  2. 训练脚本 (train.py) 核心流程: 数据加载与分词 → 数据集划分 → get_batch 函数随机取样 → 训练循环（前向传播、计算 loss、反向传播、更新权重）。
  3. 损失评估的重要性: 必须定期计算训练集和验证集的 loss，以监控模型训练进程，判断是否收敛或过拟合。

易混淆点/难点

  • x_batch 和 y_batch 的关系: y_batch 是 x_batch 向右平移一位的结果。在 Transformer 中，模型需要根据前面的 token (x) 预测紧随其后的 token (y)，这是训练自回归语言模型的基础。
  • optimizer.zero_grad() 的作用: 在 PyTorch 中，梯度是默认累加的。如果不清零，每次 backward() 计算出的梯度会叠加到之前的梯度上，导致错误的参数更新。因此，在每个训练步骤开始前都必须清零。
  • 张量维度的操作: 对于初学者，PyTorch 中各种 slicing, stacking, view 等操作可能会比较复杂，需要通过多练习来熟悉。

启发与思考

  • 代码的模块化: 视频将模型 (model.py) 和训练逻辑 (train.py) 分开，体现了良好的编程习惯。未来还可以将配置 (config.py)、数据处理等进一步模块化。
  • 超参数的调试: 视频中提到了 learning_rate, max_iters, eval_interval 等超参数，这些参数的选择对模型最终性能至关重要，需要通过实验来调整。
  • 从训练到推理的闭环: 视频完整展示了从模型定义、训练到最终推理的整个流程，为理解大模型的工作原理提供了清晰的路线图。

1. model.py 回顾与 generate 方法实现

1.1. Model 类回顾

  • 文件结构: 上一期视频中，我们手写了 Transformer 的核心组件，并将其封装在 model.py 文件中，包括：
    • FeedForward 类
    • Head (单头注意力)
    • MultiHeadAttention (多头注意力)
    • Block (Transformer 块)
    • Model (完整的 Transformer 模型)
  • Model 类的 forward 方法:
    • 输入: x_batch (输入样本) 和 y_batch (目标样本，可为空)。
    • 输出: logits (预测结果) 和 loss (损失值)。
    • 设计思想: 通过 y_batch 是否为空，实现了训练和推理逻辑的复用。
      • 训练时: x_batch 和 y_batch 都传入，计算 loss。
      • 推理时: 只传入 x_batch，loss 为 None。

1.2. generate 方法的核心逻辑

  • 目的: 实现模型的推理功能，即根据输入的 prompt (一段文本) 生成后续的文本。
  • 自回归 (Auto-regressive) 过程:
    1. 将初始 prompt (作为 x_batch) 输入模型，预测出下一个 token。
    2. 将预测出的新 token 拼接到 prompt 的末尾。
    3. 由于模型的上下文长度 (context_length) 有限，需要从序列的末尾截取固定长度的片段作为新的输入。
    4. 重复以上步骤，直到生成所需数量的 token。
  • 循环: 该方法的核心是一个 for 循环，循环次数由 max_new_tokens 参数决定，即希望生成的新 token 数量。

1.3. generate 方法代码详解

  1. 循环设置:

    • for _ in range(max_new_tokens):
    • 循环 max_new_tokens 次，每次生成一个新 token。

  2. 上下文截断 (Context Cropping):

    • x_cropped = x[:, -context_length:]
    • x 的形状为 (B, T)，其中 B 是 batch_size，T 是时间步/序列长度。
    • [:, -context_length:] 表示对第二个维度（时间步）进行切片，只保留从后往前数的 context_length 个 token。这确保了输入模型的序列长度始终是固定的。

  3. 获取预测 logits:

    • logits, loss = self(x_cropped)
    • 调用模型自身的 forward 方法进行前向传播。在推理时，loss 为 None，我们只关心 logits。
    • logits 的形状为 (B, T, vocab_size)。

  4. 聚焦最后一个时间步:

    • logits = logits[:, -1, :]
    • 我们只关心最后一个时间步的预测结果，因为它对应着下一个 token 的概率分布。
    • 切片后，logits 的形状变为 (B, vocab_size)。

  5. 应用 temperature:

    • logits = logits / temperature
    • temperature 是一个超参数，用于调节生成文本的随机性。
      • temperature < 1：使概率分布更“尖锐”，模型倾向于选择高概率的词，生成结果更确定。
      • temperature > 1：使概率分布更“平滑”，增加了低概率词被选中的机会，生成结果更多样。

  6. 计算概率 softmax:

    • probs = F.softmax(logits, dim=-1)
    • 对 logits 在最后一个维度（vocab_size 维度）上应用 softmax，将其转换为概率分布。

  7. 采样下一个 token:

    • x_next = torch.multinomial(probs, num_samples=1)
    • torch.multinomial 根据输入的概率分布 probs 进行采样，随机抽取 num_samples 个 token。
    • 相比于 argmax (总是选择概率最高的 token)，multinomial 引入了随机性，可以生成更多样化的文本。
    • x_next 的形状为 (B, 1)。

  8. 拼接新 token:

    • x = torch.cat([x, x_next], dim=1)
    • 将新生成的 token (x_next) 拼接到原始序列 x 的末尾（在 dim=1，即时间步维度上）。
    • 拼接后的 x 将在下一次循环中被再次截断和处理。

  9. 返回结果:

    • return x
    • 循环结束后，返回包含了初始 prompt 和所有新生成 token 的完整序列。

2. train.py 脚本编写

train.py 是驱动模型训练的核心文件，负责数据处理、模型调用、参数优化和状态监控。

2.1. 数据加载与预处理

  • 读取数据: 从 .csv 文件中读取原始文本数据。 
    with open('data/mingcheng.csv', 'r', encoding='utf-8') as f:
    text = f.read()
  • 初始化 Tokenizer: 使用 OpenAI 的 tiktoken 库进行分词。 
    import tiktoken
tokenizer = tiktoken.get_encoding("cl100k_base")
  • 编码与转换: 将文本编码为 token 索引，并转换为 PyTorch Tensor。 
    tokenized_text = tokenizer.encode(text)
data = torch.tensor(tokenized_text, dtype=torch.long, device=device)

2.2. 数据集划分

  • 将整个数据集按 90% 和 10% 的比例划分为训练集 (train_data) 和验证集 (val_data)。 
    n = int(0.9 * len(data))
train_data = data[:n]
val_data = data[n:]

2.3. get_batch 函数：获取数据批次

  • 功能: 从指定的数据集（训练集或验证集）中随机抽取一个批次 (batch) 的数据。
  • 实现步骤:
    1. 选择数据集: 根据传入的 split 参数（‘train’ 或 ‘val’）选择对应的数据集。
    2. 生成随机索引:
      • ix = torch.randint(len(data) - context_length, (batch_size,))
      • 生成 batch_size 个随机整数，作为每个样本在数据集中切片的起始位置。
      • 减去 context_length 是为了确保切片不会越界。
    3. 构建 x_batch 和 y_batch:
      • 使用列表推导式，根据随机索引 ix 从 data 中切片，构建 x 和 y。
      • x = [data[i:i+context_length] for i in ix]
      • y = [data[i+1:i+context_length+1] for i in ix] (y 是 x 向右平移一位的结果)
    4. 堆叠成 Tensor:
      • x_batch = torch.stack(x)
      • y_batch = torch.stack(y)
      • torch.stack 将列表中的多个 Tensor 堆叠成一个更高维度的 Tensor，形成批次。
    5. 返回批次: return x_batch, y_batch

2.4. estimate_loss 函数：损失评估

  • 目的: 在训练过程中，定期评估模型在训练集和验证集上的损失，以监控训练状态。
  • 核心逻辑:
    • 使用 @torch.no_grad() 装饰器，确保在该函数内不进行梯度计算，节省资源且不影响反向传播。
    • 循环 eval_iters 次，对训练集和验证集分别调用 get_batch 获取数据，并计算 loss。
    • 将每次计算的 loss 累加，最后求平均值，得到更稳定的损失评估。
    • 返回一个包含 train_loss 和 val_loss 的字典。

2.5. 模型实例化与优化器定义

  • 模型: 
    model = Model()
model.to(device) # 将模型移动到 GPU 或 CPU
  • 优化器:
    • 使用 AdamW，这是目前训练 Transformer 模型的常用优化器。
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
    • model.parameters(): 将模型的所有可训练参数传递给优化器。
    • lr: 设置学习率。

2.6. 训练主循环 (Training Loop)

  • 循环: for step in range(max_iters):
  • 核心步骤:
    1. 定期评估:
      • if step % eval_interval == 0:
      • 每隔 eval_interval 步，调用 estimate_loss 函数，并打印当前的 train_loss 和 val_loss。
    2. 获取批次数据:
      • x_batch, y_batch = get_batch('train')
    3. 前向传播与计算损失:
      • logits, loss = model(x_batch, y_batch)
    4. 反向传播:
      • optimizer.zero_grad(set_to_none=True): 清除上一轮的梯度。set_to_none=True 是一个性能优化选项。
      • loss.backward(): 根据 loss 计算所有参数的梯度。
    5. 更新参数:
      • optimizer.step(): 优化器根据计算出的梯度更新模型参数。

2.7. 运行与调试

  • 保存模型: 训练循环结束后，保存模型的状态字典。 
    torch.save(model.state_dict(), 'model.ckpt')
  • 运行脚本:
    • 启动训练后，控制台会按 eval_interval 的频率打印损失值。
    • 观察 train_loss 和 val_loss 的变化趋势，是判断模型训练是否正常的重要依据。

总结

本视频详细演示了如何从零开始编写一个用于训练大语言模型的 train.py 脚本。首先，视频回顾并完成了 model.py 中用于文本生成的 generate 方法，阐述了其自回归和循环采样的核心逻辑。随后，重点转向 train.py 的构建，内容涵盖了从加载原始文本数据、使用 tiktoken 进行分词、将数据切分为训练集和验证集，到实现一个高效的 get_batch 函数来随机生成数据批次。

视频的核心部分是训练流程的实现，包括：实例化模型、定义 AdamW 优化器、构建主训练循环。在循环中，详细讲解了前向传播计算损失、反向传播计算梯度、以及优化器更新参数这三大关键步骤。此外，视频还强调了通过编写 estimate_loss 函数并定期在训练和验证集上评估损失的重要性，这是监控模型性能和防止过拟合的关键实践。整个过程为学习者提供了一套完整、可复现的大模型训练代码框架。

Graph View