Hugging Face高效训练技术四:多GPU分布式训练(DP、PP、TP 、ZeRO)

参考资料:

  如果在单个 GPU 上训练模型的速度太慢,或者模型的权重在单个 GPU 上装不下,此时就需要多GPU的分布式训练。首先,PyTorch 内置的DataParallel (DP) 和 DistributedDataParallel (DDP)都可用于多 GPU 训练,相比DP,DDP更值得推荐。

一、分布式训练基础知识

1.1 集合通信、集合通信库

参考《分布式训练硬核技术——通信原语》《NCCL文档》

在计算机科学中,多个处理单元(计算机节点、线程、进程或其他通信实体)之间信息传递的模式有两种:

  • Point-to-point communication:点对点通信(P2P):一对一的通信方式,通常用于在两个不同的进程或处理单元之间传递数据,是最基本的通信模式;
  • Collective communication:集合通信,一对多或多对多或的通信方式,用于在多个处理单元之间进行协同操作,通常在并行计算环境中使用。

  在分布式系统中,各个节点间往往存在大量的集合通信需求,而我们可以用消息传递接口(Message Passing Interface, MPI)来定义一些比较底层的消息通信行为,譬如Broadcast、Reduce、Allreduce、Scatter、Gather、Allgather等。

  1. Broadcast:广播行为(对应MPI_Bcast接口)。执行Broadcast时,数据从主节点0广播至其他各个指定的节点(0~3)

  2. Scatter:和broadcast类似,都是一对多的通信方式。不同的是,Broadcast将0号节点将相同的信息发送给所有的节点,而Scatter则是将数据的不同部分,按需发送给所有的节点。
    여기에 이미지 설명을 삽입하세요.

  3. Reduce:规约运算,是一系列简单运算操作的统称,细分可以包括:SUM、MIN、MAX、PROD、LOR等类型的规约操作。Reduce意为减少/精简,因为其操作在每个节点上获取一个输入元素数组,通过执行操作后,将得到精简的更少的元素。
    여기에 이미지 설명을 삽입하세요.

  4. AllReduce:多对多的Reduce,即在所有的节点上都应用同样的Reduce操作。从下图中可以看出,all reduce操作可通过单节点上reduce+broadcast操作完成。
    여기에 이미지 설명을 삽입하세요.

  5. Gather:将多个sender上的数据收集到单个节点上,Gather可以理解为反向的Scatter。

  6. AllGather:多对多的Gather,收集所有数据到所有节点上。从最基础的角度来看,All Gather相当于一个Gather操作之后跟着一个Broadcast操作。
    여기에 이미지 설명을 삽입하세요.

  7. ReduceScatter:将各个节点的输入先进行求和,然后在第0维度按卡数切分,将数据分发到对应的卡上。
    여기에 이미지 설명을 삽입하세요.

  上世纪90年代针对HPC领域定义了一套集合通信相关的接口标准,称为 MPI,主要是被应用在科学计算,尤其是超算领域。由于容错性一般,故在机器学习场景下使用较少。

  目前流行的集合通信库 Open MPI、NCCL、Gloo等,都在MPI的基础上,对各种集合通信的模式和算法作了各自的实现。例如NCCL(NVIDIA Collective Communications Library)是基于NCIDIA-GPU的一套开源的集合通信库, 针对NVIDIA GPU进行了性能优化,实现多GPU和多节点集体通信原语,所以在在英伟达硬件上能带来更低的延迟和更高的带宽。

1.2 通信模式

参考《Framework(二):分布式训练》《MXNet之ps-lite及parameter server原理》《Ring Allreduce》

  为了保证分布式训练的结果与单机训练的结果是一致的,需要某种机制在多个机器之间同步信息。目前最流行的模式有两种:

  • 参数服务器模式(Parameter Server,PS):基于参数服务器的中心化架构模式,
  • 集合通讯模式(Collective Communication,CC):基于规约(Reduce)的去中心化架构模式。
1.2.1 Parameter Server(2014)

여기에 이미지 설명을 삽입하세요.

Parameter Server架构

  Parameter Server架构具有星状的拓扑结构,有一个或一组服务器来存储模型参数,众多worker服务器负责读取数据,执行前向和反向并计算梯度。通过网络连接,这些worker把自己的梯度上传(push)到参数服务器,参数服务器收集所有的worker的梯度并进行计算之后,各worker再下拉(pull)模型参数。

  假设每个节点的数据量是M,在一次梯度同步过程中,N台worker节点都需要和中心PS进行一次通信,则PS节点总通信量为 N×M 。可以看出,这种架构总通信量与集群规模成线性关系。因此,当集群规模较大或模型较大时,参数服务器的带宽可能会成为瓶颈。

1.2.2 Ring-AllReduce(2017)

여기에 이미지 설명을 삽입하세요.

Ring-AllReduce架构

  参数服务器的主要问题是多个 Worker 同时跟 PS 通信,PS 本身有可能成为瓶颈,随着 Worker 数量的增加,整体的通信量也线性增加,加速比可能会停滞在某个点位上。

  基于规约的架构是一种去中心化的架构,典型的有Tree All-reduceRing All-Reduce(百度于2017年提出)。在Ring All-Reduce架构下,多个worker通过网络组成了一个环,每个worker依次把自己计算出的梯度同步给下一个worker。假设集群规模为N(GPU数量),经过至多 2×(N-1) 轮同步,就可以完成所有worker的更新。

Ring AllReduce 分为3个步骤:Split, ScatterReduce, AllGather。

  • Split 阶段:根据集群的规模N,把需要同步的数据平均分成N份。
    여기에 이미지 설명을 삽입하세요.

    节点数据平均分为5个分块(Split)

  • ScatterReduce阶段:进行N-1次 ScatterReduce 迭代,每次迭代中,GPU将向其右邻居发送一个块,并从其左邻居接收一个块并累积到该块中。各个节点依次交换数据,使得每个节点只包含最终结果的一部分(1/N)。每个GPU发送和接收的块在每次迭代中都是不同的;第n个GPU从发送块N和接收块N - 1开始,然后从那里向后进行,每次迭代都发送它在前一次迭代中接收到的块。
    여기에 이미지 설명을 삽입하세요.

    第一次 ScatterReduce 操作

여기에 이미지 설명을 삽입하세요.

第二次 ScatterReduce 操作

여기에 이미지 설명을 삽입하세요.

第三次 ScatterReduce 操作

여기에 이미지 설명을 삽입하세요.

最后一次 ScatterReduce 操作

여기에 이미지 설명을 삽입하세요.

ScatterReduce完成时,每个节点都有一个数据块,它累加了其他所有worker节点相应数据块的数据

  • AllGather阶段:各个节点再次交换数据,最终得到完整的结果。此过程与scatter-reduce是相同的(发送和接收的N-1次迭代),只是gpu接收的值没有累加,而是简单地覆盖块。第n个GPU首先发送第n+1个块并接收第n个块,然后在以后的迭代中总是发送它刚刚接收到的块。
    여기에 이미지 설명을 삽입하세요.
    여기에 이미지 설명을 삽입하세요.
    여기에 이미지 설명을 삽입하세요.
    여기에 이미지 설명을 삽입하세요.
    여기에 이미지 설명을 삽입하세요.
    AllGather 操作操作完成,所有worker 节点的所有数据块就都包含了来自于其他worker节点的数据

  最终,经过N-1ScatterReduce 操作和 N-1AllGather 操作,整个集群就完成了数据同步。假设每个节点的数据量是M,每一次操作各个worker 节点发送的数据量是 M/N,接收的数据量也是 M/N,则总传输量为:
여기에 이미지 설명을 삽입하세요.

  所以在Ring All-Reduce中,通信成本是恒定的,与系统中gpu的数量无关,完全由系统中gpu之间最慢的连接决定。

  所有传输都是在离散迭代中同步进行的,因此所有传输的速度受到环中相邻GPU之间最慢(最低带宽)连接的限制.一般来说,如果一个节点上的所有GPU在环中彼此相邻,则该算法的功能最佳。

1.3 同步范式

  在基于参数服务器的架构下,多个worker 间的梯度同步需要通过一个中心节点参数服务器来完成,这不可避免要涉及多个 worker 间的合作,一般来说,模型更新有同步(sync)、异步(async)和混合三种模式。

  同步模式是指,当所有worker都完成一次梯度计算和参数更新后,才开始下一轮的迭代。这种模式会出现木桶效应,使得整个集群的速度上限受限于最慢的机器。

  异步模式则相反,每个worker只关心自己的进度,完成梯度计算后就尝试更新,至于能跟其他多少个worker“互通有无”则完全随机,其过程不可控,有可能出现无法收敛的问题。

  混合模式综合了上述两种方式,各个worker都会等待其他worker完成梯度计算和参数更新,但不是永远等待,而是通过一个超时机制来完成。混合模式虽然也带来了一定的不确定性,但影响并不大,因此其应用也最为普遍。

여기에 이미지 설명을 삽입하세요.

1.4 大模型训练的目标公式

参考《全网最全-超大模型+分布式训练架构和经典论文》

超大模型训练的总体目标就是提升总的训练速度,减少大模型的训练时间,其总的训练速度的公式为:
여기에 이미지 설명을 삽입하세요.

  1. 单卡速度:单卡速度既然是运算速度和数据IO的快慢来决定,那么就需要对单卡训练进行优化,于是主要的技术手段有精度训练、算子融合、梯度累加来加快单卡的训练性能。

  2. 加速芯片数量:理论上,AI芯片数量越多,模型训练越快。但是,随着训练数据集规模的进一步增长,加速比的增长并不明显。如数据并行就会出现局限性,当训练资源扩大到一定规模时,由于通信瓶颈的存在,增加计算资源的边际效应并明显,甚至增加资源也没办法进行加速。这时候需要通讯拓扑进行优化,例如通过ring-all-reduce的通讯方式来优化训练模式。

  3. 多卡加速比:多卡加速比既然由计算、通讯效率决定,那么就需要结合算法和集群中的网络拓扑一起优化,于是有了数据并行DP、模型并行MP、流水线并行PP相互结合的多维度混合并行策略,来增加多卡训练的效率。

  总的来说呢,超大模型训练的目标就是优化上面的公式,提升总训练速度。核心思想是将数据和计算有关的图/算子切分到不同设备上,同时尽可能降低设备间通信所需的代价,合理使用多台设备的计算资源,实现高效的并发调度训练,最大化提升训练速度。

二、数据并行

2.1 DataParallel(DP)

  最常见的并行训练方式是数据并行(DataParallel),这种方法把输入数据split到各个 workers中(每个worker拥有全部模型)做并行计算,以解决batch size过大的问题。因为求导以及加和都是线性的,所以数据并行在数学上是等价的。

  DataParallel是Pytorch中最容易实现的并行方案,只需要增加一行代码 model = nn.DataParallel(model),即可实现多卡训练。

# 数据集的长度为100,batch size为32。模型是一个简单的fc层,输入长度是5,输出是2
input_size,output_size = 5,2
batch_size,data_size = 32,100

model = Model(input_size, output_size)
if torch.cuda.device_count() > 1:
    print("Gemfield have ", torch.cuda.device_count(), "GPUs!")
    model = nn.DataParallel(model)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)
rand_loader = DataLoader(dataset=RandomDataset(input_size, data_size),batch_size=batch_size, shuffle=True)

for data in rand_loader:
    input = data.to(device)
    output = model(input)

  上述代码中,batch_size=32。但是由于使用了DataParallel,在有2个GPU时,一个batch被划分成了2份,也就是tensor.split(16),分别送往两个GPU上。

  또한 첫 번째 호출model.to(device)에서는 첫 번째 GPU 장치에 모델이 로드되고, 첫 번째 호출에서도output = model(input) 때입니다. 모델이 나머지 GPU에 복사될 때 처음으로 순방향 전파가 수행됩니다. DataParallel반복 프로세스는 다음과 같이 요약됩니다.

여기에 이미지 설명을 삽입하세요.

  • 모델은 메인 GPU에서 초기화된 후 각 GPU에 복사됩니다.
  • 메인 GPU는 데이터 배치를 읽은 다음 한 배치의 데이터를 여러 개의 작은 배치로 나누고 이를 다른 GPU에 배포합니다.
  • 각 GPU에서 순방향 계산이 완료되고 출력은gather 손실 계산을 위해 기본 GPU로 전송됩니다.
  • 손실scatter은 각 GPU에 적용되며, 각 GPU는 BP를 통해 기울기를 계산합니다.
  • 각 GPU의 기울기는reduce 기본 GPU로 전송된 다음 모델 가중치가 기본 GPU에서 업데이트됩니다.
  • 다음 반복 전에 기본 GPU는 모델 매개변수broadcast를 다른 GPU로 전송하여 가중치 매개변수 값의 동기화를 완료합니다.

DataParallel매개변수 서버 병렬 아키텍처가 채택되었으며 여기에는 다음과 같은 제한 사항이 있습니다.

  • PS 아키텍처를 사용하면 통신 오버헤드가 높음(각 작업자는 매개변수 서버와 여러 번 통신해야 함)
  • 부하가 불균형하고 메인 GPU의 부하가 높아 GPU 활용도가 부족합니다.
  • 단일 기계 다중 카드 모드만 지원하며 다중 기계 및 다중 카드 교육은 달성할 수 없습니다.

2.2 분산데이터병렬(DDP)

2.2.1 원리

  DDP에는 더 이상 메인 GPU가 없으며, 각 GPU는 동일한 작업을 수행합니다. 추론, 손실 함수 계산, 기울기 계산 등을 각 GPU에서 병렬 및 독립적으로 완료할 수 있어 훈련 효율성과 속도가 향상됩니다.

  DDP 병렬 훈련 구현의 핵심은 모델 간의 기울기 동기화에 있습니다. 이는 각 GPU가 정확히 동일한 기울기를 얻을 수 있도록 전체 감소 통신 작업을 통해 달성됩니다. 각 GPU는 또한 독립적인 최적화 프로그램을 구축합니다. 모델의 초기 상태와 후속 기울기가 동일하므로 각 반복 후 서로 다른 프로세스 간의 모델은 정확히 동일하므로 DDP의 수학적 일관성이 보장됩니다. DDP의 흐름도를 바탕으로 주요 단계를 소개하면 다음과 같습니다.

여기에 이미지 설명을 삽입하세요.

DDP 흐름도. 그 중 Construction은 훈련 시작 전에 한 번만 실행되고, 훈련 중에는 Forward와 Backward만 여러 번 반복됩니다.

  1. Construction단계: DDP에서는 통신 프로토콜 및 전체 프로세스 수를 결정하기 위해 추가 프로세스 그룹 구축 단계(구축)가 필요합니다.
    • 통신 프로토콜은 DDP의 기본 기반이며 병렬 훈련 방법을 결정합니다.
    • 총 프로세스 수(세계 규모): 즉, 훈련에 참여하는 독립적인 병렬 프로세스 수입니다.

      각 프로세스는 필요에 따라 하나 이상의 GPU를 점유할 수 있지만 여러 프로세스가 하나의 GPU를 공유하는 것은 권장되지 않습니다. 이로 인해 잠재적인 성능 손실이 발생할 수 있습니다. 이해를 돕기 위해 이 기사의 모든 예에서는 각 프로세스가 1개의 GPU만 점유한다고 가정하고, 여러 개의 GPU가 점유된 경우 GPU 매핑 관계를 조정하기만 하면 됩니다.
      Construction훈련 시작 전에만 실행되므로 추가 지연이 없습니다

  2. 초기 상태 브로드캐스트
    각 GPU(replica)에서 모델의 복사본을 생성하면 GPU-1의 모델 상태는 다음과 같습니다. 다른 모든 프로세스에 브로드캐스트하여 모든 모델이 동일한 초기 상태를 갖도록 합니다.
  3. 학습 반복: 각 GPU는 할당된 데이터 배치를 수신하고 로컬로(각 GPU 자체에서) 손실 함수의 기울기를 계산합니다(local gradients).
  4. 전체 차수 합계(All-Reduce): 모델 가중치의 일관성을 유지하고 모델 업데이트를 허용하려면 각 GPU의 local gradients가 다음과 같아야 합니다. A global gradients로 요약하면 이 작업을 all-reduce라고 하며 여기에는 reduce-scatterall-gather :
    여기에 이미지 설명을 삽입하세요.
    • Reduce-Scatter

      • 각 노드(복제본)local gradients는 서로 다른 블록 또는 샤드(블록 또는 샤드, 일반적으로 균등하게 분산됨)로 분할된 후 N 노드 간에 N이 수행됩니다. -1 라운드의 데이터 교환 이를 통해 각 노드는 다른 노드의 기울기 정보의 일부를 얻을 수 있습니다.
      • 각 교환 라운드에서는 그래디언트 데이터의 일부를 병합하여 전역 그래디언트 정보의 일부를 얻습니다. 결국 각 복제본은 다른 모든 복제본의 그래디언트 샤드 요약을 보유하게 됩니다(fully reduced data)

      reduced은 샤드를 줄이는 각 병합 작업 라운드를 나타냅니다. fully reduced은 모든 그래디언트 샤드가 결국 하나의 값으로 병합됨을 의미합니다.

    • All-gather

      • 각 복제본은 Reduce-Scatter 단계에서 얻은 완전히 축소된 데이터를 다시 여러 개의 작은 조각으로 나눈 후 다른 노드로 브로드캐스팅합니다. 또한 N-1 라운드의 데이터 교환 후에 각 노드는 다른 노드의 모든 기울기 정보를 얻었습니다.
      • 각 노드는 수집된 모든 그래디언트 정보를 바탕으로global gradients로 요약됩니다.
  5. 다음을 기반으로 모델 매개변수 업데이트global gradients (weight update)
  6. 훈련이 완료될 때까지 3~5단계를 반복합니다.

  Reduce-Scatter 단계에서 그래디언트 데이터를 먼저 여러 개의 샤드로 나눈 후 한 라운드 내에서 방송하지 않고 N-1번 방송해야 하는 이유는 다음과 같습니다. 통신 오버헤드와 메모리 사용량을 줄이는 동시에 훈련 속도와 확장성을 향상시킵니다. 대규모 훈련에서는 GPU 수가 매우 클 수 있으므로 서로 다른 GPU의 local gradients 을 모든 GPU에 동시에 브로드캐스팅하면 많은 양의 데이터 흐름 통신이 발생하고 공간을 차지할 수도 있습니다. 대용량 GPU 메모리. 버킷팅 작업을 통해 한 번에 데이터의 일부만 전송되므로 네트워크 부하와 메모리 부담을 줄일 수 있으므로 훈련을 더 많은 GPU 또는 더 큰 모델 크기로 확장할 수 있습니다. local gradients

2.2.2all-reduce백화

  성능을 최적화하기 위해 DDP는 all-reduce 운영에 대한 보다 심층적인 설계를 수행했습니다. 기울기 계산 과정과 프로세스 간 통신 과정에는 각각 일정 시간이 필요합니다. 통신하기 전에 모델의 모든 매개변수가 기울기를 계산할 때까지 기다리는 것은 분명히 최적이 아닙니다. 아래 그림과 같이 DDP의 설계는 모든 모델 매개변수를 수많은 작은 버킷으로 나누어 버킷 수준에서 설계됩니다all-reduce. 모든 프로세스에서 bucket0의 기울기 계산이 완료되면 즉시 통신이 시작되며, 이 때 bucket1의 기울기 계산은 계속 진행됩니다. 이를 통해 계산 및 통신 프로세스의 시간 중복을 달성하여 DDP 교육을 더욱 효율적으로 만들 수 있습니다.
여기에 이미지 설명을 삽입하세요.
  DDP 백엔드 통신은 CPP가 작성한 다양한 프로토콜에 의해 지원됩니다. 프로토콜마다 통신 운영자 지원이 다르며 개발 중 필요에 따라 선택할 수 있습니다.

여기에 이미지 설명을 삽입하세요.

2.2.3 DP 비교

  DDP의 장점은 각 GPU 노드에 부하가 분산되고 GPU 수에 관계없이 통신 비용이 일정하므로 훈련이 더 빠르며 단일 머신 멀티 카드 및 멀티 머신 멀티 카드도 지원합니다. . 단일 GPU에 모델을 로드할 수 없으면 처리할 수 없다는 한계가 있습니다. 다음은 DP와 DDP의 차이점을 설명하기 위한 실험입니다.

  • Hardware: 2x TITAN RTX 24GB each + NVlink with 2 NVLinks (NV2 in nvidia-smi topo -m).
  • Software: pytorch-1.8-to-be + cuda-11.0 / transformers==4.3.0.dev0.

其中一个实验使用NCCL_P2P_DISABLE=1禁用 NVLink 功能,整个测试结果为:

Type	NVlink	Time
2:DP	 Y	  110s
2:DDP	 Y	  101s
2:DDP	 N	  131s

  可以看到,DP 比使用 NVlink 的 DDP 慢 ~10%,但比不使用 NVlink 的 DDP 快 ~15%。真正的区别将取决于每个 GPU 需要与其他 GPU 同步多少数据——要同步的数据越多,slow link就越会阻碍整体runtime。更多内容详见《Efficient Training on Multiple GPUs》

  NVLink(NVIDIA NVLink)是由NVIDIA开发的一种高速互连技术,允许不同GPU之间以非常高的带宽进行直接通信,而无需经过较慢的PCI-E总线。

2.3 ZeRO Data Parallelism

参考《ZeRO & DeepSpeed: New system optimizations enable training models with over 100 billion parameters》《大模型分布式训练策略:ZeRO、FSDP》《大模型-LLM分布式训练框架总结》

  数据并行和模型并行都保持了整个训练过程中所需的所有模型状态,但并不是所有时候这都是必需的。例如,仅在某个层的正向传播和反向传播期间才需要与每个层对应的参数。

  ZeRO-DP是一种改进的数据并行性方法,它通过对参数(包括优化器状态、梯度和模型参数)进行分片来消除内存冗余,使得每个GPU仅保存部分参数及相关状态,提高了内存效率;同时还通过在训练过程中使用动态通信来保持计算和通信效率。

여기에 이미지 설명을 삽입하세요.

图1:使用Adam优化器进行混合精度训练时,ZeRO-DP优化的三个阶段中,每个设备的model states内存消耗。

参数解释:

  • Baseline:未优化的基线
  • Ψ:模型大小,上图假设模型参数为Ψ=75亿
  • K:存储优化器状态要消耗的内存倍数,上一节讲过,对于混合精度的Adam优化器而言,K=12
  • N d N_d Nd:数据并行度。基于Adam优化器的混合精度训练,数据并行度为Nd=64(即64个GPU)

上图展示了ZeRO-DP对数据并行优化的三个阶段:

  1. 优化器状态分割( P o s P_{os} Pos):
    在每个gpu中保存全部的参数和梯度,但是只保存1/Nd的优化器变量。通过将优化器状态进行分割,实现4倍的内存减少,同时保持与DP相同的通信量。

  2. 梯度分割( P o s + g P_{os+g} Pos+g):
    每个gpu中只保存1/Nd的梯度
    ,实现8倍的内存减少,并保持与DP相同的通信量。

  3. 参数分割( P o s + g + p P_{os+g+p} Pos+g+p):
    每个gpu中只保存1/Nd的参数
    ,实现64倍的内存减少,通信量会略微增加50%。作者通过用少量的计算的成本和通信成本换来了大幅的内存节省。

  Zero DP与DataParallel类似,不同之处在于,每个 GPU 不是复制完整的模型参数、梯度和优化器状态,而是只存储其中的一部分。假设有一个具有 3 层(La、Lb 和 Lc)的简单模型,其中每层有 3 个参数。例如,图层 La 的权重为 a0、a1 和 a2:

La | Lb | Lc
---|----|---
a0 | b0 | c0
a1 | b1 | c1
a2 | b2 | c2

  如果我们有 3 个 GPU,ZeRO-DP 会将模型拆分为 3 个 GPU,在某种程度上,这与张量并行性的水平切片相同,与垂直切片相反(垂直切片将整个层组放在不同的 GPU 上),如下所示:

GPU0:
La | Lb | Lc
---|----|---
a0 | b0 | c0

GPU1:
La | Lb | Lc
---|----|---
a1 | b1 | c1

GPU2:
La | Lb | Lc
---|----|---
a2 | b2 | c2

因为采用数据并行,每个GPU会获得一个mini-batch:

x0 => GPU0
x1 => GPU1
x2 => GPU2

La层:3 个 GPU 中的每一个都可以重建完整的张量,并使用自己的小批量进行前向传递。一旦计算完成,不再需要的数据就会被删除 (只在计算过程中使用)

  • 在 GPU0 上,x0 批次需要 a0、a1、a2 参数来执行其通过层的正向路径,但 GPU0 只有 a0。它将从 GPU1 获得 a1,从 GPU2 获得 a2,将模型的所有部分组合在一起。
  • GPU1将从 GPU0 和 GPU2获取参数a0,a2来计算x1 批次;
  • GPU2将从 GPU0 和 GPU1获取参数a0,a1来计算x2 批次;

然后对 Lb 层、Lc层重复整个过程完成前向计算,再反过来完成后向计算。完整工作流程如下:

  另外Zero除了Zero DP还包括Zero-R,后续改进工作还有ZeRO-OffloadZeRO-Infinity,分别实现将GPU的数据和计算卸载到CPU和NVMe内存,能在有限资源下能够训练前所未有规模的模型,而无需对模型代码进行重构。详情见另一篇博客《大模型分布式训练策略:ZeRO、FSDP》

  ZeRO实现DeepSpeed(实现Zero DP、Zero-R、ZeRO-Offload和ZeRO-Infinity的全部功能Accelerate集成transformers集成

三、流水线并行(Pipeline Parallelism)

参考:《Efficient Training on Multiple GPUs》《一文捋顺千亿模型训练技术:流水线并行、张量并行和3D并行》

3.1 模型并行(MP,Naive Model Parallelism)

  当一个模型大到单个GPU无法训练时,最朴素(Naive)的想法是对模型的层进行划分。Naive Model Parallelism(Naive MP)通过使用.to()来将特定层分配到特定的GPU上。当数据通过这些层传递时,它会被移动到与该层相同的GPU上,而其他层保持不变。

  Naive MP也被称作Vertical MP,因为模型通常是以垂直方式切分的。下面以一个4层的序列模型为例进行介绍:

o u t p u t = L 4 ( L 3 ( L 2 ( L 1 ( i n p u t ) ) ) ) output=L4(L3(L2(L1(input)))) output=L4(L3(L2(L1(input))))

将其按层划分至两个GPU上:

  • GPU1负责计算: i n t e r m e d i a t e = L 2 ( L 1 ( i n p u t ) ) intermediate=L_2(L_1(input)) intermediate=L2(L1(input))
  • GPU2负责计算: o u t p u t = L 4 ( L 3 ( i n t e r m e d i a t e ) ) output=L_4(L_3(intermediate)) output=L4(L3(intermediate))

GIF

  整个朴素层并行前向传播和后向传播的过程如上图所示。GPU1执行前向传播,并将激活(activations)缓存下来。然后将 L 2 L_2 L2 层的输出intermediate发送给GPU2GPU2完成前向传播和loss计算后,开始反向传播。当GPU2完成反向传播后,会将 L 3 L_3 L3的梯度返还给GPU1GPU1完成最终的反向传播。

通过以上过程可以发现Naive MP存在一些缺点:

  • 低GPU利用率: 在任意时刻,有且仅有一个GPU在工作,其他GPU都是空闲的。

  • 计算和通信没有重叠。在发送前向传播的中间结果(FWD)或者反向传播的中间结果(BWD)时,GPU也是空闲的。

  • 高显存占用。GPU1需要保存整个batch的所有激活,直至最后完成参数更新。另外Shared embeddings 也需要在不同GPU之间来回复制。如果batch size很大,这将对显存带来巨大的挑战。

  Pipeline Parallelism (PP)几乎与Naive MP相同,但通过将传入的batch分成micro-batches并人工创建一个pipeline来解决了GPU空闲问题,允许不同的GPU同时参与计算过程。

3.2 GPipe

论文《GPipe: Easy Scaling with Micro-Batch Pipeline Parallelism》《图解大模型训练之:流水线并行Gpipe》

3.2.1 工作原理

假设有4个GPU,并将模型按层划分为4个部分。朴素层并行的过程为:
여기에 이미지 설명을 삽입하세요.
  如图,阴影部分所表示的时间段里,总有GPU在空转。在Gpipe中,将阴影部分定义为bubble。假设有 K块GPU,而单块GPU上做一次forward和backward的时间为: t f b = ( t f + t b ) t_{fb} = (t_{f} + t_{b}) tfb=(tf+tb) 。则:

  • 图中灰色长方形的整体面积为: K ∗ K t f b K*Kt_{fb} KKtfb
  • 图中实际在做forward和backward的面积为: K t f b Kt_{fb} Ktfb
  • 图中阴影部分的面积为: ( K − 1 ) K t f b (K-1)Kt_{fb} (K1)Ktfb
  • 图像阴影部分的占比为: ( K − 1 ) / K (K-1)/K (K1)/K

  则我们定义出bubble部分的时间复杂度为: O ( K − 1 K ) O(\frac{K-1}{K}) O(KK1)当K越大,即GPU的数量越多时,空置的比例接近1,即GPU的资源都被浪费掉了。因此这个问题肯定需要解决。

여기에 이미지 설명을 삽입하세요.

  流水线并行的核心思想是:在模型并行的基础上,进一步引入数据并行的办法,即把原先的数据再划分成若干个batch,送入GPU进行训练。未划分前的数据,叫mini-batch。在mini-batch上再划分的数据,叫micro-batch

  图c中,第一个下标表示GPU编号,第二个下标表示micro-batch编号。假设我们将mini-batch划分为M个。假设单个GPU上完成前向传播或者后向传播的面积为1(也就是上图中的单个小方块面积为1)。上图中的总长度为2(M+K-1),宽度为K,则总面积为2K(M+K-1)。其中,彩色小方块占用的面积表示GPU执行的时间,为 2KM;空白处bubble面积的占比为 1 − 2 K M 2 K ( M + K − 1 ) 1-\frac{2KM}{2K(M+K-1)} 12K(M+K1)2KM= K − 1 K + M − 1 \frac{K-1}{K+M-1} K+M1K1

  PP的目标是平衡计算强度(即每个micro-batch的大小)和最小化Pipeline中的空闲时间,以提高效率。Gpipe通过实验证明,当 M > = 4 K M>=4K M>=4K 时,bubble产生的空转时间占比对最终训练时长影响是微小的,可以忽略不计,但这可能会导致更多的通信开销。

  将batch切好,并逐一送入GPU的过程,就像一个流水生产线一样(类似于CPU里的流水线),因此也被称为Pipeline Parallelism。

3.2.2 active checkpoint

  增大batch size就会线性增大需要被缓存激活的显存需求。在GPipe中,GPU需要在前向传播至反向传播这段时间内缓存激活(activations)。以GPU0为例,micro-batch1的激活需要从timestep 0保存至timestep 13。随着模型的增加,每块GPU中存储的中间结果也会越大。

  GPipe为了解决显存的问题,使用了为re-materalization技术,后称为active checkpoint。 该技术不需要缓存所有的激活,而是在反向传播的过程中重新计算激活。这降低了对显存的需求,但是增加了计算代价(时间换空间),图例如下:

  每块GPU上,我们只保存来自上一块的最后一层输入z,其余的中间结果我们算完就废。等到backward的时候再由保存下来的z重新进行forward来算出activations。

如果你使用Pytorch提供的pipeline接口,其中有一个参数叫checkpoint,就是用来做这一项的。

3.2.3 实验结果

여기에 이미지 설명을 삽입하세요.
Gpipe分别在AmoebaNet(图像)和Transformer(自然语言)两个大模型上做了实验。

  • Naive:单卡
  • Pipeline-N:re-materalization + N卡。
  • AmeobaNet-D和Trasformer-L:超参数的量
  • # of Model Parameter表示模型的参数量
  • Total Model Parameter Memory:模型参数所占内存大小
  • Peak Activation Memory:峰值时中间结果大小。可以发现,中间结果占据的内存大小是相当可观的。

从实验结果里,我们可以发现:

  • Transformer上,Gpipe基本实现了模型大小(参数量)和GPU个数之间的线性关系。例如从32卡增到128卡时,模型的大小也从21.08B增加到82.9B,约扩4倍
  • AmoebaNet而言,却没有完全实现线性增长。例如从4卡到8卡,模型大小从1.05B到1.8B,不满足2倍的关系。本质原因是AmoebaNet模型在切割时,没有办法像Transformer一样切得匀称,保证每一块GPU上的内存使用率是差不多的。因此对于AmoebaNet,当GPU个数上升时,某一块GPU可能成为木桶的短板。

여기에 이미지 설명을 삽입하세요.
  上图是使用不同数量的GPU数 K 和不同micro-batches数M 在TPUs上进行GPipe Normalized training throughput(吞吐量)。性能随M提高。当 M ≥ K 时,Transformer模型的加速器数量几乎呈线性加速。如果必要,批次大小会根据内存进行调整。

3.3 PipeDream

论文《PipeDream: Generalized Pipeline Parallelism for DNN Training》《PipeDream: Fast and Efficient Pipeline Parallel DNN Training》《PipeDream: 数据并行+流水线》

여기에 이미지 설명을 삽입하세요.

Model parallel training

여기에 이미지 설명을 삽입하세요.
GPipe’s inter-batch parallelism,频繁的管道刷新导致了 idle time 的增加

  GPipe使用 active checkpoint技术,在反向传播的过程中重新计算激活,降低了对显存的需求。在将mini-batch划分为Mmicro-batch后,如果 M 较小,则由于重新计算开销和频繁的管道刷新,其在硬件效率方面可能会受到影响。

  GPipe需要等所有的microbatch前向传播完成后,才会开始反向传播。PipeDream则是当一个microbatch的前向传播完成后,立即进入反向传播阶段。理论上,反向传播完成后就可以丢弃掉对应microbatch缓存的激活。由于PipeDream的反向传播完成的要比GPipe早,因此也会减少显存的需求。

  下图是PipeDream的调度图,4个GPU和8个microbatchs。蓝色的方块表示前向传播,绿色表示反向传播,数字则是microbatch的id。
여기에 이미지 설명을 삽입하세요.

An example PipeDream pipeline with 4 workers

  PipeDream在bubbles上与GPipe没有区别,但是由于PipeDream释放显存的时间更早,因此会降低对显存的需求。

  Transformer는 기본적으로 간단한 파이프라인 병렬 처리를 지원합니다. 이렇게 하려면 device="auto"을 사용하여 모델을 로드하면 자동으로 해당 GPU에 다양한 레이어가 배치됩니다. 추론을 위한 대형 모델 처리를 참조하세요. . 그러나 이 단순한 파이프라인 병렬 처리는 매우 효과적이지만 유휴 GPU 문제를 해결하지는 못합니다.

4. 텐서 병렬성

참고 논문:"Megatron-LM을 사용한 GPU 클러스터에 대한 효율적인 대규모 언어 모델 훈련", "파이프라인 병렬성, 텐서 병렬성, 3D 병렬성 등 수천억 개 모델의 훈련 기술을 요약한 기사가 하나 있습니다."

  모델 병렬성과 데이터 병렬성 외에도 더 작은 규모의 병렬성 방법인 텐서 병렬성도 있습니다. Transformer의 주요 구성 요소는 완전히 연결된 레이어와 어텐션 메커니즘이며, 그 핵심은 행렬 곱셈입니다. 텐서 병렬화의 핵심은 단일 레이어 매개변수가 너무 큰 문제를 해결하기 위해 행렬 곱셈을 분할하는 것입니다.

  Tensor Parallelism 에서 각 GPU는 텐서 조각을 처리하고 이를 필요로 하는 작업에 대해서만 전체 텐서를 집계합니다. 이 장의 개념에 대한 자세한 내용은Megatron-LM 논문 "Efficient Large-Scale Language Model Training on Megatron-LM을 사용하는 GPU 클러스터 》.

4.1 1D 텐서 병렬성

  완전 연결 레이어의 경우 가장 간단한 아이디어는 블록 행렬 계산 규칙을 ​​사용하여 결과의 ​​일관성을 얻는 것입니다. 다음에서는 행렬 곱셈을 사용하여 1D 텐서 병렬성을 이해하는 데 도움을 줍니다.

  • 열 평행: 가중치 행렬의 행과 열을 n 부분으로 나누고 행렬 곱셈은 다음과 같이 표현됩니다.
    X A = X [ A 1 , A 2 , … , An ] = [ X A 1 , X A 2 , … , X A n ]XA=X[A1,2,,n]=[XA1,XA2,,XAn]
  • 행 병렬성: 가중치 행렬과 입력 행렬을 모두 나누고 행렬 곱셈은 다음과 같이 표현됩니다.
    X A = [ X 1 , X 2 , … , X n ] [ A 1 A 2 … An ] = X 1 A 1 + X 2 A 2 + ⋯ + X n A n \ \dots \\ A_n \end{array} \right] = X_1A_1+X_2A_2+\dots+X_nA_n XA=[X1,엑스2,,엑스n] 12n =엑스11+엑스22++엑스nn

  행 병렬이든 열 병렬이든 각 부분을 계산한 후 한 번만 통신하면 되는 것을 알 수 있다. 단지 열 병렬 처리는 통신 결과를 연결하는 반면 행 병렬 처리는 통신 결과를 추가합니다.

4.2 MLP 병렬성

   트랜스포머의 주요 부분은 MLP 레이어입니다. 아래에서는 2계층 풀링크 계층을 시뮬레이션합니다. X와 Y가 입력과 출력이고 A와 B가 완전히 연결된 두 레이어의 가중치라고 가정하면 다음과 같습니다.

Y = G e L U ( G e L U ( X A ) B ) Y={GeLU(GeLU(XA)B)}그리고=GeLU(G eLU(XA )B)

使用列会行可电影:
Y = G e L U ( G e L U ( X A ) B ) = GeLU ( GeLU ( [ X A 1 , X A 2 , … , X A n ] ) [ B 1 B 2 … B n ] ) = GeLU ( [ GeLU ( X A 1 ) B 1 , GeLU ( X A 2 ) B 2 , … , GeLU ( X A n ) B n ] ) = [ GeLU ( GeLU ( X A 1 ) B 1 ) , … , GeLU ( GeLU ( X A n ) B n ) ] \begin{aligned} \text Y={GeLU(GeLU(XA)B)}&=\text{GeLU}\Big( \ text{GeLU}([XA_1,XA_2,\dots,XA_n]) \left[ \begin{array}{l} B_1 \\ B_2 \\ \dots \\ B_n \end{array} \right] \Big) \ \ &=\text{GeLU}([\text{GeLU}(XA_1)B_1,\text{GeLU}(XA_2)B_2,\dots,\text{GeLU}(XA_n)B_n]) \\ & = [\text{GeLU}(\text{GeLU}(XA_1)B_1),\dots,\text{GeLU}(\text{GeLU}(XA_n)B_n)] \end{aligned} 그리고=GeLU(G eLU(XA )B)=GeLU(GeLU([X1,XA2,,XAn]) 12n )=GeLU([GeLU(X1)B1,GeLU(XA2)B2,,GeLU(XAn)Bn])=[GeLU(GeLU( XA1)B1),,GeLU(GeLU(X n)Bn)]

  위의 수식을 보면 알 수 있습니다. A와 B를 미리 나누어서 독립적으로 계산하면 G e L U ( GeLU ( X A i ) B i ) {GeLU}(\text{GeLU} (XA_i) B_i) GeLU(GeLU (XA나는)B나는) 통신하기 전에. 즉, 이 예에는 완전히 연결된 두 개의 레이어가 있지만 최종 결과를 얻기 전에만 통신하면 됩니다. **완전히 연결된 레이어가 여러 개 쌓인 경우 최종 출력 중에 하나의 통신만 필요합니다. **이 원칙을 사용하면 모든 깊이의 MLP를 업데이트할 수 있습니다.

여기에 이미지 설명을 삽입하세요.

4.3 다중 헤드 어텐션 병렬성

다중 헤드 어텐션의 각 헤드는 본질적으로 독립적이므로 각 헤드는 병렬로 작동할 수 있습니다.

여기에 이미지 설명을 삽입하세요.
  TP에는 매우 빠른 네트워크가 필요하므로 여러 노드에 걸쳐 TP를 수행하는 것은 권장되지 않습니다. 실제로 노드에 GPU가 4개 있는 경우 최대 TP 수준은 4입니다. 또한 DeepSpeed에서는 TP를 텐서 슬라이싱이라고 합니다.

4.4 구현

  • 메가트론-LM
  • 병렬 형성자(현재 추론만 사용 가능)
  • SageMaker: 보다 효율적인 처리를 위해 TP와 DP를 결합하지만 AWS에서만 사용 가능
  • OSLO에는 Transformer 기반 텐서 병렬 구현이 있습니다
  • 가속화: Megatron-LM과 통합된 텐서 병렬 처리.

  Transformer는 모델 아키텍처를 작성하는 특정 방법이 필요하기 때문에 기본적으로 텐서 병렬 처리를 지원하지 않습니다. 텐서 병렬 방식으로 모델을 작성하는 데 관심이 있다면 언제든지 TGI(텍스트 생성 추론) 라이브러리를 확인해 보세요.

5. 3D 병렬성

《DeepSpeed와 Megatron을 사용하여 대규모 생성 언어 모델인 Megatron-Turing NLG 530B 훈련》

5.1 데이터 병렬성 + 파이프라인 병렬성

  DeepSpeed ​​​​파이프라인 튜토리얼에서는 DP와 PP를 사용하는 방법을 보여줍니다. 아래 그림과 같이 DP 랭크 0은 GPU2를 볼 수 없고, DP 랭크 1은 GPU3을 볼 수 없으며, DP의 경우 GPU가 2개만 있는 것처럼 GPU 0과 1만이 데이터를 제공합니다. GPU0은 PP를 사용하여 로드의 일부를 GPU2로 오프로드하고, 마찬가지로 GPU1은 PP를 사용하여 로드의 일부를 GPU3으로 오프로드합니다. 각 차원에는 최소 2개의 GPU가 필요하므로 여기서는 최소 4개의 GPU가 필요합니다.
여기에 이미지 설명을 삽입하세요.

5.2 데이터 병렬성 + 파이프라인 병렬성 + 텐서 병렬성

연결:Megatron-DeepSpeed,OSLO

  3D 병렬성은 데이터 병렬성(DP), 텐서 병렬성(TP), 파이프라인 병렬성(PP)으로 구성됩니다. 다음 그림은 《3D 병렬성: 1조 매개변수 모델로 확장》을 참조합니다. 각 차원에는 최소 2개의 GPU가 필요하므로 최소 8개가 필요합니다. 여기 GPU입니다. 이 기능은 DeepSpeed , Megatron-LM , OSLO 달성합니다. , SageMaker, Varuna

여기에 이미지 설명을 삽입하세요.
  DeepSpeed의 주요 기능 중 하나는 DP의 슈퍼 확장인 ZeRO로, 독립적인 기능으로 PP나 TP가 필요하지 않습니다. 그러나 PP 및 TP와 함께 사용할 수 있습니다.

  Zero-DP를 PP(또는 TP)와 함께 사용하는 경우 일반적으로 ZeRO 1단계(최적화 샤딩)만 활성화됩니다. 이론적으로 이 시점에서 ZeRO 2단계(그라디언트 샤딩)를 활성화할 수 있지만 이는 성능에 부정적인 영향을 미칩니다.

  • DP 자체는 배치의 데이터를 분할했으며 PP는 mini-batchmicro-batch로 추가로 분할하여 각 micro-batch그래디언트 정보를 전송하려면 추가적인 감소-산란 작업이 필요합니다. 분할이 너무 많으면 통신 오버헤드가 증가하고 컴퓨팅 효율성이 감소합니다.
  • PP 자체는 각 단계에서 그라디언트 정보의 일부만 처리하면 되므로 그라디언트의 크기를 줄였습니다. 이 경우 그라디언트 샤딩을 추가로 사용하여 그라디언트 크기를 줄이면 그라디언트가 이미 상대적으로 작기 때문에 메모리가 크게 절약되지 않을 수 있습니다.

같은 이유로 현재 ZeRO 3단계(파라미터 샤딩)를 활성화하는 것은 노드 간 통신이 더 많아지기 때문에 적절하지 않습니다.

6. 요약: 병렬 전략을 선택하는 방법

6.1 단일 노드 병렬화 전략

  1. 단일 GPU로 전체 모델에 적합
    • DDP - 분산 데이터 병렬
    • ZeRO: 사용법과 구성에 따라 이 방법이 더 빠를 수도 있고 그렇지 않을 수도 있지만 시도해 볼 가치가 있습니다.
  2. 단일 GPU는 전체 모델을 로드할 수 없습니다.
    • 파이프라인병렬(PP)
    • 텐서병렬(TP)

매우 빠른 노드 간 연결(예: NVLINK 또는 NVSwitch)의 경우 세 가지 전략(PP, ZeRO, TP) 모두 유사한 성능을 제공해야 합니다. 그러나 이것이 없으면 PP는 TP나 ZeRO보다 빠릅니다. 또한 TP는 단일 노드 내에서 사용하기에 적합합니다. 즉, TP 크기 <= 노드당 GPU 수입니다.

  1. 단일 GPU는 모델의 가장 큰 레이어를 로드할 수 없습니다.

6.2 다중 노드 병렬화 전략

  1. NVLINK 또는 NVSwitch와 같은 빠른 노드 간 연결을 사용하는 경우 다음을 사용하는 것이 좋습니다.
    • ZeRO: 모델 수정이 거의 필요하지 않습니다.
    • PP, TP, DP의 조합을 사용하면 이 접근 방식으로 의사소통이 줄어들지만 모델에 상당한 변경이 필요합니다.
  2. 노드 간 연결이 느리고 GPU 메모리가 여전히 부족한 경우 PP, TP, DP 및 ZeRO를 조합하여 사용하십시오.

Guess you like

Origin blog.csdn.net/qq_56591814/article/details/134099476