MLTalks

2022年google在GShard之后发表另一篇跟MoE相关的paper, 论文名为GLaM (Generalist Language Model), 最大的GLaM模型有1.2 trillion参数, 比GPT-3大7倍, 但成本只有GPT-3的1/3, 同时效果也超过GPT-3. 以下是两者的对比:

Switch Transformers也是google在2022年发表的一篇论文, 该论文简化了MoE的路由算法, 减少了计算量和通信量; 第一次支持bfloat16精度进行训练. 基于T5-Base和T5-Large设计的模型在相同的算力下训练速度提升了7x倍; 同时发布了1.6万亿(1.6 trillion)参数的MoE模型，相比T5-XXL模型训练速度提长了4x倍.

1. 背景说明

GShard是Google在2020年的一篇论文, 将 \(Sparsely\ Gated\ Mixture\ of\ Experts\) 与 transformer结合, 支持了600B大小的参数量, 使用了2048块TPU v3训练了4天翻译任务(100种语言翻译为英文).

1. 背景介绍

这是2017年Google Brain团队发表的一篇比较经典的论文, 可以做为MOE在大语言模型上应用的第一篇, MOE全称是混合专家(Mixture of Experts). 这里的每个Expert都是一个更小的神经网络, 比如最简单就是FC全连接网络, MOE整体不是一个完整的网络结构, 而是作为layer层嵌入到别的网络中.

首先介绍这篇论文的背景, 模型的学习能力强弱跟模型的参数量有关, 即模型越大参数量越多那么模型效果越好, 同时存在个问题就是按之前稠密模型的设计, 模型参数量越大计算量也越大, 计算越慢. 有没有一种方式既可以增大模型参数量又能使得计算量控制在一定范围内(对应有个名词叫条件计算conditional computation), 这就是这次论文讨论的重点. 本文中基于LSTM堆栈基础上应用了MOE结果使得参数量提升至137B大小, 相比之前模型容量提升了有1000倍, 但是计算量只增加了一少部分, 具体数据和效果看最后.

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.')