Zhangzhe's Blog

URL

• paper: https://arxiv.org/pdf/2101.00190
• code: https://github.com/XiangLi1999/PrefixTuning

TL;DR

• LLM 作为基础预训练模型，由于其参数量极大（以 B 计）的特点，无法像传统视觉或语言预训练模型一样，在下游任务数据集上全局微调
• 本文提出一种非常新颖的 LLM 微调方法，不对原始模型参数进行改进，而是在每个任务的 Prompt 输入之前加入不同长度的虚拟 prefix token seqence 来实现
• 其中 prefix token seqence 实际上是一个 $\mathbb{R}^{len\times dims}$ 的可学习矩阵
  • 其中，$len$ 表示 prefix token seqence 的长度
  • $dims$ 表示原始模型的词嵌入的长度
  • 虚拟 token 的含义是不是真正来自词表的 token，而是无意义的可学习的连续的任务相关的 token 表征
  • 不同任务的可学习矩阵不共享

Algorithm

传统 fine tuning 和 prefix tuning 的对比

• 传统 fine tuning 更新全部模型参数
• prefix tuning 不更新任何模型参数，只是额外加入了部分参数（虚拟 token 表征参数）

任务相关应该怎么理解？

• 任务相关是指，每一个任务都需要独立的无法共享的 prefix token 表征，而且表征的长度可能不同
• 例如，用预训练的 GPT-2 small + prefix tuning，使其可以做内容总结任务，就需要初始化一个 $\mathbb{R}^{200\times 768}$ 的矩阵，然后通过有监督训练优化这个虚拟 embedding 层相关的参数
  • 其中：200 是通过实验发现的最合适的总结任务 prefix token sequence 长度
  • 768 是 GPT-2 small 模型真实 token embedding 的维度
• 如果想用预训练的 GPT-2 small + prefix tuning，使其可以做表格转文本的任务，就需要重新初始化虚拟 token embedding 矩阵，重新训练
  

实验证明，对于总结任务，200 是个合适的 prefix token sequence 长度，对于表格转文本，10 长度更合适

prefix tuning 如何适配不同架构的预训练模型？

• 对于 GPT 系列的 Casual Decoder 架构，只需要在输入文本 token sequence 之前加入 prefix token sequence
• 对于 Encoder Decoder 架构，需要在 encoder 和 decoder 输入之前都加入 prefix，encoder 和 decoder 的 prefix token 表征可以共享，也可以不共享，具体看任务效果

和 fine tuning 相比，效果如何？

• 在数据充足且任务难度不大的情况下，prefix tuning 效果不差于 fine tuning
• 由于 prefix 没有修改模型原始参数，所以不容易出现 full fine tuning 中常出现的模型退化问题

Thought

• prefix tuning 看起来是非常有希望的 topic，prompt engineering 可挖掘的空间还很大
• 这种微调方式实际上还需要存储大模型的所有参数和推理的中间结果用于梯度回传，所以实际微调代价不如想象中小

URL

• paper: https://arxiv.org/pdf/1707.06347
• code: https://github.com/ikostrikov/pytorch-a2c-ppo-acktr-gail （民间实现）

TL;DR

• 本文提出一种强化学习算法 PPO（ Proximal Policy Optimization，近端策略优化），旨在解决传统策略梯度方法中存在的数据效率低、鲁棒性差等问题。
• PPO 对 TRPO（Trust Region Policy Optimization，信任域策略优化）进行了简化。

Algorithm

1. 算法核心思想

策略梯度方法

• PPO 基于策略梯度方法，通过估计策略的梯度并使用随机梯度上升来优化策略。

裁剪概率比

• PPO 引入了裁剪概率比（clipped probability ratio）的概念，通过对概率比进行裁剪，限制策略更新的幅度，从而避免过大的策略更新导致的性能下降。

近端目标函数

• PPO 使用近端目标函数（surrogate objective function），它是一个关于策略参数的函数，用于在策略更新时进行优化。
• 这个目标函数考虑了策略更新前后的概率比和优势函数（advantage function）的乘积，并通过裁剪概率比来调整这个乘积，使得策略更新更加稳健。

多轮更新

• 与传统的策略梯度方法每次只使用一个数据样本进行一次梯度更新不同，PPO 允许对每个数据样本进行多轮（epochs）的梯度更新，提高了数据的使用效率。

简化实现

• 相比于 TRPO 等算法，PPO 的实现更加简单直观，易于与其他算法和框架集成。

2. 伪代码实现

1. 在每次迭代中，多个 actor 并行地在环境中运行旧策略 $\pi_{\theta_{old}}$，收集数据并计算优势估计。
2. 然后，使用这些数据来优化近端目标函数，通常使用 K epochs（重复使用收集到的数据更新 K 次） 和 SGD 或 Adam 优化器来进行优化。
3. 优化完成后，更新策略参数 $\theta$，并在下一次迭代中使用新的策略。

3. 公式分析

• PPO 的目标函数公式分为三部分：
  • 裁剪的目标函数（CLIP），在时间步 t 的裁剪目标函数，用于优化策略参数 $\theta$
  • 价值函数损失（VF, value function）
  • 策略熵（S, entropy bouns），用于鼓励探索。
    

裁剪概率比的优势

• $L_t^{CLIP+VF+S}(\theta)=\hat E_t[L_t^{CLIP}(\theta)-c_1L_t^{VF}(\theta)-c_2S[\pi_\theta](s_t)]$
  • $L_t^{CLIP}(\theta)=\hat E_t[\min(r_t(\theta)\hat A_t,\ clip(r_t(\theta), 1-\epsilon,1+\epsilon)\hat A_t)]$，$r_t(\theta)=\frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)}$
    • $r(\theta)$ ：概率比率表示新旧策略在状态 s 下采取动作 a 的概率之比
    • $\pi_\theta(a|s)$ ：在状态 s 下采取动作 a 的概率，有策略参数 $\theta$ 决定的随机策略
    • $\pi_{\theta_{old}}(a|s)$ ：在策略更新前，状态 s 下采取动作 a 的概率
    • $\epsilon$ ：超参数，用于控制裁剪的严格程度
    • $\hat A$ ：优势函数（advantage function）的估计值，它估计了采取行动 a 相比于平均状况能带来多少额外回报
  • $L_t^{VF}(\theta)=(V_\theta(s_t)-V_t^{target})^2$
    • $V_\theta(s)$ ：策略参数 $\theta$ 下，状态 s 的期望回报，V 表示价值函数
    • $V^{target}$ ：折扣回报的估计
    • 价值函数损失有助于确保价值函数的预测尽可能接近真实回报，从而为策略提供准确的价值估计
  • $S[\pi_\theta](s_t)=-\mathbb{E}_{a\sim \pi_\theta(\cdot|s_t)}[\log \pi_\theta(a|s_t)]$
    • 表示在状态 s 下，根据策略 $\pi_\theta$ 采取动作的熵
    • 熵是衡量随机性的一个指标，策略的熵越高，表示策略在选择动作时越随机，这有助于算法探索更多的状态-动作空间

Thought

• 对强化学习了解不够深，无法完全 get 到 PPO 的精髓
• InstructGPT / ChatGPT 使用了 PPO 强化学习策略，所以做 LLM 需要掌握此技能

URL

• paper: https://arxiv.org/pdf/2203.02155

TL;DR

• 大模型预训练后，在很多任务上已有非常强的能力，但输出的结果可能是不可信、有害或无帮助的。
• 本文提出 InstructGPT 算法和微调流程，旨在让大模型通过人类反馈的微调方式，在大量任务上对齐人类意图。
• 目标是做到 3H:
  • Helpful
  • Honest
  • Harmless

Algorithm

• 具体来说，InstructGPT 微调分为三步：
  • 根据采集的 SFT (Supervised FineTune) 数据集对 GPT-3 进行有监督的微调；
  • 收集人工标注的 对比数据，训练奖励模型 RM (Reword Model)；
  • 使用 RM 作为强化学习的优化目标，利用 PPO 算法（Proximal Policy Optimization 最近策略优化）微调 SFT 模型。

数据集

1. SFT 数据集

• SFT 数据一部分来自使用 OpenAI 的 PlayGround 的用户，另一部分来自 OpenAI 雇佣的 40 名标注工。
• 在这个数据集中，标注工的工作是根据内容自己编写指示，并且要求编写的指示满足下面三点：
  • 简单任务：labeler 给出任意一个简单的任务，同时要确保任务的多样性；
  • Few-shot 任务：labeler 给出一个指示，以及该指示的多个查询-响应对；
  • 用户相关的：从接口中获取用例，然后让 labeler 根据这些用例编写指示。

2. RM 数据集

• RM 数据集是对比数据类型，用于训练一个奖励模型（Reward Model）。
• InstructGPT/ChatGPT 的做法是先让模型生成一批候选文本，让后通过 labeler 根据生成数据的质量对这些生成内容进行排序。
• 具体做法是对于每一个 Prompt，模型给出 K 个候选输出（$4\le K\le9$），然后将 K 个输出中挑选 2 个样本给 labeler，labeler 给出哪个更好，一个 Prompt 一共需要 labeler 标注 $C_K^2$ 次，最终得到 K 个候选输出的排序结果。

3. PPO 数据集

• InstructGPT 的 PPO 数据没有进行标注，它均来自 GPT-3 的 API 用户。
• 有不同用户提供的不同种类的生成任务，其中占比最高的包括生成任务（45.6%），QA（12.4%），头脑风暴（11.2%），对话（8.4%）等。

微调过程

1. SFT 有监督微调过程

• 和 GPT-3 的预训练过程保持一致，只是监督信号从监督一段文本中下一个 token 变成监督模型对 prompt 输出和 ground truth 之间的差异。

2. Reward Model 训练过程

• Reward Model 的输入是 Prompt + Response，输出是对应的奖励值，是个回归模型，其中 Response 是 RM 数据集中的排序的模型候选输出。
• Reward Model 实际上是原始 LLM 模型将最后一层的分类层改成回归层，用来回归奖励值。
• 损失函数是：$loss(\theta)=-\frac{1}{C_K^2}E_{(x,y_w,y_l) \sim D}[\log(\sigma(r_\theta(x,y_w), r_\theta(x,y_l)))]$，使用对比学习的方式训练。
• 其中：
  • $r_\theta$ 表示参数为 $\theta$ 的 Reward Model
  • x 表示 Prompt
  • $y_w\ ,y_l$ 分别表示对 Prompt x，labeler 喜欢和不喜欢的一对 Response

3. 使用近端策略优化强化学习算法（PPO） + Reward Model 继续优化 SFT 模型

• PPO 可参考 https://zhangzhe.space/2024/09/03/Proximal-Policy-Optimization-Algorithms/

InstructGPT/ChatGPT 的分析和计划

1. 收益

1. 效果比 GPT-3 更加真实
2. 在模型的无害性上比GPT-3 效果要有些许提升
3. 具有很强的 Coding 能力

2. 损失

1. 会降低模型在通用 NLP 任务上的效果
2. 有时候会给出一些荒谬的输出
3. 模型对指示非常敏感
4. 模型对简单概念的过分解读
5. 对有害的指示可能会输出有害的答复

3. 未来计划

1. 人工标注的降本增效：将人类表现和模型表现有机和巧妙的结合起来是非常重要的
2. 模型对指示的泛化/纠错等能力：这不仅可以让模型能够拥有更广泛的应用场景，还可以让模型变得更“智能”
3. 避免通用任务性能下降：需要方案来让生成结果的 3H 和通用 NLP 任务的性能达到平衡

Thought

• 在 ChatGPT 爆发之前，Open-AI 多年来一直在默默研究这个领域，包括 GPT 系列、RLHF、PPO 等，这些东西的有机结合就成了 InstructGPT，放大就变成了 ChatGPT（InstructGPT 是基于 GPT-3 对齐的，ChatGPT 是基于 GPT-3.5 对齐的）
• 想在一个领域做到断崖式的领先，小聪明是没用的，需要长期积累 + 修炼内功

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

公式角度理解

$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

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