大模型分布式训练策略:ZeRO、FSDP

文章目录

    • 一、ZeRO(零冗余优化器)
      • 1.1 背景
      • 1.2 深度学习内存消耗分析
      • 1.3 主要方法
        • 1.3.1 `ZeRO-DP`优化模型状态内存
        • 1.3.2 `ZeRO-R`优化残余状态内存
      • 1.4 总结
      • 1.5 官方视频:ZeRO & Fastest BERT,提高 DeepSpeed 深度学习训练的规模和速度(有空再补)
    • 二、FSDP: 使用更少的GPU进行更快的AI训练
      • 2.1 简介
      • 2.2 数据并行(Data Parallelism)
      • 2.3 sharding weight update
      • 2.4 optimizer states all-gather
      • 2.5 `FSDP`的使用
      • 2.6 FSDP和ZeRO的异同

DeepSpeed开源项目、DeepSpeed 官网 、huggingface DeepSpeed文档、知乎deepspeed入门教程

一、ZeRO(零冗余优化器)

  • 论文:ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
  • 《论文解读系列第十三篇:ZeRO——面向万亿级参数的模型训练方法》

  现有普遍的数据并行模式下的深度学习训练,每一台机器都需要消耗固定大小的全量内存,这部分内存和并不会随着数据的并行而减小,因而,数据并行模式下机器的内存通常会成为训练的瓶颈。这篇论文开发了一种新颖的解决方案Zero Redundancy Optimizer (ZeRO),主要用于解决数据并行状态下内存不足的问题,使得模型的内存可以平均分配到每个gpu上,每个gpu上的内存消耗与数据并行度成反比,而又基本不影响通信效率。

  按照作者的优化方案,在64个GPU的数据并行的情况下,我们可以将一个75亿个参数的模型的内存消耗从每GPU 120GB减少到每GPU 1.9GB,同时通信开销变为原来的1.5倍;或者每GPU 16.6GB而不增加通信开销。

  通过这样的方式,我们可以以同样的内存运行更大的模型,或者将原来必须使用模型并行才能运行的大模型使用数据并行方式进行训练,以减少模型并行的额外开销。

1.1 背景

  深度学习领域的模型越来越大,这显著地增加了模型的准确性。在NLP领域,像Bert-large(0.3B)、GPT-2(1.5B)、Megatron-LM(8.3B)、T5(11B)这样的transformer大型模型已经出现。然而,要继续扩展模型大小(从数十亿到数万亿参数),我们遇到了训练这些模型的挑战——它们无法容纳在单个设备(GPU或TPU)的内存中,而且仅仅添加更多设备也无法有效扩展训练。

  现有的解决方案存在一些限制。基本的数据并行(DP,Data Parallelism)不会减少每个设备的内存占用,对于具有超过1.4B参数的模型,它在当前32GB内存的GPU上会耗尽内存。其他解决方案,如流水线并行(PP,Pipeline Parallelism)、模型并行(MP,Model Parallelism)、CPU-Offloading等,虽然有一些作用,但它们都在功能性、可用性以及内存和计算/通信效率之间做出了权衡,而这些方面对于大规模的高速训练都是非常重要的。

  目前模型并行(MP)是其中一个有希望的解决方案,现在的文献中的大模型都采用了MP,但MP的扩展能力也有限。MP将模型在垂直方向上分割,将每层的计算和参数分配到多个设备上,这需要在每一层之间进行大量通信。因此,MP在单个节点内运行效果良好,但在单个节点之外的效率迅速下降。我们使用Megatron-LM在两个DGX-2节点上测试了一个40B参数模型,观察到每个V100 GPU的性能只有5Tflops(不到硬件峰值的5%)。

1.2 深度学习内存消耗分析

  一个15亿参数的GPT-2模型在16位训练中的权重为3GB内存。但在使用Tensorflow或Pythorch这样的框架时,它无法在具有32GB内存的单个GPU上进行训练。那么在训练过程中,这些内存都消耗在了哪里呢?主要在是以下两个部分:

  1. model states(模型状态):大部分内存用于存储这一部分对象,其中包括优化器参数(比如Adam中的动量和方差)、梯度、模型参数等。

  通常情况下,深度学习模型中的参数和激活值使用32位浮点数(fp32)来表示,这提供了高的数值精度,但也需要更多的内存和计算资源。目前在NVIDIA GPU上训练大型模型的最先进方法是使用混合精度(fp16/32)训练。即在前向传播和反向传播过程中使用fp16(半精度)表示的参数和激活值进行计算;而在计算和应用参数更新时,通常需要将这些值转换回32位浮点数来进行计算,以避免数值精度损失,所以混合精度训练中需要同时保留fp16和fp32副本。

  以Adam优化器为例,Adam需要存储动量和方差两种优化器状态,以计算参数更新;此外,还需要存储梯度和模型权重本身的信息。所以对一个具有Ψ个参数的模型进行混合精度训练时,其内存需求为:

  • fp16格式的参数和梯度,都是字节
  • fp32格式的优化器的状态信息,包括参数、动量和方差,都是字节

  我们用来表示存储优化器状态信息所需要的内存数,则对于混合精度Adam优化器而言,总的内存需求为(2+2+k)Ψ=(2+2+4+4+4)Ψ=16Ψ字节,所以对于一个拥有15亿参数的模型(GPT-2),其内存需求至少为24GB,远远超过只保存fp16参数所需的3GB内存。

  1. residual states(残余状态):剩余内存用于存储这部分对象,其中包括激活函数、临时缓冲区、不可用的碎片化内存。

  例如对于一个拥有15亿参数的GPT-2模型,在序列长度为1K和批量大小为32的情况下,即使使用激活值检查点(Activation checkpointing)技术来减少其内存占用,但最终模型的激活值内存消耗还是有8GB左右。如果是1000亿参数的模型,这部分内存消耗将达到60G。另外用于存储中间结果的临时缓冲区会消耗6GB内存(GPT-2)。

  此外,如果没有足够的连续内存来满足内存请求,即使总可用内存大于请求的内存量,请求内存的操作也会失败,因为在训练非常大的模型时,会发生显著的内存碎片化,导致在一些极端情况下,即使仍然有超过30%的内存可用,也会出现内存不足的问题。

  本文提出的ZeRO,就从这两个方面入手,使用ZeRO-DPZeRO-R方法分别优化其占用的内存。

1.3 主要方法

1.3.1 ZeRO-DP优化模型状态内存

在深度学习训练中,模型状态通常需要占用大量内存。现有的解决方法有两种:

  • 数据并行性(DP):

    • 优点:具有良好的计算和通信效率,适用于多设备训练
    • 缺点:内存冗余严重。因为它在所有数据并行进程之间复制整个模型状态,导致冗余的内存消耗。
  • 模型并行性(MP):

    • 优点:内存效率较高,通过分割模型状态来减少内存消耗。
    • 缺点:计算和通信效率不佳,因为垂直切分模型会导致计算过于细粒度化,不利于规模化训练。

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

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

大模型分布式训练策略:ZeRO、FSDP_第1张图片

图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、FSDP_第2张图片
  上图显示了在不同数据并行度下(不同数量的GPU),每个设备内存消耗情况。可见,当启用了所有三个阶段时,ZeRO可以在仅使用1024个NVIDIA GPU的情况下训练具有万亿参数的模型(使用1024个GPU训练1T Model时, P o s + g + p P_{os}+g+p Pos+g+p =15.6GB,当时的V100足有32GB)。

大模型分布式训练策略:ZeRO、FSDP_第3张图片
  上表是不同模型并行度和gpu下可实现的最大模型(以参数计量),最右边是作者的实现的测量值,左边是理论值。因此,这说明作者提出的内存计算是基本可靠的。

1.3.2 ZeRO-R优化残余状态内存

  在ZeRO-DP提高了模型状态内存效率之后,主要消耗在激活值、临时缓冲区以及无法使用的内存碎片这三个方面的剩余内存成为次要的内存瓶颈。为了解决这个问题,我们开发了ZeRO-R来进行优化:

  1. 通过激活值分区来优化激活值内存
    对于激活函数值,一般使用checkpoint技术来优化存储,但这对于大模型仍旧是不足的。ZeRO-R通过激活值分区来识别和删除现有MP方法中的激活值复制,从而优化激活值内存。同时,它还会在适当的情况下将激活值卸载到CPU上,以释放GPU内存(CPU-offload)

  2. 恒定临时缓冲区大小,以在内存和计算效率之间取得平衡。
    之前的缓冲区的大小通常会随着模型的规模变化而变化,但ZeRO-R采用了固定大小的缓冲区,这样可以防止随着模型规模的增加而导致缓冲区过大。同时ZeRO-R会保持缓冲区足够大,以保证计算效率。

  3. 根据张量的不同生命周期来管理内存,以防止内存碎片化
    训练时产生的内存碎片化,是由于不同张量的生命周期差异引起的。内存碎片化可能导致内存分配失败,即使总内存足够也可能出现问题,因为无法获得足够的连续内存。

1.4 总结

  ZeRO(Zero Redundancy Optimizer)是一种用于深度学习模型训练的内存优化方法,它的主要目标是减少模型训练过程中的内存占用,从而允许训练更大规模的模型。以下是ZeRO的具体训练方法:

  1. ZeRO-DP(Zero Redundancy Optimizer - Data Parallelism)ZeRO-DP通过分区而不是复制模型状态来消除数据并行进程之间的内存冗余。它保持了数据并行的计算粒度和通信量,并使用动态通信计划来保持计算/通信效率。ZeRO-DP有三个主要优化阶段——优化器状态分区、梯度分区和参数分区。这些阶段可以显著减少每个设备的内存占用,使其能够在有限的GPU内存上训练大规模模型。

  2. ZeRO-R(Zero Redundancy Optimizer - Residual States)ZeRO-R用于优化模型训练过程中的剩余内存占用,包括激活函数、临时缓冲区和内存碎片。它通过激活函数分区来减少激活内存的占用,还会根据不同张量的生命周期来主动管理内存,以防止内存碎片化。

  3. 混合精度训练(Mixed-Precision Training)ZeRO通常与混合精度训练结合使用,参数和激活函数以fp16(半精度)存储,以利用现代GPU上的高吞吐量张量核心单元。ZeRO通过减少模型状态的内存占用,有助于更有效地实施混合精度训练。

  4. ZeRO与模型并行(Model Parallelism)的结合:虽然ZeRO减少了数据并行的内存占用,但在某些情况下,模型并行仍然有用,它可以用于减少激活内存的占用,特别是对于非常大的模型。 二者结合可以在训练大型模型时实现更好的内存管理和训练效率。

  总之,ZeRO采用了一系列内存优化策略,以实现在有限的硬件资源下训练超大规模的深度学习模型。

  有关训练更具体的内容,可以参考论文第7章Communication Analysis of ZeRO-DP和第8章Communication Analysis of ZeRO-R

1.5 官方视频:ZeRO & Fastest BERT,提高 DeepSpeed 深度学习训练的规模和速度(有空再补)

  • YouTube视频:ZeRO & Fastest BERT: Increasing the scale and speed of deep learning training in DeepSpeed

  人工智能的最新趋势是更大的自然语言模型提供更好的准确性;然而,由于成本、时间和代码集成的难易程度,较大的模型很难训练。为了通过提高全球模型开发人员的规模、速度、成本和可用性来推进大型模型训练,微软于 2020 年 2 月开源了 DeepSpeed 库。过使用名为 ZeRO的内存优化系统,DeepSpeed 可以有效地训练具有 100-2000 亿个参数的模型,速度比最先进的技术快 10 倍。

在本次网络研讨会中,DeepSpeed 团队将讨论 :

  • DeepSpeed 功能、速度和规模优化以及未来路线图
  • 如何使用 DeepSpeed 训练您自己的模型以及其他流行模型(例如 BERT 和 GPT-2)
  • 深入探讨 ZeRO 优化器背后的技术和即将推出的功能,以及研究人员利用这些突破创建了Turing-NLG(最大的公开语言模型之一,拥有 170 亿个参数)
  • 我们如何使用这项技术创造 BERT 训练世界纪录

资源列表:

  • DeepSpeed Website 、 DeepSpeed Library (GitHub)
  • 论文: ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (publication)
  • ZeRO 和 DeepSpeed:新的系统优化支持超过 1000 亿个参数的训练模型 (blog)
  • ZeRO-2 和 DeepSpeed:打破深度学习速度和规模的障碍 (blog)
  • DeepSpeed Fastest Bert deep dive (blog)
  • Turing-NLG:微软的 170 亿参数语言模型 (blog)
  • AI at Scale 大规模人工智能 (Project Page)
  • ONNX runtime (GitHub)

二、FSDP: 使用更少的GPU进行更快的AI训练

  • 论文《Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Training》
  • Fully Sharded Data Parallel: faster AI training with fewer GPUs、如何用数据并行训练万亿参数模型

2.1 简介

  Facebook发布的FSDP(Fully Sharded Data Parallel)是一种用于分布式深度学习模型训练的并行计算策略,它对标的是微软在DeepSpeed中提出的ZeRO,可以看成PyTorch中的DDP优化版本。

  FSDP的核心思想是将模型参数分成多个小片段(shards),并在多个设备上并行处理这些参数片段。每个设备只处理模型参数的一个子集,而不是整个模型参数。FSDP的主要特点和原则包括:

  1. 参数分片:模型的每个参数都被分成多个小片段(shards),每个片段存储在不同的设备上。这样,每个设备只需要处理自己负责的参数片段,而不需要存储整个模型。

  2. 异步通信:FSDP采用异步通信模式,允许不同设备上的参数片段在不同的时间进行前向和反向传播,而无需等待其他设备。这提高了训练的效率。

  3. 内存效率:由于模型参数被分片存储,FSDP可以更好地利用每个设备的内存,允许训练非常大的模型。

  4. 通信策略:FSDP使用通信策略来控制参数片段之间的信息交换。通信可以定期发生,以确保参数片段的同步。

  5. 容错性:FSDP具有一定的容错性,即使某个设备失败,训练也可以继续进行,因为参数片段可以在其他设备上备份。

  FSDP旨在解决大规模深度学习模型训练中的内存和计算瓶颈问题,允许研究人员和工程师训练超大规模的模型,同时有效地利用多个设备的计算资源。它通常用于大型深度学习训练集群中,以加速训练过程并提高训练效率。

2.2 数据并行(Data Parallelism)

  对于典型的数据并行实现(PyTorch的DDP和TF的tf.distribute.MirroredStrategy)来说,每个GPU通常存储的是完整的模型,包括权重和参数。然后将不同的数据批次分配给不同的GPU来进行处理,接着在所有GPU之间共享梯度信息,以更新模型的权重和参数。因为多个GPU可以同时处理不同的数据批次,而无需等待其他GPU完成它们的工作,所以这种方法可以提高训练速度,整个过程如下图所示:
大模型分布式训练策略:ZeRO、FSDP_第4张图片

整个数据并行训练过程可以用以下步骤来总结:

  1. 初始化:在每个GPU上创建一个模型的副本(replica),每个副本都有相同的初始权重和参数。
  2. 训练迭代:每个GPU接收其分配的数据批次,并在本地(每个GPU本身)计算损失函数的梯度(local gradients)。
  3. 全体度求和(All-Reduce):为了保持模型权重的一致性并允许模型更新,需要将来自各个GPU的local gradients汇总为一个global gradients,这个操作称为all-reduce,其包含两个操作,reduce-scatterall-gather
    大模型分布式训练策略:ZeRO、FSDP_第5张图片
    • Reduce-Scatter:每个GPU的local gradients都被分成不同的块或分片(blocks or shards,通常是均匀分配的),然后在N个 replicas间进行 N-1 轮数据交换。每一轮的交换都会将部分梯度数据合并,最终每个 replica 都会持有来自所有其他 replica 的梯度分片的汇总(fully reduced data

      reduced是指每一轮的合并操作,使得shards减少。fully reduced就是指最终所有的梯度shards都被合并成一个值。

    • All-gather:每个 replica 将在 Reduce-Scatter 阶段中获得的自己的 fully reduced data 广播给所有的 GPU(这个广播操作可以在一轮内完成),以确保每个 replica 都具有全部的 fully reduced data,汇总后的梯度即为global gradients

  4. 基于global gradients进行模型参数更新(weight update
  5. 重复步骤2至4,直到训练完成。

  在 Reduce-Scatter阶段,之所以要先将梯度数据被分成不同的shards,再分N-1次进行广播,而不是在一轮之内广播完成,是为了减少通信开销和内存占用,同时提高训练速度和可扩展性,特别是在大规模深度学习模型和分布式训练中,这一点尤为重要。具体来说:

  1. 通信效率

    • 一次性将不同GPU的本地梯度广播给所有GPU可能导致大量的网络通信。在大规模训练中,GPU的数量可能非常大,广播操作会在网络上产生大量的数据流量,可能会成为瓶颈。
    • 通过分步的 Reduce-Scatter 阶段,每个 replica 只需将自己的梯度分片发送给其他 replica,而不是全部的梯度数据。这减少了每轮通信的数据量,减轻了网络负载。
  2. 内存效率

    • 深度学习模型通常具有大量参数,模型的梯度数据也很大。一次性将所有梯度广播给所有GPU需要存储所有梯度数据的副本,这会占用大量的GPU内存。
    • 通过 Reduce-Scatter 阶段,每个 replica 只需要存储部分梯度数据,即其本地梯度分片。这降低了每个GPU上需要存储的梯度数据的量,从而减轻了内存压力。
  3. 可扩展性

    • 一次性广播对于大规模分布式训练来说可能不够可扩展,因为网络带宽和内存资源是有限的。通过分步的方式,可以更好地控制通信和内存资源,使得训练可以扩展到更多的GPU或更大的模型规模。

2.3 sharding weight update

  在上述标准的数据并行操作中,每个分布式节点(replicas)都拥有完整的模型参数副本,并在每个训练步骤中都会做重复的update weight操作。对于大型模型,这种全局参数同步可能会成为性能瓶颈,因为需要大量的计算和通信资源。即使对于小模型而言,为防止global batch size过大,每个分布式节点都会采用较小的batch size,此时update weight也会成为训练中的重要耗时项。

为了解决这个问题,谷歌在2020年提出了sharding weight update。它的具体步骤如下:

  1. 参数分片(Sharding):模型参数被划分成多个分片(Shard),通常是均匀分布的。每个分布式节点只负责更新自己持有的参数分片。

  2. 本地更新:在每个分布式节点上,只有该节点持有的分片被用于计算梯度和更新参数。这减少了每个节点需要处理的参数数量,从而提高了计算效率。

  3. 全局同步:在一定的训练间隔之后,分布式节点之间进行全局同步。在全局同步中,各个节点交换它们的参数分片,以确保每个节点最终都具有完整的全局模型参数。

  这种方法的优点是,它减少了在每个训练步骤中进行全局参数同步的频率,因为只有在全局同步点才会进行参数交换,所以减少了计算和通信的开销,特别是在大型模型和小批次大小的情况下。同时,它允许分布式节点在本地更新参数,提高了计算效率。

  如下图所示,经过reduce-scatter后每个replica得到一个gradient shard,每个replica先更新自己的shard的weight,然后再进行all-gather,这样其实是和原始的all-reduce是等价的。但是经过这个调整,每个replica只是update weight shard,耗时就会降低了,相当于update weight也被各个replica给分担了。 ​

大模型分布式训练策略:ZeRO、FSDP_第6张图片

2.4 optimizer states all-gather

  optimizer往往包含额外的参数,比如SGD包含梯度的动量,而Adam包含动量和方差,,这些参数可以统称为optimizer states(优化器状态),它们也是需要同步更新的。如果也对optimizer states进行all-gather的话,通信成本就会比较大(原始的all-reduce并不需要)。

  optimizer states通常只在参数更新时需要,而在前向传播和反向传播中并不需要,所以只在需要时才对optimizer states进行all-gather,可以减少通信开销和内存使用,这样整个过程就变成:
大模型分布式训练策略:ZeRO、FSDP_第7张图片

图4:两种在循环中分片化优化器辅助变量(sharding auxiliary variables)的方式。左图:只在迭代之间保留分片化的auxiliary。右图:在迭代中保留auxiliary和权重分片,并在前向/后向传递之前对权重进行全局收集。

  左图weight的all-gather是在update后立即进行的,而右图是在需要的时候(forward和backward)才进行all-gather。右图方案有更大的优化空间,因为在前后向传播过程中,可以采用低精度fp16来all-gather来得到所需要的全部weight,这样就大大降低了内存使用和通信成本。另外weight和auxliary weight的生存周期也减少了。特别是optimizer的auxliary weight,这样就节省一部分内存空间,可以用来存储前后向传播中的activations和gradients。

  假定模型参数大小是W,auxliary weight大小是V,前后向传播中activations和gradients的峰值大小是P,共有N个shards,那么训练的峰值大小就从W+V+P降低为max(W+V/N+P,W+V),这带来的一个好处是Adam将和SGD一样高效(Adam比SGD要多一份方差)。
大模型分布式训练策略:ZeRO、FSDP_第8张图片
下图对比了standard DDP trainingFSDP training

大模型分布式训练策略:ZeRO、FSDP_第9张图片

  在标准数据并行训练中,每个 GPU 上都有完整模型的一份拷贝。然后,在每个 GPU 上,仅对数据的一个分片进行一系列前向传播和反向传播的计算。在本地计算完成后,每个 GPU 的参数和优化器状态需要与其他 GPU 共享,以计算全局模型的权重更新。

  而在完全分片数据并行训练中,每个 GPU 上只有模型的一个分片。然后,在本地计算前向传播时,需要从其他 GPU 收集所有分片的权重信息,以完成前向传播计算。同样,在反向传播之前,再次执行权重的收集。在完成反向传播后,本地梯度需要通过减少散布步骤进行平均并分片到各个 GPU 上,以允许每个 GPU 更新其本地权重分片。

  这两种方法的区别在于参数和权重信息的共享方式以及计算的分布方式。完全分片数据并行训练试图减少每个 GPU 上的模型复制,但需要更复杂的权重信息的收集和分布,以确保每个 GPU 具有必要的信息来执行前向和反向传播。这个方法的选择通常取决于训练环境和硬件资源的特性。

2.5 FSDP的使用

FSDP特点如下:

  • FSDP对parameters 和optimizer state进行sharding;
  • reshard_after_forward=False,和PyTorch DDP通信成本一样,类似ZeRO-DP-2(optimizer state+gradient sharding DP,即两阶段的 P o s + g P_{os+g} Pos+g);
  • reshard_after_forward=True,通信成本增加50%,类似ZeRO-DP-3 P o s + g + p P_{os+g+p} Pos+g+p),速度会慢,但是显存开销最小(ZeRO-DP详见本文1.3.1章节)。
  • FSDP用8 GPUs可以训练13B parameter models,用128 GPUs可以训练175B parameter models。当设置cpu_offload=True,可以用256 GPUs训练 1T parameter models。
  • FSDP只兼容pointwise Optimizers(Adam, AdamW, Adadelta, Adamax, SGD等),如果是non-pointwise Optimizers(Adagrad, Adafactor, LAMB等),sharding将得到稍微不一样的结果。

可以通过以下示例直接使用 FairScale 的 FSDP,只需替换 DDP(my_module):

from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
...
sharded_module = DDP(my_module)FSDP(my_module)
optim = torch.optim.Adam(sharded_module.parameters(), lr=0.0001)
for sample, label in dataload.next_batch:
  out = sharded_module(x=sample, y=3, z=torch.Tensor([1]))
  loss = criterion(out, label)
  loss.backward()
  optim.step()

2.6 FSDP和ZeRO的异同

  FSDP(Fully Sharded Data Parallel)和ZeRO(Zero Redundancy Optimizer)都是用于大规模深度学习模型训练的优化策略,都采用了内存优化,用于解决大型深度学习模型训练中的内存和计算瓶颈,以便在有限的硬件资源上训练更大规模的模型;但它们的焦点和目标略有不同:

  1. 关注点

    • FSDP:FSDP的主要关注点是模型参数的并行存储和计算,它通过将参数分成小片段并在多个设备上并行处理来降低内存占用。
    • ZeRO:ZeRO的主要关注点是优化训练中的冗余内存,特别是优化器状态、梯度和参数的内存占用。它的目标是在有限的设备内存上训练大规模模型。
  2. 通信策略

    • FSDP:FSDP使用异步通信来实现参数片段之间的信息交换,以提高训练效率。
    • ZeRO:ZeRO使用动态通信调度来减少通信量,同时保持计算粒度,以提高训练速度。它特别关注了通信中的内存重叠。
  3. 组合使用

    • FSDP:通常与数据并行一起使用,以实现参数的全分片和计算的全并行。
    • ZeRO:通常与数据并行一起使用,以减少模型参数的内存冗余,但也可以与模型并行结合使用,以降低激活内存。
  4. 适用性

    • FSDP:主要用于分布式训练和大规模模型的内存优化,适用于需要高度并行化和大规模训练的情况。
    • ZeRO:主要用于在有限设备内存上训练大规模模型,适用于需要降低内存占用的情况。

你可能感兴趣的:(LLMs,分布式,深度学习)