DL专栏3-GPU并行训练

GPU并行计算

为什么要使用多GPU并行训练

本简单来说，有两种原因：第一种是模型在一块GPU上放不下，两块或多块GPU上就能运行完整的模型（如早期的AlexNet）。第二种是多块GPU并行计算可以达到加速训练的效果。想要成为“炼丹大师“，多GPU并行训练是不可或缺的技能。

常见的多GPU训练方法：

1.模型并行方式：如果模型特别大，GPU显存不够，无法将一个显存放在GPU上，需要把网络的不同模块放在不同GPU上，这样可以训练比较大的网络。（下图左半部分）

2.数据并行方式：将整个模型放在一块GPU里，再复制到每一块GPU上，同时进行正向传播和反向误差传播。相当于加大了batch_size。（下图右半部分）

在pytorch1.7 + cuda10 + TeslaV100的环境下，使用ResNet34，batch_size=16, SGD对花草数据集训练的情况如下：使用一块GPU需要9s一个epoch，使用两块GPU是5.5s， 8块是2s。这里有一个问题，为什么运行时间不是9/8≈1.1s ? 因为使用GPU数量越多，设备之间的通讯会越来越复杂，所以随着GPU数量的增加，训练速度的提升也是递减的。

误差梯度如何在不同设备之间通信？

在每个GPU训练step结束后，将每块GPU的损失梯度求平均，而不是每块GPU各计算各的。

BN如何在不同设备之间同步？

假设batch_size=2，每个GPU计算的均值和方差都针对这两个样本而言的。而BN的特性是：batch_size越大，均值和方差越接近与整个数据集的均值和方差，效果越好。使用多块GPU时，会计算每个BN层在所有设备上输入的均值和方差。如果GPU1和GPU2都分别得到两个特征层，那么两块GPU一共计算4个特征层的均值和方差，可以认为batch_size=4。注意：如果不用同步BN，而是每个设备计算自己的批次数据的均值方差，效果与单GPU一致，仅仅能提升训练速度；如果使用同步BN，效果会有一定提升，但是会损失一部分并行速度。

下图为单GPU、以及是否使用同步BN训练的三种情况，可以看到使用同步BN（橙线）比不使用同步BN（蓝线）总体效果要好一些，不过训练时间也会更长。使用单GPU（黑线）和不使用同步BN的效果是差不多的。

两种GPU训练方法：DataParallel 和 DistributedDataParallel：

• DataParallel是单进程多线程的，仅仅能工作在单机中。而DistributedDataParallel是多进程的，可以工作在单机或多机器中。
• DataParallel通常会慢于DistributedDataParallel。所以目前主流的方法是DistributedDataParallel。

pytorch中常见的GPU启动方式：

注：distributed.launch方法如果开始训练后，手动终止程序，最好先看下显存占用情况，有小概率进程没kill的情况，会占用一部分GPU显存资源。

下面以分类问题为基准，详细介绍使用DistributedDataParallel时的过程:

首先要初始化各进程环境：

def init_distributed_mode(args):
    # 如果是多机多卡的机器，WORLD_SIZE代表使用的机器数，RANK对应第几台机器
    # 如果是单机多卡的机器，WORLD_SIZE代表有几块GPU，RANK和LOCAL_RANK代表第几块GPU
    if'RANK'in os.environ and'WORLD_SIZE'in os.environ:
        args.rank = int(os.environ["RANK"])
        args.world_size = int(os.environ['WORLD_SIZE'])
        # LOCAL_RANK代表某个机器上第几块GPU
        args.gpu = int(os.environ['LOCAL_RANK'])
    elif'SLURM_PROCID'in os.environ:
        args.rank = int(os.environ['SLURM_PROCID'])
        args.gpu = args.rank % torch.cuda.device_count()
    else:
        print('Not using distributed mode')
        args.distributed = False
        return

    args.distributed = True

    torch.cuda.set_device(args.gpu)  # 对当前进程指定使用的GPU
    args.dist_backend = 'nccl'# 通信后端，nvidia GPU推荐使用NCCL
    dist.barrier()  # 等待每个GPU都运行完这个地方以后再继续

在main函数初始阶段，进行以下初始化操作。需要注意的是，学习率需要根据使用GPU的张数增加。在这里使用简单的倍增方法。

def main(args):
        if torch.cuda.is_available() isFalse:
            raise EnvironmentError("not find GPU device for training.")

        # 初始化各进程环境
        init_distributed_mode(args=args)

        rank = args.rank
        device = torch.device(args.device)
        batch_size = args.batch_size
        num_classes = args.num_classes
        weights_path = args.weights
        args.lr *= args.world_size  # 学习率要根据并行GPU的数倍增

实例化数据集可以使用单卡相同的方法，但在sample样本时，和单机不同，需要使用DistributedSampler和BatchSampler。

#给每个rank对应的进程分配训练的样本索引
train_sampler=torch.utils.data.distributed.DistributedSampler(train_data_set)
val_sampler=torch.utils.data.distributed.DistributedSampler(val_data_set)
#将样本索引每batch_size个元素组成一个list
train_batch_sampler=torch.utils.data.BatchSampler(
train_sampler,batch_size,drop_last=True)

DistributedSampler原理如图所示：假设当前数据集有0~10共11个样本，使用2块GPU计算。首先打乱数据顺序，然后用 11/2 =6（向上取整），然后6乘以GPU个数2 = 12，因为只有11个数据，所以再把第一个数据（索引为6的数据）补到末尾，现在就有12个数据可以均匀分到每块GPU。然后分配数据：间隔将数据分配到不同的GPU中。

BatchSampler原理: DistributedSmpler将数据分配到两个GPU上，以第一个GPU为例，分到的数据是6，9，10，1，8，7，假设batch_size=2，就按顺序把数据两两一组，在训练时，每次获取一个batch的数据，就从组织好的一个个batch中取到。注意：只对训练集处理，验证集不使用BatchSampler。

接下来使用定义好的数据集和sampler方法加载数据：

train_loader = torch.utils.data.DataLoader(train_data_set,
                                               batch_sampler=train_batch_sampler,
                                               pin_memory=True,   # 直接加载到显存中，达到加速效果
                                               num_workers=nw,
                                               collate_fn=train_data_set.collate_fn)

 val_loader = torch.utils.data.DataLoader(val_data_set,
                                             batch_size=batch_size,
                                             sampler=val_sampler,
                                             pin_memory=True,
                                             num_workers=nw,
                                             collate_fn=val_data_set.collate_fn)

如果有预训练权重的话，需要保证每块GPU加载的权重是一模一样的。需要在主进程保存模型初始化权重，在不同设备上载入主进程保存的权重。这样才能保证每块GOU上加载的权重是一致的：

# 实例化模型
    model = resnet34(num_classes=num_classes).to(device)

    # 如果存在预训练权重则载入
    if os.path.exists(weights_path):
        weights_dict = torch.load(weights_path, map_location=device)
        # 简单对比每层的权重参数个数是否一致
        load_weights_dict = {k: v for k, v in weights_dict.items()
                             if model.state_dict()[k].numel() == v.numel()}
        model.load_state_dict(load_weights_dict, strict=False)
    else:
        checkpoint_path = os.path.join(tempfile.gettempdir(), "initial_weights.pt")
        # 如果不存在预训练权重，需要将第一个进程中的权重保存，然后其他进程载入，保持初始化权重一致
        if rank == 0:
            torch.save(model.state_dict(), checkpoint_path)

        dist.barrier()
        # 这里注意，一定要指定map_location参数，否则会导致第一块GPU占用更多资源
        model.load_state_dict(torch.load(checkpoint_path, map_location=device))

如果需要冻结模型权重，和单GPU基本没有差别。如果不需要冻结权重，可以选择是否同步BN层。然后再把模型包装成DDP模型，就可以方便进程之间的通信了。多GPU和单GPU的优化器设置没有差别，这里不再赘述。

# 是否冻结权重
    if args.freeze_layers:
        for name, para in model.named_parameters():
            # 除最后的全连接层外，其他权重全部冻结
            if"fc"notin name:
                para.requires_grad_(False)
    else:
        # 只有训练带有BN结构的网络时使用SyncBatchNorm采用意义
        if args.syncBN:
            # 使用SyncBatchNorm后训练会更耗时
            model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)

    # 转为DDP模型
         model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])

     # optimizer使用SGD+余弦淬火策略
        pg = [p for p in model.parameters() if p.requires_grad]
        optimizer = optim.SGD(pg, lr=args.lr, momentum=0.9, weight_decay=0.005)
        lf = lambda x: ((1 + math.cos(x * math.pi / args.epochs)) / 2) * (1 - args.lrf) + args.lrf  # cosine
    scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)

与单GPU不同的地方：rain_sampler.set_epoch(epoch)，这行代码会在每次迭代的时候获得一个不同的生成器，每一轮开始迭代获取数据之前设置随机种子，通过改变传进的epoch参数改变打乱数据顺序。通过设置不同的随机种子，可以让不同GPU每轮拿到的数据不同。后面的部分和单GPU相同。

for epoch in range(args.epochs):
            train_sampler.set_epoch(epoch)  

            mean_loss = train_one_epoch(model=model,
                                        optimizer=optimizer,
                                        data_loader=train_loader,
                                        device=device,
                                        epoch=epoch)

            scheduler.step()

            sum_num = evaluate(model=model,
                               data_loader=val_loader,
                               device=device)
            acc = sum_num / val_sampler.total_size

我们详细看看每个epoch是训练时和单GPU训练的差异（上面的train_one_epoch）

def train_one_epoch(model, optimizer, data_loader, device, epoch):
        model.train()
        loss_function = torch.nn.CrossEntropyLoss()
        mean_loss = torch.zeros(1).to(device)
        optimizer.zero_grad()

        # 在进程0中打印训练进度
        if is_main_process():
            data_loader = tqdm(data_loader)

        for step, data in enumerate(data_loader):
            images, labels = data

            pred = model(images.to(device))

            loss = loss_function(pred, labels.to(device))
            loss.backward()
            loss = reduce_value(loss, average=True)  #  在单GPU中不起作用，多GPU时，获得所有GPU的loss的均值。
            mean_loss = (mean_loss * step + loss.detach()) / (step + 1)  # update mean losses

            # 在进程0中打印平均loss
            if is_main_process():
                data_loader.desc = "[epoch {}] mean loss {}".format(epoch, round(mean_loss.item(), 3))

            ifnot torch.isfinite(loss):
                print('WARNING: non-finite loss, ending training ', loss)
                sys.exit(1)

            optimizer.step()
            optimizer.zero_grad()

        # 等待所有进程计算完毕
        if device != torch.device("cpu"):
            torch.cuda.synchronize(device)

        return mean_loss.item()

    def reduce_value(value, average=True):
        world_size = get_world_size()
        if world_size < 2:  # 单GPU的情况
            return value

        with torch.no_grad():
            dist.all_reduce(value)   # 对不同设备之间的value求和
            if average:  # 如果需要求平均，获得多块GPU计算loss的均值
                value /= world_size

        return value

接下来看一下验证阶段的情况，和单GPU最大的额不同之处是预测正确样本个数的地方。

    @torch.no_grad()
    def evaluate(model, data_loader, device):
        model.eval()

        # 用于存储预测正确的样本个数，每块GPU都会计算自己正确样本的数量
        sum_num = torch.zeros(1).to(device)

        # 在进程0中打印验证进度
        if is_main_process():
            data_loader = tqdm(data_loader)

        for step, data in enumerate(data_loader):
            images, labels = data
            pred = model(images.to(device))
            pred = torch.max(pred, dim=1)[1]
            sum_num += torch.eq(pred, labels.to(device)).sum()

        # 等待所有进程计算完毕
        if device != torch.device("cpu"):
            torch.cuda.synchronize(device)

        sum_num = reduce_value(sum_num, average=False)  # 预测正确样本个数

    return sum_num.item()

需要注意的是：保存模型的权重需要在主进程中进行保存。

if rank == 0:
            print("[epoch {}] accuracy: {}".format(epoch, round(acc, 3)))
            tags = ["loss", "accuracy", "learning_rate"]
            tb_writer.add_scalar(tags[0], mean_loss, epoch)
            tb_writer.add_scalar(tags[1], acc, epoch)
            tb_writer.add_scalar(tags[2], optimizer.param_groups[0]["lr"], epoch)

            torch.save(model.module.state_dict(), "./weights/model-{}.pth".format(epoch))

如果从头开始训练，主进程生成的初始化权重是以临时文件的形式保存，需要训练完后移除掉。最后还需要撤销进程组。

if rank == 0:# 删除临时缓存文件        
    if os.path.exists(checkpoint_path) is True:            
        os.remove(checkpoint_path)    
        dist.destroy_process_group()  # 撤销进程组，释放资源

OpenAI如何训练千亿级模型？

“炼大模型”已成为人工智能领域的主流研发趋势。从GPT-3的1750亿，到如今悟道2.0的1.75万亿，超大语言模型在 NLP 基准任务中不断刷新SOTA。

然而，参数和数据集的快速增长让 GPU 算力开始捉襟见肘。单个GPU内存已经远远不能满足大模型的需求。如，阿里用480块GPU训练千亿模型；英伟达用3072块GPU训练万亿模型；谷歌用2048块TPU训练1.6万亿模型（1 TPU约等于2～3 GPU）。

如何利用上百块GPU上训练大规模语言模型？并行计算是一种行之有效的方法。

并行训练

大规模神经网络模型对存储空间有强烈的需求，单个GPU的内存已经远远不够。究其原因，一方面模型权重有数百亿个浮点数，随机梯度下降和Adam优化需要极高的计算成本；另一方面在预训练阶段，大模型与大规模语料库的配对需要很长的时间周期。综合来看，跨GPU并行计算显得尤为重要。

并行计算在数据、模型架构和张量等不同维度上都可以操作，接下来本文将具体介绍一些主流方法：

数据并行

数据并行（ Data parallelism ，DP）最简单的方法是将相同的模型权重复制到worker节点，并分配一部分数据以同时进行处理。我们知道，如果模型的参数量大于单个GPU节点的内存，DP无法正常工作，GeePS架构（Cui等人，2016）的解决思路是使用有限的GPU内存。也就是，如果模型太大无法嵌入到一台机器，就将暂时未使用的参数卸载回CPU。

数据交换传输通常在后端进行（不干扰训练计算），在每个Mini-batch计算结束后，worker需要同步梯度或权重，以保证学习效率。现有的同步方法有两种，各自优缺点如下：

1、批量同步并行（BSP）：worker在每个Mini-batch结束时同步数据，这种方法保证了模型权重传递的及时性，但每台机器都必须排队等待其他机器发送梯度。

2、异步并行（ASP）：每个GPU采用异步方式处理数据，这种方法避免了不同机器之间的相互等待或暂停，但影响了权重传递的时效，降低了统计学习效率。而且即使增加计算时长，也不会加快训练的收敛速度。

在中间某些地方的每一次迭代（＞1）都需要同步全局梯度。自Pytorch v1.5版（Li等人，2021年）提出后，该特征在分布式数据并行（Distribution Data Parallel，DDP）中被称为“梯度累积（gradient accumulation）”。分桶梯度(bucketing gradients)避免立即执行AllReduce操作，而是将多个梯度存储到一个AllReduce中以提高吞吐量，并基于计算图优化计算和通信调度。

模型并行

模型并行（Model parallelism，MP）用于解决模型权重不能适应单个节点的情况，在这里，计算和模型参数都需要跨多台机器进行处理。在数据并行中，每个worker承载着整个模型的完整副本，而MP只在一个worker上分配部分模型参数，因此对内存和计算的需求要小很多。

深度神经网络包含一堆垂直层，如果逐层拆分将连续的小层分配到工作层分区，操作起来并不难，但通过大量具有顺序依赖性的Workers来运行每个数据batch会花费大量的时间，计算资源的利用率也严重不足。

管道并行

管道并行（Pipeline parallelism，PP）是将模型并行与数据并行结合起来，以减少低效时间“气泡”的过程。主要思想是将Mini-batch拆分为更多个微批次（microbatch），并使每个阶段worker能够同时处理。需要注意的是，每个微批次需要两次传递，一次向前，一次向后。worker分区的数量称为管道深度，不同worker分区之间的通信仅传输激活（向前）和梯度（向后）。这些通道的调度方式以及梯度的聚合方式在不同的方法中有所不同。

在GPipe（Huang et al.2019）方法中，多个微批次处理结束时会同时聚合梯度和应用。同步梯度下降保证了学习的一致性和效率，与worker数量无关。如图3所示，“气泡”仍然存在，但比图2少了很多。给定m个均匀分割的微批次和d个分区，假设每个微批次向前和向后都需要一个时间单位，则气泡的分数为：

GPipe论文表明，如果微批次的数量超过分区数量4倍（m>4d），则“气泡”开销几乎可以忽略不计。

GPipe在吞吐量方面与设备数量成线性关系，设备数量越多，吞吐量越大。不过，如果模型参数在worker中分布不均匀，这种线性关系不会稳定出现。

PipeDream（Narayanan等人，2019年）方法要求每个worker交替处理向前和向后传递的消息（1F1B）。它将每个模型分区命名为“stage”，每个stage worker可以有多个副本来并行运行数据。这个过程使用循环负载平衡策略在多个副本之间分配工作，以确保相同minibatch 向前和向后的传递发生在同一副本上。

由于PipeDream没有在所有worker batch结束时同步全局梯度，1F1B 很容易导致不同版本的模型权重的微批次向前和向后传递，降低学习效率。对此，PipeDream提供了一些解决的思路：

• 权重存储：每个worker跟踪多个模型版本，给定数据 batch 的向前和向后传递相同版本的权重。
• 垂直同步：不同模型权重版本与激活和梯度一起在全局worker之间传递，计算采用上一个worker传播的相对应的隐藏版本。这个过程确保了worker之间的版本一致性（不同于GPipe，采用异步计算）。

在训练开始时，PipeDream会先分析模型每一层的计算内存和时间成本，然后将层划分为不同的stage进行优化。

后来有学者提出了PipeDream两种变体，主要思路是通过减少模型版本来缓解内存占用（Narayanan等人，2021年）。其中，PipeDream-flush增加了定期刷新全局同步管道的功能，就像GPipe一样，这种方式虽然牺牲了一点吞吐量，但显著减少了内存占用（例如仅需要维护单一版本的模型权重）。

PipeDream-2BW维护两个版本的模型权重，“2BW”代表“双缓冲权重（double-buffered weights）”，它会在每个微批次生成一个新的模型版本K（K>d）。由于一些剩余的向后传递仍然依赖于旧版本，新的模型版本无法立即取代旧版本，但因为只保存了两个版本，内存占用的也被大大降低了。

张量并行

模型并行和管道并行都会垂直拆分模型，而张量并行（Tensor Parallelism，TP）是将张量运算的计算水平划分到多个设备上。

以Transformer为例。Transformer架构主要由多层MLP和自注意力块组成。Megatron-LM（Shoeybi et al.2020）采用了一种简单的方法来并行计算层内MLP和自注意力。

MLP层包含GEMM（通用矩阵乘法）和非线性GeLU传输。如果按列拆分权重矩阵A，可以得到：

注意力块根据上述分区并行运行GEMM的 查询（Q）、键（K）和 权重（V），然后与另一个GEMM组合以生成头注意力结果。

今年Narayanan等人将管道、张量和数据并行与新的管道调度策略相结合，提出了一种名为PTD-P的新方法。该方法不仅在设备上能够定位一组连续的层（“模型块”），该可以为每个wokers分配多个较小的连续层子集块（例如，设备1具有第1、2、9、10层；设备2具有第3、4、11、12层；每个具有两个模型块）

每个batch中，微批次的数量应精确除以wokers数量（mm）。如果每个worker有v个模型块，那么与GPipe调度相比，管道的“气泡”时间可以减少 v 倍。

混合专家（MoE）

为了突破模型大小的限制，谷歌后来提出一种混合专家（MoE）方法，其核心理念是：综合学习，它假设多个弱学习者组合起来就会拥有一个强学习者。

在深度神经网络中，混合专家（MoE）通过连接多个专家的门机制（gating mechanism）实现集成（Shazeer等人，2017）。门机制激活不同网络的专家以产生不同的输出。作者在论文将其命名为“稀疏门控专家混合层（sparsely gated MoE）”。

仅一个MoE层包含：（1）前馈网络专家n；（2）可训练的门控网络G，通过学习n个专家的概率分布，将流量路由到几个特定的专家。

根据门控输出，并非每个专家都必须进行评估。当专家的数量太大时，可以考虑使用两层MoE。

G将输入与可训练权重矩阵Gg相乘，然后执行softmax：

由于这个过程会产生密集的门控制向量，不利于节省计算资源，而且

时也不需要评估专家。所以，MoE层仅保留了顶部k值，并通过向G中添加高斯噪声改进负载平衡，这种机制被称为噪声top-k门。

为了避免门控网络可能始终偏向少数强势专家的自我强化效应，Shazeer等人（2017）提出了通过额外重要损失的软约束，以鼓励所有专家拥有相同的权重。其数值相当于每个专家的分批平均值变异系数的平方：

其中，CV是变异系数，失重的waux是可调节的超参数。由于每个专家网络只能获得小部分训练样本（“收缩批次问题”），所以在MoE中应该尽可能使用大batch，但这又会受到GPU内存的限制。数据并行和模型并行的应用可以提高模型的吞吐量。

GShard（Lepikhin等人，2020年）通过自动分片将MoE transformer 模型的参数扩展到了6000亿。MoE transformer 用MoE层取代其他每一个前馈网络层。需要说明的是，在多台机器上MoE transformer 仅在MoE层分片，其他层只是复制。

GShard中的门控功能G有几种改进设计：

• 专家容量：通过一位专家的令牌数量不应超过“专家容量”的阈值。如果令牌被路由到已达到容量的专家，则令牌将被标记为“溢出”，并且门输出将更改为零向量。
• 本地组调度：令牌被均匀地划分为多个本地组，专家能力在组水平上得到加强。
• 辅助损失：与原始MoE aux损失相似，添加辅助损失可以最小化路由到每个专家的数据的均方。
• 随机路由：以与其权重成比例的概率选择第二位最佳专家；否则GShard遵循随机路由，增加随机性。

Switch Transformer（Fedus et al.2021）用稀疏开关FFN层取代了密集前馈层（每个输入仅路由到一个专家网络），将模型规模扩展到数万亿个参数。负载平衡的辅助损失是

给定n个专家，fi是路由到第i个专家的令牌分数，pi是门控网络预测的专家i的路由概率。

为提高训练稳定性，switch transformer采用以下设计：

• 选择精度：使用FP32精度以提高模型局部的稳定性，并降低FP32张量的通信成本。FP32精度仅在路由器功能主体内使用，结果将还原到FP16。 较小的初始化：权重矩阵的初始化从平均μ=0且标准偏差

​ 的正态分布中采样，同时将Transformer初始化参数从s=1减小到s=0.1 。

• 使用更高的专家辍学率：微调通常适用于小数据集，增加每个专家的辍学率以避免过度拟合。他们发现，所有层中的辍学率增加会导致性能下降。在论文中，他们在非专家层中使用了0.1的辍学率，但在专家FF层中使用了0.4的辍学率。

switch transformer论文总结了用于训练大型模型的不同数据和模型并行策略，并给出了一个很好的示例：

其他节省内存的设计

CPU卸载

如果GPU内存已满，可以将暂时未使用的数据卸载到CPU，并在以后需要时将其读回（Rhu等人，2016）。不过，这种方法近年来并不太流行，因为它会延长模型训练的时间。

激活重新计算

激活重新计算，也称“激活检查点”或“梯度检查点”（Chen et al，2016），其核心思路是牺牲计算时间来换取内存空间。它减少了训练 ℓ 层深层神经网络到O(√ℓ)的内存开销，每个batch只消耗额外的前向传递计算。

具体来说，该方法将ℓ层网络平均划分为d个分区，仅保存分区边界的激活，并在workers之间进行通信。计算梯度仍然需要在分区内层进行中间激活，以便在向后过程中重新计算梯度。在激活重新计算的情况下，用于训练M(ℓ) 是：

激活重新计算的方法可以得出与模型大小有关次线性内存开销，如下图：

混合精度训练

此前，Narang（Narang&Micikevicius等人，2018年）介绍了一种使用半精度浮点（FP16）数训练模型而不损失模型精度的方法。

其中涉及三种关键技术：

• 全精度权重复制：保持累积梯度的模型权重的全精度（FP32）复制。对于向前和向后的传递的信息做四舍五入至半精度处理，因为每次梯度更新（即梯度X学习率）太小，可能无法完全包含在FP16范围内。
• 缩放损失：放大损失以更好地处理小幅度的梯度（见图16），放大梯度以使其向可表示范围的右侧部分（包含较大的值）移动，从而保留可能丢失的值。
• 算术精度：对于常见的网络算法（如矢量点积、矢量元素求和归约），将部分结果累加到FP32中，然后输出保存为FP16。逐点操作可以在FP16或FP32中执行。

在这项实验中，图像分类、更快的R-CNN等不需要损失缩放，但其他网络，如多盒SSD、大LSTM语言模型是需要损失缩放的。

压缩（Compression）

模型权重在向前和向后传递的过程中会消耗大量内存。考虑到这两种传递方式会花费大量时间，2018年Jain （Jain et al，2018）提出了一种数据编码策略，即在第一次传递后压缩中间结果，然后将其解码用于反向传播。

Jain和团队研发的Gist系统包含两种编码方案：一是特定于层的无损编码，包括 ReLU-Pool和 ReLU-Conv模式；二是有攻击性的有损编码，主要使用延迟精度缩减（DPR）。需要注意的是，第一次使用特征图时应保持高精度，第二次使用时要适度降低精度。这项实验表明，Gist可以在5个最佳图像分类DNN上减少2倍的内存开销，平均减少1.8倍，性能开销仅为4%。

内存高效优化器

优化器也会消耗内存。以主流的Adam优化器为例，其内部需要维护动量和方差，这两者与梯度和模型参数比例基本相同。这意味着，我们需要节省4倍模型权重的内存。

为了减少内存消耗，学术界已经提出了几款主流优化器。与Adam相比，Adafactor（Shazeer et al.2018）优化器没有存储全部动量和变化，只跟踪移动平均数的每行和每列总和，然后根据这些总和估计二阶矩。

SM3（Anil et al.2019）优化器采用了一种不同的自适应优化方法。

ZeRO（Rajbhandari et al.2019）零冗余优化器节省了大型模型训练在两方面的内存消耗：

• 大多数内存由模型状态消耗，包括优化器状态（例如Adam动量和方差）、梯度和参数。混合精度训练也需要大量内存，因为除了FP16版本之外，优化器还需要保存FP32参数和其他优化器状态的副本。
• 未被激活、临时缓冲区以及不可用的碎片内存消耗（论文中称为剩余状态）。

ZeRO结合了ZeRO-DP和ZeRO-R两种方法。ZeRO-DP是一种增强的数据并行，避免了模型状态的简单冗余。它以动态的方式跨多个并行数据划分优化器状态、梯度和参数，以最小化通信量。ZeRO-R使用分区激活二次计算、恒定缓冲区大小和动态内存碎片，以优化剩余状态的内存消耗。

参考资料：

[1] Li et al. “PyTorch Distributed: Experiences on Accelerating Data Parallel Training” VLDB 2020.

[2] Cui et al. “GeePS: Scalable deep learning on distributed GPUs with a GPU-specialized parameter server” EuroSys 2016

[3] Shoeybi et al. “Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism.” arXiv preprint arXiv:1909.08053 (2019).

[4] Narayanan et al. “Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM.” arXiv preprint arXiv:2104.04473 (2021).

[5] Huang et al. “GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism.” arXiv preprint arXiv:1811.06965 (2018).

[6] Narayanan et al. “PipeDream: Generalized Pipeline Parallelism for DNN Training.” SOSP 2019.

[7] Narayanan et al. “Memory-Efficient Pipeline-Parallel DNN Training.” ICML 2021.

[8] Shazeer et al. “The Sparsely-Gated Mixture-of-Experts Layer Noam.” arXiv preprint arXiv:1701.06538 (2017).

[9] Lepikhin et al. “GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding.” arXiv preprint arXiv:2006.16668 (2020).

[10] Fedus et al. “Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity.” arXiv preprint arXiv:2101.03961 (2021).

[11] Narang & Micikevicius, et al. “Mixed precision training.” ICLR 2018.

[12] Chen et al. 2016 “Training Deep Nets with Sublinear Memory Cost.” arXiv preprint arXiv:1604.06174 (2016).

[13] Jain et al. “Gist: Efficient data encoding for deep neural network training.” ISCA 2018.

[14] Shazeer & Stern. “Adafactor: Adaptive learning rates with sublinear memory cost.” arXiv preprint arXiv:1804.04235 (2018).

[15] Anil et al. “Memory-Efficient Adaptive Optimization.” arXiv preprint arXiv:1901.11150 (2019).

[16] Rajbhandari et al. “ZeRO: Memory Optimization Towards Training A Trillion Parameter Models Samyam.” arXiv preprint arXiv:1910.02054 (2019).