AI大语言模型之分布式训练概述

爱喝白开水a

于 2024-08-26 10:29:41 发布

阅读量246

点赞数 3

文章标签：人工智能 embedding ai大模型 LLM 分布式训练大模型训练计算机技术

本文链接：https://blog.csdn.net/python123456_/article/details/141557649

版权

一、前言

随着语言模型参数量和所需训练数据量的急速增长，单个机器上有限的资源已无法满足大语言模型训练的要求。需要设计分布式训练（Distributed Training）系统来解决海量的计算和内存资源需求问题。

在分布式训练系统环境下需要将一个模型训练任务拆分成多个子任务，并将子任务分发给多个计算设备，从而解决资源瓶颈。

但是如何才能利用包括数万计算加速芯片的集群，训练模型参数量千亿甚至是万亿的大规模语言模型？这其中涉及到集群架构、并行策略、模型架构、内存优化、计算优化等一系列的技术。

二、分布式训练概述

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

下图给出了单个计算设备和多个计算设备的示例，这里计算设备可以是中央处理器（Central Processing Unit，CPU）、图形处理器（Graphics Processing Unit，GPU）、张量处理器（Tensor Processing Unit，TPU）也可以是神经网络处理器（Neural network Processing Unit，NPU）。

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

一个模型训练任务往往会有大量的训练样本作为输入，可以利用一个计算设备完成，也可以将整个模型的训练任务拆分成子任务，分发给不同的计算设备，实现并行计算。

此后，还需要对每个计算设备的输出进行合并，最终得到与单个计算设备等价的计算结果。由于每个计算设备只需要负责子任务，并且多个计算设备可以并行执行，因此其可以更快速地完成整体计算，并最终实现对整个计算过程的加速。

促使人们设计分布式训练系统的一个最重要的原因就是单个计算设备的算力已经不足以支撑模型训练。机器学习模型快速发展，从2013年AlexNet开始，到2022年拥有5400亿参数的PalM模型，机器学习模型以每18个月增长56倍的速度发展。模型参数规模增大的同时，对训练数据量的要求也指数级增长，这更加剧了对算力的需求。

近几年CPU的算力增加已经远低于摩尔定律（Moore’s Law），虽然计算加速设备（如GPU、TPU等）为机器学习模型提供了大量的算力，但是其增长速度仍然没有突破每18个月翻倍的摩尔定律。为了能够满足机器学习模型的发展，只有通过分布式训练系统才可以匹配模型不断增长的算力需求。

2.1、分布式训练目标

分布式训练的总体目标就是提升总的训练速度，减少模型训练的总体时间。总训练速度可以用如下公式简略估计：

总训练速度 ∝ 单设备计算速度 X 计算设备总量 X 多设备加速比

单设备计算速度主要由单块计算加速芯片的运算速度和数据I/O能力来决定，对单设备训练效率进行优化，主要的技术手段有混合精度训练、算子融合、梯度累加等；

分布式训练系统中计算设备数量越多，其理论峰值计算速度就会越高，但是受到通信效率的影响，计算设备数量增大则会造成加速比急速降低；

多设备加速比则是由计算和通讯效率决定，需要结合算法和网络拓扑结构进行优化，分布式训练并行策略主要目标就是提升分布式训练系统中的多设备加速比。

2.2、典型的分布式训练模型

OPT模型训练使用了992块NVIDIA A100 80G GPU，采用全分片数据并行（Fully Sharded Data Parallel）以及Megatron-LM张量并行（Tensor Parallelism），整体训练时间将近2个月。
BLOOM模型的研究人员则公开了更多在硬件和所采用的系统架构方面的细节。该模型的训练一共花费3.5个月，使用48个计算节点。每个节点包含8块NVIDIA A100 80G GPU（总计384个GPU），并且使用4*NVLink用于节点内部GPU之间通信。节点之间采用四个Omni-Path 100 Gbps网卡构建的增强8维超立方体全局拓扑网络进行通信。
LLaMA模型训练采用NVIDIA A100-80GB GPU，LLaMA-7B模型训练需要 82432 GPU小时， LLaMA-13B模型训练需要 135168 GPU小时， LLaMA-33B模型训练花费了 530432 GPU小时，而LLaMA-65B模型训练花费则高达 1022362 GPU小时。

2.3、分布式训练需要克服的挑战

分布式训练系统仍然需要克服计算墙、显存墙、通信墙等多种挑战，以确保集群内的所有资源得到充分利用，从而加速训练过程并缩短训练周期。

计算墙：单个计算设备所能提供的计算能力与大语言模型所需的总计算量之间存在巨大差异。2022年3月发布的NVIDIA H100 SXM的单卡FP16算力也只有 2000 TFLOPs，而 GPT-3 则需要 314 ZFLOPs 的总计算量，两者相差了 8 个数量级。
显存墙：单个计算设备无法完整存储一个大语言模型的参数。GPT-3包含1750亿参数，在推理阶段如果采用FP32格式进行存储，需要700GB的计算设备内存空间，而 NVIDIA H100 GPU只有80GB显存。
通信墙：分布式训练系统中各计算设备之间需要频繁地进行参数传输和同步。由于通信的延迟和带宽限制，这可能成为训练过程的瓶颈。GPT-3训练过程中，如果分布式系统中存在128个模型副本，那么在每次迭代过程中至少需要传输 89.6TB的梯度数据。而截止2023年8月，单个InfiniBand链路仅能够提供不超过800Gb/s 带宽。