Introduction to popular LLM components

连理o

已于 2024-07-25 19:06:28 修改

阅读量591

点赞数 22

文章标签： LLM

于 2024-06-05 16:51:43 首次发布

本文链接：https://blog.csdn.net/weixin_42437114/article/details/139436443

版权

模型部署专栏收录该内容

24 篇文章 1 订阅

订阅专栏

Popular LLM Components
LLM Architecture
References

Popular LLM Components

Normalization

RMSNorm

LayerNorm v.s. RMSNorm. LayerNorm 的前向计算过程为
$y=\frac{x-\mu}{\sqrt{\sigma^2+\varepsilon}}\odot w + b$ RMSNorm 发现 LayerNorm 的中心偏移没什么用 (减去均值等操作)，将其去掉之后效果几乎不变，但是速度提升了 40%. RMSNorm 就相当于是设定均值为 0 的 LayerNorm. 最终公式为：
$y=\frac{x}{\sqrt{\sigma^2+\varepsilon}}\odot w$

class LlamaRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * hidden_states.to(input_dtype)

Why LN instead of BN?

(1) 文本序列是变长的，长度不同的序列原则上属于不同的统计对象，所以很难得到稳定的统计量，而得不到稳定的统计量，BN 就无法成立了；(3) 还可以从归一化维度上进行解释，不同文本序列相同位置上的 token 含义可能是非常不同的，因此在样本维度上做归一化可能意义不大；(2) 但即使是变长序列，我们仍然可以强行使用 BN，甚至 ViT 的输入不是变长序列也还是使用 LN，似乎 LN 就是和 Transformer 更配，对此可以结合自注意力计算进行一些解释， $q\cdot k=\|q\|\|k\|\cos\theta$ ，LN 对于 $\|q\|\|k\|$ 的控制要比 BN 更有效
不过归根到底，最大原因还是 LN 实验效果最好~

Activation Function

SwiGLU / SiLU (Sigmoid Linear Unit)

SiLU.
$\text{silu}(x)=x\odot\text{sigmoid}(x)$

MHA

bsz, q_len, _ = hidden_states.size()

query_states = self.q_proj(hidden_states)
key_states = self.k_proj(hidden_states)
value_states = self.v_proj(hidden_states)

query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)

cos, sin = self.rotary_emb(value_states, position_ids)
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
key_states = repeat_kv(key_states, self.num_key_value_groups)
value_states = repeat_kv(value_states, self.num_key_value_groups)

attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
attn_weights = attn_weights + causal_mask

attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
attn_output = torch.matmul(attn_weights, value_states)
attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.reshape(bsz, q_len, -1)

attn_output = self.o_proj(attn_output)

LLM Architecture

Why Decoder-only？

目前 LLM 的主要架构有：以 BERT 为代表的 encoder-only、以 T5 和 BART 为代表的 encoder-decoder、以 GPT 为代表的 decoder-only，还有以 UNILM 为代表的 PrefixLM (PrefixLM 就是对 decoder-only 模型的 attention mask 做了调整，类似于 encoder-decoder，前缀部分的 attention 为双向注意力)
encoder-only 其实并不太适合做 NLP 的生成式任务，目前在 CV 领域应用的比较多 (ViT)，输入/输出通常是固定像素大小的图片，token 之间没有非常强的先后依赖关系，因此先排除 encoder-only 架构；另外 PrefixLM 在预训练后泛化性并没有明显强于 decoder-only；因此下面主要比较 decoder-only 和 encoder-decoder 架构

Zero/Few-shot Generalization

paper: Wang, Thomas, et al. “What language model architecture and pretraining objective works best for zero-shot generalization?.” International Conference on Machine Learning. PMLR, 2022.
作者发现使用 autoregressive language modeling objective 训练的 causal decoder-only 模型在自监督预训练后展现出了最强的 zero-shot 能力，而使用 masked language modeling objective 预训练的 encoder-decoder 模型则需要在一定量的标注数据上做 multitask finetuning 才能超过 causal decoder-only 模型；这可能是因为 decoder-only 的预训练任务难度更大，因此模型上限更高
值得注意的是，上述实验是在 5B 参数量 + 170B token 数据量的条件下进行的，而当前的 LLM 都依赖于涌现能力，当参数量足够大时，模型能力会得到质的飞跃，因此大数据训练 + 大参数模型的涌现能力替代了 multitask finetuning

paper: Dai, Damai, et al. “Why can gpt learn in-context? language models implicitly perform gradient descent as meta-optimizers.” arXiv preprint arXiv:2212.10559 (2022).
此外也有研究表明，decoder-only 架构在 few-shot 任务上泛化性也更强

Scalability (Pipeline Parallelism)

流水并行是千卡以上分布式训练中最重要的特性. NVIDIA 在 3076 张 A100 集群上训练的 1T 参数量 LLM 使用的并行方式是：Data Parallel Size = 6, Tensor Parallel Size = 8, Pipeline Parallel Size = 64. 可见并行度最高的是流水并行，超过 DP 和 TP 十倍左右。为什么在三千卡集群上最主要的并行方式是流水并行？流水并行的核心优势就是用比较少的 Pipeline Bubble 代价 (当 gradient accumulation step 很大时可以忽略不计)，较少的 Tensor Buffer 显存代价，以及非常低的通信开销，将大模型分割在不同的 Group 中，大幅减少了单张 GPU 上的 weight tensor 大小 (数量) 和 Activation tensor 大小 (数量). 同时，跟 Tensor Parallel 相比， Pipeline Parallel 的通信代价很低且可以被 overlap， Tensor Parallel 虽然也能切分模型大小，但是需要全量的数据 (没有减少 Activation tensor 大小)，另外极高的通信频率和通信量使得 Tensor Parallel 只能在机器内 8 张卡用 NVLink 等高速互联来实现，跨机的 TP 会严重拖慢速度。不仅如此， Pipeline Parallel 还将 Data Parallel 的模型更新限定在一个很小的范围内 (比如六台机器)， DP 所需的 AllReduce 通信会随着机器数量增多而变慢。 PP 也让 DP 所需同步的模型梯度变少了，大大减缓了模型更新对于训练速度的影响。因此 Pipeline Parallel 是让模型可以达到千亿、集群可以扩充到千卡以上的一个最重要的特性。然而流水并行有很重要的约束条件：需要一个规整对称的、线性顺序的网络结构

流水并行训练时的 time line 参考如下 (反向的计算时间是前向的两倍)：整个集群最高效的训练时间段是 step 4、5、6、7 的前向和 step 0、1、2、3 的反向同时在所有 stage 上并行计算的时候，这个时候集群没有空闲，全部都在并行执行。当我们增加 acc step （比如从 8 增加到 64）时，中间部分完美并行的时间段占比就会更长， bubble time 的占比就会越来越小

Decoder-Only 架构就是这样一个典型的网络结构：完全一样的 Transformer Layer 顺序堆叠，没有分叉和不对称情况，当均匀切分 Layer 时，各个 Stage 的前向/反向计算时间均一致，非常适合做流水并行，因此 Decoder-Only 架构最核心的优势是非常方便于 Scale Up，基于 Scaling Laws 的实际训练成本最低。而反观 Encoder-Decoder 架构，整个网络分为两大块，且 Encoder 和 Decoder 的 Transformer Layer 参数大小、Attention 计算量、Context Length 等均不一致，导致 Encoder 的理论计算量要比 Decoder 大很多 (整个网络不是均匀对称的)。更要命的是， Encoder 的输出要发给每个 Decoder Layer，网络结构不是线性而是有大量的分叉，前向反向之间包含了复杂的数据依赖关系，会导致流水并行中，各个 Stage 之间会产生大量的、非对称的、间隔跨多个 Stage 的数据依赖，更加剧了流水并行的 load balance 问题。如果我们不用 Pipeline Parallelism 来训练 Encoder-Decoder 架构，那么只能借助：TP 和 ZeRO 来进行并行优化了，这就约束了 Encoder-Decoder 架构的所有 Layer 都必须放在每一个 GPU 上，这种方式在 13B 量级的模型上是 OK 的，但是再往上扩展到 100B、1T 量级就不 work 了。同时由于 TP 只能开到 8 (跨机器也会慢几倍)，在千卡 GPU 集群以上，大量的 DP 带来的通信变慢的影响也很严重 (ZeRO-2/3 会大幅加剧这种通信开销)。所以我们才说，虽然 Encoder-Decoder 架构的理论计算量相较于 Decoder-Only 架构没有增加很多，但是在千亿参数、千卡集群以上规模的时候，Encoder-Decoder 架构的实际训练效率比 Decoder-Only 架构慢很多倍。即使到现在，也没有一个超过 11B 的 T5 模型 (Encoder-Decoder 架构) 发布，而 11B 恰好是一个不借助 PP，仅通过 ZeRO + TP 就可以训练的模型大小，避免了 T5 的模型结构非对称性对于 PP 的灾难性影响
总结来说，有可能从纯算法模型结构上，Encoder-Decoder 架构是比 Decoder-Only 架构更优的结构，但是由于 Encoder-Decoder 架构的模型结构不是线性的（Decoder 和 Encoder 之间有复杂的连接关系），导致 Encoder-Decoder 架构的流水并行实际上无法真正高效的执行起来。因此，在目前已知的并行优化上， Encoder-Decoder 架构很难 Scale 到千亿参数以上的规模。针对 Decoder-Only 架构的分布式并行优化，会比 Encoder-Decoder 架构容易很多。在 LLM 时代，如果你提出的新的算法结构可能有 5% 的效果提升，但是引入了额外 50% 的训练成本（计算时间 or 通信量）的话，那这个新的算法一定是一个负优化。因为这 50% 的训练成本，基于 Scaling Laws 我可以在原模型上多训练 50% 的 tokens ，或者训练大一半的模型，带来的最终提升都远大于新算法的 5%。新的算法研究必然在探索阶段就需要引入 Infra 因素的考量。因此，我的观点是，也许不同的算法结构互有高低，但是在 LLM 时代，更容易分布式并行、更容易扩展、训练 Token 效率更高的模型，一定是更具备优势的，这就是 Infra 反过来影响算法

Why Pre-LN?

在这里插入图片描述

原始的 Transformer 架构采用 Post-LN 架构，模型在优化时往往需要进行学习率预热 (learning rate warm-up)，并且优化过程对 warm-up 的超参数 (持续轮数、增长方式、初始学习率等) 非常敏感，若调整不慎，往往会使得模型无法正常收敛，同时模型的收敛速度也很慢
对此，“On layer normalization in the transformer architecture.” 提出采用 Pre-LN 架构。由于 warm-up 发生在模型的训练初期，因此作者主要从模型的初始化阶段开始研究。作者发现，当采用 Xavier 方法对 Post-LN 进行初始化后，通过对各隐层梯度值进行分析可以证明，在初始化点附近的 Post-LN 结构最后一层的梯度值非常大，同时随着反向传播的前传会导致梯度值迅速衰减。这种在各层之间不稳定的梯度分布必然会影响优化器的收敛效果，导致训练过程初始阶段的不稳定。造成 Post-LN Transformer 梯度分布出现问题的核心原因在于各子层之后的 LN 会使得各层的输入尺度与层数 $L$ 无关，因此当 Layer Normalization 对梯度进行归一化时，也与层数 $L$ 无关。使用相同的方法对 Pre-LN 结构进行分析后，发现最后一层 LN 的输入尺寸的量级只有 Post-LN 的 $\sqrt{\frac{1}{L}}$ 倍，并且每个 LN 层都会对梯度以 $\sqrt L$ 的比例归一化。所以对于 Pre-LN 结构来说，其每层梯度范数都近似不变，这种结构更利于优化器进行优化 (直观理解就是 pre-LN 结构从输出到输入有一条残差连接组成的通路，梯度可以不受阻碍地传递下去，使得各层之间的梯度分布比较近似；而 post-LN 在梯度反传时则必须要经过 LN，使得梯度值迅速衰减，各层之间梯度分布不稳定，也就增加了优化难度)
当使用 Pre-LN 结构时，warm-up 已不再是必需，并且 Pre-LN 结构可以大幅提升 Transformer 的收敛速度