Skip to content

assignment3-scaling

Directory actions

More options

Directory actions

More options

Latest commit

 

History

History

assignment3-scaling

README.md

Assignment3

这是一个基于PyTorch从零构建Transformer语言模型并研究缩放定律的项目，通过训练探究非嵌入参数量、训练数据量与模型性能的关联，并运用DeepMind的缩放定律预测同架构不同规模模型的最终loss （但该预测仅基于非嵌入参数与训练token数量，未涵盖其他超参数设置，详见最后常见问题中的调参说明）。

代码特点

  • 🚀 从零实现的Transformer语言模型
  • 📊 基于Chinchilla缩放定律的性能预测
  • 🔬 完整的训练监控和可视化
  • 📈 FLOPs-Loss缩放关系分析
  • 🎯 支持多种模型规模配置

代码结构

├── model.py              # decoder-only模型实现
├── train.py              # 训练脚本
├── scaling.py            # 缩放定律相关函数
├── fit_scaling.py        # 缩放定律拟合
├── predict loss.py       # 损失预测
├── test1.py              # IsoFLOPs曲线分析
├── data.bin              # 训练数据（需要通过Assignment1中的`get train data.py`得到）
├── tokenizer.json        # 分词器配置
├── isoflops_curves.json  # IsoFLOPs实验数据
├── 配置.text             # 模型配置参数
└── ckpt/                 # 检查点和结果
    ├── epoch_*.pt        # 训练检查点
    ├── *.png             # 曲线图
    └── scaling_records_70M.npy  # 使用第一个配置得到的缩放数据记录

data.bin通过get train data.py得到， 需要根据训练模型的规模的大小确定训练总token数量，这里建议每一个参数至少分配到20个token训练。

核心功能

1. Transformer模型 (model.py)

实现了完整的Transformer语言模型，包括：

  • 多头自注意力机制
  • 前馈神经网络
  • 层归一化和残差连接
  • 位置编码

模型配置示例：

config = {
    "vocab_size": 32000,
    "context_length": 2048,
    "d_model": 768,
    "num_layers": 12,
    "num_heads": 12,
    "d_ff": 3072,
    "attn_pdrop": 0.1,
    "residual_pdrop": 0.1,
}

2. 训练系统 (train.py)

  • 自定义AdamW优化器实现
  • 余弦学习率调度
  • 梯度裁剪
  • 内存监控
  • 自动检查点保存
  • 训练/验证损失可视化

3. 缩放定律研究 (scaling.py)

基于Chinchilla论文实现的缩放定律：

L(N,D) = E + A/N^α + B/D^β

其中：

  • N: 模型参数量
  • D: 训练数据量
  • L: 损失值
  • E, A, B, α, β: 拟合参数(论文的实现数据为E=1.69, A=406.4, B=410.7, alpha=0.34, beta=0.28)

快速开始

环境要求

pip install torch numpy matplotlib scipy transformers psutil

训练模型

# 使用默认配置训练70M参数模型
python train.py --epochs 8 --batch_size 16 --d_model 384 --num_layers 24

# 训练130M参数模型
python train.py --d_model 512 --num_layers 36 --num_heads 8

# 训练300M参数模型  
python train.py --d_model 768 --num_layers 36 --num_heads 12

分析缩放定律

# 拟合缩放定律并生成可视化
python fit_scaling.py

# 预测特定配置的损失（通过模型的非嵌入参数、训练总token数量）
python "predict loss.py"

实验配置

项目预设了三种模型规模配置：

模型规模 参数量 d_model num_layers num_heads
小型 ~70M 384 24 6
中型 ~130M 512 36 8
大型 ~300M 768 36 12

结果分析

训练完成（实验采用第一种默认配置）后，可以得到：

  1. 训练曲线

    • 训练和验证困惑度变化ckpt/train_ppl.png、ckpt/val_ppl.png

  2. 缩放定律拟合

    • 根据已有的实验数据分析每一组FLOPs理想最优的(N, D)，拟合得到理想最优的loss-FLOPs曲线optimal_curve_70M.png
    • 在固定的有效模型参数（非嵌入参数），以及实验数据得到loss-FLOPs的拟合曲线fix_N_70M.png

  3. 缩放数据记录

    • 包含step、tokens、FLOPs、损失的详细记录ckpt/scaling_records_70M.npy

  4. IsoFLOPs分析

    • 计算预算与最优数据集大小的关系Scaling Law.png

核心算法

FLOPs估算

FLOPs ≈ 6 × N × D

其中N为非嵌入参数量，D为训练token数量。

Chinchilla缩放定律

基于DeepMind的研究，模型损失与参数量、数据量的关系为：

L(N,D) = E + A × N^(-α) + B × D^(-β)

注意事项

  1. 内存管理: 项目包含内存监控功能，建议在GPU环境下运行大模型
  2. 数据格式: 训练数据使用int32格式的二进制文件
  3. 检查点: 每个epoch自动保存检查点，支持训练中断恢复
  4. 可视化: 所有训练曲线和分析结果自动保存为PNG格式

扩展功能

  • 支持自定义tokenizer
  • 可配置的dropout和学习率策略
  • 灵活的数据集划分
  • 详细的训练日志和性能监控

常见问题

Q1：对于缩放定律所有大小的模型都适用吗？

A1：缩放定律的适用性取决于模型是否处于可缩放区间。当模型参数规模过小（例如低于约50M参数，具体阈值依赖于模型架构与任务复杂度）时，模型通常处于容量受限状态，此时即使持续增加训练数据或计算量，训练损失的下降幅度也较为有限，难以呈现稳定的幂律关系。 OpenAI或者DeepMind提出的缩放定律均建立在以下前提之上：模型规模或训练资源的增加能够带来显著且连续的性能提升。因此，只有当模型进入可观测的可缩放区间时，loss与模型规模（或FLOPs、数据量）之间的幂律关系才具有统计意义。 基于上述考虑，在本实验设置中选取约70M参数规模的模型作为最小实验模型，以尽量避免容量受限效应的干扰，确保模型行为处于可缩放区间，从而提高缩放定律拟合结果的稳定性与可信度。

在固定的模型架构、训练目标前提下，只要不同规模模型均处于可缩放区间，并在相同数据分布上训练至充分收敛，其训练损失或验证损失通常可以用统一的缩放定律进行拟合。

Q2：如何确定最优学习率？

A2：关于最优学习率的确定， $μP$ 提供了一种有效的方法。 OpenAI与微软在相关研究中指出，通过采用 $μP$ 参数化方式包括特定的权重初始化与尺度设定，可以使模型在不同参数规模下保持一致的训练动态。在该参数化下，对于同一模型结构与任务中，当小规模模型（如10M参数）在某一学习率下达到最优性能时，仅通过增加模型参数规模（如扩展至100M），而无需改变学习率设置，模型同样能够接近其对应规模下的最优性能。

因此， $μP$ 使得最优学习率在模型规模扩展过程中近似保持不变，从而显著降低了跨模型规模调参的成本。

Q3：如何确定深宽比、高宽比以及Attention计算的隐藏维度空间？

A3： 关于前馈网络宽度与模型嵌入维度之比（ $$\frac{d_{ff}}{d_{model}}$$ 通常为4）、模型嵌入维度与层数之比（ $$\frac{d_{model}}{num_{layers}}$$ ），以及每个注意力头的特征子空间维度（ $$\frac{d_{model}}{num_{head}}$$ ）的选择，已有研究表明，在固定总参数量且模型处于可缩放区间的前提下，模型最终性能对上述比例的变化通常表现出相对较低的敏感性，其影响程度显著小于模型规模或训练计算量等一阶因素。

在具体结构设计中，上述比例仍需保持在合理范围内，以避免对模型表达能力与训练稳定性产生不利影响：

  • 嵌入维度（Embedding Size）过小会形成信息瓶颈，限制输入token语义特征的表达能力；嵌入维度过大则可能导致特征表示过于分散，降低注意力机制在有限计算预算下的建模效率；
  • 注意力层数过少会限制模型对长程上下文信息的逐层整合能力，使高层语义表征不充分；注意力层数过多则可能增加优化难度，降低训练稳定性，并在数据规模受限时引入过拟合风险；
  • 前馈网络维度过小会限制非线性变换的表达能力，从而削弱模型对复杂语义模式的建模；前馈网络维度过大则可能显著增加参数冗余，出现收益递减现象，使模型更易出现收敛缓慢或过拟合等问题。

因此，在固定总参数量且模型处于可缩放区间的前提下，已有研究表明模型性能对深宽比、高宽比以及注意力头维度的变化通常呈现二阶敏感性，其影响程度显著小于模型规模与训练计算量等一阶因素。

Q4：如何确定多头注意力机制的数量？

A4：关于多头注意力机制中注意力头数量的确定，其核心作用在于提升模型对不同子空间关系的建模能力。 在多头注意力机制中，不同注意力头采用彼此独立的Q、K、V投影参数，对同一输入表示进行并行建模，从而使模型能够在不同表示子空间中关注不同类型的依赖关系。 在固定embedding维度的前提下，注意力头数量的增加并不会显著改变模型的总参数规模，而是将表示空间划分为更多低维子空间。然而，注意力头数量并非越多越好：

  • 头数过少时，模型对不同关系模式的建模能力受限；
  • 头数过多时，则会导致单个注意力头的维度过小，削弱其表示能力，并可能降低注意力计算的有效性。 在固定embedding维度的前提下，注意力头数量通常选择在8-16的范围内，以保证单个注意力头具有足够的表示维度（通常不少于32-128），从而在多样化关系建模能力与计算稳定性之间取得平衡。

多头注意力并不是增加容量，而是重新组织的容量。

Q5：为什么在缩放注意力中注意力头数基本变化不大？

A5：在缩放注意力机制的实验中，注意力头数通常保持相对固定，其原因在于多头注意力主要影响表示结构而非模型容量。 在固定embedding维度(d_model)的前提下，增加注意力头数量并不会显著提升模型的表达能力，而是将原有表示空间划分为更多低维子空间。 当注意力头数量过多时，单个注意力头的维度(d_model // num_head)随之减小，削弱了其建模能力，同时不同注意力头之间容易产生高度冗余的注意力模式。 此外，过多的注意力头还会带来额外的计算与数值稳定性开销，其性能收益往往呈现明显的递减趋势。 因此，在缩放实验中通常保持注意力头数量不变，以减少结构性因素对模型性能的干扰，使实验结果更专注于规模变量本身。

Q6：如何确定处理批次的数量？

A6：关于训练批次大小的确定，首先受限于GPU显存容量或则CPU内存，应在显存可承受范围内选择尽可能大的batch size。 在此基础上，可通过梯度累积构建更大的有效批次规模，以提升训练过程的数值稳定性。 然而，批次规模并非越大越好。尽管较大的batch size能降低梯度方差、使训练过程更加平稳，但在超过一定范围后，其对模型最终性能的提升会出现明显的收益递减，甚至可能影响泛化能力。因此，在实际训练中，通常在稳定性与泛化性能之间选择一个折中范围，并与学习率设置协同调整。

📌 常见经验范围（LLM）

  • 小模型或单机实验：有效批次规模通常在几十至数千的范围内；
  • 中大型LLM：有效批次规模通常在8k~64k之间。

当有效批次规模超过上述范围后，模型训练稳定性虽可能继续提升，但对最终性能的改善往往呈现明显的收益递减。

Q7：如何选择epoch大小？

A7：关于训练轮数的选择，其本质在于控制模型在固定数据集上对样本的重复暴露次数。 在训练初期，随着epoch的增加，模型能够更充分地利用数据分布信息，性能通常持续提升。已有研究表明，在随机梯度优化过程中即使重复使用同一数据样本，其在不同训练阶段所产生的梯度更新效果并不完全等价，因此适度增加训练轮数仍可带来性能增益。 然而，当epoch数量超过一定范围后，模型对训练数据的拟合程度不断加深，梯度多样性逐渐下降，性能提升趋于饱和，甚至可能出现过拟合或泛化能力下降的问题。因此，在实际训练中，通常通过验证集性能或损失收敛情况，选择位于性能提升与收益递减之间的合适epoch数。

💡不同规模和容量的模型，其最优训练epoch通常不同，因此epoch往往需要在训练过程中通过验证集性能动态确定。 在训练过程中，当模型在验证集上的loss下降趋势逐渐趋于平缓，或在继续训练时出现上升，而验证指标不再提升甚至下降时，说明模型已开始过拟合，此时对应的epoch可视为该模型在当前设置下的最优训练轮数。