MLTalks

1. Context Parallel并行原理介绍

megatron中的context并行(简称CP)与sequence并行(简称SP)不同点在于，SP只针对Layernorm和Dropout输出的activation在sequence维度上进行切分，CP则是对所有的input输入和所有的输出activation在sequence维度上进行切分，可以看成是增强版的SP。除了Attention模块以外，其他的模块(Layernorm、Dropout)由于没有多token的处理，在CP并行时都不用任何修改。

1. 背景介绍

Gemma模型是在2024.2.21号Google新发布的大语言模型, Gemma复用了Gemini相同的技术(Gemini也是Google发布的多模态模型)，Gemma这次发布了了2B和7B两个版本的参数，不仅提供了预训练的checkpoints，还提供了用于对话、指令跟随等fine-tune的checkpoints。在QA问答、常识。在11

1. 使用入口

• DistributedOptimizer类定义在megatron/optimizer/distrib_optimizer.py文件中。创建的入口是在megatron/optimizer/__init__.py文件中的get_megatron_optimizer函数中。根据传入的args.use_distributed_optimizer参数来判断是用DistributedOptimizer还是Float16OptimizerWithFloat16Params。

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

def get_megatron_optimizer(model,
                           no_weight_decay_cond=None,
                           scale_lr_cond=None,
                           lr_mult=1.0):
        ...
        # Megatron optimizer.
        opt_ty = DistributedOptimizer \
            if args.use_distributed_optimizer else \
            Float16OptimizerWithFloat16Params
        return opt_ty(optimizer,
                      args.clip_grad,
                      args.log_num_zeros_in_grad,
                      params_have_main_grad,
                      args.use_contiguous_buffers_in_local_ddp,
                      args.fp16,
                      args.bf16,
                      args.params_dtype,
                      grad_scaler,
                      model)

• 相关的Optimizer的使用参考【Megatron-LM源码系列(六)：Distributed-Optimizer分布式优化器实现Part1】

1. 使用说明

在megatron中指定--use-distributed-optimizer就能开启分布式优化器, 参数定义在megatron/arguments.py中。分布式优化器的思路是将训练中的优化器状态均匀地分布到不同数据并行的rank结点上，相当于开启ZERO-1的训练。

1
2

group.add_argument('--use-distributed-optimizer', action='store_true',
                   help='Use distributed optimizer.')

在使用--use-distributed-optimizer, 同时会check两个参数 args.DDP_impl == 'local'(默认开启)和args.use_contiguous_buffers_in_local_ddp(默认开启)。

1
2
3
4
5

# If we use the distributed optimizer, we need to have local DDP
# and we should make sure use-contiguous-buffers-in-local-ddp is on.
if args.use_distributed_optimizer:
    assert args.DDP_impl == 'local'
    assert args.use_contiguous_buffers_in_local_ddp

分布式优化器节省的理论显存值依赖参数类型和梯度类型，以下是每一个parameter对应占用的理论字节数(d表示数据并行的size大小，也就是一个数据并行中的卡数, 等于 \(TP \times PP\) )：

1. 浮点格式说明

浮点数的格式通常由三部分组成：符号位(Sign bit)、指数部分(Exponent)和尾数部分(Significand/Fraction)。整个浮点数占用的位数取决于不同的浮点数格式。例如，IEEE 754标准的单精度浮点数（float）有32位，双精度浮点数（double）有64位。参考：Floating-point arithmetic

最终的浮点表示如下，s是significand；p是precision精度(significand中的数字的个数)；b是base，这里base用的是10或者2。 \[ \frac{s}{b^{p-1}} \times b^e\]

一个具体的示例如下：

float32/bfloat16/float16/tf32三种格式的比较:

Megatron-LM代码仓：Megatron-LM

1. FP16参数指定

• 训练模型要使用fp16时，训练启动参数中指定--fp16, 对应megatron/arguments.py中的定义如下：

1
2

group.add_argument('--fp16', action='store_true',
                   help='Run model in fp16 mode.')

• 在计算lm-cross-entropy时默认是使用fp32来计算的，在开启--fp16选项的前提下可以通过指定--fp16-lm-cross-entropy来使用fp16计算lm-loss-entropy，对应megatron/arguments.py中的定义如下：

1
2
3

group.add_argument('--fp16-lm-cross-entropy', action='store_true',
                   help='Move the cross entropy unreduced loss calculation'
                   'for lm head to fp16.')

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