JRY's digital garden

第6课 手写 Inference.py 大模型代码逻辑

17 min read

目录

1. 训练代码修正与概念补充

1.1. 验证过程中的梯度问题

在 estimate_loss 函数中，模型会进行一次前向传播来计算验证集上的损失。这个过程仅用于评估，不应计算和累积梯度。

  • 问题: 如果不进行特殊处理，即使在验证阶段，PyTorch 默认也会跟踪计算图以备反向传播，消耗不必要的计算资源。
  • 解决方案: 使用 torch.no_grad() 上下文管理器。
    • 代码示例： 
      @torch.no_grad()
def estimate_loss():
    # ... a series of calculations ...
    return losses
    • 作用: 在 with torch.no_grad(): 块内的所有计算都不会被跟踪，从而阻止梯度累积。

1.2. 模型状态切换：训练与评估

模型在训练和推理（或验证）时，行为应该有所不同，特别是对于包含 Dropout 或 BatchNorm 等层的模型。

  • model.train(): 将模型设置为 训练模式。在这种模式下，Dropout 层会随机丢弃神经元，BatchNorm 层会使用当前批次的均值和方差。

  • model.eval(): 将模型设置为 评估模式。在这种模式下，Dropout 层会失效（即所有神经元都参与计算），BatchNorm 层会使用在整个训练集上学习到的均值和方差。

  • 正确实践: 在计算验证损失之前，调用 model.eval()；计算完毕后，调用 model.train() 将其切换回训练模式，以便继续训练。

    # In estimate_loss function or validation loop
model.eval()
# ... calculate validation loss ...
model.train()

Screenshot-[01:10]

1.3. Epoch 概念

  • Epoch: 指的是将整个训练数据集完整地过一遍的过程。如果一个 epoch 完成后模型还未充分学习，可以进行多个 epoch 的训练。
  • 实现方式: 通常在 max_iters 的外层再套一个 for 循环，即大循环 (epoch) 套小循环 (iterations/steps)。本课程的代码为了简化，暂时没有引入 epoch 的概念。

2. 初步训练与Inference脚本编写

2.1. 本地小模型训练

为了快速验证代码逻辑，首先在本地 CPU 环境下使用较小的参数进行训练。

  • 训练参数:
    • batch_size: 4
    • d_model: 64
    • n_head: 2
    • max_iters: 200 (后增加到 2000)
  • 流程:
    1. 运行 train.py 脚本。
    2. 脚本会定期打印 train_loss 和 validation_loss。
    3. 训练结束后，会将模型的状态字典（state_dict）保存到 model.ckpt 文件中。 Screenshot-[03:50]

2.2. 编写 inference.py 推理脚本

推理脚本的核心任务是加载训练好的模型并用它来生成文本。

  • 基本步骤:
    1. 导入依赖: 导入 torch、模型类 (GPT) 和 Tokenizer。
    2. 加载模型:
      • 实例化模型 model = GPT(...)。
      • 加载检查点文件 checkpoint = torch.load('model.ckpt')。
      • 将状态字典加载到模型实例中 model.load_state_dict(checkpoint)。
    3. 准备输入:
      • 定义一个输入提示词（prompt），例如：“农夫山泉”。
      • 使用 tokenizer.encode() 将提示词转换为 token 序列。
      • 将 token 序列转换为 torch.Tensor。
    4. 生成文本:
      • 调用 model.generate() 方法，传入输入 Tensor、最大生成长度 max_new_tokens 等参数。
    5. 解码输出:
      • generate() 方法返回的是 token 序列的 Tensor。
      • 使用 tokenizer.decode() 将 token 序列转换回人类可读的文本。
    6. 打印结果。 Screenshot-[04:51]

3. 初步推理结果分析与代码重构

3.1. 推理结果分析

在 d_model=64, n_head=2 (即每个头的维度仅为32) 的小模型上训练 2000 次后，推理结果并不理想。

  • 现象: 模型能生成一些在训练数据中出现过的词组，但无法构成有意义的、连贯的句子。例如，输入 “德氏”，模型可能续写出 “玉米”，但 “德氏玉米” 这个组合在数据集中可能并不存在。
  • 原因:
    1. 模型容量过小: d_model 和每个头的维度太小，限制了模型学习复杂语义关系的能力。
    2. 训练不充分: 训练迭代次数相对较少。
    3. 数据集特性: 训练数据是独立的 SKU 列表，模型倾向于“过拟合”这些短语，这对于 SKU 补全任务来说是期望的行为，但也需要更大的模型容量来学习更准确的组合。 Screenshot-[07:03]

3.2. 代码重构计划

为了便于进行不同超参数的实验，并使代码更具可维护性，需要进行重构。

  • 目标:
    1. 集中管理超参数: 将所有超参数（如 d_model, n_head, learning_rate 等）统一放在一个配置文件或一个配置对象中，而不是散落在代码各处。
    2. 集成实验跟踪工具: 引入 Weights & Biases (W&B)，用于可视化损失曲线、监控系统资源，并对比不同实验的效果。

4. 代码重构：引入配置文件管理

4.1. model.py 的改动

  • 将模型的 __init__ 方法修改为接收一个配置对象 config。
  • 模型内部的所有超参数都通过 config 对象来获取，例如 config.d_model, config.n_head。
  • generate 方法中增加了 top_k 采样参数，使其更加灵活。top_k 是一种限制采样范围的技术，只在概率最高的 k 个词中进行采样，可以增加生成文本的相关性。

4.2. train.py 的改动

这是改动的核心，train.py 成为配置和启动训练的中心。

  • 超参数配置块: 在文件顶部定义所有超参数，方便统一修改和管理。 Screenshot-[12:22]
  • 关键超参数解释:
    • d_model: 模型的嵌入维度，决定了模型参数量的关键因素。
    • batch_size: 每批次训练的样本数，影响训练速度和GPU显存占用。
    • context_length (block_size): 模型能处理的上下文窗口大小。
    • n_blocks (n_layer): Transformer 堆叠的层数。
    • n_head: 多头注意力机制中的头的数量。
    • d_model / n_head: 每个注意力头的实际工作维度，代表了其学习的语义空间大小。
    • learning_rate: 学习率。可以设置为动态衰减，但目前是固定的。

5. 实验跟踪与可视化：集成Weights & Biases (W&B)

W&B 是一个广泛使用的机器学习实验跟踪平台，可以帮助开发者记录和分析训练过程。

  • 集成步骤:

    1. 安装与登录: pip install wandb 并在命令行 wandb login。
    2. 初始化: 在训练脚本开始时调用 wandb.init()，并指定项目名称 (project) 和运行名称 (run_name)。
    3. 记录配置: 将超参数配置字典传给 wandb.config，W&B 会自动保存每次运行的配置。
    4. 记录指标: 在训练循环中，使用 wandb.log() 记录损失、准确率等关键指标，例如 wandb.log({'train_loss': train_loss, 'val_loss': val_loss})。 Screenshot-[14:36]

  • 模型保存的改进:

    • 在保存模型检查点 (.ckpt 文件) 时，不仅保存模型的权重 (state_dict)，还要将 超参数配置 一同保存进去。
    • 格式: 
      checkpoint = {
    'model': model.state_dict(),
    'model_args': gpt_config, # 包含d_model等超参数的配置对象
    'max_token_value': max_token_value
}
torch.save(checkpoint, 'model.ckpt')
    • 好处: 在推理或微调时，可以直接从模型文件中读取原始训练配置，无需手动重新指定，确保了环境的一致性。

6. 重构后的推理与分布式训练概念

6.1. 重构后的 inference.py

  • 加载流程:
    1. 加载检查点文件 checkpoint = torch.load('model.ckpt')。
    2. 从检查点中提取超参数配置 model_args = checkpoint['model_args']。
    3. 使用提取的配置来实例化模型 model = GPT(model_args)。
    4. 从检查点中提取模型权重 state_dict = checkpoint['model']。
    5. 将权重加载到模型中 model.load_state_dict(state_dict)。
  • 优势: 推理脚本变得非常简洁，不再需要重复定义模型的各种超参数，降低了出错的风险。 Screenshot-[17:18]

6.2. 分布式训练概念

当模型或数据量非常大，单块 GPU 无法满足需求时，就需要分布式训练。

  • 数据并行 (Data Parallelism / DDP):
    • 思想: 将同一份模型复制到多个 GPU 上，然后将训练数据切分成多份，每个 GPU 处理一部分数据。
    • 适用场景: 模型可以被单个 GPU 容纳，但数据集非常大。
  • 模型并行 (Model Parallelism):
    • 思想: 将模型本身的不同部分（例如，不同的层）切分到不同的 GPU 上。
    • 适用场景: 模型本身巨大，单块 GPU 的显存无法容纳。

7. 深入探讨：W&B结果分析与微调 (Fine-tuning) 策略

7.1. W&B 结果分析

W&B 仪表盘提供了丰富的监控信息：

  • 损失曲线: 直观地看到 train_loss 和 validation_loss 的下降趋势，判断模型是否收敛或过拟合。 Screenshot-[20:02]
  • 梯度与权重分布: 可以监控特定层的权重和梯度的直方图，诊断梯度消失或爆炸等问题。
  • 系统监控: 实时查看 GPU/CPU 使用率、显存占用、磁盘 I/O 等，帮助优化硬件资源利用。 Screenshot-[21:11]

7.2. 微调 (Fine-tuning) 策略

微调是基于一个已经预训练好的基础模型，在特定的下游任务和数据集上进行少量额外训练，使其适应新任务。微调策略因任务而异。

  • 场景1：指令遵循/问答 (Instruction Following / QA)

    • 任务: 让模型按特定格式回答问题。例如，输入公司名，返回其地址和成立时间。
    • 数据准备: 创建一个包含几百到几千条 (问题, 回答) 对的数据集。
    • 格式: {"prompt": "问题文本", "completion": "期望的回答文本"}。
    • 目标: 教会模型理解指令并以结构化的方式输出。

  • 场景2：文本补全 (Text Completion)

    • 任务: 补全 SKU 名称或特定领域的短语。
    • 数据准备: 将数据处理成 BOS (Beginning-of-Sentence) + 文本 + EOS (End-of-Sentence) 的格式。
    • 目标: 让模型学习特定语料的语言风格和模式。

  • 场景3：分类/情感分析 (Classification / Sentiment Analysis)

    • 任务: 判断一段文本的情感是积极还是消极。
    • 数据准备: 标注好的数据集，每条文本都有一个分类标签（如 positive, negative）。
    • 目标: 在模型末尾增加一个分类头，微调整个模型或仅微调分类头，使其学会对文本进行分类。

8. GPU实战与项目总结

8.1. 在 GPU 上进行训练

  • 环境: 使用 Lambda Labs 提供的 GPU 服务。
  • 参数设置:
    • d_model: 1024
    • n_head: 8 (每个头维度为 128)
    • context_length: 256
  • 结果:
    • 训练速度显著提升。
    • 损失值下降到更低的水平 (约 1.8)。
    • 最终模型参数量达到 1.5 - 3.5 亿级别。
    • W&B 监控显示 GPU 显存占用约 10GB / 21GB，说明超参数设置在硬件承受范围内。 Screenshot-[27:24]
  • 挑战: 即使模型变大，对于中文的 sub-word 分词，如果训练数据不够庞大，模型仍然可能无法很好地学习到 token 之间的关联，导致生成不完整的汉字或词语。

8.2. 项目总结

  • 本系列课程从零开始，手写了构建一个 GPT 模型的核心代码，包括 Transformer Block、多头自注意力、前馈网络等。
  • 最终实现了一个完整的训练和推理流程，并引入了配置管理和实验跟踪等工程实践。
  • 提供的代码是一个很好的起点，用户可以在此基础上进行修改、调参，并尝试在自己的数据集上进行训练和微调。
  • 更高级的微调技术和应用将在后续课程中展开。

AI 总结

该视频是手写大模型代码系列的收官之作，主要围绕 代码优化、推理实现和工程实践 展开。首先，视频修正了前期训练代码中的两个关键问题：在验证阶段需使用 torch.no_grad() 避免梯度计算，并正确切换 model.eval() 和 model.train() 模式。接着，视频演示了如何编写一个独立的 inference.py 脚本来加载训练好的模型并进行文本生成。

视频的核心内容在于 代码重构。为了便于实验和维护，所有超参数被整合到统一的配置中，并与模型权重一同保存在检查点文件里，极大地简化了推理和微调时的模型加载流程。此外，视频引入了强大的实验跟踪工具 Weights & Biases (W&B)，用于实时监控和可视化损失曲线、系统资源占用等关键指标，为模型调优提供了数据支持。

最后，视频展望了 微调 (Fine-tuning) 的多种策略，如问答、文本补全和情感分析，并展示了在专业 GPU 平台上训练更大模型（亿级参数）的实践过程，强调了超参数与硬件资源的匹配关系。本视频为学习者提供了一套从理论到实践、从粗糙到规范的大模型训练与推理代码框架。

Graph View