MLTalks

1. 背景介绍

Causal Attention论文是一篇因果推断(causal inference)和注意力(attention)结合的一篇文章，主要用在视觉和文本结合的领域，如VQA(Visual Question Answering)视觉问答。

VQA(Visual Question Answering)视觉问答的一个基本流程如下，对输入图进行self-attn编程得到K和V的向量，从文本得到Q的向量进行Attn计算，得到填空的结果(riding)。这个过程可以看成是一个因果推断的过程，对应的示意图如下X->Z->Y，X是输入，Z是模型过程，Y是输出，箭头表示相互依赖的关系。

github: https://github.com/NVIDIA/Megatron-LM

1. recompute参数配置

在megatron/arguments.py中有重计算的参数配置如下：

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

group.add_argument('--recompute-activations', action='store_true',
                   help='recompute activation to allow for training '
                   'with larger models, sequences, and batch sizes.')
group.add_argument('--recompute-granularity', type=str, default=None,
                   choices=['full', 'selective'],
                   help='Checkpoint activations to allow for training '
                   'with larger models, sequences, and batch sizes. '
                   'It is supported at two granularities 1) full: '
                   'whole transformer layer is recomputed, '
                   '2) selective: core attention part of the transformer '
                   'layer is recomputed.')
group.add_argument('--distribute-saved-activations',
                   action='store_true',
                   help='If set, distribute recomputed activations '
                   'across model parallel group.')
group.add_argument('--recompute-method', type=str, default=None,
                   choices=['uniform', 'block'],
                   help='1) uniform: uniformly divide the total number of '
                   'Transformer layers and recompute the input activation of '
                   'each divided chunk at specified granularity, '
                   '2) recompute the input activations of only a set number of '
                   'individual Transformer layers per pipeline stage and do the '
                   'rest without any recomputing at specified granularity'
                   'default) do not apply activations recompute to any layers')
group.add_argument('--recompute-num-layers', type=int, default=1,
                   help='1) uniform: the number of Transformer layers in each '
                   'uniformly divided recompute unit, '
                   '2) block: the number of individual Transformer layers '
                   'to recompute within each pipeline stage.')

说明：

• --recompute-activations: 设置recompute_activations等同于recompute_granularity为selective；selective运行效率更高，大部分场景只设置这个就可以。如果显存更紧张时，再通过recompute-granularity来进行full的设置。
• --recompute-granularity: 支持不同颗粒度的重计算，设为full会重计算整个transformer层，设为selective只会重算transformer中的core_attention部分。
• --distribute-saved-activations: 按TP并行度分开存储activation。
• --recompute-method: uniform计算会把所有的transformer layer分为若干组，分别把每组的input activation保存在内存中, GPU显存不足时，可通过设大每个组内的layer数来运行更大的model；block是针对pipeline并行的每个stage，checkpoint部分transformer layer的input activation, 剩余部分不进行checkpoint缓存，对于一个pipeline stage中有8层的来说，当设为5时，前5层中每一层的input activation都会被缓存，后3层在反向的时候正常计算。
• --recompute-num-layers: 对于uniform类型，表示设置在每个重计算的transformer layer group中的层数, 默认为1表示对每一层transformer layer都分别进行checkpoint；对于block类型，设为N表示单个pipeline stage中的前N个layers会缓存input activation。

1. LayerNorm使用介绍

pytorch中的函数定义如下：

1

torch.nn.LayerNorm(normalized_shape, eps=1e-05, elementwise_affine=True, device=None, dtype=None)

函数参数说明如如下： * normalized_shape: 进行LayerNorm的维度定义，对于一个多维矩阵[N, C, H, W]来说，这里的normalized_shape定义都是要和矩阵最后几个维度保持一致的，这里就是[C, H, W]。对比数学公式，其中的 \(\gamma\) 和 \(\beta\) 的维度都是[C, H, W]，\(x\) 和 \(y\) 的维度都是[N, C, H, W]。 * eps：为了防止计算公式中的分母为0，加上一个极小的数，默认值: 1e-5 * elementwise_affine：设为True的时候，进行elementwise的仿射变换, \(\gamma\) 和 \(\beta\) 才会生效，在训练过程中做为参数会被学习更新，为False的话不生效。\(\gamma\) 所有元素初始为1， \(\beta\) 所有元素初始为0的。\(\gamma\) 在代码实现中对应 \(gamma\), \(\beta\) 在代码实现中对应 \(beta\)。

LayerNorm的数学公式定义如下：

\[\begin{align*} Y &= \frac{X - E[X]}{\sqrt{Var[X] + \epsilon}} * \gamma + \beta \end{align*}\]

论文：GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

1. 背景介绍

Google在2023年发表的一篇关于Transformer Attention的论文，整体论文写的清晰易读，思想简单但很好用。论文名字简写是GQA，但实际分别代表了两种缩写： 1. Generalized Multi Query Attention 2. Grouped Query Attention

1. 基本介绍

LLaMA-2是2023年7月24日Meta发布的LLaMA第二代，跟LLaMA-1几个显著区别:

• 免费可商用版本的大模型
• context上下文增加了一倍，从2K变为了4K
• 训练的总token数从1.0T/1.4T增加为2.0T(\(2 \times 10^{12}\)), 在1.4T基础上增加40%
• 对于最大的模型参数量65B也增加到了70B(\(70 \times 10^{9}\))，并在34B和70B两个版本上使用了 \(Group-Query-Attention(GQA)\) 的方法

github: https://github.com/NVIDIA/Megatron-LM

在【Megatron-LM源码系列(二)：Tensor模型并行和Sequence模型并行训练】基础上增加了Pipeline模型并行训练的介绍，对于Pipeline模型并行思路可参考【详解MegatronLM流水线模型并行训练(Pipeline Parallel)】。pipeline并行中网络是按层的粒度进行纵向切分，在通信组通信上中在pipeline的不同stage中进行横向通信。如下图中2机16卡每个色块就是一个pipeline通信组，训练前向通信的顺序是从左向右。

代码库地址: https://github.com/NVIDIA/Megatron-LM/tree/23.05

1. 整体介绍

模型并行训练实现的核心代码在megatron/core/目录下，按README.md介绍来说，Megatron Core是一个专门针对transformer类模型、效率高、可扩展的计算库。

1. 使用KV缓存(KV Cache)

1. 背景介绍

全切片数据并行(Fully Sharded Data Parallel，简称为FSDP)是数据并行的一种新的方式，FSDP最早是在2021年在FairScale-FSDP中提出的，后来合入了PyTorch 1.11版本中。微软之前Deepspeed框架中提出过三种级别的ZERO算法，FSDP可以看成是ZERO-3的实现。

1. 背景介绍

MegatronLM的第三篇论文【Reducing Activation Recomputation in Large Transformer Models】是2022年出的。在大模型训练过程中显存占用过大往往成为瓶颈，一般会通过recomputation重计算的方式降低显存占用，但会带来额外的计算代价。这篇论文提出了两种方法，分别是sequece parallel和selective activation recomputation，这两种方法和Tensor并行是可以相结合的，可以有效减少不必要的计算量。

下图中绿色部分表示不同模型中需要用于保存activation需要的显存大小，蓝色部分表示不同模型中需要用于保存parameter和optimizer state需要的显存大小。红色线表示A100的显存大小80G。