Transformers 4.37 中文文档（八）（5）

Transformers 4.37 中文文档（八）（4）https://developer.aliyun.com/article/1563211

高效的训练技巧

在单个 GPU 上进行高效训练的方法和工具

原文链接：huggingface.co/docs/transformers/v4.37.2/en/perf_train_gpu_one

本指南演示了您可以使用的实用技术，通过优化内存利用率、加快训练速度或两者兼而有之来提高模型训练的效率。如果您想了解 GPU 在训练过程中的使用情况，请先参考模型训练解剖概念指南。本指南侧重于实用技术。

如果您可以访问具有多个 GPU 的计算机，这些方法仍然有效，此外，您还可以利用多 GPU 部分中概述的其他方法。

在训练大型模型时，应同时考虑两个方面：

• 数据吞吐量/训练时间
  
• 模型性能
  

最大化吞吐量（样本/秒）可以降低训练成本。通常通过尽可能充分利用 GPU 并将 GPU 内存填满来实现这一目标。如果所需的批处理大小超出了 GPU 内存的限制，可以使用内存优化技术，如梯度累积，来帮助。

然而，如果首选的批处理大小适合内存，就没有理由应用内存优化技术，因为它们可能会减慢训练速度。仅仅因为可以使用大批处理大小，并不一定意味着应该这样做。作为超参数调整的一部分，您应该确定哪种批处理大小产生最佳结果，然后相应地优化资源。

本指南涵盖的方法和工具可以根据它们对训练过程的影响进行分类：

注意：当使用混合精度与小模型和大批处理大小时，会有一些内存节省，但对于大模型和小批处理大小，内存使用量会更大。

您可以结合上述方法以获得累积效果。这些技术对您都是可用的，无论您是使用 Trainer 训练模型，还是编写纯 PyTorch 循环，在这种情况下，您可以通过使用🤗 Accelerate 配置这些优化。

如果这些方法没有带来足够的收益，您可以探索以下选项：

• 考虑构建自己的自定义 Docker 容器，其中包含高效的软件预构建
  
• 考虑使用专家混合（MoE）的模型
  
• 将您的模型转换为 BetterTransformer 以利用 PyTorch 原生注意力
  

最后，如果即使切换到像 A100 这样的服务器级 GPU 后仍然不够，考虑切换到多 GPU 设置。所有这些方法在多 GPU 设置中仍然有效，此外，您还可以利用多 GPU 部分中概述的其他并行技术。

批处理大小选择

为了实现最佳性能，首先要确定适当的批处理大小。建议使用大小为 2^N 的批处理大小和输入/输出神经元计数。通常是 8 的倍数，但可以更高，具体取决于所使用的硬件和模型的数据类型。

有关参考，请查看 NVIDIA 关于全连接层（涉及 GEMMs（通用矩阵乘法））的输入/输出神经元计数和批次大小的建议。

Tensor Core 要求根据数据类型和硬件定义乘数。例如，对于 fp16 数据类型，推荐使用 8 的倍数，除非是 A100 GPU，此时使用 64 的倍数。

对于较小的参数，还要考虑维度量化效应。这是瓦片化发生的地方，正确的乘数可以显著加快速度。

梯度累积

梯度累积方法旨在以较小的增量计算梯度，而不是一次为整个批次计算梯度。这种方法涉及通过模型执行前向和反向传递，并在过程中累积梯度来迭代计算较小批次的梯度。一旦积累了足够数量的梯度，就会执行模型的优化步骤。通过使用梯度累积，可以将有效批次大小增加到 GPU 内存容量所施加的限制之外。然而，需要注意的是，梯度累积引入的额外前向和反向传递可能会减慢训练过程。

您可以通过向 TrainingArguments 添加gradient_accumulation_steps参数来启用梯度累积：

training_args = TrainingArguments(per_device_train_batch_size=1, gradient_accumulation_steps=4, **default_args)

在上述示例中，您的有效批次大小变为 4。

或者，使用🤗 Accelerate 来完全控制训练循环。在本指南的后面找到🤗 Accelerate 示例。

尽可能地最大化 GPU 使用率是建议的，但是高数量的梯度累积步骤可能会导致训练减速更加明显。考虑以下示例。假设per_device_train_batch_size=4，没有梯度累积达到了 GPU 的极限。如果您想要使用大小为 64 的批次进行训练，请不要将per_device_train_batch_size设置为 1，并将gradient_accumulation_steps设置为 64。相反，保持per_device_train_batch_size=4，并设置gradient_accumulation_steps=16。这样可以实现相同的有效批次大小，同时更好地利用可用的 GPU 资源。

有关更多信息，请参考RTX-3090和A100的批次大小和梯度累积基准。

梯度检查点

即使将批次大小设置为 1 并使用梯度累积，一些大型模型仍可能面临内存问题。这是因为还有其他组件也需要内存存储。

保存前向传递中的所有激活以便在反向传递期间计算梯度可能会导致显着的内存开销。另一种方法是在反向传递期间丢弃激活并在需要时重新计算它们，这将引入相当大的计算开销并减慢训练过程。

梯度检查点提供了这两种方法之间的折衷方案，并在计算图中保存了策略性选择的激活，因此只需重新计算一小部分激活以获得梯度。有关梯度检查点的深入解释，请参阅这篇很棒的文章。

要在 Trainer 中启用梯度检查点，请将相应的标志传递给 TrainingArguments：

training_args = TrainingArguments(
    per_device_train_batch_size=1, gradient_accumulation_steps=4, gradient_checkpointing=True, **default_args
)

或者，使用🤗 Accelerate-在本指南中找到🤗 Accelerate 示例（#使用加速）。

虽然梯度检查点可能提高内存效率，但会使训练速度减慢约 20%。

混合精度训练

混合精度训练是一种旨在通过利用较低精度数值格式来处理某些变量来优化训练模型的计算效率的技术。传统上，大多数模型使用  32 位浮点精度（fp32 或  float32）来表示和处理变量。然而，并非所有变量都需要这种高精度级别才能获得准确的结果。通过将某些变量的精度降低到较低的数值格式，如 16  位浮点（fp16 或  float16），我们可以加快计算速度。因为在这种方法中，一些计算是以半精度进行的，而一些仍然是以全精度进行的，所以这种方法被称为混合精度训练。

最常见的混合精度训练是通过使用 fp16（float16）数据类型来实现的，但是一些 GPU 架构（例如 Ampere 架构）提供了 bf16 和 tf32（CUDA 内部数据类型）数据类型。查看NVIDIA 博客以了解这些数据类型之间的区别。

fp16

混合精度训练的主要优势来自于将激活保存在半精度（fp16）中。尽管梯度也是以半精度计算的，但它们在优化步骤中被转换回全精度，因此在这里没有节省内存。虽然混合精度训练可以加快计算速度，但也可能导致更多的  GPU 内存被利用，特别是对于小批量大小。这是因为模型现在以 16 位和 32 位精度（GPU 上原始模型的 1.5 倍）存在于 GPU 上。

要启用混合精度训练，请将fp16标志设置为True：

training_args = TrainingArguments(per_device_train_batch_size=4, fp16=True, **default_args)

如果您更喜欢使用🤗 Accelerate，请在本指南中找到🤗 Accelerate 示例（#使用加速）。

BF16

如果您可以访问 Ampere 或更新的硬件，您可以使用 bf16 进行混合精度训练和评估。虽然 bf16 的精度比 fp16 差，但它具有更大的动态范围。在 fp16 中，您可以拥有的最大数字是65535，任何超过这个数字的数字都会导致溢出。bf16 的数字可以达到3.39e+38（！），这大约与 fp32 相同-因为两者都使用了 8 位用于数值范围。

您可以在🤗 Trainer 中启用 BF16：

training_args = TrainingArguments(bf16=True, **default_args)

TF32

Ampere 硬件使用一种名为 tf32 的神奇数据类型。它具有与 fp32 相同的数值范围（8 位），但是精度只有 10 位（与  fp16 相同），总共只使用了 19 位。它在这种意义上是“神奇的”，即您可以使用正常的 fp32 训练和/或推理代码，并通过启用 tf32  支持，您可以获得高达 3 倍的吞吐量改进。您只需要将以下内容添加到您的代码中：

import torch
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

CUDA 将在可能的情况下自动切换到使用 tf32 而不是 fp32，假设所使用的 GPU 来自 Ampere 系列。

根据NVIDIA 研究，大多数机器学习训练工作负载显示出与 fp32 相同的困惑度和收敛性。如果您已经在使用 fp16 或 bf16 混合精度，这也可能有助于提高吞吐量。

您可以在🤗 Trainer 中启用此模式：

TrainingArguments(tf32=True, **default_args)

tf32 无法通过tensor.to(dtype=torch.tf32)直接访问，因为它是内部 CUDA 数据类型。您需要torch>=1.7才能使用 tf32 数据类型。

有关 tf32 与其他精度的更多信息，请参考以下基准测试：RTX-3090和A100。

Flash Attention 2

您可以通过在 transformers 中使用 Flash Attention 2 集成来加快训练吞吐量。查看单 GPU 部分中的适当部分，了解如何加载带有 Flash Attention 2 模块的模型的更多信息。

优化器选择

用于训练变压器模型的最常用优化器是 Adam 或 AdamW（带有权重衰减的 Adam）。Adam  通过存储先前梯度的滚动平均值实现良好的收敛；然而，它会增加与模型参数数量相同数量级的额外内存占用。为了解决这个问题，您可以使用另一种优化器。例如，如果您在  NVIDIA GPU 上安装了NVIDIA/apex，或者在 AMD GPU 上安装了ROCmSoftwarePlatform/apex，adamw_apex_fused将为您提供所有支持的 AdamW 优化器中最快的训练体验。

Trainer 集成了各种可立即使用的优化器：adamw_hf、adamw_torch、adamw_torch_fused、adamw_apex_fused、adamw_anyprecision、adafactor或adamw_bnb_8bit。更多优化器可以通过第三方实现插入。

让我们更仔细地看看两种替代 AdamW 优化器：

1. adafactor可在 Trainer 中使用
   
2. adamw_bnb_8bit也可在 Trainer 中使用，但以下提供了第三方集成以供演示。
   

以 3B 参数模型“t5-3b”为例进行比较：

• 标准的 AdamW 优化器将需要 24GB 的 GPU 内存，因为它为每个参数使用 8 字节（8*3 => 24GB）
  
• Adafactor 优化器将需要超过 12GB。它为每个参数使用略多于 4 字节，因此 4*3，然后再加一些。
  
• 8 位 BNB 量化优化器将仅使用（2*3）6GB，如果所有优化器状态都被量化。
  

Adafactor

Adafactor 不会为权重矩阵中的每个元素存储滚动平均值。相反，它保留聚合信息（按行和列的滚动平均和），显著减少了内存占用。然而，与 Adam 相比，Adafactor 在某些情况下可能收敛较慢。

您可以通过在 TrainingArguments 中设置optim="adafactor"来切换到 Adafactor：

training_args = TrainingArguments(per_device_train_batch_size=4, optim="adafactor", **default_args)

结合其他方法（梯度累积、梯度检查点和混合精度训练），您可以在保持吞吐量的同时实现高达 3 倍的改进！然而，如前所述，Adafactor 的收敛性可能比 Adam 更差。

8 位 Adam

与 Adafactor 等聚合优化器状态不同，8 位 Adam 保留完整状态并对其进行量化。量化意味着以较低精度存储状态，并仅在优化时对其进行反量化。这类似于混合精度训练的思想。

要使用adamw_bnb_8bit，您只需在 TrainingArguments 中设置optim="adamw_bnb_8bit"：

training_args = TrainingArguments(per_device_train_batch_size=4, optim="adamw_bnb_8bit", **default_args)

然而，我们也可以使用第三方实现的 8 位优化器进行演示，以了解如何集成。

首先，按照 GitHub repo中的安装指南安装实现 8 位 Adam 优化器的bitsandbytes库。

接下来需要初始化优化器。这涉及两个步骤：

• 首先，将模型的参数分为两组 - 一组应用权重衰减，另一组不应用。通常，偏置和层归一化参数不会被权重衰减。
  
• 然后进行一些参数整理，以使用与先前使用的 AdamW 优化器相同的参数。
  

import bitsandbytes as bnb
from torch import nn
from transformers.trainer_pt_utils import get_parameter_names
training_args = TrainingArguments(per_device_train_batch_size=4, **default_args)
decay_parameters = get_parameter_names(model, [nn.LayerNorm])
decay_parameters = [name for name in decay_parameters if "bias" not in name]
optimizer_grouped_parameters = [
    {
        "params": [p for n, p in model.named_parameters() if n in decay_parameters],
        "weight_decay": training_args.weight_decay,
    },
    {
        "params": [p for n, p in model.named_parameters() if n not in decay_parameters],
        "weight_decay": 0.0,
    },
]
optimizer_kwargs = {
    "betas": (training_args.adam_beta1, training_args.adam_beta2),
    "eps": training_args.adam_epsilon,
}
optimizer_kwargs["lr"] = training_args.learning_rate
adam_bnb_optim = bnb.optim.Adam8bit(
    optimizer_grouped_parameters,
    betas=(training_args.adam_beta1, training_args.adam_beta2),
    eps=training_args.adam_epsilon,
    lr=training_args.learning_rate,
)

最后，将自定义优化器作为参数传递给Trainer：

trainer = Trainer(model=model, args=training_args, train_dataset=ds, optimizers=(adam_bnb_optim, None))

结合其他方法（梯度累积、梯度检查点和混合精度训练），您可以期望获得大约 3 倍的内存改进，甚至比使用 Adafactor 时的吞吐量稍高。

multi_tensor

pytorch-nightly 引入了torch.optim._multi_tensor，应该显着加快具有大量小特征张量的优化器的速度。最终应该成为默认设置，但如果您想更早尝试它，请查看这个 GitHub 问题。

数据预加载

达到良好训练速度的一个重要要求是能够以 GPU 能够处理的最大速度提供数据。默认情况下，所有操作都在主进程中进行，可能无法快速从磁盘读取数据，从而导致瓶颈，导致 GPU 利用率不足。配置以下参数以减少瓶颈：

• DataLoader(pin_memory=True, ...) - 确保数据预加载到 CPU 上的固定内存中，通常会导致从 CPU 到 GPU 内存的传输速度更快。
  
• DataLoader(num_workers=4, ...) - 生成几个工作进程以更快地预加载数据。在训练过程中，观察 GPU 利用率统计数据；如果远离 100％，尝试增加工作进程的数量。当然，问题可能出在其他地方，因此许多工作进程不一定会导致更好的性能。
  

在使用 Trainer 时，相应的 TrainingArguments 是：dataloader_pin_memory（默认为True），和dataloader_num_workers（默认为0）。

DeepSpeed ZeRO

DeepSpeed 是一个开源的深度学习优化库，与🤗 Transformers 和🤗 Accelerate 集成。它提供了各种功能和优化，旨在提高大规模深度学习训练的效率和可扩展性。

如果您的模型适合单个 GPU 并且有足够的空间来容纳小批量大小，则不需要使用  DeepSpeed，因为它只会减慢速度。但是，如果模型无法适应单个 GPU 或无法容纳小批量，则可以利用 DeepSpeed ZeRO +  CPU Offload，或 NVMe Offload 来处理更大的模型。在这种情况下，您需要单独安装该库，然后按照指南之一创建配置文件并启动  DeepSpeed：

• 有关 DeepSpeed 与 Trainer 集成的详细指南，请查看相应文档，特别是单个 GPU 部署部分。在笔记本中使用 DeepSpeed 需要进行一些调整；请查看相应指南。
  
• 如果您更喜欢使用🤗 Accelerate，请参考🤗 Accelerate DeepSpeed 指南。
  

使用 torch.compile

PyTorch 2.0 引入了一个新的编译函数，不需要对现有的 PyTorch 代码进行任何修改，只需添加一行代码即可优化您的代码：model = torch.compile(model)。

如果使用 Trainer，您只需要在 TrainingArguments 中传递torch_compile选项：

training_args = TrainingArguments(torch_compile=True, **default_args)

torch.compile使用 Python 的帧评估 API 自动从现有的 PyTorch 程序创建图。在捕获图之后，可以部署不同的后端以将图降低到优化引擎。您可以在PyTorch 文档中找到更多详细信息和基准测试。

torch.compile有一个不断增长的后端列表，可以通过调用torchdynamo.list_backends()找到，每个后端都有其可选依赖项。

通过在 TrainingArguments 中指定torch_compile_backend来选择要使用的后端。一些最常用的后端包括：

调试后端：

• dynamo.optimize("eager") - 使用 PyTorch 运行提取的 GraphModule。这在调试 TorchDynamo 问题时非常有用。
  
• dynamo.optimize("aot_eager") - 使用 AotAutograd 而不使用编译器，即仅使用 PyTorch eager 进行 AotAutograd 的提取前向和反向图。这对调试很有用，但不太可能提供加速。
  

训练和推理后端：

• dynamo.optimize("inductor") - 使用 TorchInductor 后端，通过利用 codegened Triton 内核实现 AotAutograd 和 cudagraphs。阅读更多
  
• dynamo.optimize("nvfuser") - nvFuser 与 TorchScript。阅读更多
  
• dynamo.optimize("aot_nvfuser") - nvFuser 与 AotAutograd。阅读更多
  
• dynamo.optimize("aot_cudagraphs") - 使用 AotAutograd 的 cudagraphs。阅读更多
  

仅推理后端：

• dynamo.optimize("ofi") - 使用 Torchscript 的 optimize_for_inference。阅读更多
  
• dynamo.optimize("fx2trt") - 使用 NVIDIA TensorRT 进行推理优化。阅读更多
  
• dynamo.optimize("onnxrt") - 使用 ONNXRT 进行 CPU/GPU 推理。阅读更多
  
• dynamo.optimize("ipex") - 使用 IPEX 进行 CPU 推理。阅读更多
  

使用torch.compile与🤗 Transformers 的示例，请查看这篇博文，介绍如何使用最新的 PyTorch 2.0 功能微调 BERT 模型进行文本分类

使用🤗 PEFT

参数高效微调（PEFT）方法在微调期间冻结预训练模型参数，并在其上添加少量可训练参数（适配器）。

因此，与优化器状态和梯度相关的内存大大减少。

例如，对于普通的 AdamW，优化器状态的内存需求将是：

• fp32 参数的副本：4 字节/参数
  
• 动量：4 字节/参数
  
• 方差：4 字节/参数
  

假设一个具有 70 亿参数和 2 亿参数注入低秩适配器的模型。

普通模型的优化器状态的内存需求将为 12 * 7 = 84 GB（假设有 7B 可训练参数）。

添加 Lora 会略微增加与模型权重相关的内存，并大幅减少优化器状态的内存需求至 12 * 0.2 = 2.4GB。

在PEFT 文档或PEFT 存储库中详细了解 PEFT 及其详细用法。

使用🤗 Accelerate

使用🤗 Accelerate可以在完全控制训练循环的同时使用上述方法，并且基本上可以使用纯 PyTorch 编写循环并进行一些微小修改。

假设您已将 TrainingArguments 中的方法组合如下：

training_args = TrainingArguments(
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    gradient_checkpointing=True,
    fp16=True,
    **default_args,
)

使用🤗 Accelerate 的完整示例训练循环只有几行代码：

from accelerate import Accelerator
from torch.utils.data.dataloader import DataLoader
dataloader = DataLoader(ds, batch_size=training_args.per_device_train_batch_size)
if training_args.gradient_checkpointing:
    model.gradient_checkpointing_enable()
accelerator = Accelerator(fp16=training_args.fp16)
model, optimizer, dataloader = accelerator.prepare(model, adam_bnb_optim, dataloader)
model.train()
for step, batch in enumerate(dataloader, start=1):
    loss = model(**batch).loss
    loss = loss / training_args.gradient_accumulation_steps
    accelerator.backward(loss)
    if step % training_args.gradient_accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

首先，我们将数据集包装在DataLoader中。然后，我们可以通过调用模型的 gradient_checkpointing_enable()方法启用梯度检查点。当我们初始化Accelerator时，我们可以指定是否要使用混合精度训练，并且它将在prepare调用中为我们处理。在prepare调用期间，如果我们使用多个 GPU，数据加载器也将分布在工作进程之间。我们使用与之前示例相同的 8 位优化器。

最后，我们可以添加主要的训练循环。请注意，backward调用由🤗 Accelerate 处理。我们还可以看到梯度累积的工作原理：我们规范化损失，因此在累积结束时获得平均值，一旦我们有足够的步骤，我们就运行优化。

使用🤗 Accelerate 实现这些优化技术只需要几行代码，并且在训练循环中具有更大的灵活性。要查看所有功能的完整文档，请查看Accelerate 文档。

高效的软件预构建

PyTorch 的pip 和 conda 构建预先构建了 cuda 工具包，足以运行 PyTorch，但如果需要构建 cuda 扩展，则不足。

有时，可能需要额外的努力来预构建一些组件。例如，如果您使用的是未经预编译的库，如apex。在其他情况下，弄清楚如何在系统范围内安装正确的 cuda 工具包可能会很复杂。为了解决这些情况，PyTorch 和 NVIDIA 发布了一个新版本的 NGC docker 容器，其中已经预先构建了一切。您只需在其中安装您的程序，它就可以立即运行。

这种方法在您想要调整 pytorch 源代码和/或制作新的定制构建时也很有用。要找到您想要的 docker 镜像版本，请从PyTorch 发布说明开始，选择最新的一个月发布之一。进入所需发布的发布说明，检查环境的组件是否符合您的需求（包括 NVIDIA 驱动程序要求！），然后在该文档的顶部转到相应的 NGC 页面。如果由于某种原因您迷失了方向，这里是所有 PyTorch NGC 镜像的索引。

接下来按照说明下载和部署 docker 镜像。

专家混合

一些最近的论文报告了 4-5 倍的训练加速和将 Mixture of Experts（MoE）集成到 Transformer 模型中以实现更快的推理。

由于发现更多的参数会导致更好的性能，这种技术允许将参数数量增加一个数量级，而不增加训练成本。

在这种方法中，每个其他的 FFN 层都被一个 MoE 层替换，该层由许多专家组成，具有一个门控函数，根据输入令牌在序列中的位置平衡地训练每个专家。

(来源：GLAM)

您可以在本节末尾列出的论文中找到详尽的细节和比较表。

这种方法的主要缺点是它需要大量的 GPU 内存 - 几乎比其密集等价物大一个数量级。提出了各种蒸馏和方法，以克服更高的内存需求。

然而，存在直接的权衡，您可以使用少量专家和 2-3 倍较小的基础模型，而不是数十或数百个专家，从而导致 5 倍较小的模型，因此适度增加训练速度，同时适度增加内存需求。

大多数相关论文和实现都是围绕 Tensorflow/TPUs 构建的：

• GShard: 使用条件计算和自动分片扩展巨型模型
  
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
  
• GLaM: Generalist Language Model (GLaM)
  

对于 Pytorch，DeepSpeed 也构建了一个：DeepSpeed-MoE: 推进混合专家推理和训练以支持下一代 AI 规模，Mixture of Experts - 博文：1，2以及大型基于 Transformer 的自然语言生成模型的特定部署：博文，Megatron-Deepspeed 分支。

使用 PyTorch 原生注意力和 Flash Attention

PyTorch 2.0 发布了一个原生的torch.nn.functional.scaled_dot_product_attention（SDPA），允许使用融合的 GPU 内核，如内存高效注意力和闪存注意力。

安装optimum包后，可以替换相关的内部模块以使用 PyTorch 的原生注意力，方法如下：

model = model.to_bettertransformer()

转换后，像往常一样训练模型。

PyTorch 原生的scaled_dot_product_attention操作符只有在没有提供attention_mask时才能分派到 Flash Attention。

默认情况下，在训练模式下，BetterTransformer 集成取消了掩码支持，只能用于不需要填充掩码的批量训练。例如，在掩码语言建模或因果语言建模期间。BetterTransformer 不适用于需要填充掩码的任务的微调模型。

查看这篇博文，了解有关 SDPA 加速和节省内存的更多信息。

-in-context-learning-with.html))

您可以在本节末尾列出的论文中找到详尽的细节和比较表。

这种方法的主要缺点是它需要大量的 GPU 内存 - 几乎比其密集等价物大一个数量级。提出了各种蒸馏和方法，以克服更高的内存需求。

然而，存在直接的权衡，您可以使用少量专家和 2-3 倍较小的基础模型，而不是数十或数百个专家，从而导致 5 倍较小的模型，因此适度增加训练速度，同时适度增加内存需求。

大多数相关论文和实现都是围绕 Tensorflow/TPUs 构建的：

• GShard: 使用条件计算和自动分片扩展巨型模型
  
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
  
• GLaM: Generalist Language Model (GLaM)
  

对于 Pytorch，DeepSpeed 也构建了一个：DeepSpeed-MoE: 推进混合专家推理和训练以支持下一代 AI 规模，Mixture of Experts - 博文：1，2以及大型基于 Transformer 的自然语言生成模型的特定部署：博文，Megatron-Deepspeed 分支。

使用 PyTorch 原生注意力和 Flash Attention

PyTorch 2.0 发布了一个原生的torch.nn.functional.scaled_dot_product_attention（SDPA），允许使用融合的 GPU 内核，如内存高效注意力和闪存注意力。

安装optimum包后，可以替换相关的内部模块以使用 PyTorch 的原生注意力，方法如下：

model = model.to_bettertransformer()

转换后，像往常一样训练模型。

PyTorch 原生的scaled_dot_product_attention操作符只有在没有提供attention_mask时才能分派到 Flash Attention。

默认情况下，在训练模式下，BetterTransformer 集成取消了掩码支持，只能用于不需要填充掩码的批量训练。例如，在掩码语言建模或因果语言建模期间。BetterTransformer 不适用于需要填充掩码的任务的微调模型。

查看这篇博文，了解有关 SDPA 加速和节省内存的更多信息。