zero-offload论文阅读记录

LogicOverlord
619 阅读8分钟

注：新手文章。

1.摘要

    大模型训练需要大量GPU，zero-offload可以增加10倍的模型参数量且不牺牲性能及无需修改模型。zero-offload通过把数据和计算下发到cpu来实现大模型训练。

2.介绍

    随着深度学习的发展，模型越来越大，目前一些分布式训练方法，如模型并行、zero等依然需要大量GPU，以至于普通人玩不起。

    异构计算是个解决上述问题的好法子，但是目前的异构计算大多数集中在解决CNN网络的问题，这些CNN网络的瓶颈在于中间状态内存而非模型参数量，而近期的基于注意力机制的网络的瓶颈在于网络参数量而非中间状态，而且这些研究基本上都在使用CPU的内存而没有用CPU来计算，其次他们往往基于单个GPU。而本文提出的方法就可以解决以上问题。

    zero-offload基于三方面来设计：效率、可扩展性以及易用性。实现上，zero-offload在CPU上的计算需求远低于GPU、需要的CPU和GPU之间的通信也极少，这防止了CPU侧出现瓶颈。为了实现上述目标，我们必须把梯度、优化器状态以及优化过程offload到CPU，同时将模型、前向和反向计算保持在GPU进行，这让我们训练的模型参数量大了十倍。

    offloading要CPU进行O(M)的计算量，需要GPU进行O(MB)的计算量，其中M是模型参数，B是batch size，对于大多数情况，batch size都很大，因此CPU计算不是瓶颈，但是对于小batch size，CPU就会成为瓶颈，本文采用两个优化方法：①一个有效的CPU优化器。②One-step delayed parameter update使得CPU的优化步骤和GPU的计算一起进行（overlap）。

    可扩展性对于利用多GPU至关重要，深度学习的社区中，数据并行是使用最广泛的，由于数据并行会把优化器状态、参数和梯度都复制一遍，以及在每个GPU都进行优化计算，因此，在这种情况下使用offload会导致显著的通信量增加以及CPU计算量的增加。

    为了解决这些问题，本文的方法使用基于ZeRO的数据并行方法，这样在CPU内存中就只需要保存一份优化器状态，同时，通信量也不会线性增长，这使得本文的方法可以扩展至128卡。

    易用性：ZeRO-Offload是deepspeed的一部分，使用起来很方便，如下图所示:

image.png

2.背景和相关工作

①模型训练过程中的内存消耗

    模型训练内存消耗包含两部分：模型状态（model states）和剩余状态（residual states）。模型状态包括模型参数、梯度和优化器状态；剩余状态包括激活值、缓存和内存碎片等。大模型训练的主要瓶颈来自模型状态，如Megatron-LM、T5等，他们主要通过FP16 amp和ADAM来训练。

    FP16的混合精度训练需要16×M字节的内存量，分别是2Bytes的FP16类型的模型参数和梯度（共4Bytes）、4Bytes的FP32类型的adam参数以及模型参数（共12Bytes），这导致模型过大而无法训练。

  目前有两个主流方向来解决这个问题:scale-out training以及scale-up training。

②Scale out large model training.

    scale-out training的方法是用多个GPU的总体显存来进行大模型训练，比如模型并行和流水线并行。模型并行垂直方向切割模型，从而让模型的每一层进入到不同的GPU，而流水线并行则是水平方向切割模型，从而让模型的不同层进入不同的GPU。这两种方法都会改变模型的工作模式，从而限制其使用。

    ZeRO提供了另外一种方法来训练大模型，其将大模型分散到不同的GPU，并且提供良好的计算效率和可扩展性。但是上述方法都要求多个GPU的总显存要够大，而ZeRO-Offload则可以利用CPU内存从而训练大模型，而且在多GPU训练的时候还可以和ZeRO或者模型并行一起使用从而提供出色的可扩展性。

③Scale up large model training.

    现有的scale up的方法有三种：recompute、低精度或者混合精度、用CPU内存作为GPU显存的扩展。我们的方法属于第三种。

④ZeRO powered data parallel training.

    ZeRO-Offload可以和ZeRO-2很好的协作。在ZeRO-2中，每个GPU都保存整个模型参数，但是每个step只更新一部分模型参数，因此每个GPU只保存一部分优化器状态和梯度，即优化optimizer state和gradient，之后再把更新后的参数传播出去，其计算和通信步骤为:在前向过程中，每个GPU计算不同的mini-batch所对应的loss，在反向过程中，不同部分的梯度被计算并同步（注意ZeRO-2除了模型参数之外都是只保存1/N），反向传播后每个GPU更新自己负责的模型参数，最后再同步整个模型。

3.ZeRO-Offload是独一无二的最优offload策略

    本文把深度学习的中的数据流做成图，并用第一原则分析法来将图切分到CPU和GPU，ZeRO-Offload的切分方法在以下三个方面是最优的：GPU的计算量是CPU的几何倍，从而防止CPU成为瓶颈、ZeRO-Offload保证了CPU和GPU的通信最少、理论上是内存和通信最优的。

3.1数据流图

image.png

    圆形节点代表model states、方形节点代表计算过程、有向边代表数据流。

3.2减轻CPU的计算负担

  前向和反向传播的计算量由于和batch size以及模型大小成正比，因此放在GPU，而范数计算和模型参数更新等则只和模型大小成正比，因此放在CPU。

3.3减少通信开销

    反向传播后GPU发送2M的梯度给CPU，以及CPU更新参数后发送2M的参数给GPU。之所以这样做，是因为模型参数更新在CPU，所以CPU必须有梯度，而更新后的参数又要参与下一次GPU计算，所以传播至少需要两次，而2M+2M就是最小的。

3.4减少GPU显存占用

穷举法证明这样做GPU显存占用最少。

4.ZeRO-Offload调度策略

4.1单卡训练

    如上节讨论的那样，ZeRO-Offload中把fp16的模型参数保存在GPU显存中，而把fp16的梯度和fp32的模型参数、优化器参数都保存在CPU。在反向传播的过程中，梯度算出来之后可以立刻被传到CPU内存或者分组传到CPU内存，因此GPU只需要很小的显存来缓存这些梯度。同时，梯度的传输可以和反向传播同时进行，从而进一步减少通信开销。反向传播结束后，CPU进行参数更新并将更新后的参数传给GPU，整个过程如下图所示（感觉p swap应该是CPU->GPU?）:

image.png

4.2多卡训练

    ZeRO-Offload可以和ZeRO-2很好的结合起来。在多卡训练中，ZeRO-Offload同样是把每张卡自己负责的部分Adam参数和梯度以及FP32格式的模型参数offload到CPU里面，仅在GPU显存保存完整的FP16格式的模型参数。反向传播的时候，照样是通过reduce-scatter先同步梯度，然后负责该区域的GPU把这部分梯度offload到CPU，然后就可以在CPU更新相应的参数了，然后该GPU负责的参数再被copy回GPU，然后全部GPU再同步整个模型参数。

伪代码:

for_parallel rank in range(world_size):
    initialize_layers()
    for batch in dataset:
        x = forward(batch)
        compute_loss(x,batch).backward()
        backward(x.grad)
        step()

def _is_owner(i):
    return True if rank owns i else False

 def initialize_layers():
    for i in range(num_layers):
        l = layers[i]
        allocate_on_gpu l.param_fp16            #将该层的fp16格式的参数复制到GPU（毕竟stage2所有的GPU都要保存完整的fp16格式的参数）
        if _is_owner(i):                        #如果本GPU负责该层，则在CPU内存中分配相应的空间来保存FP32格式的模型参数、优化器参数、梯度     
            allocate_on_cpu l.param_fp32        
            allocate_on_cpu l.optim_states_fp32
            allocate_on_cpu l.cpu_grad

 def forward(x): #stage2可以无脑forward
    for i in range(num_layers):
        x = layers[i].forward(x)
    return x

 def backward(dx):
    for i in range(num_layers, 0, -1):
        dx=layers[i].backward(dx)
        reduce(layers[i].grad, dest_rank
            = _owner_rank(i)) #计算梯度并传给相关的GPU
        if _is_owner(i) l.cpu_grad.copy(l.grad) #把梯度复制到CPU
        else pass
        del layers[i].grad  #释放梯度占用的空间

 def step():
    for i in range(num_layers):
        l=layers[i]
        if _is_owner(i):
            #CPU中更新自己负责的参数并把这部分参数COPY到GPU
            update_in_cpu(l.optim_states_fp32,
                            l.cpu_grad,
                            l.param_fp32)
            l.param_fp16.copy(l.param_fp32)
        #同步全部模型参数
        BROADCAST(l.param_fp16, src=_owner_rank(i))
avatar
文章
阅读
粉丝