cs336Assignment1 训练组件

模型训练

对于一个神经网络模型，其训练流程基本如上图所示，本笔记旨在在从零搭建Transformer-CS336作业所搭建的Transformer模型基础上，搭建模型训练所需要的一些基本模块，并最后使用这些模块实现对上述模型的训练与推理验证。依旧感谢【斯坦福 CS336】零基础手搓大模型：分词器、transformer 结构、数据与训练、文本生成这个视频的讲解，帮助了很多。

在上述的训练步骤中，主要要自己实现的部分有：

• 损失函数与Loss的计算，训练中使用交叉熵(CrossEntropyLoss)损失函数
• 优化器的实现，包括SGD与AdamW优化器
• 梯度调优的策略，如梯度裁剪(Gradient Clipping)与学习率调度器(Learning Rate Scheduler)

CrossEntropyLoss

在实现交叉熵损失函数之前，我们先来了解一下什么是信息熵。

信息熵

对于随机事件\(X\)，其可能的取值为\(x_1,x_2,...,x_n\)，其概率分布为\(P(X=x_i)=p_i\)，不妨此时先关注其中的一个事件\(X=x_i\)，我们希望去描述下面的几个属性：

• 对于事件\(X=x_i\)，如何表述该结果发生之后，我们获得了多少信息量？
• 对于随机事件X，如何表述整个事件的平均信息量？

我们做出以下假设：

1. 一个事件发生的概率越小，我们获得的信息量就越大；反之，一个事件发生的概率越大，我们获得的信息量就越小。因此信息量与概率之间应该是一个反比的关系。
   • 比如，抛硬币。对于普通的硬币，正反面出现的概率都是\(0.5\)，那么我们抛一次硬币的结果难以预测，当该结果发生后，给我们的信息量是比较大的；对于小说电影中常见的“作弊硬币”，总是会固定偏向某一面，此时我们几乎有着100%的概率预测到结果，那么当该结果发生后，给我们的信息量是比较小的，这是当然的，因为即使该事件不发生，我们也会有着很大的概率预测到结果。
2. 对于两个独立的\(X=x_i\)与\(X=x_j\)同时发生时，其概率有\(p_{ij}=p_i*p_j\)，而这两个事件同时发生后，我们获得的信息量应该是两个事件各自信息量的和，即我们使用\(h(x)\)代表信息量，则有\(h(x_i,x_j)=h(x_i)+h(x_j)\)，不难看出事件发生的概率与信息量之间可以使用对数去描述这种乘积变成加和的关系。

综上，我们可以选取对数函数去描述一个事件发生的概率与信息量之间的关系，假设我们选取以\(2\)为底的对数函数，则有：

\[h(x_i)=-\log_2(p_i) \qquad p_i \uparrow \Rightarrow h(x_i) \downarrow \\ p_i \in (0,1] \qquad h(x_i) \in [0,+\infty) \]

而对于整个随机事件，我们想要去描述对于其包含的所有样本在发生之前，平均有多少的不确定性，或者说发生之后，平均获得多少信息量。显然这是一个期望：

\[ H(X)=-\sum_x P(x)\log_2 P(x) \]

交叉熵

对于交叉熵，其数学定义如下：

\[H(P,Q)=-\sum_x P(x)\log Q(x) \]

其中：

• \(P\)为真实分布，\(Q\)为预测分布

在这个公式中，当真实分布与预测分布接近时，我们不妨对于一个二分类任务，取0-1分布:

\[P(x = 1) = 1 \qquad P(x = 0) = 0 \quad \]

然后我们通过一些方法(如模型推理或者数理统计的估计)得出了预测的分布Q，取值如下：

\[Q(x = 1) = 0.9 \qquad Q(x = 0) = 0.1 \]

则交叉熵为：

\[H(P,Q)=-\sum_x P(x)\log Q(x) = -\log(0.9)-0\log(0.1) = -\log(0.9) \approx 0.10536 \]

上方的\(log() = ln()\)

显然，对于一个分类任务，一般都是使用一个one hot的向量标签去表示真实分布，向量中的目标元素位置上会是1，其余位置都是0，而模型的输出一般是一个概率分布向量，表示每个类别的预测概率。而由上面的信息熵的定义，我们指定目标元素的下标为\(i\)，不难想象，如果模型的预测比较正确，此时\(Q(x_i)\)会接近1，而此时其对应的信息熵\(-\log(Q(x_i))\)会接近0。因此最后的结果会很小；否则这样计算出的交叉熵会很大。

同时，因为在我们期望的真实分布中，仅有目标元素的位置上是1，其余位置都是0，因此交叉熵在对于序列预测进行计算的时候可以简化为：

\[L_t=-\ln Q(y_t\mid x_{\leq t}) \]

• \(x_{\leq t}\)：当前位置及之前的输入 token。
• \(y_t\)：正确的下一个 token。
• \(Q(y_t\mid x_{\leq t})\)：模型给这个正确 token 分配的概率。

我们给出一个简单的例子，在训练的时候，我们的训练文本为：

我 喜欢 吃 苹果

此时模型已经看过了前三个token，接下来我们希望模型预测下一个token为苹果，那么在我们期望的真实分布中，词表中苹果这个token的下标位置上是1，其余位置都是0，假设模型预测的概率中，苹果这个token的预测概率为0.9，即

我	...	...	苹果
0.05	...	...	0.9
则交叉熵为：

\[L_t = -\ln(0.9) \approx 0.10536 \]

当然，在训练的时候，模型的输出一般形状为：

\[X.shape = (Batch\_Size, Seq\_Len, Vocab\_Size) \]

Batch_Size即为批大小，Seq_Len为输入的token序列长度，以及Vocab_Size为词表大小。

我们关注最后一维，对于固定一个batch的样本中，一共Seq_Len个token，模型会吃进去一个token，然后输出一个长度为Vocab_Size的预测分数，即对于词表中的所有词进行打分，分数越高，表示模型越认为下一个词应该是这个词。而吃进去所有token之后，自然会一共有着Seq_Len个这样的Vocab_Size分数向量。而对于每个batch中的样本中所有的token，我们都可以计算一个交叉熵损失，而总的交叉熵损失就是：

\[L = \frac{1}{Batch * Seq\_Len}\sum_{b=1}^{Batch}\sum_{i=1}^{Seq\_Len} -lnQ(x_{i+1} | x_1, x_2, ..., x_i) \]

此处\(-lnQ(x_{i+1} | x_1, x_2, ..., x_i)\)的含义与上文中\(L_t=-\ln Q(y_t\mid x_{\leq t})\)的含义是一样的。都是取出预测词表中，对于真实标签的token的预测概率，然后取负对数。

当然，在实际工程中，我们不妨对于第i个token的输出记作\(o_i\)，然后上述的\(x_i\)表示的是该位置对应token的token_id，那么模型吞入前i个向量后输出的打分表中，我们期待的第\(x_{i+1}\)位置上的分数为：

\[o_i[x_{i+1}] \]

其对应的概率为：

\[Q(x_{i+1} | x_1, x_2, ..., x_i) = softmax(o_i)[x_{i+1}] = \frac{exp(o_i[x_{i+1}])}{\sum_{j=1}^{Vocab\_Size} exp(o_i[j])} \]

同样，为了防止我们在从零搭建Transformer-CS336作业中关于SoftMax的实现中会出现的以e为底的指数运算溢出的问题，我们可以在计算softmax的时候，先减去该向量的最大值，然后再进行指数运算，这样就不会出现溢出的问题了。

即

\[m_i = max(o_i) \\ 则参与计算的向量：\\ o_i\prime = o_i - m_i \]

所以新的公式即为

\[Q(x_{i+1} | x_1, x_2, ..., x_i) = softmax(o_i\prime)[x_{i+1}] = \frac{exp(o_i[x_{i+1}] - m_i)}{\sum_{j=1}^{Vocab\_Size} exp(o_i[j] - m_i)} \]

因此此时的交叉熵损失：

\[L_i = -lnQ(x_{i+1} | x_1, x_2, ..., x_i) = -ln\frac{exp(o_i[x_{i+1}] - m_i)}{\sum_{j=1}^{Vocab\_Size} exp(o_i[j] - m_i)} = -(o_i[x_{i+1}] - m_i) + ln\sum_{j=1}^{Vocab\_Size} exp(o_i[j] - m_i) \]

故总的交叉熵损失为：

\[L = \frac{1}{Batch * Seq\_Len}\sum_{b=1}^{Batch}\sum_{i=1}^{Seq\_Len} -(o_i[x_{i+1}] - m_i) + ln\sum_{j=1}^{Vocab\_Size} exp(o_i[j] - m_i) \]

pytorch实现的代码如下：

import torch

def cross_entropy(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:

    '''
    logits: (batch_size, seq_len, vocab_size)
    targets:(batch_size, seq_len)，存放的是每个位置上token的token_id
    '''

    max_logits = logits.max(dim=-1)[0]
    # (沿着最后一维取max) ，然乎返回元组(最大值Tensor， 最大值对应的下标)
    # 我们取第一个元素，即最大值列表，此时max_logits.shape = (B, S)(没有规约维度)

    target_logit = torch.gather(logits, -1, targets[..., None]).squeeze(-1)
    # targets[..., None] -> (B, S, 1)
    # .squeeze(-1) : 抹掉最后一个维度
    # 最终取得的是每个(B, S)维度上的对应的那个logits

    term1 = -(target_logit - max_logits) # 删除最大分数
    exp_sum = torch.exp(logits - max_logits[..., None]).sum(dim=-1) 
    # (B, S, V) - (B, S, 1)广播，最后沿着最后一个维度加和(V -> 1)，不保留规约维度 :(B, S)
    return term1 + torch.log(exp_sum) # 交叉熵损失，各个位置上的(B, S)

其中，torch.gather函数的作用是从logits中取出每个位置上对应的真实标签的预测分数。

我们关注其前三个参数：

• input：输入的张量
• dim：指定沿着哪个维度进行索引
• index：指定索引的张量

需要注意的几点：

1. index与input的维数必须相同，即若input是一个三维张量，则index也必须是一个三维张量，但是维度上的大小可以不同。
2. index中每个元素的取值范围必须在input指定维度的大小范围内，否则会报错。
3. 在dim指定的维度上取值

其取值的方式可以使用下面的公式简单描述，对于我们的代码中，输入的logits是一个三维张量，形状为(B, S, V)，其中B表示批大小，S表示序列长度，V表示词表大小。而targets[..., None]的形状为(B, S, 1)，表示每个位置上对应的真实标签的token_id。

\[when \qquad dim = -1, \quad \\ output[B, S, k] = logits[B, S, targets[B, S, k]] \\ when \qquad dim = 0, \quad \\ output[k, S, V] = logits[targets[k, S, V], S, V] \\ when \qquad dim = 1, \quad \\ output[B, k, V] = logits[B, targets[B, k, V], V] \\ \]

不难看出，这个函数实际上就是在固定其他维度的情况下，我们期望index上对应位置存储的为input对应位置上的索引，然后用这个索引去寻值。

最后的最后，输出的是一个维度为(B, S)的张量，表示每个位置上对应的交叉熵损失:

可以看作原先的一行经过计算后变为新矩阵中的一个Loss元素，而对于一个Batch上的所有Sequence，一共有Seq_Len个元素，因此一个Batch上所有的Loss对应着新矩阵中的一行，所有Batch的Loss对应着整个新矩阵。

优化器

在确定好损失函数之后，训练的主要目的就是通过一个优化策略，使得在不断的训练过程中，让模型的各个参数想向着使得损失函数取值最小的方向去更新。

SGD

当然，最简单的想法就是让各个参数沿着负梯度的方向去更新，因为负梯度方向一定是损失函数下降最快的方向。

\[\theta_{t+1} = \theta_t - \eta \nabla L(\theta_t) \]

其中

• \(\theta_t\)：表示第t次迭代的参数
• \(\eta\)：表示学习率
• \(\nabla L(\theta_t)\)：表示第t次迭代的参数对应的梯度

当然，对于不同的batch的不同预测位置，由于其输入输出的影响，其梯度计算有所不同，因此其梯度大小也会有所不同，因此我们去取所有batch的梯度的平均值，作为最终的梯度去更新参数。这就是全批量的梯度下降算法：

\[\theta \leftarrow \theta - \eta \frac{1}{N}\sum_{i=1}^{N} \nabla L_i(\theta) \qquad N = Batch\_Size * Seq\_Len \]

以一个简单的Sequence：我 爱吃 苹果 <EOS>为例，忽略Batch维度，需要优化的参数为\(\theta\)，则：

模型能看到的内容	要预测的 token	对应损失
我	爱吃	\(L_1(\theta)\)
我 爱吃	苹果	\(L_2(\theta)\)
我 爱吃 苹果	<EOS>	\(L_3(\theta)\)

那么这个样本所提供的梯度为：

\[\nabla_\theta L = \frac{ \nabla_\theta L_1 + \nabla_\theta L_2 + \nabla_\theta L_3 }{3} \\ \qquad \\ \nabla_\theta L_{b} = \frac{1}{Seq\_Len} \sum_{i=1}^{\text{Seq\_Len}} \nabla_\theta L_i(\theta) \]

而所有Batch的样本对于参数\(\theta\)的梯度再次加和平均:

\[\nabla_\theta L = \frac{1}{B} \sum_{b=1}^{B} \nabla_\theta L_b(\theta) \]

但是，当样本过多的时候，每进行一次优化，都要计算所有样本的所有参数的梯度，这样的计算量是非常大的，我们可以在所有Batch的样本中，随机选择一部分样本，使用这些样本的平均梯度作为全局的梯度去更新参数，这就是随机梯度下降算法(Stochastic Gradient Descent, SGD)：

\[\theta \leftarrow \theta - \eta_t \frac{1}{|\mathcal B_t|} \sum_{i\in\mathcal B_t} \nabla_\theta L_i(\theta_t) \]

其中：

• \(\mathcal B_t\)：表示第t次迭代中，随机选择的样本集合
• \(|\mathcal B_t|\)：表示第t次迭代中，随机选择的样本集合的大小
• \(\eta_t\)：表示第t次迭代的学习率

可以看到，不仅样本是随机选择的，学习率实际上也可以随着迭代次数的变化而变化，这使得训练过程：

1. 计算的参数更少了，计算速度加快了
2. 训练初期使用较大的学习率，快速下降到一个较优的区域；训练后期使用较小的学习率，慢慢收敛到局部最优，收敛速度加快。

下面给出一个pytorch实现的SGD优化器的代码：

import torch
import math
from collections.abc import Callable, Iterable
from typing import Optional

class SGD(torch.optim.Optimizer):
    '''
    自定义优化器类必须继承torch.optim.Optimizer
    需要重写构造函数以及step(self)函数
    '''

    def __init__(self, params, lr=1e-3):
        '''
        params: 待优化的参数集合，可以是全部参数，也可以是分组参数
        lr: 学习率
        '''

        if lr < 0:
            raise ValueError(f"Invalid learning rate: {lr}")
        defaults = {"lr": lr} # 存储外部参数
        super().__init__(params, defaults)

    def step(self, closure:Optional[Callable] = None):

        loss = None if closure is None else closure()
        for group in self.param_groups: # 遍历不同的参数组，参数组中存放的是需要优化的参数以及对应的超参数
            lr = group["lr"] # 取出学习率
            for p in group["params"]: # 遍历当前参数组的参数值的所有参数
                if p.grad is None:    # 如果当前参数没有梯度，则不进行更新
                    continue

                state = self.state[p] # state字典，保存每个参数对应的状态信息，如迭代步数t等
                t = state.get("t", 0) # 读取迭代步数 t, 若t开始是不存在，则会返回0
                grad = p.grad.data # 获取损失loss关于参数p的梯度张量
                p.data -= lr / math.sqrt(t + 1) * grad # 带入公式进行参数更新
                state["t"] = t+1 # 更新迭代步数
        return loss

在上面的代码中，对应的公式实际上为：

\[\theta \leftarrow \theta - \frac{\alpha}{\sqrt{t+1}} \nabla_\theta L(\theta, \mathcal B_t) \]

但是上面实际上并没有抽样这一过程，但是优化器确实也是不用管这一过程，采样过程通常是由DataLoader去实现的：

loader = torch.utils.data.DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,  # 随机打乱，每次取出32批样本
)

AdamW

在实现AdamW优化器之前，我们先来了解一下作为其前置的四个优化器：

• SGDM
• Adagrad
• RMSProp
• Adam

1. SGDM(Stochastic Gradient Descent with Momentum)

SGDM实际上也就是SGD with Momentum，即带有动量的随机梯度下降算法。其公式为：

\[m_t = \beta m_{t-1} + g_t \qquad where \quad g_t = \nabla_\theta L(\theta, \mathcal B_t) \\ \theta \leftarrow \theta - \eta m_t \]

其中，\(m_t\)被称为动量，Momentum，此时更新的时候，我们并没有直接在负梯度方向进行下降，而是将当前的梯度\(g_t\)与上一次更新的动量\(m_{t-1}\)进行加权平均，得到当前的动量\(m_t\)，然后沿着这个动量方向去更新参数。我们不妨从t=0开始迭代几次：

\[m_0 = 0 \\ m_1 = \beta m_0 + g_1 = g_1 \qquad \theta_1 = \theta_0 - \eta m_1 = \theta_0 - \eta g_1 \\ m_2 = \beta m_1 + g_2 = \beta g_1 + g_2 \qquad \theta_2 = \theta_1 - \eta m_2 = \theta_1 - \eta (\beta g_1 + g_2) \\ \dots \]

可以看到，当前更新的方向其实并不是直接为当前的负梯度方向，以t=2更新为例，更新方向为\(m_2 = \beta m_1 + g_2\)，而\(m_1\)实际上是上一次的参数更新方向。换言之，本次参数更新的推力实际上是由当前的负梯度方向与上一次的更新方向的向量加权来提供，而这个\(m_t\)使得本次移动似乎掺入了上次移动的惯性，当：

• 参数更新的方向与上一次的更新方向一致时，当前的更新方向会被放大，此时下降的速度会加快
• 参数更新的方向与上一次的更新方向相反时，当前的更新方向会被抵消，此时下降的速度会减慢

因此参数更新会变得更加平滑，同时如果下降到了局部的鞍点附近，此时\(g_t = 0\)，而由于上次更新提供的动量\(m_{t-1}\)，此时参数更新的方向仍然会沿着上次的更新方向继续下降，从而有可能跳出局部的最优点，使得参数可以继续探索更优化的位置。

同时，不难看出，随着更新轮数的增加，越早期的\(g_t\)前所乘的系数会越多，也就是其占比会越小，而越靠近当前的\(g_t\)，其占比会越大，这也符合我们对于参数更新的直觉，越靠近当前的梯度，其对参数更新的影响应该越大，而对于之前的梯度，就应当随着更新的推进逐渐忽略。

2. Adagrad

在SGDM中，动量的引入确实很神奇，但是在公式中的学习率\(\eta\)仍然是一个固定的值，而在训练的过程中，实际上训练初期为了快速下降，我们希望学习率可以大一些，而随着训练的进行，我们希望学习率可以小一些，去精调参数。Adagrad提出了一种自适应学习率的方法：

\[v_t = v_{t-1} + g_t^2 \\ \theta \leftarrow \theta - \frac{\eta}{\sqrt{v_t + \epsilon}} g_t \]

可以看到Adagrad是在SGD的基础上进行的优化，主要是将学习率从\(\eta \rightarrow \frac{\eta}{\sqrt{v_t + \epsilon}}/\)，我们可以暂时忽略\(\epsilon\)，其主要作用是防止\(v_t\)为0时，分母为0而报错。

不难看出，上述的公式实际上也就是：

\[\theta \leftarrow \theta - \eta \frac{1}{\sqrt{\sum_{i=1}^t g_i^2 + \epsilon}} g_t \]

学习率会除以之前所有梯度的平方和的平方根，而这个\(v_t\)会越来越大，对应着的\(\eta\)会越来越小。因此使用Adagrad这种自适应的学习率的时候，在训练初期，学习率会比较大，而随着训练的进行，学习率会越来越小。

但是对于极端情况，如开始训练的时候，梯度非常大，而后面的梯度非常小，那么学习率在训练的整个过程中实际上不有太大的变化，导致自适应学习率的设计失效。

迭代 \(t\)	当前梯度 \(g_t\)	累积梯度平方 \(v_t\)	有效学习率 \(\eta_t^{\mathrm{eff}}\)
1	\(10^6\)	\(10^{12}\)	\(\eta\times10^{-6}\)
2	\(10^{-4}\)	\(10^{12}+10^{-8}\)	\(\approx\eta\times10^{-6}\)
3	\(10^{-4}\)	\(10^{12}+2\times10^{-8}\)	\(\approx\eta\times10^{-6}\)
4	\(10^{-4}\)	\(10^{12}+3\times10^{-8}\)	\(\approx\eta\times10^{-6}\)

3. RMSProp

RMSProp是对Adagrad的改进，其更新公式如下：

\[v_t = \beta v_{t-1} + (1-\beta) g_t^2 \\ \theta \leftarrow \theta - \frac{\eta}{\sqrt{v_t + \epsilon}} g_t \]

主要的改进在于如何计算这个衰减项\(v_t\)，不同于Adagrad的直接梯度平方相加，RMSProp使用一个参数\(\beta\)，对之前的梯度平方进行不同比例的加权。当前的\(v_t\)是由之前的\(v_{t-1}\)与当前的梯度平方\(g_t^2\)进行加权平均得到的。当然\(\beta \in [0, 1)\)。

我们不妨迭代几轮

\[v_0 = 0 \\ v_1 = \beta v_0 + (1-\beta) g_1^2 = (1-\beta) g_1^2 \\ v_2 = \beta v_1 + (1-\beta) g_2^2 = \beta (1-\beta) g_1^2 + (1-\beta) g_2^2 \\ v_3 = \beta v_2 + (1-\beta) g_3^2 = \beta^2 (1-\beta) g_1^2 + \beta (1-\beta) g_2^2 + (1-\beta) g_3^2 \\ \]

不难看出，随着梯度更新的次数增加，\(v_t\)中越是之前的梯度，其对应的系数会乘上越多的\(\beta\)，比如\(g_1\)的系数就从\((1-\beta) \rightarrow \beta(1-\beta) \rightarrow \beta^2(1-\beta)\)，这使得在本次的学习率更新中，靠的越近的梯度对当前的影响越大，同时解决了初期梯度过大，导致学习率一致无法变化的问题。

4. Adam

Adam优化器实际上将SGDM的动量参与更新的机制与RMSProp的自适应学习率结合在了一起，其更新公式如下：

\[\theta \leftarrow \theta - \frac{\eta}{\sqrt{\hat{v_t} + \epsilon}} \hat{m_t} \\ \hat{m_t} = \frac{m_t}{1-\beta_1^t} \qquad \hat{v_t} = \frac{v_t}{1-\beta_2^t} \]

我们回顾一下SGDM与RMSProp的公式：

• SGDM:

\[\theta \leftarrow \theta - \eta m_t \qquad where \quad m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t \]

• RMSProp:

\[v_t = \beta v_{t-1} + (1-\beta) g_t^2 \\ \theta \leftarrow \theta - \frac{\eta}{\sqrt{v_t + \epsilon}} g_t \]

非常直观的可以看到Adam在下降的时候，使用了SGDM的动量\(m_t\)去更新参数，同时使用了RMSProp的自适应学习率\(\frac{\eta}{\sqrt{v_t + \epsilon}}\)去更新学习率。但是在其中使用的是\(\hat{m_t}\)与\(\hat{v_t}\)，他们是上面两者分别都除了一个衰减项\(1-\beta_1^t\)与\(1-\beta_2^t\)。作用是在训练初期，防止更新的参数会由于\(m_t = 0\) or \(v_t = 0\)而导致更新参数过小。

在更新初期，\(m_0=0\)，一般取\(\beta_1=0.9\)，则第一次更新下：

\[m_1 = \beta_1 m_0 + (1-\beta_1) g_1 = (1-\beta_1) g_1 = 0.1 g_1 \]

这是非常不合理的，第一次更新的值直接被缩小的了十倍，或者说第一次更新时，所含有的实际上的梯度贡献仅有\(10\%\)。为了让动量中的数值完全来自于实际上的梯度，而非来自于初始的0值，我们需要对其进行修正，使用\(\hat{m_t} = \frac{m_t}{1-\beta_1^t}\)，则第一次更新时：

\[\hat{m_1} = \frac{m_1}{1-\beta_1^1} = \frac{0.1 g_1}{0.1} = g_1 \]

同理我们模拟第二次更新：

\[m_2 = \beta_1 m_1 + (1-\beta_1) g_2 = \beta_1 (1-\beta_1) g_1 + (1-\beta_1) g_2 \\ \hat{m_2} = \frac{m_2}{1-\beta_1^2} = \frac{\beta_1 (1-\beta_1) g_1 + (1-\beta_1) g_2}{1-\beta_1^2} = \frac{\beta_1 g_1 + g_2}{\beta_1 + 1} = \frac{0.9 g_1 + g_2}{1.9} = 0.4737 g_1 + 0.5263 g_2 \]

不难看出，修正之后的更新项几乎全部都来自于实际的梯度，而非初始的0值。

对于AdamW优化器，看名字来说，应该是Adam优化器的改进版本。对于Adam优化器，我们考虑这样一个问题：

在训练的过程中，为了防止某个参数过大，我们希望通过在损失函数中加入一个正则化项去约束参数的大小，如L2正则化项：

\[L = L(\theta) + \frac{\lambda}{2} ||\theta||^2 \qquad where \quad \lambda \in [0, 1] \]

即对于参数\(\theta\)，此时其梯度为：

\[\frac{\partial L}{\partial \theta} = \frac{\partial L(\theta)}{\partial \theta} + \lambda \theta \]

则在更新参数的时候，实际上是：

\[\theta - \eta \frac{\partial L}{\partial \theta} = \theta - \eta (\frac{\partial L(\theta)}{\partial \theta} + \lambda \theta) = (1 - \eta \lambda) \theta - \eta \frac{\partial L(\theta)}{\partial \theta} \\ 即 \\ \theta \leftarrow (1 - \eta \lambda) \theta - \eta \frac{\partial L(\theta)}{\partial \theta} \]

对比一下，\(\theta\)参数不仅沿着负梯度方向下降，同时对于\(\theta\)本身，会乘以一个小于1的系数\((1 - \eta \lambda)\)，这使得参数\(\theta\)会逐渐减小，由于是乘的系数，因此对于比较大的参数，其下降的幅度会比较大，而对于比较小的参数，其下降的幅度会比较小，这样就达到了约束参数大小的目的，尽量避免出现某些过大的参数主导网络的推理。

在使用Adam优化的时候，我们考虑加入L2正则化项，那么梯度项实际上相较于原来，会加上一个\(\lambda \theta\)的正则项，我们期望正则项将参数修正为\((1 - \eta \lambda) \theta\)。但是实际上，根据Adam的更新公式：

\[\theta \leftarrow \theta - \frac{\eta}{\sqrt{\hat{v_t} + \epsilon}} \hat{m_t} \\ \qquad \\ 当前梯度 \hat{g_t} = g_t + \lambda \theta \\ 则 \\ m_t = \beta_1 m_{t-1} + (1-\beta_1) \hat{g_t} = \beta_1 m_{t-1} + (1-\beta_1) (g_t + \lambda \theta) \\ v_t = \beta_2 v_{t-1} + (1-\beta_2) \hat{g_t}^2 = \beta_2 v_{t-1} + (1-\beta_2) (g_t + \lambda \theta)^2 \]

最后上述的\(\hat{m_t}\)与\(\hat{v_t}\)会经过修正，得到\(\hat{m_t} = \frac{m_t}{1-\beta_1^t}\)与\(\hat{v_t} = \frac{v_t}{1-\beta_2^t}\)，然后带入更新公式中：

\[\theta \leftarrow \theta - \frac{\eta}{\sqrt{\hat{v_t} + \epsilon}} \hat{m_t} \]

在这个更新过程中，显然出现以下问题：

我们本来期望正则项的作用是对于任何参数\(\theta\)，都可以让其衰减为\((1 - \eta \lambda) \theta\)，但是实际上，由于Adam的自适应学习率的存在，这套规则不仅经过了动量以及自适应学习率历史累计的影响，而且还经过了修正\(\hat{m_t}\)与\(\hat{v_t}\)的影响，这使得正则项的作用不再是单纯的让参数衰减，而是会受到梯度以及历史梯度的影响，对于不同参数，其衰减的幅度会有所不同，这显然不是我们期望的。

因此在AdamW中，提出了一个新的更新公式：

\[\theta \leftarrow (1 - \eta \lambda) \theta - \frac{\eta}{\sqrt{\hat{v_t} + \epsilon}} \hat{m_t} \]

也就是说，在更新的时候，我们直接去将正则项的作用单独拿出来，直接让参数衰减为\((1 - \eta \lambda) \theta\)，而不再将正则项加入到梯度中去计算，同时在更新AdamW的\(\hat{m_t}\)和\(\hat{v_t}\)时，不在将正则项加入到梯度中去计算，而是仅仅使用原始的梯度\(g_t\)去计算\(\hat{m_t}\)和\(\hat{v_t}\)。

下面给出一个使用python实现的AdamW优化器的代码：

class AdamW(torch.optim.Optimizer):
    def __ini__(self, params, lr: float = 1e-3, betas:tuple[float, float] = (0.9, 0.999), eps:float = 1e-8, weight_decay:float =0.0):
        '''
        lr： float 学习率
        betas: 计算动量m_t与累计梯度v_t的衰减系数
        eps: epsilon, 防止除0的小量
        weight_decay 权重衰减系数，用于L2正则化
        '''

        if lr < 0:
            raise ValueError(f"Invalid learning rate: {lr}")

        # 将超参数保存到字典中
        defaults = {"lr":lr, "betas":betas, "eps":eps, "weight_decay":weight_decay}

        super().__init__(params, defaults)

    def step(self, closure: Optional[Callable] = None):
        '''
        执行单步的AdamW参数更新
        '''

        loss = None if closure is None else closure()

        # 遍历每一组参数的group
        for group in self.param_groups:
            lr = group["lr"]
            beta1, beta2 = group["betas"]
            eps = group["eps"]
            weight_decay = group["weight_decay"]

            for p in group["params"]:
                if p.grad is None:
                    continue

                state = self.state[p]
                # 从字典中取出上一次的
                # t:更新迭代步数
                # m: 一阶矩或者说动量项
                # v: 二阶矩
                t = state.get("t", 0)
                m = state.get("m", 0)
                v = state.get("v", 0)
                grad = p.grad.data

                t+=1

                p.data -= lr * weight_decay * p.data # 正则项对参数进行衰减

                m_t = beta1 * m + (1 - beta1) * grad
                v_t = beta2 * v + (1 - beta2) * grad ** 2 # 更新m_t与v_t
                m_hat = m_t / (1 - beta1 ** t)
                v_hat = v_t / (1 - beta2 ** t) # 修正
                p.data -= lr * m_hat / (torch.sqrt(v_hat) + eps) # 下降


                # 将更新后的数值存入字典
                state["t"] = t 
                state["m"] = m_t
                state["v"] = v_t

        return loss

Learning Rate Scheduler

之前我们也提到过，我们在训练的不同阶段中，其实是希望以不同的学习率去训练模型的参数的；

在训练的初期，我们希望使用较大的学习率，让模型快速下降到一个较优的区域；而在训练的后期，我们希望使用较小的学习率，让模型慢慢收敛到局部最优。

在前面的Adagrad，RMSProp，Adam以及AdamW中，其实都去实现了可变的自适应学习率，当然，这个自适应的实现是对我们给定的学习率除以一个自适应的衰减项。

而学习率调优(Learning Rate Scheduler)则是我们在训练的不同阶段，去直接指定学习率的大小，也就是\(\eta\)的值。当然，优化器的自适应与学习率调优是可以同时使用的。

Cosine Annealing Schedule

余弦退火调度（Cosine Annealing Schedule）是一种学习率调优的策略，其公式如下:

\[\eta_t = \frac{t}{T_w}\eta_{max} \quad where \quad t < T_w \\ \qquad \\ \eta_t = \eta_{min} + \frac{1}{2}(1 + \cos(\frac{t-T_w}{T_c-T_w}\pi))(\eta_{max} - \eta_{min}) \quad where \quad T_w \leq t \leq T_c \\ \qquad \\ \eta_t = \eta_{min} \quad where \quad t > T_c \]

其中：

• \(\eta_t\)：表示第t次迭代的学习率
• \(\eta_{max}\)：表示学习率的最大值，超参数，自己设置
• \(\eta_{min}\)：表示学习率的最小值，最终的学习率会收敛到这个值，超参数，自己设置
• \(T_w\)：表示学习率的预热步数，前\(T_w\)步的学习率会从0线性增加到\(\eta_{max}\)，超参数，自己设置
• \(T_c\)：表示学习率的退火步数，前\(T_c\)步的学习率会从\(\eta_{max}\)以余弦退火方式下降到\(\eta_{min}\)，超参数，自己设置
• \(t<T_w\): 预热阶段，此时学习率从0线性增加到\(\eta_{max}\)
• \(T_w \leq t \leq T_c\): 退火阶段，此时学习率从\(\eta_{max}\)以余弦退火方式下降到\(\eta_{min}\)
• \(t > T_c\): 收敛阶段，此时学习率保持在\(\eta_{min}\)

ps: t与T都是迭代步数，t是当前迭代步数，T是总的迭代步数，不是时间！

公式中的第一步很好理解，就是在前\(T_w\)步中，学习率随着训练步数从0线性增加到\(\eta_{max}\)，我们关注于第二步，也就是退火阶段的公式：

\[\eta_t = \eta_{min} + \frac{1}{2}(1 + \cos(\frac{t-T_w}{T_c-T_w}\pi))(\eta_{max} - \eta_{min}) \quad where \quad T_w \leq t \leq T_c \]

其中，有\(p = \frac{t-T_w}{T_c-T_w}\)，这部分将从\(T_c\rightarrow T_w\)的退火迭代步数总数转换到一个百分比形式，因为在退火阶段，有\(t = T_w + t\prime\)，需要退火的总迭代步数为\(T_c - T_w\)，因此\(p = \frac{t\prime}{T_c-T_w}\)，当\(t\prime = 0\)时，\(p=0\)，当\(t\prime = T_c-T_w\)时，\(p=1\)，也就是\(p \in [0, 1]\)。

然后，我们将上述的进度\(p\)使用余弦函数进行映射，映射到一条从1降低到0的曲线上：

\[p \in [0, 1] \rightarrow p\pi \in [0, \pi] \rightarrow \cos(p\pi) \in [1, -1] \\ \qquad \\ 1 + \cos(p\pi) \in [2, 0] \rightarrow \frac{1}{2}(1 + \cos(p\pi)) \in [1, 0] \\ \qquad \\ \omega = \frac{1}{2}(1 + \cos(p\pi)) \in [1, 0] \]

通过这种构造，使得退火的步数实际上与一个以余弦形状从1下降到0的权重\(\omega\)映射上，即\(\omega = \omega(t)\)，最后的退火公式可以写为：

\[\eta_t = \eta_{min} + \omega(\eta_{max} - \eta_{min}) \]

当开始退火，即\(t = T_w\)时，\(\omega = 1\)，此时\(\eta_t = \eta_{max}\)；当退火结束，即\(t = T_c\)时，\(\omega = 0\)，此时\(\eta_t = \eta_{min}\)。

我们去设定\(T_w\) = 10， \(T_c\) = 100，\(\eta_{max}\) = 0.1，\(\eta_{min}\) = 0.01，绘制出学习率随迭代次数的变化曲线：

当然，虽然我们期望前期使用较大的学习率，后期使用较小的学习率，但是在训练刚刚开始的时候，模型的参数是随机初始化的，AdamW的各个历史量还在累计，直接使用较大的学习率可能会导致模型的参数更新过大，导致损失震荡甚至训练不稳定。因此才加入了预热阶段，让学习率从0线性增加到\(\eta_{max}\)，让模型的参数慢慢适应较大的学习率。

下面给出一个使用python实现的Cosine Annealing Schedule的代码：

import math

def learing_rate_schedule(
        it: int,
        max_learning_rate: float,
        min_learning_rate: float,
        warmup_iters:int,
        cosine_cycle_iters:int,
):
    '''
    it: 当前迭代步数t
    max_learning_rate 最大学习率
    min_learning_rate 最小学习率
    warmup_iters 预热迭代数T_w
    cosine_cycle_iters 余弦退火的结束迭代步数T_c
    return 第it步使用的学习率
    '''

    if it < warmup_iters:
        return it / warmup_iters * max_learning_rate
    elif warmup_iters <= it <= cosine_cycle_iters:
        return min_learning_rate + 0.5 * (1 + math.cos(it - warmup_iters) / (cosine_cycle_iters - warmup_iters) * math.pi) * (max_learning_rate - min_learning_rate)
    else:
        return min_learning_rate

代码很简单，都是简单的加减乘除套一个cos函数，然后直接分支选择不同的阶段，返回对应的学习率即可。

梯度裁剪

在更新参数的时候，以简单的梯度下降为例，参数的更新公式为：

\[\theta \leftarrow \theta - \eta g_t \]

当某一步的梯度\(g_t\)过大时，参数的更新幅度就会过大，此时参数可能会直接掠到不是我们期望的最优位置，此时的损失值会突然升高，甚至可能会出现梯度爆炸的情况，导致训练不稳定。

为了防止出现梯度过大，我们所采用的策略就是进行梯度裁剪，也就是将较大梯度给他"剪"成一个较小的梯度，保证参数的更新幅度不会过大。一个简单的梯度裁剪方法是按照全局的梯度范数进行裁剪：

\[\tilde g_t = g_t\min\left(1,\quad\frac{M}{\|g_t\|_2}\right) \\ \qquad \\ \|g_t\|_2 = \sqrt{\sum_{i=1}^{N} g_{t,i}^2} \]

其中：

• \(g_t\)：表示第t次迭代的全量梯度
• \(\|g_t\|_2\)：表示第t次迭代的全量梯度的L2范数，也就是其所有元素平方和的平方根
• \(\tilde g_t\)：表示第t次迭代的裁剪后的梯度
• \(M\)：阈值，当\(\|g_t\|_2 > M\)时，才会进行裁剪
• \(\min(1, \frac{M}{\|g_t\|_2})\)：表示裁剪的系数，当\(\|g_t\|_2 > M\)时，裁剪系数为\(\frac{M}{\|g_t\|_2}\)，否则为1，也就是不进行裁剪。

不难看出，我们所作的操作就是计算全量梯度的L2范数，如果其大于阈值\(M\)，那么我们对于所有的梯度，都乘以一个系数\(\frac{M}{\|g_t\|_2} \in [0,1]\)，使得裁剪后的梯度的L2范数为\(M\)，否则不进行裁剪。

当然，有时候为了防止公式中出现除0的情况，我们会在分母中加入一个小量\(\epsilon\)，使得公式变为：

\[\tilde g_t = g_t\min\left(1,\quad\frac{M}{\|g_t\|_2 + \epsilon}\right) \]

下面给出一个使用python实现的梯度裁剪的代码：

import torch
from typing import Iterable
import math
def gradient_clipping(parameters: Iterable[torch.nn.Parameter], max_l2_norm: float, eps:float = 1e-6) -> None :

    parameters = list(parameters)
    l2_norm = math.sqrt(sum(p.grad.pow(2).sum().item() for p in parameters if p.grad is not None))
    # p.pow(2).sun()->所有梯度平方和
    # for p in parameters if p.grad is not None 列表推导，最终会有一个所有参数p梯度平方和的列表
    # 最后再对这个列表求和，就是全量的平方和，对其开方就是其L2范数
    if l2_norm > max_l2_norm:
        for p in parameters:
            with torch.no_grad():
                if p.grad is not None:
                    p.grad *= max_l2_norm / (l2_norm + eps)

数据准备

在训练Transformer的时候，我们准备的数据集实际上是一段或者很多文本，我们需要将文本进行分词，然后使用分词后的token去构建输入以及输入对应的输出所构成的训练样本。

对于一段已经分词结束的文本，我们对其做如下操作：

• 在原始token序列中随机采样Batch_size个长度为seq_len的子序列，作为一个batch的训练样本
• 对于每一个子序列，我们将其作为输入序列，同时，我们期望的是模型吃进去前seq_len个token后，能够预测出后面第seq_len+1个token的值，因此我们会向后偏移一个token，同时保持seq_len的长度，作为输出序列

在上图的例子中，我们从原始序列中抽取了一个长度为3的子序列\([x_2, x_3, x_4]\)作为输入序列，在训练时，输入输出时这样对应的：

模型一次接收整段输入：
input  = [x_2, x_3, x_4]
target = [x_3, x_4, x_5]

各时间位置的预测目标对应为：
x_2 所在位置 → 预测 x_3
x_3 所在位置 → 预测 x_4
x_4 所在位置 → 预测 x_5

箭头表示时间位置与下一 token 标签的对应关系。
每个位置可以利用输入序列中该位置及之前的上下文。

注意，在选取子序列的时候，我们只需要在不越界的前提下，随机采样起始点，然后指定子序列长度即可。无需担心各个输入会重叠，因为我们训练的目的是输入一段序列并让模型预测出下一步的token，因此即使输入序列有重叠，模型也会学习到不同的上下文信息。如输入选取\([x_2, x_3, x_4]\)同时另一个batch选取\([x_3, x_4, x_5]\)，虽然有重叠，但是对于\(x_3\)来说，其上下文信息是不同的，因此模型会学习到不同的上下文信息。

给出一个简单的python实现的DataLoader即数据加载器，用于从原始token序列中采样出训练样本：

import numpy as np
import torch

def get_batch(
        x: np.ndarray,
        batch_size: int,
        context_length: int,
        device: str
) -> tuple[torch.Tensor, torch.Tensor]:
    '''
    x : 一维整数数组，元素为token_id
    batch_size: 批大小
    context_length: 每条句子的上下文长度
    device: 设备字符串
    retruns [输入序列，输出序列]，每个形状都是(batch_size, context_length)
    '''

    indices = np.random.randint(0, len(x) - context_length, size = batch_size)
    # 随机生成一个数组，其中元素在[0, len(x) - context_length)大小之间，数组大小为batch_size
    # len(x) - context_length + 1 -> 文本token长度 - 每个句子的token长度 + 1。
    # 这样保证每个随机数取到的下标开始向后多取一个token不会越界（注意右边界是取不到的）
    # 最大下标为len(x) - context_length - 1
    # indices 一维向量(batch_size,)

    input_seq = x[indices[:, None] + np.arange(0, context_length)]
    # indices[:, None] -> (batch_size, 1)
    # np.arange(0, context+length)，生成一个数据为[0, 1, 2, ..., context_length -1]的数组
    # 相加(B, 1) + (Seq,) -> (B, 1) + (1, Seq) -> (B, Seq) + (1, Seq) -> (B, Seq) + (B, Seq)
    # 得到的是每个样本完整的token下标

    indices += 1
    # 所有的头索引向后一个单位，准备取target

    target_seq = x[indices[:, None] + np.arange(0, context_length)] # 如法炮制

    input_seq_tensor = torch.tensor(input_seq, device=device) # 将np数组转换为torch中的Tensor

    target_seq_tensor = torch.tensor(target_seq, device=device)

    return (input_seq_tensor, target_seq_tensor)

几点说明：

• 这里的x是一个一维的整数数组，表示文本中每个token对应的id
• context_length表示每个训练样本的长度，也就是我们前面提到的seq_len

在开始的时候我们使用np.random.randint随机来生成一个起始索引数组indices，其大小为batch_size，每个元素的取值范围为[0, len(x) - context_length)，因此最大索引为len(x) - context_length - 1，这样保证了我们在后续取子序列的时候不会越界，因为如果此时取到最大下标，则输入的最后一个元素的下标为len(x) - context_length - 1 + context_length - 1 = len(x) - 2，此时取到的子序列的最后一个元素的下标为len(x) - 2，而我们在取target的时候会向后偏移一个token，因此target的最后一个元素的下标为len(x) - 1，此时不会越界。注意，x的下标是从0开始的，因此最大下标为len(x) - 1。

在后面的操作中，我们使用indices[:, None] + np.arange(0, context_length)来生成每个样本的完整下标：

• indices[:, None]将indices从一维数组(batch_size,)转换为二维数组(batch_size, 1)
  • 如[1, 2 ,3] -> [[1], [2], [3]]
• np.arange(0, context_length)生成一个一维数组，表示每个样本的相对下标，从0到context_length-1，维度(context_length,)
  • np.arange(0, 3) -> [0, 1, 2]
• 两者相加，有如下广播

   ① 原始形状：(B, 1) + (Seq,)

   [[1],
    [2],       +    [0, 1, 2]
    [3]]


② 从右向左对齐维度：(B, 1) + (1, Seq)

   [[1],
    [2],       +    [[0, 1, 2]]
    [3]]


③ 左侧沿列方向广播：(B, Seq) + (1, Seq)

   [[1, 1, 1],
    [2, 2, 2], +    [[0, 1, 2]]
    [3, 3, 3]]


④ 右侧沿行方向广播：(B, Seq) + (B, Seq)

   [[1, 1, 1],      [[0, 1, 2],
    [2, 2, 2], +     [0, 1, 2],
    [3, 3, 3]]       [0, 1, 2]]


⑤ 逐元素相加，结果形状为 (B, Seq)

   [[1+0, 1+1, 1+2],      [[1, 2, 3],
    [2+0, 2+1, 2+2],  =    [2, 3, 4],
    [3+0, 3+1, 3+2]]       [3, 4, 5]]

经过这种操作，就有了一个以随机起始索引为头的长度为context_length的子序列的下标数组，形状为(batch_size, context_length)，然后我们直接使用这个下标数组去索引原始token序列x，就得到了每个样本的输入序列，我们不妨记运算后的结果为input_seq_indices，则有：

\[input[i][j] = x[input\_seq\_indices[i][j]] \quad where \quad i \in [0, batch\_size), j \in [0, context\_length) \]

此时这个矩阵中，每一行就是一个batch的输入序列，所有行是所有batch的输入序列。

同理，让头索引indices向后偏移一个单位，得到target的下标数组，做相同的操作就有output_seq。最后将np数组转换到torch的Tensor并将结果返回即可。

Checkpoint

在训练的过程中，如果突然断电或者程序崩溃，那么我们之前训练的所有参数都将丢失，因此我们需要在训练的过程中，定期保存模型的参数，这就是所谓的Checkpoint。同时主动去保存检查点，也可以让我们了解在不同的训练阶段，模型的参数的变化情况，方便我们去分析模型的训练情况。

保存检查点的步骤只需要使用torch提供的一些api，将模型的参数作为字典量存入对应的文件中即可，当我们想要载入检查点的时候，只需要将文件中的字典量载入到模型中即可。

python实现如下：

from torch import Module, optim
import torch
import os
import typing
def save_checkpoint(
        model: Module,
        optimizer: optim.Optimizer,
        iteration: int,
        out:str | os.PathLike | typing.BinaryIO | typing.IO[bytes]
) -> None:
    '''
    model : 模型
    optimizer : 优化器
    iteration : 当前迭代步数
    out: 保存路径或者文件对象
    '''

    torch.save(
        {
            "model" : model.state_dict(),
            "optim" : optimizer.state_dict(),
            "iter_num": iteration,
        }
    , out)
    # torch.save, 将当前的参数状态以键值对的方式保存到对应的文件中
    # 常见的保存文件的后缀名为.pt，表示这是一个pytorch的专属格式

def load_checkpoint(
                    src: str | os.PathLike | typing.BinaryIO | typing.IO[bytes],
                    model:Module,
                    optimizer:optim.Optimizer = None
                    ) -> int:
    '''
    src: 检查点文件路径或者文件对象
    model: torch.Module，模型
    optimizer: 优化器
    returns: 检查点中保存的迭代步数
    '''

    device = next(model.parameters()).device
    # 读取模型第一个参数的设备，作为参数设备
    state_dic = torch.load(src, map_location=device, weights_only=True)
    # 使用torch.load去加载src指定的键值对参数

    # 下面就是使用键分别加载对应的状态到对应的工具中
    model.load_state_dict(state_dic["model"])

    if optimizer is not None:
        optimizer.load_state_dict(state_dic["optim"])

    return state_dic["iter_num"]

训练脚本

综合前面的所有模块，最终给出一个实现了参数读取以及训练循环以及保存检查点的训练脚本：

import torch
import numpy as np
import time
import argparse
import os
from torch.nn import Module
from torch.optim import Optimizer
from Transformer import TransformerLM
from train1_crossEntropy import cross_entropy
from train4_GradientClipping import gradient_clipping
from train2_optimizer import AdamW
from train5_GetBatch import get_batch
from train3_CosineAnnealingSchedule import learing_rate_schedule
from train6_CheckPoint import load_checkpoint,save_checkpoint

DATA_DIR = 0

def train_step(
        model: Module,
        optimizer: Optimizer, 
        x: torch.Tensor,
        y: torch.Tensor,
        lr: float,
        max_l2_norm: float =1.0
) -> float:
    '''
    model: 模型
    optimizer：优化器
    x: 输入tokenID，需要为long类型，维度(B, S)
    y: 目标tokenID
    lr： 学习率
    max_l2_norm: 梯度裁剪允许的最大L2范数
    '''

    optimizer.param_groups[0]["lr"] = lr # 在step之前对优化器写入学习率，lr有余弦退火给出

    logits = model(x) # 前向传播
    loss = cross_entropy(logits, y).mean() # 得到所有位置上的总平均损失

    optimizer.zero_grad() # 情况上一步梯度
    loss.backward()       # 反向传播计算梯度
    # 在pytorch中，所有的运算都会记录在计算图中，因此loss即使没有显示的告诉他我们自定义的损失函数长什么样子，他也已经将对应参数的梯度求出来了
    gradient_clipping(model.parameters(), max_l2_norm) # 梯度裁剪
    optimizer.step() # 参数更新

    return loss.item()

def train_llm(  model: torch.nn.Module,
                optimizer: torch.optim.Optimizer,
                train_data: np.ndarray,
                val_data: np.ndarray,
                args
                ) -> None:
    '''
    训练循环: 反复调用train_step,并按间隔做验证，保存检查点与记录日志
    train_data/val_data:np.ndarray, token_ID数组
    args:超参数集合
    '''
    start_time = time.time()

    for i in range(1, args.total_iters + 1):
        model.train() # 切换到训练模式

        # 取出一批训练数据并计算本步的学习率
        x, y = get_batch(train_data, args.batch_size, args.context_length, args.device)
        lr = learing_rate_schedule(i, args.max_lr, args.min_lr, args.warmup_iters, args.cosine_cycle_iters)

        # 单步训练
        loss = train_step(model, optimizer, x.long(), y.long(), lr=lr, max_l2_norm=args.max_l2_norm)

        # 记录指标

        print(f"iter{i:5d} train_loss{loss:.4f} lr{lr: .6f} wall_time={time.time() - start_time:.1f}s")

        # 定期在验证集上评估，此时要torch.no_grad()

        if i % args.val_interval == 0:
            model.eval() # 切换到评估模式
            with torch.no_grad():
                val_x, val_y = get_batch(val_data, args.batch_size, args.context_length, args.device)
                val_loss = cross_entropy(model(val_x.long()), val_y.long()).mean().item()
            print(f"iter{i:5d} val_loss = {val_loss:.4f}")

        # 定期保存检查点
        if i % args.save_interval == 0:
            save_checkpoint(model, optimizer, i, args.save_path)
            print(f"iter {i:5d} 已保存检查点 -> {args.save_path}")




# -------------------- 超参与入口 --------------------
def parse_args():
    parser = argparse.ArgumentParser(description="训练一个 Transformer 语言模型")

    # 数据
    parser.add_argument("--train_data_path", type=str, default=os.path.join(DATA_DIR, "TinyStoriesV2-GPT4-train.npy"))
    parser.add_argument("--val_data_path", type=str, default=os.path.join(DATA_DIR, "TinyStoriesV2-GPT4-valid.npy"))

    # 优化器超参
    parser.add_argument("--max_lr", type=float, default=6e-4)
    parser.add_argument("--min_lr", type=float, default=6e-5)
    parser.add_argument("--weight_decay", type=float, default=0.01)
    parser.add_argument("--beta1", type=float, default=0.9)
    parser.add_argument("--beta2", type=float, default=0.95)

    # 训练超参
    parser.add_argument("--batch_size", type=int, default=32)
    parser.add_argument("--context_length", type=int, default=256)
    parser.add_argument("--total_iters", type=int, default=5000)
    parser.add_argument("--warmup_iters", type=int, default=500)
    parser.add_argument("--cosine_cycle_iters", type=int, default=5000)
    parser.add_argument("--max_l2_norm", type=float, default=1.0)
    parser.add_argument("--val_interval", type=int, default=100)
    parser.add_argument("--save_interval", type=int, default=500)
    parser.add_argument("--save_path", type=str, default="checkpoint.pt")
    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")

    # 模型超参
    parser.add_argument("--vocab_size", type=int, default=10000)
    parser.add_argument("--d_model", type=int, default=512)
    parser.add_argument("--num_layers", type=int, default=4)
    parser.add_argument("--num_heads", type=int, default=16)
    parser.add_argument("--rope_theta", type=float, default=10000.0)

    return parser.parse_args()



def main():

    args = parse_args()
    # mmap_mode = "r" : 以内存映射方式打开，不会一次性将整个数据集读入内存
    train_data  = np.load(args.train_data_path, mmap_mode="r")
    val_data = np.load(args.val_data_path, mmap_mode="r")
    print(f"训练集{train_data.shape}, 验证集{val_data.shape}, 设备{args.device}")

    d_ff = int(args.d_model * 8 / 3 +63) // 64 *64 # 向上取整到64整数倍

    model = TransformerLM(
        vocab_size=args.vocab_size,
        context_length=args.context_length,
        d_model=args.d_model,
        num_layers=args.num_layers,
        num_heads=args.num_heads,
        d_ff=d_ff,
        rope_theta=args.rope_theta
    ).to(args.device)

    optimizer = AdamW(
        model.parameters(),
        lr = args.max_lr,
        betas=(args.beta1, args.beta2),
        weight_decay=args.weight_decay
    )

    train_llm(model, optimizer, train_data, val_data, args)

if __name__ == "__main__":
    main()

参考

• 跟着 stanford 动手学大模型博客
• 跟着 stanford 动手学大模型b站
• Transformer 详解：从架构到训练
• 知乎回答：信息熵是什么
• 机器学习2 -- 优化器（SGD、SGDM、Adagrad、RMSProp、Adam）
• 深度学习中正则化与L2正则化详解