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。

1. 背景介绍

MegatronLM的第一篇论文【Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism】是2020年出的，针对billion级别的模型进行训练，例如具有38亿参数的类GPT-2的transformer模型和具有39亿参数的BERT模型。

分布式训练的模型并行有两种方式，一种是层间并行(inter-layer)，也就是Pipeline流水线并行，相当于下图对整个模型竖切后每个device各保存3个layer(0,1,2和3,4,5)；一种是层内并行(intra-layer)的方式进行，也就是Tensor模型并行，相当于下图横切后每个device各保留6个layer的一半。

1. 背景介绍

MegatronLM的第二篇论文【Efficient Large-Scale Language Model Training on GPU ClustersUsing Megatron-LM】是2021年出的，同时GPT-3模型参数已经达到了175B参数，GPU显存占用越来越多，训练时间也越来越长。

在本系列中，我们将探讨Megatron-LM的源代码。Megatron-LM是由Nvidia开发的一个大规模语言模型训练框架，它采用模型并行的方式实现分布式训练。在本篇文章中，我们将关注模型并行初始化的过程。

1. pretrain

在Megatron中pretrain函数是框架执行的入口，定义在megatron/training.py文件中。

1
2
3
4
5
6
7
8

def pretrain(train_valid_test_dataset_provider,
             model_provider,
             model_type,
             forward_step_func,
             process_non_loss_data_func=None,
             extra_args_provider=None,
             args_defaults={}):
    ...

以gpt训练为例，在pretrain_gpt.py中使用如下：

1
2
3
4
5
6

if __name__ == "__main__":
    pretrain(train_valid_test_datasets_provider,
             model_provider,
             ModelType.encoder_or_decoder,
             forward_step,
             args_defaults={'tokenizer_type': 'GPT2BPETokenizer'})

1. 背景介绍

因为Transformer的自注意力机制(self-attention)的计算的时间复杂度和空间复杂度都与序列长度有关，所以在处理长序列的时候会变的更慢，同时内存会增长更多。通常的优化是针对计算复杂度(通过\(FLOPs\) 数衡量), 优化会权衡模型质量和计算速度。

在FlashAttention中考虑到attention算法也是IO敏感的，通过对GPU显存访问的改进来对attention算法的实现进行优化。如下图，在GPU中片上存储SRAM访问速度最快，对应的HBM(high bandwidth memory)访问速度较慢，为了加速要尽量减少HBM的访问次数。