关于人工智能: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 是一个根底多模态模型，由基于 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,
           )

如上，两头激活存储 (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：不同配置下可能的最大本地批尺寸

从中能够察看到：

扩大模型尺寸：

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

对于较小的模型，当批尺寸为 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 内存限度下运行，导致老本昂扬的碎片整顿操作。然而，对于这些大模型，批尺寸的减少足以抵销这种开销。

批尺寸：

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

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

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

PyTorch 官网博客、教程

最新进展、最佳实际

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

扩大 FLAVA 概览

优化概览

分布式数据并行

齐全分片式数据并行

activation checkpointing

试验

论断