《大规模语言模型：从理论到实践》-LLM集群训练阅读分享

ljg2np

《大规模语言模型：从理论到实践》-LLM集群训练阅读分享 [复制链接]

本文分享对书中第4章“分布式训练”的阅读理解，介绍LLM集群训练的相关理论和实践内容。

一、引言

LLM的主要成果和技术核心是预训练模型，它是LLM发展的平台和支撑。预训练模型的实现，主要包括预训练数据集和集群训练两部分，前文介绍了预训练数据集，下面介绍集群训练的相关内容。

二、集群训练

LLM预训练对算力和内存资源需求庞大，目前主要通过设计分布式训练系统来实现，涉及集群架构、并行策略、模型架构、内存优化、计算优化等技术。

1、分布式训练

分布式训练（Distributed Training）指将机器学习或深度学习模型训练任务分解成多个子任务，并在多个计算设备上并行地进行训练。

由于同一个服务器内部的多个计算设备之间内存也可能并不共享，因此无论这些计算设备是否处于一个服务器还是多个服务器中，其系统架构都属于分布式系统范畴。

上图给出了单个计算设备和多个计算设备的示例，这里计算设备可以是中央处理器（CPU）、图形处理器（GPU）、张量处理器（TPU）或神经网络处理器（NPU）。

分布式训练的总体目标是提升总的训练速度，减少模型训练的总体时间。总训练速度按下式估计：

单设备计算速度由单块计算加速芯片的运算速度和数据I/O能力决定。

2、并行策略

单个计算设备模型训练系统的流程如下图所示：

从数据和模型两方面着手考虑并行加速策略，包括数据并行、模型并行、混合并行，设备内存优化。

（1）数据并行

对数据进行切分，将同一模型复制到多个设备上，并行执行不同的数据分片。

（2）模型并行

对模型划分，将模型中的算子分发到多个设备上分别完成处理。

模型并行按照层间运行和层内运行方式的不同分为流水线并行和张量并行。

流水线并行：

流水线并行将模型的各个层分段处理，将每个段分布在不同的计算设备上，使前后段能够流水式、分批工作。这种并行计算策略要注意在设计时尽量避免流水线气泡导致资源利用率降低的情况发生。

GPipe方法采用流水线并行策略来有效减少并行气泡，如下图所示。

Megatron-LM采用1F1B流水线策略，通过引入任务调度机制来减少流水线气泡。1F1B 交错式调度模式要求微批次（Micro-batch）的数量是流水线阶段的整数倍。每个设备不仅负责连续多个层的计算，还可以处理多个层的子集，这些子集被称为模型块。如下图所示，设备1 可以处理层1、2、9、10，设备2 处理层3、4、11、12，依此类推。在这种模式下，每个设备在流水线中被分配到多个阶段，不仅在内存消耗方面表现出色，还能提高计算效率。

张量并行：

张量并行需要根据模型的具体结构和算子类型，解决如何将参数切分到不同设备，以及如何保证切分后数学一致性这两个问题。

大语言模型都是以Transformer 结构为基础，Transformer 结构主要由嵌入式表示（Embedding）、矩阵乘（MatMul）和交叉熵损失（Cross Entropy Loss）计算构成。

这三种类型的算子有较大的差异，都需要设计对应的张量并行策略才可以实现将参数切分到不同的设备。

①嵌入层：对于嵌入表示层的参数，按照词维度切分，每个设备存储部分词向量，最后通过汇总得到完整词向量，如图所示：

②矩阵乘：按照行、列切分方式不同，执行相应的处理来保证数学等价，如图所示：

其在前向网络层的应用，如下图所示：

在多头注意力计算中的应用，如下图所示：

③交叉熵损失：

分类网络最后一层一般会选用Softmax 和Cross_entropy 算子来计算交叉熵损失（Cross Entropy Loss）。如果类别数量非常大，则会导致单计算设备内存无法存储和计算logit 矩阵。针对这一类算子，可以按照类别维度切分，同时通过中间结果通信，得到最终的全局交叉熵损失。

计算Softmax值，公式如下：

其中，p表示张量并行的设备号。

得到Softmax计算结果后，对标签Target按类别切分，每个设备得到部分损失，最后一次通信，得到所有类别的损失。整个过程需要三次小量的通信，就可以完成交叉熵的计算。

（3）混合并行

对数据和模型同时进行切分实现更高程度的并行，即将多种并行策略如数据并行、流水线并行和张量并行等混合使用。

通过结合不同的并行策略，混合并行可以充分发挥各种并行策略的优点，最大程度地提高计算性能和效率。

BLOOM 模型使用Megatron-DeepSpeed框架进行训练，主要包含两个部分：Megatron-LM提供张量并行能力和数据加载原语；DeepSpeed提供ZeRO 优化器、模型流水线及常规的分布式训练组件。通过这种方式可以实现数据、张量和流水线三维并行。下图是BLOOM模型训练时采用的并行计算结构：

（4）计算设备内存优化

大语言模型训练通常采用Adam优化算法，除了需要每个参数梯度，还需要一阶动量（Momentum）和二阶动量（Variance）。虽然Adam 优化算法相较SGD 算法效果更好也更稳定，但是对计算设备内存的占用显著增大。为了降低内存占用，大多数系统采用了混合精度训练方式，即同时存在FP16（16 位浮点数）或者 BF16（Bfloat16）和FP32（32 位浮点数）两种格式的数值。

混合精度优化的过程如下图所示：

3、集群架构

包括参数服务器架构和去中心化架构。

（1）参数服务器（Parameter Server，PS）架构的分布式训练系统中有两种服务器角色：训练服务器和参数服务器。

（2）去中心化（Decentralized Network）架构采用集合通信实现分布式训练系统。

在去中心化架构中，没有中央服务器或控制节点，而是由节点之间进行直接通信和协调。这种架构的好处是可以减少通信瓶颈，提高系统的可扩展性。

书中介绍的大语言模型训练所使用的分布式训练并行策略，大都是使用去中心化架构，并利用集合通信进行实现。

在分布式训练过程中，节点之间需要周期性地交换参数更新和梯度信息。可以通过集合通信（Collective Communication，CC）技术实现，常用通信原语包括Broadcast、Scatter、Reduce、All Reduce、Gather、All Gather、Reduce Scatter、All to All 等。

三、DeepSpeed实践

分布式框架DeepSpeed、Megatron-LM、Colossal-AI都可以匹配支持HuggingFace transformers，书中以DeepSpeed为例进行了实践讲解。

DeepSpeed是Microsoft 公司开发的开源深度学习优化库，旨在提高大语言模型训练的效率和可扩展性，使研究人员和工程师能够更快地迭代和探索新的深度学习模型和算法。它采用了多种技术手段来加速训练，包括模型并行化、梯度累积、动态精度缩放、本地模式混合精度等。

此外，DeepSpeed 还提供了一些辅助工具，例如分布式训练管理、内存优化和模型压缩，以帮助开发者更好地管理和优化大规模深度学习训练任务。

DeepSpeed可以实现三种并行方法（ZeRO支持的数据并行、流水线并行和张量并行）的灵活组合来应对不同工作负载的需求。

DeepSpeed 基于PyTorch 构建，将PyTorch 训练代码迁移到DeepSpeed 上通常仅需简单修改，这使开发者可以快速利用DeepSpeed 的优化功能来加速训练任务。

DeepSpeed 的分布式计算框架，主要包括API接口、RunTime运行空间和Ops内核三个部分，如下图所示：

几个重要的基础概念：

（1）主节点（master_ip+master_port）负责协调所有其他节点和进程的工作，由主节点所在服务器的IP 地址和主节点进程的端口号来确定主节点。主节点还负责监控系统状态、处理任务分配、结果汇总等任务，因此是整个系统的关键部分。

（2）节点编号（node_rank）是系统中每个节点的唯一标识符，用于区分不同计算机之间的通信。

（3）全局进程编号（rank）是整个系统中的每个进程的唯一标识符，用于区分不同进程之间的通信。

（4）局部进程编号（local_rank）是单个节点内的每个进程的唯一标识符，用于区分同一节点内的不同进程之间的通信。

（5）全局总进程数（world_size）是整个系统中运行的所有进程的总数，用于确定可以并行完成多少工作及完成任务所需的资源数量。

在网络通信策略方面，DeepSpeed 提供了MPI、GLOO、NCCL 等选项，可以根据具体情况进行选择和配置。

DeepSpeed 支持多种类型ZeRO 的分片机制，包括ZeRO-0、ZeRO-1、ZeRO-2、ZeRO-3以及ZeRO-Infinity，其中ZeRO-Infinity允许通过使用NVMe固态硬盘扩展GPU和CPU内存来训练LLM，极大程度上降低了训练门槛。

通信策略和分片机制在DeepSpeed配置文件中的optimizer部分进行配置。

书中讲解了使用DeepSpeed框架训练LLaMA模型的基本过程（附有代码）：

（1）训练数据配置
使用PyTorch 和transformers 库来设置预训练模型的数据加载器，以实现在单机或多机分布式训练环境中对数据的加载和采样。
args.local_rank为-1时，进行单机采样；其它值，则进行分布式采样。

（2）模型载入
使用transformers 库加载和配置LLaMA 模型及其相关的分词器。从transformers 库中导入LLaMA模型、相应的分词器和模型配置后，使用from_pretrained 方法加载预训练的LLaMA 模型、分词器和配置。

（3）优化器设置
DeepSpeed 库提供了高效的优化器算法，如DeepSpeedCPUAdam 和FusedAdam，这些算法经过特殊优化以提高在大规模数据和模型上的训练速度。

（4）DeepSpeed 设置
训练配置函数get_train_ds_config 主要包括：ZeRO优化配置、混合精度训练、梯度裁剪、混合引擎配置、TensorBoard配置、验证集配置函数。

（5）DeepSpeed初始化
按照初始化流程进行模型训练的初始化。

（6）模型训练

借助DeepSpeed 框架实现对模型的训练，训练步骤分为以下几个阶段：

①训练前准备：使用print_rank_0 函数输出当前的训练状态，在训练前对模型进行一次评估，计算模型的困惑度。

②训练循环：每个周期的开始，打印当前周期和总周期数。在每次迭代中，数据批次被移动到相应的GPU 设备，然后对这个批次进行前向传播计算损失。使用model.backward(loss)计算梯度，并使用model.step() 更新模型参数。对于主进程，使用print_throughput 函数打印吞吐量，以了解模型的训练速度和效率。

③保存模型：如果指定了输出目录，模型的状态和配置将会被保存。模型可以保存为不同的格式，例如HuggingFace 格式或DeepSpeed 的Zero Stage 3 特定格式。

下面是部分的训练代码：

#模型训练部分
print_rank_0(“***** Running training *****”, args.global_rank)
print_rank_0(
	f“***** Evaluating perplexity , \
	Epoch {0}/{args.num_train_epochs} *****”,
	args.global_rank)
perplexity=evaluation(model, eval_dataloader)
print_rank_0(f“ppl: {perplexity}”,args.global_rank)

for epoch in range(args.num_train_epochs):
	print_rank_0(
		f“Beginning of Epoch {epoch+1}/{args.num_train_epochs},\
		Total Micro Batches {len(train_dataloader)}”,
		args.global_rank)
	model.train()
	import time
	for step,batch in enumerate(train_dataloader):
	start=time.time ()
	batch=to_device(batch,device)
	outputs=model(**batch,use_cache=False)
	loss=outputs.loss
	if args.print_loss:
		print(
			f“Epoch: {epoch}, Step: {step}, \
			Rank: {torch.distributed.get_rank()}, loss={loss}”
		)
		model.backward(loss)
		model.step()
		end=time.time()
	if torch.distributed.get_rank()==0:
		print_throughput(model.model, args, end-start, args.global_rank)

if args.output_dir is not None:
	print_rank_0(‘saving the final model …’, args.global_rank)
	model=convert_lora_to_linear_layer(model)

if args.global_rank==0:
	save_hf_format(model,tokenizer,args)

if args.zero_stage==3:
	save_zero_three_model(model,args.global_rank,args.output_dir,zero_stage=args.zero_stage)

def print_rank_0(msg,rank=0):
	if rank<=0:
		print(msg)

def print_throughput(hf_model,args,e2e_time,rank=0):
……
def save_hf_format(model,tokenizer,args,sub_folder=””):
……
def save_zero_three_model(model_ema,global_rank,save_dir,zero_stage=0):
……

四、结语

分布式训练这一章主要介绍了并行策略、集群架构等基础理论，提出分布式训练的总体目标就是提升总的训练速度，减少模型训练的总体时间，针对这一目标在训练设计的各个环节进行切分、通信和并行优化处理，最后使用DeepSpeed对LLaMA模型进行了分布式训练的代码实践。