大模型学习笔记 - LLM 预训练

LLM 模型预训练

[LLM 模型预训练](#LLM 模型预训练)
- [1. 预训练任务](#1. 预训练任务)
- [2. 优化参数设置](#2. 优化参数设置)
  - [2.1 基于批次数据的训练](#2.1 基于批次数据的训练)
  - [2.2 学习率](#2.2 学习率)
  - [2.3 优化器](#2.3 优化器)
  - [2.4 稳定优化技术](#2.4 稳定优化技术)
- [3. 可扩展的训练技术](#3. 可扩展的训练技术)
  - [3.1 3D并行训练](#3.1 3D并行训练)
  - [3.2 零冗余优化器](#3.2 零冗余优化器)
  - [3.3 激活重计算](#3.3 激活重计算)
  - [3.4 混合精度训练](#3.4 混合精度训练)
- [4. 模型参数量计算与效率分析](#4. 模型参数量计算与效率分析)
  - [4.1 参数量估计](#4.1 参数量估计)
  - [4.2 训练运算量估计](#4.2 训练运算量估计)
  - [4.3 训练时间估计](#4.3 训练时间估计)
  - [4.4 训练显存估计](#4.4 训练显存估计)
- [5. 预训练代码实践](#5. 预训练代码实践)

1. 预训练任务

常用的预训练任务主要分为三类：语言建模(Language Modeling,LM)，去噪自编码(Denoisiting Autoencoding,DAE),混合去噪器(Mixture-of-Denosisers,MoD).

语言建模：语言建模任务是目前绝大部分大语言模型广泛采用的预训练任务。给定一个词元序列𝒖 = {𝑢1, . . . , 𝑢𝑇 }，语言建模任务的目标定义为词元的预测任务：基于序列中当前位置之前的词元序列𝒖<𝑡，采用自回归的方式对于目标词元𝑢𝑖 进行预测。 L l m ( u ) = ∑ t = 1 T l o g P ( u t ∣ u < t L_{lm}(u)=\sum_{t=1}^T log P(u_t|u_{<t} Llm(u)=t=1∑TlogP(ut∣u<t

从本质上看，基于语言建模的预训练还可以看做是一种多任务学习过程。比如预测情感词语、数量词语、打分词语时其实针对是情感分析任务、数学算术任务、分类任务等。因此，基于大规模文本语料的预训练任务能够潜在地学习到解决众多任务的相关知识与能力。

语言建模有一些变种，最重要的是 前缀语言建模 ，这个专门为前缀解码架构设计的。在训练阶段，每个文本序列𝒖 会根据随机选择的位置𝑘 (1 ≤ 𝑘 ≤ 𝑇) 切分为前缀𝒖prefix = {𝑢1, · · · , 𝑢𝑘 } 和后缀𝒖suffix = {𝑢𝑘+1, · · · , 𝑢𝑇 } 两个部分。与标准语言建模任务不同，在前缀解码器中，仅后缀中的词元损失会被计入总损失。 L p r e f i x ( u ) = l o g P ( u s u f f i x ∣ u p r e f i x ) = ∑ t = k + 1 T l o g P ( u t ∣ u < t ) L_{prefix}(u) = log P(u_{suffix}| u_{prefix}) = \sum_{t=k+1}^T log P(u_t|u_{<t}) Lprefix(u)=logP(usuffix∣uprefix)=t=k+1∑TlogP(ut∣u<t)

前缀语言建模任务本质上是基于前缀信息来预测后缀的词元，这与自然语言处理任务中常见的基于输入来预测输出的模式十分相似。然而，在模型预训练阶段的损失函数中，由于并未将所有词元的损失都纳入计算，当使用相同规模的数据集进行训练时，采用前缀语言建模训练的模型在性能上通常会稍逊于使用标准语言建模任务训练的模型。

语言建模的另一个重要变种是中间填充任务,通过重新调整输入序列的顺序，旨在训练模型对于中间缺失信息的填充能力。具体来说，一个输入序列u 被划分为三个部分：前缀u_{prefix}、中间部分u_{middle} 和后缀u_{suffix}。随后，中间部分被移至序列末尾。因此，模型需要自回归地对新序列u_{prefix}⊕ u_{suffix}⊕ u_{middle}进行预测。通过这种方式，模型能够学习填充中间缺失信息的能力。这种任务的训练函数可表示如下： L F I M ( u ) = l o g P ( u p r e f i x ) + l o g P ( u s u f f i x ∣ u p r e f i x ) + l o g P ( u m i d d l e ∣ u p r e f i x , u s u f f i x ) L_{FIM} (u) = log P(u_{prefix}) + log P(u_{suffix} |u_{prefix}) + log P(u_{middle} |u_{prefix}, u_{suffix}) LFIM(u)=logP(uprefix)+logP(usuffix∣uprefix)+logP(umiddle∣uprefix,usuffix).通常来说，中间填充任务被用作标准语言建模方法的辅助任务。在保留预测下一个词能力的同时，这种方法使得模型具备对于文本中间部分内容的恢复能力。这种预训练任务经常被用于训练代码预训练模型，从而提升模型在代码补全等实际应用场景中的表现。
去噪自编码:在去噪自编码任务中，输入文本经过一系列随机替换或删除操作，形成损坏的文本𝒖\𝒖'。模型的目标是根据这些损坏的文本恢复出被替换或删除的词元片段𝒖'。去噪自编码器的训练目标可以用以下数学公式表示： L D A E ( u ) = l o g P ( u ′ ∣ u u ′ ) L_{DAE}(u)=log P(u'|u_{u'}) LDAE(u)=logP(u′∣uu′).与语言建模相比，去噪自编码任务的实现更为复杂，需要设定额外的优化策略，如词元替换策略、替换片段长度、替换词元比例等。这些策略的选择会直接影响模型的训练效果。尽管去噪自编码任务在许多预训练语言模型中得到了广泛应用。然而，相比于语言建模任务，目前完全使用去噪自编码进行预训练的大语言模型还较为有限。代表性的模型包括FLAN-T5。
混合去噪器: 又称UL2 损失,通过将语言建模和去噪自编码的目标均视为不同类型的去噪任务，对于预训练任务进行了统一建模。具体来说，混合去噪器定义了三种去噪器：S-去噪器、R-去噪器和X-去噪器。

S-去噪器与前缀语言建模的目标相同，旨在训练模型学习基于给定前缀信息生成合理的后缀文本的能力。相比之下，R-去噪器和X-去噪器与去噪自编码任务的优化目标更为相似。二者仅仅在被掩盖片段的跨度和损坏比例上有所区别。R-去噪器屏蔽序列中约15% 的词元，且每个被屏蔽的片段仅包含3 到5 个词元。而X-去噪器则采用更长的片段（12 个词元以上）或更高的损坏比例（约50%），进而要求模型能够精准还原原始信息。这种设置增加了任务难度，迫使模型学习到更全面的文本表示。为了引导模型针对不同类型的输入选择相应的去噪器，输入句子会以特殊词元（如[R], [S], [X]）作为开头。这种标记方式引导模型识别输入中使用的去噪器，并对该去噪器损坏的词元进行还原。混合去噪器被应用于训练UL2 和PaLM2等大语言模型。

2. 优化参数设置

与传统神经网络的优化类似，通常使用批次梯度下降算法来进行模型参数的调优。同时，通过调整学习率以及优化器中的梯度修正策略，可以进一步提升训练的稳定性。为了防止模型对数据产生过度拟合，训练中还需要引入一系列正则化方法。

2.1 基于批次数据的训练

在大模型预训练中，通常将批次大小（Batch Size）设置为较大的数值，例如1M 到4M 个词元，从而提高训练的稳定性和吞吐量。

现在很多工作都采用了动态批次调整策略，即在训练过程中逐渐增加批次大小，最终达到百万级别。例如，GPT-3 的批次大小从32K 个词元逐渐增加到3.2M个词元；PaLM-540B 的批次大小从1M 个词元逐渐增加到4M 个词元。相关研究表明，动态调整批次大小的策略可以有效地稳定大语言模型的训练过程。这是因为较小的批次对应反向传播的频率更高，训练早期可以使用少量的数据让模型的损失尽快下降；而较大的批次可以在后期让模型的损失下降地更加稳定，使模型更好地收敛。

2.2 学习率

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

2.3 优化器

大语言模型的训练通常采用Adam 及其变种AdamW 作为优化器。Adam 优化器使用梯度的"动量"作为参数的更新方向，它使用历史更新步骤中的梯度加权平均值来代替当前时刻的梯度，从而缓解样本随机性带来的损失震荡。进一步，Adam 使用自适应的学习率方法，通过梯度的加权"二阶矩"对梯度进行修正（可以看做使用"标准差"进行"归一化"），从而防止梯度过小导致模型难以优化。Adam 在优化中引入了三个超参数，在大模型训练中通常采用以下设置：𝛽1 = 0.9，𝛽2 = 0.95 和𝜖 = 10−8。此外，谷歌的研究者提出了Adafactor优化器，它是Adam 优化器的一个变种，通过引入了特殊设计可以在训练过程中节省显存，被用于PaLM 和T5 等大语言模型的训练。Adafactor 常见的超参数设置如下：𝛽1 = 0.9，𝛽2 = 1.0 − 𝑘−0.8，其中𝑘 表示训练步数。

2.4 稳定优化技术

梯度裁剪:训练中一种常见的现象是损失的突增。为了解决这一问题，可以采取梯度裁剪（Gradient Clipping）的方法，把梯度限制在一个较小的区间内：当梯度的模长超过给定的阈值后，便按照这个阈值进行截断。在大模型训练中，这个阈值通常设置为1.0
训练恢复:每隔固定的步数设置一些模型存档点。当模型发生了训练异常时（例如损失激增），便可以选择前一个存档点重启训练过程，并跳过可能导致问题的数据。
权重衰减:在模型的训练过程中也通常会引入正则化技术来稳定训练过程，提高模型的泛化能力。AdamW 中采用了权重衰减（Weight Decay）方法，在每次更新模型参数的时候引入衰减系数，这个系数通常设置为0.1。
Dropout:传统深度学习通常采用Dropout 技术来避免模型过拟合，即在训练中随机将一些神经元的输出值置零来避免过拟合。但是在大模型训练中，考虑到大规模的训练数据和模型中存在的归一化结构，已有工作很少使用Dropout技术。

3. 可扩展的训练技术

LLM面临的技术问题：一是如何提高训练效率；二是如何将庞大的模型有效地加载到不同的处理器中。

3.1 3D并行训练

3D 并行策略实际上是三种常用的并行训练技术的组合，即数据并行（DataParallelism）、流水线并行（Pipeline Parallelism）和张量并行（Tensor Parallelism）。有的工作也会使用模型并行一词，它同时包括了张量并行和流水线并行。

数据并行：数据并行是一种提高训练吞吐量的方法，它将模型参数和优化器状态复制到多个GPU 上，然后将训练数据平均分配到这些GPU 上。这样，每个GPU 只需要处理分配给它的数据，然后执行前向传播和反向传播以获取梯度。当所有GPU 都执行完毕后，该策略会将不同GPU 的梯度进行平均，以得到整体的梯度来统一更新所有GPU 上的模型参数。
流水线并行：将大语言模型不同层的参数分配到不同的GPU 上。在实践中，可以将Transformer 连续的层加载到同一GPU 上，以减少GPU 之间传输隐藏状态或梯度的成本。流水线并行通常需要配合梯度累积（Gradient Accumulation）技术进行优化。该技术的主要思想是，计算一个批次的梯度后不立刻更新模型参数，而是累积几个批次后再更新，这样便可以在不增加显存消耗的情况下模拟更大的批次。在流水线并行中使用了梯度累积后，1 号卡前向传播完第一个批次后，便可以不用等待，继续传播第二个和后续的批次，从而提高了流水线的效率。
张量并行:与流水线并行是两种将大模型参数加载到多个GPU 上的训练技术。流水线并行侧重于将模型的不同层分配到不同的GPU 上。相较之下，张量并行的分配粒度更细，它进一步分解了模型的参数张量（即参数矩阵），以便更高效地利用多个GPU 的并行计算能力。，对于大语言模型中的某个矩阵乘法操作𝑾𝑯，参数矩阵𝑾 可以按列分成两个子矩阵𝑾1 和𝑾2，进而原式可以表示为[𝑾1𝑯,𝑾2𝑯]。然后，可以将参数矩阵𝑾1 和𝑾2 放置在两张不同的GPU上，然后并行地执行两个矩阵乘法操作，最后通过跨GPU 通信将两个GPU 的输出组合成最终结果。常见的张量并行策略是分解模型注意力层的𝑾𝑄，𝑾𝐾，𝑾𝑉，𝑾𝑂 矩阵参数和前馈网络层的𝑾𝑈，𝑾𝐷 矩阵参数.

3.2 零冗余优化器

零冗余优化器（Zero Redundancy Optimizer, ZeRO）技术由DeepSpeed 代码库提出，主要用于解决数据并行中的模型冗余问题，即每张GPU 均需要复制一份模型参数。数据并行时每个GPU 都需要存储大语言模型的相同副本，包括模型参数和优化器参数等。对于每个GPU，在模型传播到某一层时，其他层的模型和优化器参数并不参与计算，这导致了严重的显存冗余现象，同时也限制了每个GPU 可以支持的前向传播数据量，降低了训练效率。为了解决这个问题，ZeRO 技术仅在每个GPU 上保留部分模型参数和优化器参数，当需要时再从其它GPU 中读取。，模型被均分在两张GPU 上，当需要使用第一层计算时，两张卡分别从对方获取相应的模型参数进行计算，使用完之后便可以释放相应显存，从而降低了显存冗余度。PyTorch 中也实现了与ZeRO 相似的技术，称为完全分片数据并行（Fully Sharded Data Parallel, FSDP）。

3.3 激活重计算

激活重计算（Activation Recomputation），也称为梯度检查点（Gradient Checkpointing），是一种用于优化反向传播时显存占用的技术。具体来说，给定一个待优化函数𝒀 = 𝑿𝑾，在反向传播时需要𝑿 的值才能计算𝑾 的导数，所以在前向传播时需要保留这些𝑿（通常被称为激活值）。然而，保存每一层所有的激活值需要占用大量的显存资源。因此，激活重计算技术在前向传播期间仅保留部分的激活值，然后在反向传播时重新计算这些激活值，以达到节约显存的目的，但是同时也会引入额外的计算开销。在大语言模型的训练过程中，激活重计算的常见方法是将Transformer 的每一层的输入保存下来，然后在反向传播时计算对应层内的激活值。

3.4 混合精度训练

早期的预训练语言模型（例如BERT）主要使用单精度浮点数（FP32）表示模型参数并进行优化计算。近年来，为了训练超大规模参数的语言模型，研发人员提出了混合精度训练（Mixed Precision Training）技术，通过同时使用半精度浮点数（2 个字节）和单精度浮点数（4 个字节）进行运算，以实现显存开销减半、训练效率翻倍的效果。具体来说，为了保证表示精度，需要保留原始32 位模型的参数副本。但在训练过程中，会先将这些32 位参数转换为16 位参数，随后以16 位精度执行前向传播和反向传播等操作，最后在参数更新时再对32 位模型进行优化。由于在模型训练中前向传播和反向传播占用了绝大部分优化时间，混合精度训练因而能够显著提升模型的训练效率。常见的半精度浮点数表示方式为FP16，其包含1 位符号位、5 位指数位和10 位尾数位，表示范围为−65504 到65504。进一步，谷歌的研究人员深度学习场景提出了新的半精度浮点数表示BF16，其包含1 位符号位、8 位指数位和7 位尾数位，表示范围可以达到1038 数量级。相比于FP16，BF16 有着更大的数值范围，在大模型训练中被广泛使用。值得一提的是，目前较为主流的GPU （例如英伟达A100）都支持16 位计算单元运算，因此混合精度训练能够被硬件很好地支持。

4. 模型参数量计算与效率分析

4.1 参数量估计

由于当前主流的大模型普遍采用因果解码器架构，因此下面以LLaMA 模型为范例，深入剖析其参数数量计算方式。对于其他模型，其参数量计算算法可参照此方法计算。首先，假设词表大小为𝑉，模型包含𝐿 层解码器，中间状态的维度大小为𝐻，前馈网络层的中间状态维度大小为𝐻′。我们主要关注计算以下几个部分的参数量：

输入嵌入维度：输入嵌入层（𝑬 ∈ R𝑉×𝐻）将词表中的每个单词映射到一个𝐻 维的向量，因此输入编码层有𝑉𝐻 个参数。
多头注意力层：传统的注意力机制部分包含查询（𝑾𝑄 ∈ R𝐻×𝐻）、键（𝑾𝐾 ∈R𝐻×𝐻）和值（𝑾𝑉 ∈ R𝐻×𝐻）的线性变换矩阵，每个变换矩阵都包含𝐻^2 个参数，所以这部分需要 3 x H 2 3xH^2 3xH2个参数。同时还需要一个额外的线性变换来将多头注意力机制的输出拼接后映射成最终输出（𝑾𝑂 ∈ R𝐻×𝐻），这又需要𝐻2 个参数。因此，多头注意力层总共需要 4 x H 2 4xH^2 4xH2 个参数。
前馈网络层:LLaMA 的前馈网络层由三个线性变换组成，中间有一个非线性激活函数。前两个线性变换（𝑾𝑈 ∈ R𝐻×𝐻′ 和𝑾𝐺 ∈ R𝐻×𝐻′）将输入从𝐻 维映射到𝐻′ 维，需要2 × 𝐻𝐻′ 个参数；最后一个线性变换（𝑾𝐷 ∈ R𝐻′×𝐻）将输出从𝐻′维映射回𝐻 维，需要𝐻𝐻′ 个参数。因此，前馈网络层总共需要3 × 𝐻𝐻′ 个参数。
归一化层:每一层解码器还包含两个RMSNorm 操作，分别用于对多头注意力层和前馈网络层的输入进行归一化处理，共需要2 × 𝐻 个参数。此外，最后一层的输出也需要进行归一化处理，这又需要额外的𝐻 个参数.
输出层:最后，LLaMA 的输出层包含一个线性变换（𝑾𝐿 ∈ R𝐻×𝑉），将解码器的输出映射到词表大小𝑉 的维度上，使用softmax 归一化后预测下一个单词的概率分布。这个线性变换需要𝑉𝐻 个参数。

综上所述，累积输入嵌入层、输出层和𝐿 层解码器每层的多头注意力层、前馈网络层和归一化层，LLaMA 模型的参数量计算公式为：
参数量 = 2 V H + H + L ∗ ( 4 H 2 + 3 H H ′ + 2 H ) 参数量= 2VH+H+L*(4H^2+3HH'+2H) 参数量=2VH+H+L∗(4H2+3HH′+2H)

以LLaMA (7B) 为例计算其参数量，给定𝑉 = 32000, 𝐿 = 32, 𝐻 = 4096, 𝐻′ = 11008，将这些值代入上述公式中:参数量= 2 × 32000 × 4096 + 4096 + 32 × (4 × 40962 + 3 × 4096 × 11008 + 2 × 4096)= 6, 738, 415, 616.计算得到的参数量与LLaMA (7B) 模型的实际参数量完全一致。

4.2 训练运算量估计

模型训练运算量指的是模型在训练过程中，需要进行的浮点运算次数（FloatingPoint Operations, FLOP） 。这里的浮点运算包括浮点数的加减乘除运算，以及浮点数的指数函数，对数函数，三角函数等运算操作。使用Transformer 架构进行训练的运算量主要集中在多头注意力计算和线性变换计算 。相比之下，归一化、输出映射和旋转位置编码计算所需的运算量较少，而输入编码层则无需计算，因此后续的分析中省略了这些部分。在分析多头注意力和线性变换的运算量时，我们进一步设定以下参数：模型总参数量为𝑃，批处理大小为𝐵，输入序列长度为𝑇，因此训练词元总数为𝐶 = 𝐵𝑇；多头注意力机制包含𝑁 个头，每个头的维度为𝐷，因此和中间状态维度𝐻 满足关系𝐻 = 𝑁𝐷

注意：矩阵𝑨 ∈ R𝑛×𝑚 和矩阵𝑩 ∈ R𝑚×𝑝 相乘所需的运算量为2𝑛𝑚𝑝。

多头注意力：首先分析多头注意力机制一次计算的运算量。对于批次化的数据，计算得到相应的查询、键和值张量，𝑸, 𝑲,𝑽 ∈ R𝐵×𝑇×𝐻，考虑到需要进行多头计算，这些张量需要进行拆分和转置 ，得到𝑸′, 𝑲′,𝑽′ ∈ R𝐵×𝑁×𝑇×𝐷。在注意力计算中 𝑸′𝑲′⊺ 的矩阵乘法需要 2 B T 2 N D 2BT^2 ND 2BT2ND 次浮点运算；接着，进行标准化操作（ D \sqrt{D} D 放缩）需要 2 B T 2 N 2BT^2 N 2BT2N次浮点运算，softmax 操作需要进行指数、加和、归一化操作，总计3𝐵𝑇^2𝑁 次浮点运算；最后结果与𝑽′ 进行矩阵乘法，再需要2𝐵𝑇2𝑁𝐷次浮点运算。因此，一次多头注意力计算总的浮点运算量为 4 B T 2 N D + 4 B T 2 N 4BT^2 ND + 4BT^2N 4BT2ND+4BT2N 考虑到后向传播的运算量大致为前向传播的两倍 ，整个模型中多头注意力运算量可表达为：运算量 = 12 ∗ ( B T 2 N D + B T 2 N ) ∗ L = 12 C T L ∗ ( H + N ) 运算量= 12*(BT^2ND+BT^2N)*L = 12CTL*(H+N) 运算量=12∗(BT2ND+BT2N)∗L=12CTL∗(H+N).
线性变换：以前馈网络层中的上映射操作𝑿′𝑾𝑈 为例，中间状态𝑿′ ∈ R𝐵×𝑇×𝐻，上映射矩阵𝑾𝑈 ∈ R𝐻×𝐻′，因此其前向传播需要2𝐵𝑇𝐻𝐻′ 次浮点运算，反向传播则需要4𝐵𝑇𝐻𝐻′ 次浮点运算，总计需要6𝐶𝐻𝐻′ 次浮点运算，其中𝐶 = 𝐵𝑇 为训练的词元总数。可以看到，线性变换的运算量与矩阵参数量相关，因此Transformer中所有线性变换部分的运算量公式可表达为：运算量= 6𝐶 · 线性变换的参数量.
若训练过程中采用了激活重计算技术，反向传播时需要额外进行一次前向传播，则总运算量将变换：运算量= 8𝐶 · 线性变换的参数量.
多头注意力的运算量约为线性变换运算量的𝑇/6𝐻 ，考虑到大模型训练场景下序列长度𝑇 小于等于中间状态维度𝐻，因此多头注意力运算量最多为线性变换运算量的1/6，其影响相对较小.
参数量为𝑃 的模型在𝐶 个词元上进行预训练的总体运算量可以按照下式进行估计:运算量≈ 6𝐶𝑃.如果使用了激活重计算技术，则运算总量约为8𝐶𝑃.