Zhangzhe's Blog

InstructGPT: Training language models to follow instructions with human feedback

Posted on 2024-08-29 Edited on 2025-03-03 In LLM Valine:

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. 收益

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

2. 损失

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

3. 未来计划

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

Thought

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

ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

Posted on 2024-08-21 Edited on 2025-03-03 In Data Parallelism Valine:

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` 想要解决哪些问题？

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

`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

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

Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

Posted on 2024-08-21 Edited on 2025-03-03 In Tensor Parallelism Valine:

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 函数，下面是其实现：

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，一声叹息…

PipeDream: Fast and Efficient Pipeline Parallel DNN Training

Posted on 2024-08-20 Edited on 2025-03-03 In Pipeline Parallelism Valine:

URL

TL;DR

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

Algorithm

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

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

`weigt stashing`

从上图可以看出 Machine 1 在 forward 5 时，使用的是更新过 1 次的参数（forward 5 之前只做了 backward 1）；但当 backward 5 时，参数已经被更新了 4 次，如果直接使用会出现 forward 和 backward 参数不一致问题
因此，PipeDream 提出了 weight stashing (权重存储)，即保存多组 weights，每次 forward 过程使用最新版本的 weights，当 backward 时，使用 forward 时刻对应的 weights
使用 weight stashing 之前， $w^{(t+1)} = w^{(t)}+\gamma \cdot \triangle f(w_1^{(t)},w_2^{(t)},...,w_n^{(t)})$ $w^{(t + 1)} = w^{(t)} + γ \cdot △ 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$ 表示梯度
使用 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`

weight stashing 解决了 forward 和 backward 使用的 weight 版本不一致问题，但依然存在另外一个问题：同一个 mini-batch 在不同 stage 上 forward 使用的 weight 版本不一致
例如，对于 mini-batch 4，在 Machine 2 上 forward 使用的参数版本是未更新版本，在 Machine 3 上 forward 使用的版本是更新过一次的版本
为解决这个问题，作者提出 vertical sync 概念，即：对于一个 mini-batch 在每个 stage 上 forward 都使用最旧版本的参数，stage forward and backward 结束后统一更新参数
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)})$
事实上实验证明，weight stashing 对训练精度提升明显，vertical sync 提升较小

Thought

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

GPipe: Easy Scaling with Micro-Batch Pipeline Parallelism

Posted on 2024-08-20 Edited on 2025-03-03 In Pipeline Parallelism Valine:

URL

TL;DR

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

Algorithm

Gpipe 方案主要包含三个部分

划分 `Stage`

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

划分 `Micro-Batch`

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

重计算

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

Thought

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

SpecInfer: Accelerating Large Language Model Serving with Tree-based Speculative Inference and Verification

Posted on 2024-08-15 Edited on 2025-03-03 In LLM Valine:

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 对候选树进行验证

如何对树式结构并行解码

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

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

总体算法流程

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

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

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

贪心验证

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

随机验证

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

Thought

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

Fast Inference from Transformers via Speculative Decoding

Posted on 2024-08-14 Edited on 2025-03-03 In LLM Valine:

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

算法流程

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

代码实现

@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

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

Locally Typical Sampling

Posted on 2024-08-09 Edited on 2025-03-03 In decoding method Valine:

URL

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) 之间

代码角度理解

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

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

Truncation Sampling as Language Model Desmoothing

Posted on 2024-08-09 Edited on 2025-03-03 In decoding method Valine:

URL

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 表示概率

代码表示

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

A Contrastive Framework for Neural Text Generation

Posted on 2024-08-08 Edited on 2025-03-03 In decoding method Valine:

URL

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 阶段再推理一次，这一步比较麻烦

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

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

# 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()

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

# 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])

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

# 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])

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

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

TL;DR

Algorithm

数据集

1. SFT 数据集

2. RM 数据集

3. PPO 数据集

微调过程

1. SFT 有监督微调过程

2. Reward Model 训练过程

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

InstructGPT/ChatGPT 的分析和计划

1. 收益

2. 损失

3. 未来计划

Thought

URL

TL;DR

Algorithm

ZeRO-DP

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

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

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

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

ZeRO-R

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

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

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

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

Thought

URL

TL;DR

Algorithm

Tensor Parallelism

Thought

URL

TL;DR

Algorithm

weigt stashing

vertical sync

Thought

URL

TL;DR

Algorithm

划分 Stage

划分 Micro-Batch

重计算

Thought

URL

TL;DR

Algorithm

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

候选树生成和验证流程

如何对树式结构并行解码

总体算法流程

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

贪心验证

随机验证

Thought

URL

TL;DR

Algorithm

算法流程

代码实现

收益分析

Thought

URL

TL;DR

Algorithm

公式角度理解

代码角度理解

Thought

URL

TL;DR

Algorithm

公式表示

代码表示

Thought

URL

TL;DR

1. `SFT` 数据集

2. `RM` 数据集

3. `PPO` 数据集

1. `SFT` 有监督微调过程

2. `Reward Model` 训练过程

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

`InstructGPT/ChatGPT` 的分析和计划

`ZeRO-DP`

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

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

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

`ZeRO-R`

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

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

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

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

`Tensor Parallelism`

`weigt stashing`

`vertical sync`

划分 `Stage`

划分 `Micro-Batch`