Zhangzhe's Blog

URL

• paper: https://arxiv.org/pdf/2104.14294
• code: https://github.com/facebookresearch/dino

TL;DR

• DINO（Distillation with No Labels） 是一种自监督学习方法，主要用于 Vision Transformer (ViT) 的训练
• 在无标签的图片数据上训练模型，让模型学习图像的表示意义
• 利用 MoCo 提出的 Momentum Teacher 算法做蒸馏

Algorithm

训练流程

1. 创建两个完全一样的网络，命名为教师 teacher 网络和学生 student 网络
2. 对同一个输入 x，进行不同的数据增强，得到 x1 和 x2
3. 交叉计算对比损失，再求均值得到 loss for student
4. 只对 student 网络进行反向传播和梯度更新
5. 基于 student 网络的参数更新 teacher 的参数，更新方式是 EMA (exponential moving average)，即：$\theta_t=\lambda \theta_t+(1-\lambda)\theta_s$
6. 更新 teacher 网络输出的中心点：$C = m*C + (1 - m)*mean(t1, t2)$

中心化和锐化

• 两种操作本质上是互补的，防止模型训练崩溃

中心化（centering）

• 中心化的目的是防止特征向量的某个维度占主导地位，从而导致模型输出分布过于集中
• 本质就是一种均值为 0 的归一化，可以提高模型训练的稳定性

锐化（Sharpening）

• 锐化操作的目的是增加教师网络输出的概率分布的锐度，使得输出的概率更加集中在少数几个维度上
• 实现上，锐化通过修改蒸馏温度系数实现

模型效果

• 比一众视觉自监督模型效果都好，比如：MoCo v1/v2，SimCLR v1/v2 等

Thought

• 感觉是 MoCo 系列的升级，框架本身不变，加了数据，稳定了训练过程，增加了些许 trick

URL

• paper: https://arxiv.org/pdf/2110.07602
• code: https://github.com/THUDM/P-tuning-v2

TL;DR

• 本文提出一种提示词微调的方法，是对 P-Tuning 的升级
• P-Tuning v2 要解决的问题是：对于所有（或大多数）任务类型和模型参数规模，将提示词微调的精度达到和整体参数微调同样的效果，这也是这篇论文的题目

Algorithm

对 P-Tuning 的优化

1. P-Tuning 只在模型输入层添加可学习的连续嵌入，P-Tuning v2 在模型的每一层都添加

1. Prefix Tuning / Prompt Tuning / P-Tuning 三种方法都是在模型输入中加入连续嵌入
   • 添加方式可能是前缀，也可能是其他 Concat Pattern
   • 通过 Self-Attention 的 embedding 间信息融合机制让虚拟连续嵌入影响整个模型
2. P-Tuning V2 的做法则完全不同，是对模型的 每一层 都添加了可学习的虚拟连续嵌入
   • 具体来说是通过初始化虚拟 past_key_values 来实现的
   • 用 GPT2-small 来举例（12 层 transformer，每层 12 个头，每个头的 dim = 64）
     • 假设 virtual_prompt_seq_len=3，input_prompt_seq_len=10
     • 那么需要先初始化 past_key_values，shape = [12, 2, 12, 3, 64]，分别表示：
       • num_layers
       • key and value
       • num_heads
       • virtual_prompt_seq_len
       • dim
       • shape = [12, 2, 12, 3, 64] 以及修改后的输出层参数是可训练的所有参数
   • 然后将序列长度为 10 的 input token embeddings 输入模型，第一层输出长度还是 10
   • 第二层以及之后的每一层都将上一层输出的长度为 10 的序列和长度为 3 的 virtual_prompt_key_values 合并计算，并输出长度为 10 的序列

2. 不再使用 Verbalizer，而是使用 class head

1. 什么是 Verbalizer ?
   • 传统的预训练语言模型（例如 Bert）的输入是一个 token 序列，输出是一个 token，也就是说词表中每个词都有可能输出
   • 现在有个下游任务需要用 Bert 做情感分类，输入是一段话，输出是：{正面，负面，中性} 中的一种，而且用 P-Tuning 方法微调，那么直接把输入附加上一些虚拟连续提示嵌入，输出的结果还是整个词表，不是三分类
   • 这时候就需要 Verbalizer 的存在了，它的作用是将 Bert 模型的输出从词表空间映射到三分类空间，它的实现形式可以是规则，也可以是深度学习模型
2. P-Tuning V2 如何抛弃 Verbalizer?
   • 抛弃 Verbalizer 的方式很简单，就是打破 Prompt Tuning 模型时不应修改模型参数和结构 的限制
   • 直接删除预训练模型输出层，改成任务相关的层并随机初始化，然后微调

Thought

• 看起来比 P-Tuning v2 更优雅，和 kv cache attention 结合起来，推理耗时增加较小
• 据说对大模型来讲，这种方法和 Prompt Tuning 相比并没有显著精度优势（模型参数量小时，设计很重要；模型参数量大时，参数量几乎可以弥补一切设计上的非最优）

URL

• paper: https://arxiv.org/pdf/2103.10385
• code: https://github.com/THUDM/P-tuning

TL;DR

• 本文提出一种 Prompt Tuning 的方法 P-Tuning
• 和 Prefix Tuning 与 Prompt Tuning 这种连续词嵌入作为前缀的方法不同， P-Tuning 把连续词嵌入分段插入到 输入 和 标签 之间

Algorithm

提出问题

• 模型对人工设计的 Prompt 很敏感（指预训练模型，非大模型），同一个模型同一个数据集，只要稍微改变问题的问法，评测指标就差非常多，如图：
  
• 用 P-Tuning 可解决此类问题

解决问题

• 使用连续词嵌入（可训练）和离散词嵌入（不可训练）相结合的方法，做 Prompt Tuning 微调
  
• 上图左侧是传统全部用离散词嵌入 Prompt 过程
• 上图右侧是离散词嵌入和连续词嵌入相结合的方法，其中 captical 和 Britain 两个问题中最关键的词使用离散词嵌入（来自于词表，固定不可训练），并在离散词嵌入周围插入若干连续词嵌入（可通过反向传播梯度下降训练）

数学表述

• P-Tuning 中输入序列为 $T=\{[P_{0:i}],x,[P_{(i+1):j}],y,[P_{(j+1),k}]\}$，其中：
  • x 表示原始输入的离散词文本（还没有变成词向量）
  • y 表示原始的 label 文本
  • $[P]$ 表示连续词向量
• 输入序列 T 需要通过一种特殊的 Prompt Encoder 变成真实的词嵌入输入 $\{h_0,...,h_i,e(x),h_{i+1},...,h_j,e(y),h_{j+1},...,h_k\}$，其中：
  • $e(x),e(y)$ 是通过查词表得到的离散词嵌入
  • $h$ 是通过 MLP/LSTM 等方法得到的连续向量的词嵌入，向量的长度和离散词嵌入一致

Thought

• 比 Prefix Tuning 插入连续词嵌入的自由度更高，因此理应效果更好，但总感觉解决问题的方法不优雅，因为离散和连续嵌入结合的模板是人为规定的，包含了较多先验知识在里面

URL

• paper: https://arxiv.org/pdf/2104.08691
• code: https://github.com/google-research/prompt-tuning

TL;DR

• 本文提出的 prompt tuning 和 prefix tuning 非常相似，是一种通过给不同任务输入前添加不同前缀，同时冻结原预训练模型参数的微调方式
• 和 prefix tuning 区别主要在前缀词向量的设置和初始化方式方面

Algorithm

Prompt tuning 的前缀词向量长度应该设置多少？

• 作者实验了 {1, 5, 20, 100, 150} 等长度的前缀长度，结论是 20 最合适，超过 20 收益可忽略

Prompt tuning 的前缀初始化方式

• 作者实验了三种前缀初始化方式：
  1. 随机初始化（和 prefix tuning 一致）
  2. 从词表中随机选择常见词初始化
  3. 用自然语言描述任务，并将其根据词表转化为词向量
• 实验结论是：第三种方式最优

其他部分

• 我没看出来和 prefix tuning 有任何不同，甚至本文对 prefix tuning 的理解都是错的

Thought

• 我认为这篇论文在 prefix tuning 的基础上改动较小，比较水

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，听上去也是几倍的存储开销以及吞吐量）