乐趣区

关于人工智能:PyTorch-新库-TorchMultimodal-使用说明将多模态通用模型-FLAVA-扩展到-100-亿参数

先前的文章中,咱们介绍了 TorchMultimodal,明天咱们将从一个具体案例登程,演示如何在 Torch Distributed 技术加持下,在 TorchMultimodal 库中扩大多模态根底模型。

近年来,大模型已成为一个备受关注的钻研畛域。以自然语言解决为例,语言模型曾经从几亿参数(BERT)倒退到了几千亿参数(GPT-3),对上游工作的性能晋升显示出重大作用。

业界对大规模语言模型如何扩大进行了宽泛的钻研。在视觉畛域也能够察看到相似的趋势,越来越多的开发者开始转向基于 transformer 的模型(如 Vision Transformer、Masked Auto Encoders)。

显然,因为大规模模型的倒退,单模态(如文本、图像、视频)相干钻研不断改进,框架也迅速适应了更大的模型。

同时,随着图像 - 文本检索、视觉问答、视觉对话和文本到图像的生成等工作在事实世界中的利用,多模态越来越受到重视。

接下来就是训练大规模多模态模型。该畛域也有了一些致力成绩,如 OpenAI 的 CLIP,谷歌的 Parti 和 Meta 的 CM3。

本文将通过一个案例钻研,展现如何应用 PyTorch Distributed 技术将 FLAVA 扩大到 100 亿参数。

补充浏览:HyperAI 超神经:Meta 外部都在用的 FX 工具大起底:利用 Graph Transformation 优化 PyTorch 模型



FLAVA 是一个视觉和语言根底模型在 TorchMultimodal 中可用

FLAVA 在单模态和多模态 Benchmark 中都体现出了十分突出的性能劣势。本文将联合相干代码示例,演示如何扩大 FLAVA。

代码详见:

https://github.com/facebookresearch/multimodal/tree/main/examples/flava/native

扩大 FLAVA 概览

FLAVA 是一个根底多模态模型,由基于 transformer 的图像和文本编码器以及基于 transformer 的多模态交融模块组成。

FLAVA 在单模态和多模态数据上都进行了预训练,且这些数据的损失 (loss) 各不相同,包含掩码的语言、图像和多模态模型 loss,要求模型从其上下文中重建原始输出(自监督学习)。

此外,它还应用了图像文本匹配损失 (image text matching loss),包含对齐图像 - 文本对的 positive 和 negative 示例,以及 CLIP 格调的比照损失。

除了多模态工作(如图像 - 文本检索),FLAVA 在单模态 Benchmark(如 NLP 的 GLUE 工作和视觉的图像分类)上也体现出极佳的性能。


最后 FLAVA 模型约有 3.5 亿参数,并应用 ViT-B16 配置,用于图像和文本编码器。

Reference:https://arxiv.org/pdf/2010.11…

多模态交融 transformer 沿用了单模态编码器,但层数只有之前的 1/2。PyTorch 开发团队始终在摸索减少编码器的尺寸,以适应更大的 ViT 变量 (variant)。

扩大 FLAVA 的另一个方面,就是减少批尺寸。FLAVA 奇妙利用了 in-batch negative 的比照损失,这通常只在少量尺寸中才有。

Reference:https://openreview.net/pdfid=…

一般来说,当操作靠近最大可能的批尺寸时,也能实现最大的训练效率或吞吐量,这由可用的 GPU 内存数量决定(参见试验局部)。

下表演示了不同模型配置的输入,试验中已确定每个配置可能适应内存的最少量尺寸。

优化概览

PyTorch 提供了几种原生技术来无效地扩大模型。在上面的章节中会具体介绍三种办法,并演示如何利用这些技术,将 FLAVA 模型扩大到 100 亿参数。

分布式数据并行

分布式训练的一个常见终点是数据并行。数据并行在 GPU 之间复制模型,并进行数据集划分。不同的 GPU 会并行地解决不同的数据分区,并在模型权重更新前同步其梯度(通过 all reduce)。

下图展现了解决一个数据并行(正向迭代、反向迭代和权重更新步骤)的流程:


为了实现数据并行,PyTorch 提供了一个原生 API,即 DistributedDataParallel(DDP),它能够作为一个模块封装器 (module wrapper) 应用,如下所示:

from torchmultimodal.models.flava.model import flava_model_for_pretraining
import torch
import torch.distributed as dist

model = flava_model_for_pretraining().cuda()
# Initialize PyTorch Distributed process groups
# Please see https://pytorch.org/tutorials/intermediate/dist_tuto.html for details
dist.init_process_group(backend=”nccl”)
# Wrap model in DDP
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[torch.cuda.current_device()])

齐全分片式数据并行

训练应用程序的 GPU 内存应用能够大抵细分为模型输出、两头激活存储(intermediate activation,梯度计算须要用到)、模型参数、梯度和优化器状态。

扩大模型时通常会将这些元素同时减少。当单个 GPU 内存不足时,应用 DDP 扩大模型可能导致内存不足,因为它会在所有 GPU 上复制参数、梯度和优化器状态。

为了缩小这种复制并节俭 GPU 内存,能够将模型参数、梯度和优化器状态分片给所有 GPU,每个 GPU 只治理一个分片。这个办法参照了微软提出的 ZeRO-3。

这种办法的 PyTorch 原生实现可作为 FullyShardedDataParallel(FSDP)API,已在 PyTorch 1.12 中作为 beta 版性能公布。

在模块的正向和反向迭代过程中,FSDP 会依据计算须要对模型参数进行整合(应用 all-gather),并在计算后从新分片。它应用散射规约汇合来同步梯度,以确保分片的梯度是全局均匀的。FSDP 中模型的正向迭代和反向迭代流程如下:

应用 FSDP 时要用 API 封装模型的子模块,从而管制某一特定子模块何时被分片或不分片。FSDP 提供了一个开箱即用的 auto-wrapping API、几个封装策略 (wrapping policy) 以及编写策略的能力。

以下示例演示了如何用 FSDP 封装 FLAVA 模型。指定主动封装策略为:transformer_auto_wrap_policy。这将把单个 transformer 层 (TransformerEncoderLayer)、图像 transformer (ImageTransformer)、文本编码器 (BERTTextEncoder) 和多模态编码器 (FLAVATransformerWithoutEmbeddings) 封装为单个 FSDP 单元。

这采纳了一种递归封装的办法来进行无效的内存治理。例如,在单个 transformer 层的正向或反向迭代实现后,删除参数、开释内存从而缩小了峰值内存应用。

FSDP 还提供了一些可配置的选项来调整应用程序的性能,如本例中 limit_all_gathers 的应用。它能够避免过早地收集所有模型参数,加重应用程序的内存压力。

import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from torchmultimodal.models.flava.model import flava_model_for_pretraining
from torchmultimodal.models.flava.text_encoder import BertTextEncoder
from torchmultimodal.models.flava.image_encoder import ImageTransformer
from torchmultimodal.models.flava.transformer import FLAVATransformerWithoutEmbeddings
from torchmultimodal.modules.layers.transformer import TransformerEncoderLayer

model = flava_model_for_pretraining().cuda()
dist.init_process_group(backend=”nccl”)

model = FSDP(
               model,
               device_id=torch.cuda.current_device(),
               auto_wrap_policy=partial(
                   transformer_auto_wrap_policy,
                   transformer_layer_cls={
                       TransformerEncoderLayer,
                       ImageTransformer,
                       BERTTextEncoder,
                       FLAVATransformerWithoutEmbeddings
                   },
               ),
               limit_all_gathers=True,
           )

activation checkpointing

如上,两头激活存储 (intermediate activation)、模型参数、梯度和优化器状态会影响 GPU 内存的应用。FSDP 能够缩小后三者带来的内存耗费,但不能缩小激活所耗费的内存。激活所应用的内存随着批尺寸或暗藏层数量的减少而减少。

activation checkpointing 通过在反向迭代过程中从新计算激活,而非将其保留在特定 checkpointed 模块的内存中,来缩小内存的应用。

例如,通过对 27 亿参数模型利用 activation checkpointing,正向迭代后的流动内存峰值缩小了 4 倍。

PyTorch 提供了一个基于 wrapper 的 activation checkpointing API。且 checkpoint_wrapper 容许用户通过 check 封装单个模块,apply_activation_checkpointing 容许用户指定策略在整个模块中用 checkpointing 封装模块。

这两个 API 能够利用于大多数模型,因为它们不须要对模型定义代码进行任何批改。

然而,如果须要对 checkpointed segment 进行更细化的管制,如对模块内的特定性能进行 checkpointing,能够利用 torch.utils.checkpoint API,这须要批改模型代码。

activation checkpointing wrapper 对单个 FLAVA transformer 层(用 TransformerEncoderLayer 示意)的利用如下所示:

from torchmultimodal.models.flava.model import flava_model_for_pretraining
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import apply_activation_checkpointing, checkpoint_wrapper, CheckpointImpl
from torchmultimodal.modules.layers.transformer import TransformerEncoderLayer

model = flava_model_for_pretraining()
checkpoint_tformer_layers_policy = lambda submodule: isinstance(submodule, TransformerEncoderLayer)

apply_activation_checkpointing(
               model,
               checkpoint_wrapper_fn=checkpoint_wrapper,
               check_fn=checkpoint_tformer_layers_policy,
           )

如上所示,用 activation checkpointing 封装 FLAVA transformer 层,用 FSDP 封装整体模型,能够将 FLAVA 扩大到 100 亿参数。

试验

对于上文提到的不同优化办法,咱们将进一步试验其对系统性能的影响。

背景介绍:

  • 应用含 8 个 A100 40 GB GPU 的单节点
  • 运行 1000 次迭代预训练
  • 应用 bfloat16 数据类型的 PyTorch 混合精度训练 (automatic mixed precision)
  • 启用 TensorFloat32 格局,进步 A100 上的 matmul 性能
  • 将吞吐量定义为每秒解决的均匀项目数(测量吞吐量时疏忽前 100 次迭代)
  • 训练收敛及其对上游工作指标的影响,会作为将来钻研的新方向

图 1 显示了每个模型配置和优化的吞吐量,local batch size 为 8,在 1 个节点上可能的最大 batch size。优化的模型变体 (model variant) 没有数据点,阐明该模型无奈在单个节点上训练。

​图 1:不同配置下的训练吞吐量

图 2 展现了所有 GPU 在每个优化中可能的最少量尺寸。

​图 2:不同配置下可能的最大本地批尺寸

从中能够察看到:

  1. 扩大模型尺寸:

    DDP 只能在一个节点上适应 350M 和 900M 的模型。应用 FSDP 能够节俭内存,所以可能训练比 DDP 大 3 倍的模型(即 1.8B 和 2.7B 的变体)。将激活检查点(AC)与 FSDP 联合起来,能够训练更大的模型,约为 DDP 的 10 倍(如 4.8B 和 10B 变体)。

  2. 吞吐量:
  • 对于较小的模型,当批尺寸为 8 时,DDP 的吞吐量略高于或等于 FSDP,能够解释为 FSDP 须要额定的通信。FSDP 和 AC 联合在一起的吞吐量最低。这是因为 AC 在反向迭代的过程中,从新运行 checkpointed 正向迭代通道,为了节俭内存就义了额定的计算。然而,对于 2.7B 模型,与独自的 FSDP 相比,FSDP + AC 实际上具备更高的吞吐量。这是因为带有 FSDP 的 2.7B 模型即便在批处尺寸为 8 的状况下也靠近内存的极限,会触发 CUDA malloc retry,从而导致训练速度减慢。AC 有助于缩小内存压力导致 no retry。
  • 对于 DDP 和 FSDP + AC,模型的吞吐量会随着批尺寸的减少而减少。FSDP 对较小的变体也是如此。然而,对于 1.8B 和 2.7B 参数模型,当减少批尺寸时,吞吐量降落。一个潜在的起因是,在内存极限时,PyTorch 的 CUDA 内存治理可能不得不重试 cudaMalloc 调用或运行老本昂扬的碎片整顿 (defragmentation),以找到闲暇的内存块来解决工作负载的内存需要,这可能导致训练速度减慢。
  • 对于只能用 FSDP 训练的大模型 (1.8B,2.7B,4.8B) 而言,最高吞吐量的设置是用 FSDP+AC 扩大到最少量尺寸。对于 10B,能够察看到小批尺寸和最少量尺寸的吞吐量简直相等。这是因为 AC 会导致计算量减少,而最少量尺寸可能会因为在 CUDA 内存限度下运行,导致老本昂扬的碎片整顿操作。然而,对于这些大模型,批尺寸的减少足以抵销这种开销。
  1. 批尺寸:

    与 DDP 相比,独自应用 FSDP 能够实现略高的批尺寸。对于 350M 参数模型,应用 FSDP+AC 能够实现比 DDP 高 3 倍的批尺寸,对于 900M 参数模型,能够实现 5.5 倍的批尺寸。即便是 10B,最大的批尺寸也约是 20,这相当不错。FSDP+AC 基本上能够用较少的 GPU 实现较大的全局批尺寸,对比照学习工作 (contrastive learning task) 特地无效。

论断

随着多模态根底模型的倒退,扩大模型参数和高效训练正在成为一个重点畛域。PyTorch 生态系统旨在通过提供不同的工具,减速训练和扩大多模态模型。

将来,PyTorch 将减少对其余类型模型的反对,比方多模态生成模型,以及晋升相干技术的自动化。欢送大家继续关注 PyTorch 开发者社区公众号,你也能够扫码备注「PyTorch」,退出 PyTorch 社群。

PyTorch 官网博客、教程

最新进展、最佳实际

扫码备注退出讨论组

退出移动版