Zhangzhe's Blog

URL

• paper: https://arxiv.org/pdf/1910.02054
• code: https://github.com/microsoft/DeepSpeed

TL;DR

• 传统的 DDP 训练大模型有个很大的问题是：由于 DDP 训练中每个 GPU 上都需要存储模型状态信息和其他状态信息：
  • 模型状态信息包括：Prameter + Gradient + Optimizer states
  • 其他状态信息包括：激活值（Activation）、各种临时缓冲区（buffer）以及无法使用的显存碎片（fragmentation）
    所以显存占用较大，使得不能通过 DDP 训练较大的模型。
• 本文提出来 Zero-DP 和 Zero-R 两种大模型训练优化策略，分别解决模型状态信息过大和其他状态信息过大的问题。
• Zero-DP 全称是 Zero Redundancy Optimizer Data Parallelism，主要优化点是将 GPU 间冗余存储的模型状态信息（Parmeter + Gradient + Optimizer states）删除，每个 GPU 只保留一小部分信息，所有 GPU 上的信息汇聚之后才是完整模型。
• Zero-R 的目的是解决 Tensor Parallelism 训练时，Activation 信息占用过大的问题。

Algorithm

ZeRO-DP

ZeRO-DP 想要解决什么问题？

• 解决 DDP 训练过程中，Parmeter + Gradient + Optimizer states 对显存的占用过大，导致无法训练很大的模型的问题。

没有其他方法可以降低显存占用吗？

• 有的，比如混合精度训练 (https://arxiv.org/pdf/1710.03740) 就是一个常用的有效的手段
  
• 已知 LLM 通常使用 Adam / AdamW 优化器进行优化，假设模型参数量为 $\Phi$，那么用混合精度训练过程中，模型状态信息实际的显存占用是:
  • Parameter 占用 $2\Phi$ 字节（fp16 类型存储）
  • Gradient 占用 $2\Phi$ 字节（fp16 类型存储）
  • Adam Optimizer state 占用：
    • fp32 parameter 备份占用 $4\Phi$ 字节
    • fp32 momentum 一阶动量占用 $4\Phi$ 字节
    • fp32 variance 二阶动量占用 $4\Phi$ 字节
  • 共计 $16\Phi$ 字节，其中 Optimizer states 占用 75%，因此 Optimizer states 是本文主要想解决的问题

ZeRO-DP 是如何解决模型状态信息占用过大问题的？

• ZeRO-DP 的核心思想是 Data Parallelism 时，每个模型只保留 $\frac{1}{N}$ 的模型状态信息，N 表示 GPU 数量，在需要信息聚合时使用 Ring Reduce 来聚合。
• 有三种程度的优化：
  • $P_{os}$ 只删除冗余的优化器状态，所有优化器的信息加起来才是完整优化器状态。
  • $P_{os+g}$ 删除冗余的优化器状态和梯度。
  • $P_{os+g+p}$ 删除冗余优化器状态、梯度和权重。

分块存储 + Ring Reduce 会带来额外的通信代价吗？

• All Reduce 运算来计算设备间梯度均值（求和）
  
• Reduce Scatter 来求和，All Gather 做 broadcast 同步到所有设备上
  
• 使用 Ring Reduce 实现 Reduce Scatter
  

颜色逐渐累积代表梯度累加

• 使用 Ring Reduce 实现 All Gather 做 broadcast
  

深色逐渐蔓延到所有设备，代表 broadcast

• $P_{os}$ 和 $P_{os+g}$ 不会带来额外的通信代价，$P_{os+g+p}$ 会
• 证明：
  • 传统 Data Parallelism 在计算梯度后，需要进行一次 All Reduce 来计算设备间的梯度均值，All Reduce 可分为 Reduce Scatter 和 All Gather 两步，每一步实际上都是一次 Ring Reduce
    • Reduce Scatter 如图三所示，把所有设备梯度累加，分散到各个设备上
    • All Gather 如图四所示，把分散到各个设备上的梯度累加值同步到所有设备上
    • Ring Reduce 是一种理论最优通信算法，每个周期内每个设备的发送和接收都被占用，因此梯度总字节数为 $2\Phi$ (fp16)，由于做了两次 Ring Reduce，所以每个设备的通信总字节数为 $4\Phi$（发送和接收算一次）
  • ZeRO-DP $P_{os+g}$ 也需要通过 All Reduce 计算设备间梯度均值，也需要 Reduce Scatter 和 All Gather 两个步骤，也是通过两次 Ring Reduce 实现，具体做法和传统 Data Parallelism 没有什么区别，因此每个设备的通信总字节数也是 $4\Phi$
  • ZeRO-DP $P_{os+g+p}$ 在前向传播时，也需要用 All Reduce 计算得到全量 Parameter，这一步在传统 Data Parallelism 中并不需要，因此会有额外通信量

ZeRO-R

ZeRO-R 想要解决哪些问题？

1. Tensor Parallelism 中 Activation 冗余存储问题
2. 临时缓冲区占用问题
3. 显存碎片导致无法使用的问题

ZeRO-R 如何解决 Tensor Parallelism 中 Activation 冗余存储问题？

• 是通过激活值检查点分片存储（Partitioned Activation Checkpointing，简称 $P_A$）来解决 Tensor Parallelism 中 Activation 冗余存储问题。
• 其中的 分片 思想和 ZeRO-DP 是一样的，N 个设备每个设备只存储 $\frac{1}{N}$ 的信息，通过 All Reduce 进行信息聚合
• 激活值检查点 思想是来自 https://arxiv.org/pdf/1604.06174 这篇论文，实际上是传统前向反向传播过程和 ReCompute 前向反向传播过程的折衷，也是空间和时间的折衷。
• 传统前向反向传播过程：存储每个中间算子的前向传播计算结果，用于反向传播计算梯度
  
• ReCompute 前向反向传播过程：不存储任何中间算子的前向计算结果，反向传播需要用时重新计算
  
• Activation Checkpointing 前向反向传播过程：存储部分中间算子前向计算结果，反向传播需要用时，从拓扑上游最近的检查点开始重新计算
  

ZeRO-R 如何解决临时缓存区空间占用问题？

• 通常情况下，临时缓存区是用来做数据聚合的，例如 All Reduce 操作，模型越大，需要的临时缓存区也越大。
• ZeRO-R 用了一个非常简单的策略：在模型大到一定程度后，使用固定尺寸的融合缓存区，即常数尺寸缓存区（Constant Size Buffers，简称 $C_B$）

ZeRO-R 如何解决显存碎片化问题？

• 问：为什么会产生显存碎片化问题？
• 答：因为 Activation Checkpointing 的存在，不同的 Activation 有不同的生命周期，非 checkpoint 的 Activation 用完即弃，checkpoint 的 Activation 需要反向传播之后才可以丢弃，不同生命周期的存储空间交织就会出现显存碎片
• 问：显存碎片化会带来什么问题？
• 答：显存碎片化通常会有两个问题：
  1. 大量的碎片无法使用，使得实际空间占用不多，但出现 OOM
  2. 大量存储碎片会导致系统在分配内存时需要大量的搜索来找到合适的空间，会带来额外耗时
• 问：那 ZeRO-R 如何解决显存碎片化问题？
• 答：预先给不同生命周期的 Activation 划分不同的存储空间，避免产生交织。例如：把显存划分成两段连续的存储空间，其中一段存储 Checkpoint Activation 另外一段存储 Temporary Activation。

Thought

• 当大模型时代到来，算力和存储不够用了，大家也开始审视之前做的东西是不是时间/空间低效的，是不是还能榨点油水。

URL

• paper: https://arxiv.org/pdf/1909.08053
• code: https://github.com/NVIDIA/Megatron-LM

TL;DR

• 本文是英伟达 Megatron-LM 大模型分布式并行训练框架的一篇简介，大概阐述了 Megatron-LM 的机制和效果，对于实现细节着墨不多。
• 从代码实现角度看，Megatron-LM 是对 PyTorch 进行了二次封装。
• 基本思想是 Tensor Parallelism，拆分一层的权重到不同 GPU 上并行运算。

Algorithm

Tensor Parallelism

• 重点是这里的 f 和 g 函数，下面是其实现：

1
2
3
4
5
6
7
8
9
10
11
12
13

class f(torch.autograd.Function):
  def forward(ctx, x):
    return x
  def backward(ctx, gradient):
    all_reduce(gradient)
    return gradient
  
class g(torch.autograd.Function):
  def forward(ctx, x):
    all_reduce(x)
    return x
  def backward(ctx, gradient):
    return gradient

• 在本例子中 All-reduce 实际上是 ReduceSum，实际上聚合运算（包括 forward and backward）需要 GPU 间同步，存在通信代价
• Self-Attention 可以 tensor parallelism 的一个重要原因是 softmax 是逐行做的而不是全局做的
  
• 由图 3 可以得出结论，一个 transformer layer 只有 4 个同步算子（两个 LayerNorm 和两个 Dropout），其他算子都可以被分解为并行算子

Thought

• 给我的感觉有点类似 MLC 或者机器学习编译中的 Tiling（只是升级为训练及推理全流程 tiling）
• tensor parallelism 这种策略很考验对计算流程的抽象，除非有大量人力把框架搭建好且把所有算子实现写好，否则会很难用起来（当然对于 Nvidia 这些都不是事），突然想起了 Neuwizard，一声叹息…

URL

• paper: https://arxiv.org/pdf/1806.03377
• code: https://github.com/msr-fiddle/pipedream（民间实现）

TL;DR

• 本文提出一种和 Gpipe 类似的 pipeline parallelism 并行训练机制 PipeDream，和 Gpipe 可以任意划分模型不同，PipeDream 自动化划分网络，以实现负载均衡和最小通信开销。
• 与 Gpipe 的同步更新机制不同，PipeDream 采用异步参数更新机制。

Algorithm

backward 的过程实际同时包含了 参数更新

• 与 Gpipe 使用的 Micro-Batch + 同步参数更新机制不同，PipeDream 使用异步更新机制，异步更新主要包含以下几点：

weigt stashing

1. 从上图可以看出 Machine 1 在 forward 5 时，使用的是更新过 1 次的参数（forward 5 之前只做了 backward 1）；但当 backward 5 时，参数已经被更新了 4 次，如果直接使用会出现 forward 和 backward 参数不一致问题
2. 因此，PipeDream 提出了 weight stashing (权重存储)，即保存多组 weights，每次 forward 过程使用最新版本的 weights，当 backward 时，使用 forward 时刻对应的 weights
3. 使用 weight stashing 之前，$w^{(t+1)} = w^{(t)}+\gamma \cdot \triangle f(w_1^{(t)},w_2^{(t)},...,w_n^{(t)})$，其中：
   1. $\gamma$ 表示 learning rate
   2. $n$ 表示 stage 的数量，即 GPU 数量
   3. $w_i^{(t)}$ 表示 weights of stage i after t mini-batch，即当运行了 t 个 mini-batch 之后的 stage i 的权重
   4. $\triangle f$ 表示梯度
4. 使用 weight stashing 之后，$w^{(t+1)} = w^{(t)}+\gamma \cdot \triangle f(w_1^{(t-n+1)},w_2^{(t-n+2)},...,w_n^{(t)})$

vertical sync

1. weight stashing 解决了 forward 和 backward 使用的 weight 版本不一致问题，但依然存在另外一个问题：同一个 mini-batch 在不同 stage 上 forward 使用的 weight 版本不一致
2. 例如，对于 mini-batch 4，在 Machine 2 上 forward 使用的参数版本是未更新版本，在 Machine 3 上 forward 使用的版本是更新过一次的版本
3. 为解决这个问题，作者提出 vertical sync 概念，即：对于一个 mini-batch 在每个 stage 上 forward 都使用最旧版本的参数，stage forward and backward 结束后统一更新参数
4. vertical sync 应用之后，$w^{(t+1)}=w^{(t)}+\gamma \triangle f(w_1^{(t-n+1)},w_2^{(t-n+1)},...,w_n^{(t-n+1)})$
5. 事实上实验证明，weight stashing 对训练精度提升明显，vertical sync 提升较小

Thought

• 感觉 PipeDream 为了充分排流水，减小 bubble，没有等 forward 全部结束就可以 backward，然后发现和单 GPU 更新机制不等价，然后用了 weight stashing 和 vertical sync 去找补
• 和 Gpipe 相比，灵活性更差，实现难度更高（需要缓存各个版本的 weights，听上去也是几倍的存储开销以及吞吐量）

URL

• paper: https://arxiv.org/pdf/1811.06965
• code: https://github.com/tensorflow/lingvo/blob/master/lingvo/core/gpipe.py

TL;DR

• 随着模型和数据量的增长，单 GPU 已经无法满足需求，所以需要多 GPU 并行
• Data parallelism 是最常用的，即将数据划分到多个 GPU 上，每个 GPU 单独前向反向传播，仅在梯度更新时聚合，例如 pytorch ddp
• 本论文提出一种新的并行方式 Pipeline parallelism，这种方式是将模型划分为多段子图，每个设备上加载一段，各个子图直接串联
• 直接实现的 Pipeline parallelism 的 GPU 利用率较低，本文通过一些方法可大幅提高 GPU 利用率

Algorithm

Gpipe 方案主要包含三个部分

划分 Stage

1. 模型切成若干 stage，每个 GPU 只加载一段
2. 每个 stage 串起来形成一个 pipeline

划分 Micro-Batch

1. 传统模型训练过程中使用 Mini-Batch，Mini-Batch 在 pipeline parallelism 中会出现大量气泡，因此 Gpipe 提出 Micro-Batch 概念
2. Micro-Batch 是将 Mini-Batch 再划分成多份，用于排计算流水，Mini-Batch 变成最小计算单元，有利于减小气泡，上图 2.c 中横向就是一个 Mini-Batch 划分得到的一组 Micro-Batch
3. 为了保证梯度更新和 Mini-Batch 的一致性，Gpipe 会将 Micro-Batch 梯度累积，在一个 Mini-Batch 的全部 Micro-Batch 计算结束后再更新（图 2.c 中一行只有一次 update）

重计算

1. 传统模型训练过程中，计算的中间结果都需要保存，用于反向传播计算梯度，非常消耗显存，同时会导致 pipeline parallelism 的数据依赖问题，上图 2.a 中的横向连接
2. 重计算是一种用计算换空间的操作，在前向传播过程中，丢弃 stage 内所有中间结果，只保留 stage 间中间结果；在反向传播时，stage 内中间结果会再做一次前向传播计算得到
3. 重计算可大幅降低显存占用，使得可以放更大的 batch size 和更大的模型，且可以改善数据依赖问题

Thought

• 这种方案前瞻性比较强，提出时还没有 LLM，且拥有和传统非并行训练在数学上一致的参数更新机制
• 给并行训练提供了一种新思路，非常灵活

URL

• paper: https://arxiv.org/pdf/2305.09781
• code: https://github.com/flexflow/FlexFlow/

TL;DR

• 本文利用了 Speculative decoding 思想实现了一个 SpecInfer 大语言模型推理加速系统
• 和标准 Speculative decoding 使用 seq 组织预测的 token 不同，SpecInfer 使用 tree 数据结构组织预测的 token
• SpecInfer 中的 target 模型可以并行验证整个 token tree 的正确性，而不是 token seq，且可以保证和原模型输出完全一致

Algorithm

增量式解码、序列式投机解码、树式投机解码三者对比

左图表示传统增量式推理 LLM 的解码方式
右图从上到下依次表示：增量式解码、序列式投机解码、树式投机解码

候选树生成和验证流程

SSM (small speculative model) 给出候选树
target model 对候选树进行验证

如何对树式结构并行解码

注意：这的 “并行” 是指可以一次解码一棵树而不需要将树深度优先搜索得到多组序列然后一一解码

1. 使用 kv cache 对解码过程进行加速是一个常见的方法，但树式结构无法直接使用 kv cache（因为要动态剪枝，如图左侧所示）
2. 因此本文提出一种新颖的方法使得树式结构可以并行解码，这个方法包括两个改进：
   1. 对树进行 深度优先搜索，使得树变成一个序列
   2. 将传统的上三角 attention mask 变成一个 topology-aware mask，即通过 mask 来确定 token 两两之间的可见性（如图右侧所示），这种可见性包含：
      1. 每个 token 看不到之后的 token（矩阵下三角置 $-\infty$）
      2. 每个 token 看不到不在其拓扑路径上的 token

总体算法流程

1. 先用 small speculate model 推理 sequence 得到候选树 $N$
2. 在使用 LLM 树式并行解码 sequence 得到大模型预测结果 $O$
3. 用贪心算法或随机匹配算法对候选树 $N$ 逐个节点验证（接受或拒绝）
4. 重复直到 <EOS> 出现

贪心验证算法和随机验证算法

• 这里的两种算法和 https://zhangzhe.space/2024/08/14/Fast-Inference-from-Transformers-via-Speculative-Decoding/ 里的流程基本一致，不一样的地方仅仅在于 tree / sequence 两种候选 token 的组织形式

贪心验证

1. $u$ 表示当前节点，从根节点开始广度优先遍历
2. $p_v=u$ 表示 $v$ 是 $u$ 的子节点
3. $t_v=O(u)$ 表示大模型和投机模型预测结果一致

随机验证

1. $u$ 表示当前节点，从根节点开始广度优先遍历
2. 当在 [0, 1] 之间随机均匀采样得到的随机数 $r \lt \frac{P_{LLM}}{P_{SSM}}$ 时，接受这个节点，否则拒绝

Thought

• 这篇论文站在 speculate decoding 的肩膀上，加入了树结构的候选列表，算是 speculate decoding 的超集
• topology-aware mask 和 decoder attention mask 结合的方法还是挺巧妙的
• 由于是树结构存储候选词，那么多个 SSM 是天然适配的，算是个工程角度很好用的 trick

URL

• paper: https://arxiv.org/pdf/2211.17192
• code: https://github.com/feifeibear/LLMSpeculativeSampling （民间实现）

TL;DR

• 本文提出一种新颖的大模型推理加速的新范式，叫做投机解码，用一个参数量很小的近似模型预测得到一段序列，让大模型评判 接受 还是 拒绝 这段序列。
• 从数学上可以证明投机解码的结果和大模型逐个 token 预测的结果完全相同。
• 可以加速推理 2 - 3 倍，据说 GPT-4 使用了这种方法。
• 近似小模型的参数量一般要比大模型小 2 个数量级。
  $E(\textit{generated tokens}) = \frac{1-\beta^{\gamma+1}}{1-\beta}$

Algorithm

算法流程

1. 给定 prefix_seq 序列，用近似小模型（approx model， M_q）预测 $\gamma$ 长度的后续序列，记做 predict_seq，并记录 predict_seq 中每一个 token 的预测概率，记作 q
2. 大模型（target model，M_p）输入 prefix_seq + predict_seq，得到 predict_seq 每个 token 的概率（记作 p）和下一个 token（记作 t）
3. 对于每一个 token，如果 p >= q，则表示大模型接受小模型对这个 token 的提议；如果 p < q，则表示 大模型以 $\frac{p}{q}$ 的概率接受这个 token，用随机均匀采样来确定是否接受这个 token
4. 添加大模型连续接受的 predict_seq 到 prefix_seq 之后，如果 predict_seq 每个 token 都接受了，则 prefix_seq = prefix_seq + predict_seq + t。重复步骤 1

代码实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52

@torch.no_grad()
def speculative_sampling(
    prefix: torch.Tensor,
    approx_model: torch.nn.Module,
    target_model: torch.nn.Module,
    max_len: int,
    gamma: int = 4,
    temperature: float = 1,
    top_k: int = 0,
    top_p: float = 0,
) -> torch.Tensor:
    seq_len = prefix.shape[1]
    T = seq_len + max_len
    approx_model_cache = KVCacheModel(approx_model, temperature, top_k, top_p)
    target_model_cache = KVCacheModel(target_model, temperature, top_k, top_p)
    resample_count = 0
    target_sample_count = 0
    accepted_count = 0
    while prefix.shape[1] < T:
        prefix_len = prefix.shape[1]
        x = approx_model_cache.generate(prefix, gamma)
        _ = target_model_cache.generate(x, 1)
        n = prefix_len + gamma - 1
        for i in range(gamma):
            r = torch.rand(1, device=target_model_cache.device)
            j = x[:, prefix_len + i]
            if r > (target_model_cache._prob_history[:, prefix_len + i - 1, j]) / (
                approx_model_cache._prob_history[:, prefix_len + i - 1, j]
            ):
                # reject
                n = prefix_len + i - 1
                break
            accepted_count += 1
        prefix = x[:, : n + 1]
        approx_model_cache.rollback(n + 1)
        if n < prefix_len + gamma - 1:
            # reject someone, sample from the pos n
            t = sample(
                max_fn(
                    target_model_cache._prob_history[:, n, :]
                    - approx_model_cache._prob_history[:, n, :]
                )
            )
            resample_count += 1
            target_model_cache.rollback(n + 1)
        else:
            # all approx model decoding accepted
            assert n == target_model_cache._prob_history.shape[1] - 1
            t = sample(target_model_cache._prob_history[:, -1, :])
            target_sample_count += 1
        prefix = torch.cat((prefix, t), dim=1)
    return prefix

收益分析

• 大模型推理过程中，预填充阶段计算访存比较高，增量生成阶段计算访存比较低，因此增量生成阶段是主要优化方向
• 本论文的投机解码算法在某种程度上是将 target model 的增量生成阶段变成 approx model 的增量生成 + target model 的预填充，收益主要来自于这里
• 且投机解码算法和其他大模型推理算法（例如 kv cache、flash attention、MQA / GQA）并不冲突，可以合并使用达到最优加速效果

Thought

• 从另外一种新颖且符合逻辑的方面优化了大模型推理

URL

• paper: https://arxiv.org/pdf/2202.00666
• code: https://github.com/huggingface/transformers/blob/main/src/transformers/generation/logits_process.py#L608

TL;DR

• 本文提出来一种 Typical sampling 的解码策略，是对 top-p sampling 的改进，启发了后续的 $\eta-sampling$ 解码策略，旨在生成更自然、更符合人类语言使用习惯的文本。
• 论文使用信息论的观点，局部典型性（local typicality）采样的依据是条件熵和信息量的接近程度
  • 条件熵：$(-log(p)*p).sum()$
  • 信息量：$-log(p)$

Algorithm

公式角度理解

• 在时间步 $t$，候选词集合为：

$$L_{\epsilon}^{(T)}=\{y=y_0,...,y_T| \forall 1\le t \le T, | \log p(y_t|y_{\lt t}) + H(Y_t|Y_{\lt t=y \lt t}) | \lt \epsilon \}$$

• 其中
  • $| \log p(y_t|y_{\lt t}) + H(Y_t|Y_{\lt t=y \lt t}) | \lt \epsilon$ 表示信息量和条件熵最大相差不大于 $\epsilon$，写成 $|H(Y_t|Y_{\lt t=y \lt t}) - (-\log p(y_t|y_{\lt t}) ) | \lt \epsilon$ 更容易理解
  • $\epsilon$ 是超参数，需要在 (0, 1) 之间

代码角度理解

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

class TypicalLogitsWarper(LogitsWarper):
    def __init__(self, mass: float = 0.9, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):
        mass = float(mass)
        if not (mass > 0 and mass < 1):
            raise ValueError(f"`typical_p` has to be a float > 0 and < 1, but is {mass}")
        if not isinstance(min_tokens_to_keep, int) or (min_tokens_to_keep < 1):
            raise ValueError(f"`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}")
        self.filter_value = filter_value
        self.mass = mass
        self.min_tokens_to_keep = min_tokens_to_keep
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        # calculate entropy
        normalized = torch.nn.functional.log_softmax(scores, dim=-1)
        p = torch.exp(normalized)
        ent = -(normalized * p).nansum(-1, keepdim=True)
        # shift and sort
        shifted_scores = torch.abs((-normalized) - ent)
        sorted_scores, sorted_indices = torch.sort(shifted_scores, descending=False)
        sorted_logits = scores.gather(-1, sorted_indices)
        cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
        # Remove tokens with cumulative mass above the threshold
        last_ind = (cumulative_probs < self.mass).sum(dim=1)
        last_ind.clamp_(max=sorted_scores.shape[-1] - 1)
        sorted_indices_to_remove = sorted_scores > sorted_scores.gather(1, last_ind.view(-1, 1))
        sorted_indices_to_remove[..., : self.min_tokens_to_keep] = 0
        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
        scores_processed = scores.masked_fill(indices_to_remove, self.filter_value)
        return scores_processed

Thought

• 单看计算过程实际上比较符合直觉，用信息论的角度理解反而会云里雾里…

URL

• paper: https://arxiv.org/pdf/2210.15191
• code: https://github.com/john-hewitt/truncation-sampling/blob/main/src/TruncationVisualization.ipynb

TL;DR

• 常用的 top-p random sampling decoding method 可能会导致输出的长文本质量较差
• 本文设计了一种基于熵的动态概率阈值优化 top-p 随机采样算法，叫做 $\eta-sampling$

Algorithm

公式表示

$$A_{x \lt i}=\{x \in V | P_{\theta}(x | x \lt i) \gt \eta \}$$

$$\eta=min(\epsilon,\alpha(-h_{\eta,x<i}))$$

其中第一行公式表示随机采样概率值截断到 $\eta$
第二行中 $\epsilon,\alpha$ 是超参数，通常 $\alpha=\sqrt{\epsilon}$，h 表示输出的熵，$h=-(p*log(p)).sum()$，p 表示概率

代码表示

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

class EtaWarper(transformers.LogitsWarper):
  """Our proposed eta sampling warper."""
  def __init__(self, epsilon):
    self.epsilon = epsilon
    self.filter_value = -float("Inf")
  def __call__(self, input_ids, scores) -> torch.FloatTensor:
    probabilities = scores.softmax(dim=-1)
    entropy = - (probabilities * torch.log(probabilities)).sum()
    eta = min(self.epsilon, torch.sqrt(torch.tensor(self.epsilon))*torch.exp(-entropy))
    indices_to_remove = probabilities < eta
    max_word = torch.argmax(scores,dim=-1)
    # 防止把候选词全删了，至少留一个
    indices_to_remove[...,max_word.squeeze()] = 0
    new_scores = scores.masked_fill(indices_to_remove, self.filter_value)
    return new_scores

Thought

• 是 top-p 的小升级，计算量较小，算是一个小 trick

URL

• paper: https://arxiv.org/pdf/2202.06417
• code: https://github.com/yxuansu/SimCTG/tree/main/contrastive_search_explanation

TL;DR

• 本文提出一种叫做 Contrastive search 的解码方法，比 top-k / top-p 等随机解码方法和 argmax 这种贪心解码方法更优，模型输出内容语意更连贯，模型更不容易退化
• Contrastive search 是在 top-k 基础上增加了 degeneration penalty（退化惩罚项），算是 random sampling 类算法的升级版
• 但需要将候选的 top-k 输出追加到序列末尾作为输入重新输入到模型中输出 top-k 对应的 hidden state，因此时间复杂度较高

Algorithm

总体架构

其中 $\alpha$ 是超参数，用来平衡模型预测结果和退化惩罚项
h 表示 hidden state，用 $h_v$ 和序列中每一个 token 的 hidden state 计算相似度
由于 $h_v$ 只能在下一次 inference 时产生，所以需要在 decoding 阶段再推理一次，这一步比较麻烦

用一段代码来描述一下解码过程

1. 首先在 t-1 时刻输入 seq_len 长度的序列，输出 vocab_size 长度的 logits，选 top-k

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

# compute the logits in this step
prev_hidden_states, logits = model.compute_logits_and_hidden_states(input_ids)
_, seqlen, embed_dim = prev_hidden_states.size()
_, _, vocab_size = logits.size()
logit_for_next_step = logits[:,-1,:]
assert logit_for_next_step.size() == torch.Size([1, vocab_size])
# normalize with the softmax function
next_probs = F.softmax(logit_for_next_step, dim = -1)
assert next_probs.size() == logit_for_next_step.size()
# collecte top-k candidate tokens and their logits (model confidence)
_, next_top_k_ids = torch.topk(logit_for_next_step, dim = -1, k = beam_width)
assert next_top_k_ids.size() == torch.Size([1, beam_width])
        
next_top_k_probs = torch.gather(next_probs, dim = 1, index=next_top_k_ids)
assert next_top_k_probs.size() == next_top_k_ids.size()

2. 将 top-k 预测的每一种可能追加到原始序列，并在 batch size 维度 concat 成一个 shape = [beam_width, seq_len + 1, dim] 的输入

1
2
3
4
5
6
7

# concatenate each candidate token with prefix
expanded_context = [input_ids for _ in range(beam_width)]
expanded_context = torch.cat(expanded_context, dim = 0)
assert expanded_context.size() == torch.Size([beam_width, seqlen])
top_k_ids = top_k_ids.view(beam_width, 1)
next_input_ids = torch.cat([expanded_context, top_k_ids], dim = -1)
assert next_input_ids.size() == torch.Size([beam_width, seqlen+1])

3. 将 shape = [beam_width, seq_len + 1, dim] 输入到模型中，计算得到每个 token 的 hidden state

1
2
3
4
5
6
7

# feed these candidates into next round to get their hidden states
new_hidden_states, next_logits = model.compute_logits_and_hidden_states(next_input_ids)
assert new_hidden_states.size() == torch.Size([beam_width, seqlen+1, embed_dim])
context_hidden = new_hidden_states[:,:seqlen,:]
assert context_hidden.size() == torch.Size([beam_width, seqlen, embed_dim])
next_hidden = new_hidden_states[:,seqlen:,:]
assert next_hidden.size() == torch.Size([beam_width, 1, embed_dim])

4. 计算每一个 next hidden state 和 context hidden state 之间的相关性，相关性越高，说明重复性越高，惩罚越重，选择此 token 作为 next input id 的可能性越低

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

def ranking(context_hidden, next_hidden, next_top_k_ids, next_top_k_probs, alpha):
    '''
        context_hidden: beam_width x context_len x embed_dim
        next_hidden: beam_width x 1 x embed_dim
        next_top_k_ids: beam_width x 1
    '''
    beam_width, context_len, embed_dim = context_hidden.size()
    assert next_hidden.size() == torch.Size([beam_width, 1, embed_dim])
    norm_context_hidden = context_hidden / context_hidden.norm(dim=2, keepdim=True)
    norm_next_hidden = next_hidden / next_hidden.norm(dim=2, keepdim=True)
    cosine_matrix = torch.matmul(norm_context_hidden, norm_next_hidden.transpose(1,2)).squeeze(-1)
    assert cosine_matrix.size() == torch.Size([beam_width, context_len])
    scores, _ = torch.max(cosine_matrix, dim = -1)
    assert scores.size() == torch.Size([beam_width])
    next_top_k_probs = next_top_k_probs.view(-1)
    scores = (1.0 - alpha) * next_top_k_probs - alpha * scores 
    _, selected_idx = torch.topk(scores, k = 1)
    assert selected_idx.size() == torch.Size([1])
    selected_idx = selected_idx.unsqueeze(0)
    assert selected_idx.size() == torch.Size([1,1])
    next_id = torch.gather(next_top_k_ids, dim = 0, index=selected_idx)
    assert next_id.size() == torch.Size([1,1])
    return next_id

Thought

• 道理上讲得通，但多推理的一次是实实在在的计算代价，总计算量直接翻倍，真的会有大模型选择这种解码方式吗？

URL

• paper: https://arxiv.org/pdf/2108.12409
• code: https://github.com/ofirpress/attention_with_linear_biases/blob/a35aaca144e0eb6b789dfcb46784c4b8e31b7983/fairseq/models/transformer.py#L742

TL;DR

• 本文提出一种比 T5 bias 更简单的 position embedding 方法叫做 ALiBi (Attention with Linear Bias)，简单好用
• 可以在短数据集上训练，在长数据集上测试，即具有外推性

Algorithm

T5 bias

• 先讲一下 T5 bias 是如何实现 position embedding 的，主要分三步：
  1. 计算 query / key 的 n * n 相对位置矩阵，形如：

     1 2 3 4 5 6 7 8 9 10

     [[ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [-1, 0, 1, 2, 3, 4, 5, 6, 7, 8], [-2, -1, 0, 1, 2, 3, 4, 5, 6, 7], [-3, -2, -1, 0, 1, 2, 3, 4, 5, 6], [-4, -3, -2, -1, 0, 1, 2, 3, 4, 5], [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4], [-6, -5, -4, -3, -2, -1, 0, 1, 2, 3], [-7, -6, -5, -4, -3, -2, -1, 0, 1, 2], [-8, -7, -6, -5, -4, -3, -2, -1, 0, 1], [-9, -8, -7, -6, -5, -4, -3, -2, -1, 0]]

  2. 将相对位置矩阵分桶（超过 num_buckets 的饱和到 num_buckets）

     1 2 3 4 5 6 7 8 9 10

     [[ 0, 17, 18, 19, 20, 21, 22, 23, 24, 24], [ 1, 0, 17, 18, 19, 20, 21, 22, 23, 24], [ 2, 1, 0, 17, 18, 19, 20, 21, 22, 23], [ 3, 2, 1, 0, 17, 18, 19, 20, 21, 22], [ 4, 3, 2, 1, 0, 17, 18, 19, 20, 21], [ 5, 4, 3, 2, 1, 0, 17, 18, 19, 20], [ 6, 5, 4, 3, 2, 1, 0, 17, 18, 19], [ 7, 6, 5, 4, 3, 2, 1, 0, 17, 18], [ 8, 7, 6, 5, 4, 3, 2, 1, 0, 17], [ 8, 8, 7, 6, 5, 4, 3, 2, 1, 0]]

     这里上三角和下三角都有值是因为 encoder bidirection=True，如果是 decoder，则如下：

     1 2 3 4 5 6 7 8 9 10

     [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 0, 0, 0, 0, 0, 0, 0, 0, 0], [2, 1, 0, 0, 0, 0, 0, 0, 0, 0], [3, 2, 1, 0, 0, 0, 0, 0, 0, 0], [4, 3, 2, 1, 0, 0, 0, 0, 0, 0], [5, 4, 3, 2, 1, 0, 0, 0, 0, 0], [6, 5, 4, 3, 2, 1, 0, 0, 0, 0], [7, 6, 5, 4, 3, 2, 1, 0, 0, 0], [8, 7, 6, 5, 4, 3, 2, 1, 0, 0], [9, 8, 7, 6, 5, 4, 3, 2, 1, 0]]

  3. 最后是将此 n * n 的 relative position bucket 通过可学习的 embedding 函数变成 n * n * num_heads 的向量，和每个头的 attention score（softmax 之前） 相加，然后通过逐行 softmax 得到 attention weight

ALiBi

• 用数学公式表示：$softmax(q_iK^T+m\cdot[-(i-1),...,-2,-1,0])$
• ALiBi 的计算和 T5 bias 的前两步几乎一模一样
• 第三步不再使用可学习的 embedding 函数映射到每个头上，而是将距离矩阵的值和每个头独立的 不可学习的 常量 m 值相乘，然后和 attention score 相加
• $m_h = \frac{b}{(2^{(8/H)} \cdot b)^h}$
  • b 是一个基数
  • H 是注意力头的数量
  • h 是注意力头的索引（从 0 到 H-1）

Thought

• 标准 attention 的 $pe \in \mathbb{R}^{n\times d}$ 已经慢慢被淘汰了，不管是 RoPE / T5 Bias / ALiBi 都已经逐渐演变成 $pe \in \mathbb{R}^{n\times n}$ 直接作用在 attention score 上了
• ALiBi 的外推性其实本质是强行饱和掉远距离，有点过于粗暴了…