细嗦Transformer（三）：准备训练，讲解及代码实现优化器、学习率调整策略、正则化和KL散度损失

文章目录

• 关注我：细嗦大模型
• 批处理对象 / Batches and Masking
• 训练循环主函数 / Training Loop
• 优化器 / Optimizer
• 学习率调整策略 / Learning rate adjustment strategy
• • 样例测试
• 正则化 / Regularization
• • Label smoothing 标签平滑
  • KL散度损失
  • 样例测试
• Github完整代码----求求了给个star和关注吧
• 参考资料
• 求求了，给个star和关注吧
• 求求了，给个star和关注吧
• 求求了，给个star和关注吧

关注我：细嗦大模型

我会定期更气LLM算法系列文章，旨在结合代码去理解各知识点。可以观众我的公众号，第一时间收到更新。关注公众号回复”思维导图“还可以获得LLM基础知识、预训练、微调、人类对齐RLHF、DPO、推理加速5个部分的思维导图。

前面我们讲已经讲完了如何构建一个Transformer模型，在接下来的文章中，我们将继续讲解如何去训练模型。本文中，则首先讲解Transformer的训练机制。

批处理对象 / Batches and Masking

在训练模型的时候，数据都是分批次的给到模型去训练，而不是一条一条或者一次全部给到模型；并且还需要对数据的padding部分和掩码部分做mask操作。因此，我们需要定义一个批处理对象，其中包含了用于训练的src和target数据，并对数据进行mask操作。

class Batch:
    """Object for holding a batch of data with mask during training."""

    def __init__(self, src, tgt=None, pad=2):  # 2 = <blank>
        self.src = src
        self.src_mask = (src != pad).unsqueeze(-2)
        if tgt is not None:
            self.tgt = tgt[:, :-1]
            self.tgt_y = tgt[:, 1:]
            self.tgt_mask = self.make_std_mask(self.tgt, pad)
            self.ntokens = (self.tgt_y != pad).data.sum()

    @staticmethod
    def make_std_mask(tgt, pad):
        "Create a mask to hide padding and future words."
        tgt_mask = (tgt != pad).unsqueeze(-2)
        tgt_mask = tgt_mask & subsequent_mask(tgt.size(-1)).type_as(
            tgt_mask.data
        )
        return tgt_mask

训练循环主函数 / Training Loop

下面创建训练一个Epoch的核心函数，遍历每个批次的数据，前向传播，计算损失，反向传播以及更新梯度。核心为这四步，不管什么模型都一样。剩下的就是记录日志和训练状态。下面给出的代码中TrainState就是记录训练

import time
class TrainState:
    """Track number of steps, examples, and tokens processed"""

    step: int = 0  # 在当前epoch中训练的步数
    accum_step: int = 0  # 进行梯度累计的次数
    samples: int = 0  # 使用的样本数量
    tokens: int = 0  # 已经处理的token数量
        
def run_epoch(
    data_iter, # 数据迭代器
    model, # 模型
    loss_compute, # 计算损失的函数
    optimizer, # 优化器
    scheduler, # 学习率调度器
    mode="train", # 模式
    accum_iter=1, # 梯度累计的次数
    train_state=TrainState(), # 训练状态
):
    """Train a single epoch"""
    start = time.time()
    total_tokens = 0 # 记录总的token数量
    total_loss = 0 # 记录总的损失
    tokens = 0 # 记录当前epoch已经处理的token数量
    n_accum = 0 # 记录当前epoch已经进行的梯度累计次数
    for i, batch in enumerate(data_iter): # 遍历数据迭代器
        out = model.forward(
            batch.src, batch.tgt, batch.src_mask, batch.tgt_mask # 前向传播
        )
        loss, loss_node = loss_compute(out, batch.tgt_y, batch.ntokens) # 计算损失
        # loss_node = loss_node / accum_iter
        if mode == "train" or mode == "train+log":
            loss_node.backward() # 反向传播
            train_state.step += 1 # 更新训练步数
            train_state.samples += batch.src.shape[0] # 更新样本数量
            train_state.tokens += batch.ntokens # 更新token数量
            if i % accum_iter == 0: # 每accum_iter步更新一次参数
                optimizer.step() # 更新参数
                optimizer.zero_grad(set_to_none=True) # 梯度清零
                n_accum += 1 # 更新梯度累计次数
                train_state.accum_step += 1 # 更新累计步数
            scheduler.step()

        total_loss += loss # 更新总的损失
        total_tokens += batch.ntokens # 更新总的token数量
        tokens += batch.ntokens # 更新当前epoch
        if i % 40 == 1 and (mode == "train" or mode == "train+log"): # 每40步打印一次信息
            lr = optimizer.param_groups[0]["lr"] # 获取当前学习率
            elapsed = time.time() - start # 计算这40步使用的时间
            print( # 打印信息
                (
                    "Epoch Step: %6d | Accumulation Step: %3d | Loss: %6.2f "
                    + "| Tokens / Sec: %7.1f | Learning Rate: %6.1e"
                )
                % (i, n_accum, loss / batch.ntokens, tokens / elapsed, lr)
            )
            start = time.time() # 重置开始实践
            tokens = 0 # 重置token数量
        del loss
        del loss_node
    return total_loss / total_tokens, train_state

优化器 / Optimizer

使用Adam优化器，Adam是一种基于一阶梯度的优化算法，结合了动量和RMSprop思想，能够自适应的调整每个参数的学习率，适用于处理大规模数据的参数优化问题。

Adam更新参数的公式为：

θ t = θ t − 1 − η v ^ t + ϵ m ^ t \theta_{t} = \theta_{t-1} - \frac{\eta}{\sqrt{\hat{v}_{t}} + \epsilon} \hat{m}_{t} θt​=θt−1​−v^t​ ​+ϵη​m^t​
其中, θ \theta θ是待优化参数， η \eta η是学习率， ϵ \epsilon ϵ是一个极小数，防止除以0，其余计算公式为：

m ^ t = m t 1 − β 1 t \hat{m}_{t} = \frac{m_{t}}{1 - \beta_1^t} m^t​=1−β1t​mt​​
v ^ t = v t 1 − β 2 t \hat{v}_{t} = \frac{v_{t}}{1 - \beta_2^t} v^t​=1−β2t​vt​​
这主要是为了消除初期估计的偏差， t t t是当前迭代次数，通常， β 1 = 0.9 \beta_1=0.9 β1​=0.9， β 2 = 0.999 \beta_2=0.999 β2​=0.999， β 1 t , β 2 t \beta_1^t,\beta_2^t β1t​,β2t​则为相乘k次：

β 1 t = β 1 × β 1 × … × β 1 \beta_1^t = \beta_1 \times \beta_1 \times \ldots \times \beta_1 β1t​=β1​×β1​×…×β1​
β 2 t = β 2 × β 2 × … × β 2 \beta_2^t = \beta_2 \times \beta_2 \times \ldots \times \beta_2 β2t​=β2​×β2​×…×β2​
m t , v t m_{t},v_t mt​,vt​则分别是Adam的一阶矩估计和二阶矩估计， g t g_t gt​则是时间步 t t t的梯度。 m t , v t m_{t},v_t mt​,vt​计算公式为:

m t = β 1 m t − 1 + ( 1 − β 1 ) ⋅ g m_{t} = \beta_1 m_{t-1} + (1 - \beta_1) \cdot g mt​=β1​mt−1​+(1−β1​)⋅g
v t = β 2 v t − 1 + ( 1 − β 2 ) ⋅ g 2 v_{t} = \beta_2 v_{t-1} + (1 - \beta_2) \cdot g^2 vt​=β2​vt−1​+(1−β2​)⋅g2

"""Para
params(iterable)：可用于迭代优化的参数或者定义参数组的dicts。
lr (float, optional) ：学习率(默认: 1e-3)
betas (Tuple[float, float], optional)：用于计算梯度的平均和平方的系数(默认: (0.9, 0.999))
eps (float, optional)：为了提高数值稳定性而添加到分母的一个项(默认: 1e-8)
weight_decay (float, optional)：权重衰减(如L2惩罚)(默认: 0)
"""
torch.optim.Adam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0)

优化器都存储哪些参数？以及优化器占用的显存数量如何计算？
从公式中就可以看出来，在时间步 t t t，计算 θ t \theta_t θt​， m t m_t mt​， v t v_t vt​时，都需要用到前一个时间步的数据，因此优化器中是必须存储 θ \theta θ， m m m， v v v这三个参数的。假设模型的参数量为1B，那么优化器就需要存储模型参数量本身 θ \theta θ，以及一节动量矩估计 m m m和二阶动量矩估计 v v v，存储的总参数量为3B，每个参数用FP32表示，一个参数占4个字节，总字节数为 3 ∗ 4 ∗ 1 0 9 3*4*10^9 3∗4∗109，所以，大约需要11GB的内存。

学习率调整策略 / Learning rate adjustment strategy

大语言模型在预训练阶段都通常都采用学习率调整策略，分为预热阶段和衰减阶段。预热阶段一般占整个训练步骤的 0.1% 至 0.5%，然后学习率便开始进行衰减。

本文中，学习率的调整公式为：

l r a t e = d model − 0.5 ⋅ min ⁡ ( s t e p _ n u m − 0.5 , s t e p _ n u m ⋅ w a r m u p _ s t e p s − 1.5 ) lrate = d_{\text{model}}^{-0.5} \cdot \min({step\_num}^{-0.5}, {step\_num} \cdot {warmup\_steps}^{-1.5}) lrate=dmodel−0.5​⋅min(step_num−0.5,step_num⋅warmup_steps−1.5)
这个公式可以改写成：

l r a t e = { d model − 0.5 ⋅ s t e p _ n u m ⋅ w a r m u p _ s t e p s − 1.5 ( s t e p _ n u m < w a r m u p _ s t e p ) d model − 0.5 ⋅ s t e p _ n u m − 0.5 ( s t e p _ n u m > = w a r m u p _ s t e p ) lrate =\begin{cases} d_{\text{model}}^{-0.5} \cdot {step\_num} \cdot {warmup\_steps}^{-1.5} (step\_num < warmup\_step) \\ d_{\text{model}}^{-0.5} \cdot {step\_num}^{-0.5} (step\_num >= warmup\_step) \end{cases} lrate={dmodel−0.5​⋅step_num⋅warmup_steps−1.5(step_num<warmup_step)dmodel−0.5​⋅step_num−0.5(step_num>=warmup_step)​
所以， 在小于 w a r m u p _ s t e p s warmup\_steps warmup_steps的step中学习率是线性增加的，然后按步数的平方根的倒数比例降低学习率。本文中， w a r m u p _ s t e p s = 4000 warmup\_steps=4000 warmup_steps=4000。
还需要注意的是，当 s t e p step step为0时，学习率不能为0。

def rate(step, model_size, factor, warmup):
    if step == 0:
        step = 1
    return factor * (
        model_size ** (-0.5) * min(step ** (-0.5), step * warmup ** (-1.5))
    )

为什么要采用学习率调整策略？
在模型训练的初始阶段，由于参数是随机初始化的，梯度通常也比较大，因此需要使用较小的学习率使得训练较为稳定。训练中通常采用线性预热策略来逐步调整学习率。具体来说，学习率将从一个非常小的数值（例如 0 或者 1 × 1 0 − 8 1×10^{−8} 1×10−8）线性平稳增加，直到达到预设的最大阈值。模型在学习率较大时可以加快收敛速度，这个最大阈值通常设定在 5 × 1 0 − 5 5×10^{−5} 5×10−5 到 1 × 1 0 − 4 1×10^{−4} 1×10−4 之间。达到最大阈值之后学习率会开始逐渐衰减，以避免在较优点附近来回震荡。最后，学习率一般会衰减到其最大阈值的 10%。常见的衰减策略有线性衰减，余弦衰减，平方根倒数衰减。

样例测试

下面，我们通过三个样例来看学习率的调整过程，具体过程看代码和注释。这里只对LambdaLR进行解读。

"""学习率调整器
LambdaLR的作用就是自定义一个函数，然后根据这个函数来调整优化器的学习率，从下面代码可以看出，Adam优化器的学习率初始值为1，lr_lambda函数输出是一个学习率的倍数，这个倍数会乘以当前的学习率，然后更新优化器的学习率。

参数：
optimizer：被调整学习率的优化器
lr_lambda：用户自定义的学习率调整规则。可以是lambda表达式，也可以是函数
last_epoch：当前优化器的已迭代次数，后文我们将其称为epoch计数器。默认是-1，字面意思是第-1个epoch已完成，也就是当前epoch从0算起，从头开始训练。如果是加载checkpoint继续训练，那么这里要传入对应的已迭代次数
verbose：是否在更新学习率时在控制台输出提醒
"""
torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_epoch=-1, verbose=False)

import altair as alt
from torch.optim.lr_scheduler import LambdaLR
import torch
import pandas as pd
def example_learning_schedule():
    opts = [
        [512, 1, 4000],  # model_size=512, factor=1, warmup_steps=4000
        [512, 1, 8000],  # model_size=512, factor=1, warmup_steps=8000
        [256, 1, 4000],  # model_size=256, factor=1, warmup_steps=4000
    ]

    dummy_model = torch.nn.Linear(1, 1) # 定义一个简单的线性模型
    learning_rates = []

    # we have 3 examples in opts list.
    for idx, example in enumerate(opts):
        # run 20000 epoch for each example
        optimizer = torch.optim.Adam( # 定义优化器
            dummy_model.parameters(), lr=1, betas=(0.9, 0.98), eps=1e-9
        )
        lr_scheduler = LambdaLR( # 定义学习率调整器
            optimizer=optimizer, lr_lambda=lambda step: rate(step, *example)
        )
        tmp = []
        for step in range(20000): # 进行20000步训练, 记录每一步的学习率
            tmp.append(optimizer.param_groups[0]["lr"]) # 记录当前学习率
            optimizer.step() # 更新参数
            lr_scheduler.step() # 更新学习率
        learning_rates.append(tmp) # 记录当前example的学习率

    learning_rates = torch.tensor(learning_rates) # 转换为tensor

    # Enable altair to handle more than 5000 rows
    alt.data_transformers.disable_max_rows() # 禁用最大行数限制，使altair可以处理超过5000行的数据

    opts_data = pd.concat(
        [
            pd.DataFrame(
                {
                    "Learning Rate": learning_rates[warmup_idx, :],
                    "model_size:warmup": ["512:4000", "512:8000", "256:4000"][
                        warmup_idx
                    ],
                    "step": range(20000),
                }
            )
            for warmup_idx in [0, 1, 2]
        ]
    )

    return (
        alt.Chart(opts_data)
        .mark_line()
        .properties(width=600)
        .encode(x="step", y="Learning Rate", color="model_size:warmup:N")
        .interactive()
    )


example_learning_schedule()

运行之后，可以看到学习率的调整曲线。

正则化 / Regularization

什么是正则化？
训练模型的本质是可以看成寻找一个函数 H 0 ( x ) H_0(x) H0​(x)来你和数据集，而正则化的目的就是防止模型过拟合，增强模型的泛化能力。简单来说，正则化就是给损失函数增加一个正则项来限制损失函数的增加。正则化相关问题后续单独开一篇文章来讲。

Label smoothing 标签平滑

标签平滑（Label Smoothing）是一种在机器学习中常用的正则化技术，特别是在分类任务中。它的核心思想是将硬标签（hard labels）转换为软标签（soft labels），以此来防止模型在训练过程中对某一类别的预测过于自信，从而提高模型的泛化能力。

在传统的分类任务中，我们通常使用one-hot编码来表示标签，即目标类别的概率为1，其他类别的概率为0。这种表示方法称为硬标签。然而，硬标签可能导致模型在训练数据上过拟合，特别是当训练数据无法覆盖所有情况时。为了解决这个问题，标签平滑通过在真实标签的概率上引入一个小的噪声 ε \varepsilon ε，将其从1降低到 1 − ε 1-\varepsilon 1−ε，同时将其他标签的概率从0提高到 ε / K \varepsilon/K ε/K，其中K是类别总数。这样，每个标签的概率分布变得更加平滑，模型对于每个类别的预测不再那么绝对。

在本文中，则是使用了KL散度损失来实现标签平滑。

KL散度损失

KL散度可以用来衡量两个概率分布之间的相似度，KL散度越小，说明两个概率分布的距离越近，越相似。

计算公式为：

L ( y p r e d , y t r u e ) = y t r u e ⋅ l o g y t r u e y p r e d = y t r u e ⋅ ( log ⁡ y t r u e − log ⁡ y p r e d ) L(y_{pred}, y_{true}) = y_{true} \cdot log{\frac{y_{true}}{y_{pred}}} = y_{true} \cdot (\log y_{true} - \log y_{pred}) L(ypred​,ytrue​)=ytrue​⋅logypred​ytrue​​=ytrue​⋅(logytrue​−logypred​)

"""Para
size_average与reduce 已被弃用，具体功能由参数reduction代替
reduction：指定损失输出的形式，有四种选择：none|mean|batchmean|sum。none：损失不做任何处理，直接输出一个数组；mean：将得到的损失求平均值再输出，会输出一个数；batchmean：将输出的总和除以batchsize；sum：将得到的损失求和再输出，会输出一个数
log_target：指定是否对输入的y使用log操作
"""
torch.nn.KLDivLoss(size_average=None, reduce=None, reduction='mean', log_target=False)

import torch
import torch.nn as nn
import pandas as pd
import altair as alt
class LabelSmoothing(nn.Module):
    def __init__(self, size, padding_idx, smoothing=0.0): # size是词表大小，padding_idx是padding的索引，smoothing是平滑参数
        super(LabelSmoothing, self).__init__()
        self.criterion = nn.KLDivLoss(reduction="sum")
        self.padding_idx = padding_idx
        self.confidence = 1.0 - smoothing
        self.smoothing = smoothing
        self.size = size
        self.true_dist = None

    def forward(self, x, target):
        assert x.size(1) == self.size
        true_dist = x.data.clone() # kldivloss的输入是log_softmax
        true_dist.fill_(self.smoothing / (self.size - 2)) # 除了padding_idx和正确的标签，其他的概率都是smoothing/size-2
        """scatter_函数是一个用于在特定索引处更新张量的原地操作。这个函数接受三个参数：dim、index和src。dim参数指定了要沿着哪个维度进行索引，index是一个包含索引的张量，而src是包含要散布的值的张量。
        """
        true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence) 
        print(f"true_dist:{true_dist}")
        # 将正确的标签的概率设置为confidence
        true_dist[:, self.padding_idx] = 0 # 将padding_idx的概率设置为0
        mask = torch.nonzero(target.data == self.padding_idx) # 获取padding的位置
        print(f"mask:{mask}")
        if mask.dim() > 0:
            true_dist.index_fill_(0, mask.squeeze(), 0.0) # 将padding的概率设置为0
        self.true_dist = true_dist
        print(f"true_dist:{true_dist}")
        return self.criterion(x, true_dist.clone().detach()) # 计算损失

一定要理解这里padding的含义，词表大小为 v o c a b s i z e vocab_{size} vocabsize​，那么padding_idx就是对第padding_idx个token进行mask操作，也就是one-hot编码中的下标。对padding_idx位置的token不作预测，所以需要把每一个真实值 y t r u e y_{true} ytrue​的one-hot向量中的padding_idx位置的概率设为0，同时，真实值为padding_idx的样例也要mask掉。

具体过程可以通过下面的测试样例来debug一遍，看true_dist的变化。

样例测试

下面使用一个样例来测试Label smoothing，下面这个代码主要是展示了label 平滑之后的结果，label就是真实值。每个 y t r u e y_{true} ytrue​用one-hot表示

def example_label_smoothing():
    crit = LabelSmoothing(5, 0, 0.4) # 定义一个LabelSmoothing对象
    predict = torch.FloatTensor( # 定义一个预测值
        [
            [0, 0.2, 0.7, 0.1, 0],
            [0, 0.2, 0.7, 0.1, 0],
            [0, 0.2, 0.7, 0.1, 0],
            [0, 0.2, 0.7, 0.1, 0],
            [0, 0.2, 0.7, 0.1, 0],
        ]
    )
    print(predict.log())
    crit(x=predict.log(), target=torch.LongTensor([2, 1, 0, 3, 3]))
    LS_data = pd.concat(
        [
            pd.DataFrame(
                {
                    "target distribution": crit.true_dist[x, y].flatten(),
                    "columns": y,
                    "rows": x,
                }
            )
            for y in range(5)
            for x in range(5)
        ]
    )
    return (
        alt.Chart(LS_data)
        .mark_rect(color="Blue", opacity=1)
        .properties(height=200, width=200)
        .encode(
            alt.X("columns:O", title=None),
            alt.Y("rows:O", title=None),
            alt.Color(
                "target distribution:Q", scale=alt.Scale(scheme="viridis")
            ),
        )
        .interactive()
    )

example_label_smoothing()

输出结果为：

图中，横轴表示one-hot向量的下标，纵轴表示第i个token的one-hot向量，黄色位置就是真实值 y t r u e y_{true} ytrue​，蓝色位置是smooth之后的值，为0.1333。每个one-hot向量的0位置，以及 y t r u e = 0 y_{true}=0 ytrue​=0的向量都被mask了。
下面再通过一个示例来看Label smoothing对loss的影响。

def loss(x, crit):
    d = x + 3 * 1 # x/d = x/(x+3) =1-3/(x+3)
    # 所以，当x越大，x/d越接近1，那么predict就越接近[0, 1, 0, 0, 0]，损失越小
    predict = torch.FloatTensor([[0, x / d, 1 / d, 1 / d, 1 / d]])
    return crit(predict.log(), torch.LongTensor([1])).data


def penalization_visualization():
    crit = LabelSmoothing(5, 0, 0.1)
    loss_data = pd.DataFrame(
        {
            "Loss": [loss(x, crit) for x in range(1, 100)], # 计算每一步的损失
            "Steps": list(range(99)),
        }
    ).astype("float")

    return (
        alt.Chart(loss_data)
        .mark_line()
        .properties(width=350)
        .encode(
            x="Steps",
            y="Loss",
        )
        .interactive()
    )

penalization_visualization()

结果为：

Github完整代码----求求了给个star和关注吧

Transformer系列的完整代码和可运行notebook见https://github.com/Ace-bb/Transformer

参考资料

1. The Annotated Transformer （代码来源）
2. 深度学习各类优化器详解（动量、NAG、adam、Adagrad、adadelta、RMSprop、adaMax、Nadam、AMSGrad）_动量优化器-CSDN博客
3. 算法岗常见面试题（六）：优化器
4. RUCAIBox/LLMSurvey: The official GitHub page for the survey paper “A Survey of Large Language Models”.
5. PyTorch学习笔记：nn.KLDivLoss——KL散度损失-CSDN博客

求求了，给个star和关注吧

求求了，给个star和关注吧

求求了，给个star和关注吧