前往小程序,Get更优阅读体验!
立即前往
首页
学习
活动
专区
工具
TVP
发布
社区首页 >专栏 >Transformers 4.37 中文文档(九)

Transformers 4.37 中文文档(九)

作者头像
ApacheCN_飞龙
发布2024-06-26 14:48:39
700
发布2024-06-26 14:48:39
举报
文章被收录于专栏:信数据得永生信数据得永生

原文:huggingface.co/docs/transformers

多 GPU 上的高效训练

原始文本:huggingface.co/docs/transformers/v4.37.2/en/perf_train_gpu_many

如果在单个 GPU 上训练模型太慢或者模型的权重无法适应单个 GPU 的内存,则过渡到多 GPU 设置可能是一个可行的选择。在进行此过渡之前,彻底探索在单个 GPU 上进行高效训练的方法和工具中涵盖的所有策略,因为它们普遍适用于任意数量的 GPU 上的模型训练。一旦您采用了这些策略并发现它们在单个 GPU 上不足以满足您的情况时,请考虑转移到多个 GPU。

从单个 GPU 过渡到多个 GPU 需要引入某种形式的并行性,因为工作负载必须分布在资源之间。可以采用多种技术来实现并行性,例如数据并行性,张量并行性和管道并行性。重要的是要注意,没有一种大小适合所有的解决方案,最佳设置取决于您正在使用的特定硬件配置。

本指南提供了对各种并行性类型的深入概述,以及有关如何组合的指导

技术和选择适当的方法。有关分布式训练的逐步教程,请参考🤗加速文档

尽管本指南讨论的主要概念可能适用于各种框架,但在这里我们专注于基于 PyTorch 的实现。

在深入研究每种技术的具体细节之前,让我们回顾一下在大型基础设施上训练大型模型时的大致决策过程。

可扩展性策略

首先估算训练模型所需的 vRAM 量。对于托管在🤗 Hub 上的模型,请使用我们的模型内存计算器,该计算器可以在几个百分点的范围内给出准确的计算结果。

单节点/多 GPU 设置的并行化策略

在单节点上使用多个 GPU 训练模型时,您选择的并行化策略可能会显著影响性能。以下是您的选择:

情况 1:您的模型适合单个 GPU

如果您的模型可以轻松适应单个 GPU,则有两个主要选项:

  1. DDP - 分布式数据并行
  2. ZeRO - 根据所使用的情况和配置,这种方法可能会更快,但是,值得尝试一下。

情况 2:您的模型不适合单个 GPU:

如果您的模型太大而无法适应单个 GPU,您有几种替代方案可考虑:

  1. PipelineParallel(PP)
  2. ZeRO
  3. TensorParallel(TP)

具有非常快的节点间连接(例如 NVLINK 或 NVSwitch)时,所有三种策略(PP,ZeRO,TP)应该会产生类似的性能。但是,如果没有这些,PP 将比 TP 或 ZeRO 更快。 TP 的程度也可能会有所不同。最好根据您的特定设置进行实验,以确定最合适的策略。

TP 几乎总是在单个节点内使用。即 TP 大小<=每个节点的 GPU 数。

情况 3:您的模型最大层不适合单个 GPU

  1. 如果您没有使用 ZeRO,则必须使用 TensorParallel(TP),因为仅使用 PipelineParallel(PP)将无法容纳大层。
  2. 如果您使用 ZeRO,另外采用来自在单个 GPU 上进行高效训练的方法和工具的技术。

多节点/多 GPU 设置的并行化策略

  • 当您具有快速的节点间连接(例如 NVLINK 或 NVSwitch)时,请考虑使用以下选项之一:
    1. ZeRO - 因为它几乎不需要对模型进行修改
    2. 将 PipelineParallel(PP)与 TensorParallel(TP)和 DataParallel(DP)结合使用-这种方法将减少通信量,但需要对模型进行重大更改
  • 当您的节点间连接速度慢且 GPU 内存不足时:
    1. 使用 DataParallel(DP) 与 PipelineParallel(PP)、TensorParallel(TP) 和 ZeRO 的组合。

在本指南的后续部分中,我们将深入探讨这些不同的并行方法是如何工作的。

数据并行

即使只有 2 个 GPU,您也可以充分利用 PyTorch 内置功能提供的加速训练能力,例如 DataParallel (DP) 和 DistributedDataParallel (DDP)。请注意,PyTorch 文档建议优先选择 DistributedDataParallel (DDP) 而不是 DataParallel (DP) 用于多 GPU 训练,因为它适用于所有模型。让我们看看这两种方法是如何工作的以及它们之间的区别。

DataParallel vs DistributedDataParallel

为了了解这两种方法之间的 GPU 间通信开销的关键差异,让我们回顾每批的过程:

DDP:

  • 在开始时,主进程将模型从 GPU 0 复制到其他 GPU
  • 然后对于每批:
    1. 每个 GPU 直接消耗其数据的小批量。
    2. backward 过程中,一旦本地梯度准备就绪,它们将在所有进程中进行平均。

DP:

对于每批:

  1. GPU 0 读取数据批次,然后将一个小批量发送到每个 GPU。
  2. 最新的模型从 GPU 0 复制到每个 GPU。
  3. forward 被执行,每个 GPU 的输出被发送到 GPU 0 来计算损失。
  4. 损失从 GPU 0 分发到所有 GPU,并运行 backward
  5. 每个 GPU 的梯度被发送到 GPU 0 并进行平均。

关键差异包括:

  1. DDP 每批只执行一次通信 - 发送梯度,而 DP 每批执行五次不同的数据交换。DDP 使用 torch.distributed 复制数据,而 DP 通过 Python 线程在进程内复制数据(这会引入与 GIL 相关的限制)。因此,DistributedDataParallel (DDP) 通常比 DataParallel (DP) 更快,除非您的 GPU 卡之间的连接速度很慢。
  2. 在 DP 下,GPU 0 执行的工作明显多于其他 GPU,导致 GPU 利用率不足。
  3. DDP 支持跨多台机器的分布式训练,而 DP 不支持。

这并不是 DP 和 DDP 之间的所有差异的详尽列表,然而,其他细微差别超出了本指南的范围。您可以通过阅读这篇文章来更深入地了解这些方法。

让我们通过一个实验来说明 DP 和 DDP 之间的差异。我们将基准测试 DP 和 DDP 之间的差异,并增加 NVLink 存在的背景:

  • 硬件:2x TITAN RTX 24GB 每个 + 带有 2 个 NVLink 的 NVlink(nvidia-smi topo -m 中的 NV2)。
  • 软件:pytorch-1.8-to-be + cuda-11.0 / transformers==4.3.0.dev0

为了在其中一个基准测试中禁用 NVLink 功能,我们使用 NCCL_P2P_DISABLE=1

这里是基准测试代码和输出:

DP

代码语言:javascript
复制
rm -r /tmp/test-clm; CUDA_VISIBLE_DEVICES=0,1 \
python examples/pytorch/language-modeling/run_clm.py \
--model_name_or_path gpt2 --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 \
--do_train --output_dir /tmp/test-clm --per_device_train_batch_size 4 --max_steps 200

{'train_runtime': 110.5948, 'train_samples_per_second': 1.808, 'epoch': 0.69}

DDP w/ NVlink

代码语言:javascript
复制
rm -r /tmp/test-clm; CUDA_VISIBLE_DEVICES=0,1 \
torchrun --nproc_per_node 2 examples/pytorch/language-modeling/run_clm.py \
--model_name_or_path gpt2 --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 \
--do_train --output_dir /tmp/test-clm --per_device_train_batch_size 4 --max_steps 200

{'train_runtime': 101.9003, 'train_samples_per_second': 1.963, 'epoch': 0.69}

DDP w/o NVlink

代码语言:javascript
复制
rm -r /tmp/test-clm; NCCL_P2P_DISABLE=1 CUDA_VISIBLE_DEVICES=0,1 \
torchrun --nproc_per_node 2 examples/pytorch/language-modeling/run_clm.py \
--model_name_or_path gpt2 --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 \
--do_train --output_dir /tmp/test-clm --per_device_train_batch_size 4 --max_steps 200

{'train_runtime': 131.4367, 'train_samples_per_second': 1.522, 'epoch': 0.69}

这里是为方便起见在表格中收集的相同基准测试结果:

类型

NVlink

时间

2:DP

Y

110 秒

2:DDP

Y

101 秒

2:DDP

N

131 秒

如您所见,在这种情况下,DP 比 DDP with NVlink 慢约 10%,但比 DDP without NVlink 快约 15%。真正的差异取决于每个 GPU 需要与其他 GPU 同步多少数据 - 同步的数据越多,慢速链接就会阻碍整体运行时间。

ZeRO 数据并行

ZeRO 动力数据并行(ZeRO-DP)在这篇博文中有详细说明。

DeepSpeed-Image-1
DeepSpeed-Image-1

虽然看起来复杂,但这与DataParallel(DP)非常相似。不同之处在于,每个 GPU 只存储其一部分,而不是复制完整的模型参数、梯度和优化器状态。然后,在运行时,当需要完整的层参数时,所有 GPU 会同步以互相提供它们缺少的部分。

为了说明这个想法,考虑一个具有 3 层(La,Lb 和 Lc)的简单模型,其中每层有 3 个参数。例如,层 La 有权重 a0,a1 和 a2:

代码语言:javascript
复制
La | Lb | Lc
---|----|---
a0 | b0 | c0
a1 | b1 | c1
a2 | b2 | c2

如果我们有 3 个 GPU,ZeRO-DP 将模型分割到 3 个 GPU 上:

代码语言:javascript
复制
GPU0:
La | Lb | Lc
---|----|---
a0 | b0 | c0

GPU1:
La | Lb | Lc
---|----|---
a1 | b1 | c1

GPU2:
La | Lb | Lc
---|----|---
a2 | b2 | c2

在某种程度上,这与张量并行性相同的水平切片,与垂直切片相对,其中将整个层组放在不同的 GPU 上。现在让我们看看这是如何工作的:

这些 GPU 中的每一个都将像 DP 中那样获得通常的小批量:

代码语言:javascript
复制
x0 => GPU0
x1 => GPU1
x2 => GPU2

输入被传递而不进行修改,就好像原始模型会处理它们一样。

首先,输入到达层La。在这一点上会发生什么?

在 GPU0 上:x0 小批量需要 a0,a1,a2 参数通过层进行前向路径,但 GPU0 只有 a0。它将从 GPU1 获取 a1,从 GPU2 获取 a2,将模型的所有部分汇集在一起。

同时,GPU1 获得另一个小批量-x1。GPU1 具有 a1 参数,但需要 a0 和 a2,因此它从 GPU0 和 GPU2 获取这些。GPU2 也发生同样的情况,它获得小批量 x2。它从 GPU0 和 GPU1 获取 a0 和 a1。

这样,每个 3 个 GPU 都会得到完整的张量重建,并使用自己的小批量进行前向传递。一旦计算完成,不再需要的数据将被丢弃-它只在计算过程中使用。重建通过预取高效地完成。

然后整个过程会为层 Lb 重复,然后是 Lc 的前向,然后是 Lc -> Lb -> La 的后向。

这种机制类似于一种高效的团体背包策略:A 人携带帐篷,B 人携带炉灶,C 人携带斧头。每晚他们都分享自己拥有的东西,并从他人那里得到他们没有的东西,早上他们收拾好自己分配的装备类型,继续前进。这就是 ZeRO DP/Sharded DDP。将这种策略与每个人都必须携带自己的帐篷、炉灶和斧头的简单策略进行比较(类似于 PyTorch 中的 DataParallel(DP 和 DDP)),后者将更加低效。

在阅读这个主题的文献时,您可能会遇到以下同义词:分片,分区。如果您仔细注意 ZeRO 如何分割模型的权重-它看起来非常类似于张量并行性,稍后将对此进行讨论。这是因为它分割/分片每个层的权重,与接下来将讨论的垂直模型并行性不同。

实施:

从天真的模型并行性到管道并行性

解释管道并行性,我们首先将研究天真的模型并行性(MP),也称为垂直 MP。这种方法涉及通过使用.to()将模型层组分配到多个 GPU 上,将特定层分配给特定 GPU。当数据流经这些层时,它被移动到与该层相同的 GPU 上,而其他层保持不变。

我们将这种模型并行性称为“垂直”,因为模型通常是如何可视化的。例如,以下图表显示将 8 层模型垂直分割成两个切片,将层 0-3 放在 GPU0 上,将层 4-7 放在 GPU1 上:

代码语言:javascript
复制
===================  ===================
| 0 | 1 | 2 | 3 |  | 4 | 5 | 6 | 7 |
===================  ===================
        GPU0                 GPU1

在这个例子中,当数据从第 0 层移动到第 3 层时,与常规前向传递没有区别。然而,将数据从第 3 层传递到第 4 层需要将其从 GPU0 移动到 GPU1,引入了通信开销。如果参与的 GPU 在同一计算节点上(例如同一台物理机器),这种复制是快速的,但如果 GPU 分布在不同的计算节点上(例如多台机器),通信开销可能会大大增加。

接下来,第 4 到第 7 层的工作方式与原始模型中的工作方式相同。在完成第 7 层后,通常需要将数据发送回第 0 层,那里有标签(或者将标签发送到最后一层)。现在可以计算损失并让优化器开始工作。

天真的模型并行存在一些缺点:

  • 除了一个 GPU 外,任何时刻都有空闲的 GPU:如果使用 4 个 GPU,几乎等同于将单个 GPU 的内存量增加四倍,并忽略其他硬件。
  • 设备之间数据传输的开销:例如,使用天真的 MP,4 张 6GB 的卡片可以容纳与 1 张 24GB 卡片相同大小的模型,但是单张 24GB 卡片会更快地完成训练,因为它没有数据复制的开销。但是,假设您有 40GB 的卡片,需要适应一个 45GB 的模型,您可以使用 4 张 40GB 的卡片(但是可能很勉强,因为梯度和优化器状态)。
  • 复制共享的嵌入:共享的嵌入可能需要在 GPU 之间来回复制。

现在您已经了解了模型并行的天真方法以及其缺点,让我们来看看管道并行(PP)。PP 几乎与天真的 MP 相同,但通过将传入的批次分成微批次并人为地创建一个管道来解决 GPU 空闲问题,这允许不同的 GPU 同时参与计算过程。

以下来自GPipe 论文的插图显示了顶部的天真 MP 和底部的 PP:

MP vs PP
MP vs PP

在图表底部,您可以观察到管道并行(PP)方法最小化了空闲 GPU 区域的数量,称为“气泡”。图表的两部分显示了并行度为 4 的水平,意味着有 4 个 GPU 参与了管道。您可以看到有一个由 4 个管道阶段(F0、F1、F2 和 F3)组成的前向路径,然后是一个反向路径,顺序相反(B3、B2、B1 和 B0)。

PP 引入了一个新的超参数来调整 - chunks,它确定通过同一管道阶段连续发送多少数据块。例如,在底部图表中,您可以看到chunks=4。GPU0 在块 0、1、2 和 3(F0,0、F0,1、F0,2、F0,3)上执行相同的前向路径,然后等待其他 GPU 完成它们的工作。只有当其他 GPU 开始完成它们的工作时,GPU0 才开始再次工作,为块 3、2、1 和 0(B0,3、B0,2、B0,1、B0,0)执行反向路径。

请注意,这与梯度累积步骤的概念相同。PyTorch 使用chunks,而 DeepSpeed 将同一超参数称为梯度累积步骤。

由于 chunks,PP 引入了微批次(MBS)的概念。DP 将全局数据批次大小分成小批次,因此如果 DP 度为 4,则全局批次大小为 1024 将分成 4 个每个 256 的小批次(1024/4)。如果chunks(或 GAS)的数量为 32,则我们最终得到一个微批次大小为 8(256/32)。每个管道阶段一次处理一个微批次。要计算 DP + PP 设置的全局批次大小,请使用公式:mbs * chunks * dp_degree8 * 32 * 4 = 1024)。当chunks=1时,您将得到天真的 MP,这是低效的。当chunks值很大时,您将得到微小的微批次大小,这也是低效的。因此,我们鼓励尝试不同的chunks值,以找到导致最有效的 GPU 利用率的值。

您可能会注意到图表上的“死”时间泡泡,因为最后的forward阶段必须等待backward完成管道。找到chunks的最佳值的目的是实现所有参与 GPU 之间的高并发 GPU 利用率,从而最小化泡泡的大小。

Pipeline API 解决方案已在以下平台中实施:

  • PyTorch
  • DeepSpeed
  • Megatron-LM

这些存在一些缺点:

  • 他们必须对模型进行相当大的修改,因为 Pipeline 要求将模块的正常流重写为相同的nn.Sequential序列,这可能需要对模型的设计进行更改。
  • 目前,Pipeline API 非常受限制。如果在管道的第一个阶段传递了一堆 Python 变量,您将不得不找到解决方法。目前,管道接口要求作为唯一输入和输出的是单个张量或张量元组。这些张量的批次大小必须作为第一个维度,因为管道将小批次分成微批次。这里正在讨论可能的改进github.com/pytorch/pytorch/pull/50693
  • 在管道阶段的条件控制流不可能-例如,编码器-解码器模型(如 T5)需要特殊的解决方案来处理条件编码器阶段。
  • 他们必须安排每一层,以便一层的输出成为另一层的输入。

更近期的解决方案包括:

  • Varuna
  • Sagemaker

我们尚未尝试 Varuna 和 SageMaker,但他们的论文报告称,他们已经克服了上述问题列表,并且需要对用户模型进行较小的更改。

实施:

  • PyTorch(在 pytorch-1.8 中提供初始支持,并在 1.9 中逐渐改进,在 1.10 中更是如此)。一些示例
  • DeepSpeed
  • Megatron-LM具有内部实现-没有 API。
  • Varuna
  • SageMaker - 这是一种专有解决方案,只能在 AWS 上使用。
  • OSLO - 这是基于 Hugging Face Transformers 实现的。

🤗 Transformers 状态:截至目前,没有模型支持完整的 PP。GPT2 和 T5 模型具有天真的 MP 支持。主要障碍是无法将模型转换为nn.Sequential并使所有输入为张量。这是因为当前模型包含许多使转换非常复杂的特性,并且需要将其删除才能实现转换。

DeepSpeed 和 Megatron-LM 集成可在🤗 Accelerate中找到

其他方法:

DeepSpeed、Varuna 和 SageMaker 使用交错管道的概念

交错管道执行
交错管道执行

在这里,通过优先考虑反向传递来进一步减少气泡(空闲时间)。Varuna 通过使用模拟来发现最有效的调度来尝试改进时间表。

OSLO 具有基于 Transformers 的管道并行实现,无需nn.Sequential转换。

张量并行

在张量并行中,每个 GPU 处理张量的一个切片,并且仅在需要时聚合完整的张量进行操作。为了描述这种方法,本指南的这一部分依赖于Megatron-LM论文中的概念和图表:在 GPU 集群上进行高效的大规模语言模型训练

任何 transformer 的主要构建块是一个完全连接的nn.Linear,后面跟着一个非线性激活GeLU。它的点积部分,根据 Megatron 的论文符号表示法,可以写成Y = GeLU(XA),其中X是输入向量,Y是输出向量,A是权重矩阵。

如果我们以矩阵形式看计算,您可以看到矩阵乘法如何在多个 GPU 之间分割:

并行 GEMM
并行 GEMM

如果我们将权重矩阵AN个 GPU 上按列划分,并在并行执行矩阵乘法XA_1XA_n,那么我们将得到N个输出向量Y_1, Y_2, ..., Y_n,可以独立地输入到GeLU中:

独立 GeLU
独立 GeLU

使用这个原则,我们可以更新任意深度的多层感知器,而无需在 GPU 之间进行任何同步,直到最后,在那里我们需要从碎片中重建输出向量。Megatron-LM 论文的作者为此提供了一个有用的插图:

并行碎片处理
并行碎片处理

并行化多头注意力层甚至更简单,因为它们已经天生是并行的,由于具有多个独立的头!

并行自注意力
并行自注意力

特殊考虑:TP 需要非常快的网络,因此不建议在多个节点之间进行 TP。实际上,如果一个节点有 4 个 GPU,则最高的 TP 度数为 4。如果您需要 8 的 TP 度数,则需要使用至少有 8 个 GPU 的节点。

本节基于原始的更详细的TP 概述。by @anton-l

替代名称:

实施:

  • Megatron-LM具有内部实现,因为它非常特定于模型。
  • parallelformers(目前仅推理)
  • SageMaker - 这是一个专有解决方案,只能在 AWS 上使用。
  • OSLO具有基于 Transformers 的张量并行实现。

SageMaker 将 TP 与 DP 结合起来,以实现更高效的处理。

🤗 Transformers 状态:

  • 核心:核心尚未实现
  • 但是如果您需要推理parallelformers为我们大多数模型提供了支持。因此,在核心实现之前,您可以使用它们的支持。希望训练模式也会得到支持。
  • Deepspeed-Inference 还支持我们的 BERT、GPT-2 和 GPT-Neo 模型,采用超快的基于 CUDA 内核的推理模式,更多信息请参见这里

🤗 Accelerate 集成了Megatron-LM 的 TP

数据并行 + 管道并行

来自 DeepSpeed pipeline 教程的以下图表演示了如何将 DP 与 PP 结合使用。

DP + PP-2d
DP + PP-2d

在这里,重要的是看到 DP 等级 0 看不到 GPU2,DP 等级 1 看不到 GPU3。对于 DP 来说,只有 GPU 0 和 1,它们像只有 2 个 GPU 一样提供数据。GPU0“秘密”地将一些负载转移到 GPU2 上,使用 PP。GPU1 也通过将 GPU3 列入其援助来做同样的事情。

由于每个维度至少需要 2 个 GPU,所以这里至少需要 4 个 GPU。

实现:

🤗 Transformers 状态:尚未实现

数据并行 + 流水线并行 + 张量并行

为了获得更高效的训练,使用了 3D 并行,其中 PP 与 TP 和 DP 结合使用。可以在以下图表中看到。

dp-pp-tp-3d
dp-pp-tp-3d

这个图表来自一篇博文3D 并行:扩展到万亿参数模型,也是一篇很好的阅读。

由于每个维度至少需要 2 个 GPU,所以这里至少需要 8 个 GPU。

实现:

🤗 Transformers 状态:尚未实现,因为我们没有 PP 和 TP。

ZeRO 数据并行 + 流水线并行 + 张量并行

DeepSpeed 的主要特点之一是 ZeRO,它是 DP 的一个超可扩展扩展。已经在 ZeRO 数据并行中讨论过。通常它是一个独立的功能,不需要 PP 或 TP。但它可以与 PP 和 TP 结合使用。

当 ZeRO-DP 与 PP(和可选的 TP)结合时,通常只启用 ZeRO 阶段 1(优化器分片)。

虽然理论上可以使用 ZeRO 阶段 2(梯度分片)与流水线并行,但会对性能产生负面影响。每个微批次都需要额外的 reduce-scatter 集合来聚合梯度,然后再进行分片,这会增加潜在的显著通信开销。由于流水线并行的特性,使用小微批次,而重点是尝试平衡算术强度(微批次大小)和最小化流水线气泡(微批次数量)。因此,这些通信成本将影响性能。

此外,由于 PP 已经减少了梯度大小1/PP,所以梯度分片的节省并不像纯 DP 那样显著。

ZeRO 阶段 3 也不是一个好选择,原因是需要更多的节点间通信。

由于我们有 ZeRO,另一个好处是 ZeRO-Offload。由于这是阶段 1 优化器状态可以转移到 CPU。

实现:

重要论文:

  • 使用 DeepSpeed 和 Megatron 训练 Megatron-Turing NLG 530B,一个大规模生成式语言模型

🤗 Transformers 状态:尚未实现,因为我们没有 PP 和 TP。

FlexFlow

FlexFlow 也以略有不同的方式解决了并行化问题。

论文:“Beyond Data and Model Parallelism for Deep Neural Networks” by Zhihao Jia, Matei Zaharia, Alex Aiken

它执行一种 4D 并行化,涵盖样本-操作-属性-参数。

  1. 样本 = 数据并行化(样本方向并行)
  2. 操作 = 将单个操作并行化为多个子操作
  3. 属性 = 数据并行化(长度方向并行)
  4. 参数 = 模型并行化(无论维度是水平还是垂直)

示例:

  • 样本

让我们拿 10 批次的序列长度为 512。如果我们按样本维度将它们并行化为 2 个设备,我们得到 10 x 512,这将变为 5 x 2 x 512。

  • 操作

如果我们执行层归一化,首先计算标准差,然后计算均值,然后我们可以对数据进行归一化。操作并行性允许并行计算标准差和均值。因此,如果我们按操作维度将它们并行化为 2 个设备(cuda:0,cuda:1),首先将输入数据复制到两个设备中,cuda:0 同时计算标准差,cuda:1 计算均值。

  • 属性

我们有 10 批次,每个长度为 512。如果我们按属性维度将它们并行化为 2 个设备,10 x 512 将变为 10 x 2 x 256。

  • 参数

这与张量模型并行化或天真的逐层模型并行化类似。

flex-flow-soap
flex-flow-soap

这个框架的重要性在于,它以算法方式占用资源,如(1)GPU/TPU/CPU vs.(2)RAM/DRAM vs.(3)快速内部连接/慢速外部连接,并自动优化所有这些,决定在哪里使用哪种并行化。

一个非常重要的方面是,FlexFlow 专为优化具有静态和固定工作负载的 DNN 并行化而设计,因为具有动态行为的模型可能会在迭代中更喜欢不同的并行化策略。

因此,这个承诺非常吸引人 - 它在所选集群上运行 30 分钟的模拟,并提出了最佳策略来利用这个特定环境。如果添加/删除/替换任何部分,它将运行并重新优化该计划。然后您可以进行训练。不同的设置将有自己的定制优化。

🤗 Transformers 状态:Transformers 模型可以通过transformers.utils.fx进行 FX-trace-able,这是 FlexFlow 的先决条件,但需要在 FlexFlow 方面进行更改以使其与 Transformers 模型配合使用。

GPU 选择

在多个 GPU 上训练时,您可以指定要使用的 GPU 数量和顺序。例如,当您有计算能力不同的 GPU 并希望首先使用速度更快的 GPU 时,这可能很有用。选择过程适用于DistributedDataParallelDataParallel,以仅使用可用 GPU 的子集,您不需要 Accelerate 或 DeepSpeed integration。

GPU 的数量

例如,如果您有 4 个 GPU,但只想使用前 2 个:

torchrunAccelerateDeepSpeed

使用--nproc_per_node来选择使用多少个 GPU。

代码语言:javascript
复制
torchrun --nproc_per_node=2  trainer-program.py ...
GPU 的顺序

现在,要选择要使用的 GPU 及其顺序,您将使用CUDA_VISIBLE_DEVICES环境变量。最简单的方法是在~/bashrc或其他启动配置文件中设置环境变量。CUDA_VISIBLE_DEVICES用于映射要使用的 GPU。例如,如果您有 4 个 GPU(0、1、2、3),但只想运行 GPU 0 和 2:

代码语言:javascript
复制
CUDA_VISIBLE_DEVICES=0,2 torchrun trainer-program.py ...

只有 2 个物理 GPU(0 和 2)对 PyTorch 是“可见的”,它们分别映射到cuda:0cuda:1。您还可以颠倒 GPU 的顺序以先使用 2 个。现在,映射是 GPU 0 为cuda:1,GPU 2 为cuda:0

代码语言:javascript
复制
CUDA_VISIBLE_DEVICES=2,0 torchrun trainer-program.py ...

您还可以将 CUDA_VISIBLE_DEVICES 环境变量设置为空值,以创建一个没有 GPU 的环境。

代码语言:javascript
复制
CUDA_VISIBLE_DEVICES= python trainer-program.py ...

与任何环境变量一样,它们可以被导出,而不是添加到命令行中。然而,这并不推荐,因为如果您忘记了环境变量的设置方式,最终使用了错误的 GPU,会让人感到困惑。相反,通常的做法是在同一命令行上为特定的训练运行设置环境变量。

CUDA_DEVICE_ORDER 是一个替代环境变量,您可以使用它来控制 GPU 的顺序。您可以按照以下方式对它们进行排序:

  1. nvidia-smirocm-smi 分别匹配的 PCIe 总线 ID,用于 NVIDIA 和 AMD GPU
代码语言:javascript
复制
export CUDA_DEVICE_ORDER=PCI_BUS_ID
  1. GPU 计算能力
代码语言:javascript
复制
export CUDA_DEVICE_ORDER=FASTEST_FIRST

如果您的训练设置包括一台较旧和一台较新的 GPU,其中较旧的 GPU 显示在前,但您无法物理交换卡片使较新的 GPU 显示在前,那么 CUDA_DEVICE_ORDER 就特别有用。在这种情况下,设置 CUDA_DEVICE_ORDER=FASTEST_FIRST,始终使用较新和更快的 GPU(nvidia-smirocm-smi 仍然按照 PCIe 顺序报告 GPU)。或者您也可以设置 export CUDA_VISIBLE_DEVICES=1,0

完全分片数据并行

原始文本:huggingface.co/docs/transformers/v4.37.2/en/fsdp

完全分片数据并行(FSDP)是一种数据并行方法,它将模型的参数、梯度和优化器状态分片到可用 GPU(也称为工作器或rank)的数量上。与分布式数据并行(DDP)不同,FSDP 减少了内存使用,因为模型在每个 GPU 上都有副本。这提高了 GPU 内存效率,并允许您在较少的 GPU 上训练更大的模型。FSDP 与 Accelerate 集成,Accelerate 是一个用于轻松管理分布式环境中训练的库,这意味着可以从 Trainer 类中使用它。

在开始之前,请确保已安装 Accelerate,并且至少安装了 PyTorch 2.1.0 或更新版本。

代码语言:javascript
复制
pip install accelerate

FSDP 配置

首先,运行accelerate config命令,为您的训练环境创建一个配置文件。Accelerate 使用此配置文件根据您在accelerate config中选择的训练选项自动设置正确的训练环境。

代码语言:javascript
复制
accelerate config

当您运行accelerate config时,您将被提示一系列选项来配置您的训练环境。本节涵盖了一些最重要的 FSDP 选项。要了解更多关于其他可用的 FSDP 选项,请查看fsdp_config参数。

分片策略

FSDP 提供了许多分片策略可供选择:

  • FULL_SHARD - 在工作器之间对模型参数、梯度和优化器状态进行分片;选择1作为此选项
  • SHARD_GRAD_OP- 在工作器之间对梯度和优化器状态进行分片;选择2作为此选项
  • NO_SHARD - 不对任何内容进行分片(相当于 DDP);选择3作为此选项
  • HYBRID_SHARD - 在每个工作器内对模型参数、梯度和优化器状态进行分片,每个工作器也有完整副本;选择4作为此选项
  • HYBRID_SHARD_ZERO2 - 在每个工作器内对梯度和优化器状态进行分片,每个工作器也有完整副本;选择5作为此选项

这是通过fsdp_sharding_strategy标志启用的。

CPU 卸载

当参数和梯度不在使用时,您还可以将它们卸载到 CPU 以节省更多 GPU 内存,并帮助您适应大型模型,即使 FSDP 可能不足。这可以通过在运行accelerate config时设置fsdp_offload_params: true来启用。

包装策略

FSDP 通过包装网络中的每一层来应用。包装通常以嵌套方式应用,其中完整权重在每次前向传递后被丢弃,以节省内存供下一层使用。自动包装策略是实现这一点的最简单方法,您无需更改任何代码。您应该选择fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP来包装一个 Transformer 层,并选择fsdp_transformer_layer_cls_to_wrap来指定要包装的层(例如BertLayer)。

否则,您可以选择基于大小的包装策略,如果某一层的参数超过一定数量,则应用 FSDP。这可以通过将fsdp_wrap_policy: SIZE_BASED_WRAPmin_num_param设置为所需的大小阈值来启用。

检查点

中间检查点应该使用fsdp_state_dict_type: SHARDED_STATE_DICT保存,因为在 rank 0 上使用 CPU 卸载保存完整状态字典需要很长时间,并且经常由于广播期间的无限挂起而导致NCCL Timeout错误。您可以使用load_state`方法恢复训练。

代码语言:javascript
复制
# directory containing checkpoints
accelerator.load_state("ckpt")

然而,当训练结束时,您希望保存完整状态字典,因为分片状态字典仅与 FSDP 兼容。

代码语言:javascript
复制
if trainer.is_fsdp_enabled:
    trainer.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")

trainer.save_model(script_args.output_dir)
TPU

PyTorch XLA支持 TPU 的 FSDP 训练,可以通过修改accelerate config生成的 FSDP 配置文件来启用。除了上面指定的分片策略和包装选项之外,您还可以向文件中添加下面显示的参数。

代码语言:javascript
复制
xla: True # must be set to True to enable PyTorch/XLA
xla_fsdp_settings: # XLA-specific FSDP parameters
xla_fsdp_grad_ckpt: True # use gradient checkpointing

xla_fsdp_settings允许您配置 FSDP 的额外 XLA 特定参数。

启动训练

一个示例的 FSDP 配置文件可能如下所示:

代码语言:javascript
复制
compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch_policy: BACKWARD_PRE
  fsdp_cpu_ram_efficient_loading: true
  fsdp_forward_prefetch: false
  fsdp_offload_params: true
  fsdp_sharding_strategy: 1
  fsdp_state_dict_type: SHARDED_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_transformer_layer_cls_to_wrap: BertLayer
  fsdp_use_orig_params: true
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

启动训练,请运行accelerate launch命令,它将自动使用您之前使用accelerate config创建的配置文件。

代码语言:javascript
复制
accelerate launch my-trainer-script.py
代码语言:javascript
复制
accelerate launch --fsdp="full shard" --fsdp_config="path/to/fsdp_config/ my-trainer-script.py

下一步

FSDP 可以是训练非常大模型的强大工具,如果您有多个 GPU 或 TPU 可用。通过对模型参数、优化器和梯度状态进行分片,甚至在它们不活动时将它们卸载到 CPU 上,FSDP 可以减少大规模训练的高成本。如果您有兴趣了解更多,以下内容可能有所帮助:

在 CPU 上高效训练

原文:huggingface.co/docs/transformers/v4.37.2/en/perf_train_cpu

本指南侧重于在 CPU 上高效训练大型模型。

使用 IPEX 的混合精度

IPEX 优化了 AVX-512 或更高版本的 CPU,并且在仅具有 AVX2 的 CPU 上也可以正常工作。因此,预计在具有 AVX-512 或更高版本的英特尔 CPU 世代中,IPEX 将带来性能优势,而仅具有 AVX2 的 CPU(例如 AMD CPU 或较旧的英特尔 CPU)在 IPEX 下可能会获得更好的性能,但不能保证。IPEX 为使用 Float32 和 BFloat16 进行 CPU 训练提供了性能优化。以下部分的主要重点是使用 BFloat16。

低精度数据类型 BFloat16 已经在第三代 Xeon® Scalable Processors(又称 Cooper Lake)上本地支持,具有 AVX512 指令集,并将在下一代英特尔® Xeon® Scalable Processors 上支持,该处理器具有 Intel® Advanced Matrix Extensions(Intel® AMX)指令集,进一步提升性能。自 PyTorch-1.10 以来,已启用了 CPU 后端的自动混合精度。同时,Intel® Extension for PyTorch 大规模启用了 CPU 和 BFloat16 优化运算符的自动混合精度,并部分上游到 PyTorch 主分支。用户可以通过 IPEX 自动混合精度获得更好的性能和用户体验。

查看更多关于 自动混合精度 的详细信息。

IPEX 安装:

IPEX 发布遵循 PyTorch,可以通过 pip 安装:

PyTorch 版本

IPEX 版本

2.1.x

2.1.100+cpu

2.0.x

2.0.100+cpu

1.13

1.13.0+cpu

1.12

1.12.300+cpu

请运行 pip list | grep torch 以获取您的 pytorch_version,这样您就可以获得 IPEX version_name

代码语言:javascript
复制
pip install intel_extension_for_pytorch==<version_name> -f https://developer.intel.com/ipex-whl-stable-cpu

如果需要,您可以在 ipex-whl-stable-cpu 中查看最新版本。

查看更多关于 IPEX 安装 的方法。

Trainer 中的用法

要在 Trainer 中启用 IPEX 的自动混合精度,用户应该在训练命令参数中添加 use_ipexbf16no_cuda

Transformers 问答 用例为例

使用 BF16 自动混合精度在 CPU 上进行 IPEX 训练:

代码语言:javascript
复制
 python run_qa.py \
--model_name_or_path bert-base-uncased \
--dataset_name squad \
--do_train \
--do_eval \
--per_device_train_batch_size 12 \
--learning_rate 3e-5 \
--num_train_epochs 2 \
--max_seq_length 384 \
--doc_stride 128 \
--output_dir /tmp/debug_squad/ \
--use_ipex \
--bf16 \
--use_cpu

如果要在脚本中启用 use_ipexbf16,请像这样将这些参数添加到 TrainingArguments 中:

代码语言:javascript
复制
training_args = TrainingArguments(
    output_dir=args.output_path,
+   bf16=True,
+   use_ipex=True,
+   use_cpu=True,
    **kwargs
)
实践示例

博客:使用英特尔 Sapphire Rapids 加速 PyTorch Transformers

多 CPU 高效训练

原始文本:huggingface.co/docs/transformers/v4.37.2/en/perf_train_cpu_many

当在单个 CPU 上训练速度太慢时,我们可以使用多个 CPU。本指南侧重于基于 PyTorch 的 DDP,可以在 bare metal 和 Kubernetes 上有效地启用分布式 CPU 训练。

Intel® oneCCL 绑定的 PyTorch

Intel® oneCCL(集体通信库)是一个用于实现 allreduce、allgather、alltoall 等集体通信的高效分布式深度学习训练的库。有关 oneCCL 的更多信息,请参考oneCCL 文档oneCCL 规范

模块oneccl_bindings_for_pytorch(在 1.12 版本之前为torch_ccl)实现了 PyTorch C10D ProcessGroup API,并且可以作为外部 ProcessGroup 动态加载,目前仅在 Linux 平台上可用。

查看更多关于oneccl_bind_pt的详细信息。

Intel® oneCCL 绑定的 PyTorch 安装

以下 Python 版本的 Wheel 文件可用:

扩展版本

Python 3.6

Python 3.7

Python 3.8

Python 3.9

Python 3.10

2.1.0

2.0.0

1.13.0

1.12.100

1.12.0

请运行pip list | grep torch以获取您的pytorch_version

代码语言:javascript
复制
pip install oneccl_bind_pt=={pytorch_version} -f https://developer.intel.com/ipex-whl-stable-cpu

其中{pytorch_version}应该是您的 PyTorch 版本,例如 2.1.0。查看更多关于oneccl_bind_pt 安装的方法。oneCCL 和 PyTorch 的版本必须匹配。

oneccl_bindings_for_pytorch 1.12.0 预构建的 Wheel 文件与 PyTorch 1.12.1 不兼容(适用于 PyTorch 1.12.0)。PyTorch 1.12.1 应该与 oneccl_bindings_for_pytorch 1.12.100 兼容。

Intel® MPI 库

使用这个基于标准的 MPI 实现来在 Intel®架构上提供灵活、高效、可扩展的集群消息传递。这个组件是 Intel® oneAPI HPC Toolkit 的一部分。

oneccl_bindings_for_pytorch 是与 MPI 工具集一起安装的。在使用之前需要设置环境。

对于 Intel® oneCCL >= 1.12.0

代码语言:javascript
复制
oneccl_bindings_for_pytorch_path=$(python -c "from oneccl_bindings_for_pytorch import cwd; print(cwd)")
source $oneccl_bindings_for_pytorch_path/env/setvars.sh

对于版本< 1.12.0 的 Intel® oneCCL

代码语言:javascript
复制
torch_ccl_path=$(python -c "import torch; import torch_ccl; import os;  print(os.path.abspath(os.path.dirname(torch_ccl.__file__)))")
source $torch_ccl_path/env/setvars.sh
Intel® PyTorch 扩展安装

Intel PyTorch 扩展(IPEX)为使用 Float32 和 BFloat16 进行 CPU 训练提供了性能优化(请参考单 CPU 部分以了解更多信息)。

以下“Trainer 中的用法”以 Intel® MPI 库中的 mpirun 为例。

在 Trainer 中的用法

要在 Trainer 中使用 ccl 后端启用多 CPU 分布式训练,用户应在命令参数中添加**--ddp_backend ccl**。

让我们看一个例子,使用问答示例

以下命令在一个 Xeon 节点上启用了使用 2 个进程进行训练,每个进程在一个插槽上运行。OMP_NUM_THREADS/CCL_WORKER_COUNT 变量可以调整以获得最佳性能。

代码语言:javascript
复制
 export CCL_WORKER_COUNT=1
 export MASTER_ADDR=127.0.0.1
 mpirun -n 2 -genv OMP_NUM_THREADS=23 \
 python3 run_qa.py \
 --model_name_or_path bert-large-uncased \
 --dataset_name squad \
 --do_train \
 --do_eval \
 --per_device_train_batch_size 12  \
 --learning_rate 3e-5  \
 --num_train_epochs 2  \
 --max_seq_length 384 \
 --doc_stride 128  \
 --output_dir /tmp/debug_squad/ \
 --no_cuda \
 --ddp_backend ccl \
 --use_ipex

以下命令在两个 Xeon 上(node0 和 node1,以 node0 为主进程)启用了总共四个进程的训练,ppn(每个节点的进程数)设置为 2,每个插槽上运行一个进程。OMP_NUM_THREADS/CCL_WORKER_COUNT 变量可以调整以获得最佳性能。

在 node0 中,您需要创建一个包含每个节点的 IP 地址的配置文件(例如 hostfile),并将该配置文件路径作为参数传递。

代码语言:javascript
复制
 cat hostfile
 xxx.xxx.xxx.xxx #node0 ip
 xxx.xxx.xxx.xxx #node1 ip

现在,在 node0 中运行以下命令,将在 node0 和 node1 中启用 4DDP,并使用 BF16 自动混合精度:

代码语言:javascript
复制
 export CCL_WORKER_COUNT=1
 export MASTER_ADDR=xxx.xxx.xxx.xxx #node0 ip
 mpirun -f hostfile -n 4 -ppn 2 \
 -genv OMP_NUM_THREADS=23 \
 python3 run_qa.py \
 --model_name_or_path bert-large-uncased \
 --dataset_name squad \
 --do_train \
 --do_eval \
 --per_device_train_batch_size 12  \
 --learning_rate 3e-5  \
 --num_train_epochs 2  \
 --max_seq_length 384 \
 --doc_stride 128  \
 --output_dir /tmp/debug_squad/ \
 --no_cuda \
 --ddp_backend ccl \
 --use_ipex \
 --bf16

在 Kubernetes 中的用法

可以使用Kubeflow PyTorchJob 训练操作符将前一节中的相同分布式训练作业部署到 Kubernetes 集群。

设置

此示例假定您已经:

  • 访问已安装Kubeflow 的 Kubernetes 集群
  • 已安装并配置kubectl以访问 Kubernetes 集群
  • 可以用于存储数据集和模型文件的Persistent Volume Claim (PVC)。设置 PVC 的多种选项,包括使用 NFS storage class或云存储桶。
  • 一个包含您的模型训练脚本和运行脚本所需的所有依赖项的 Docker 容器。对于分布式 CPU 训练作业,这通常包括 PyTorch、Transformers、Intel Extension for PyTorch、Intel oneCCL Bindings for PyTorch 和 OpenSSH 以在容器之间进行通信。

下面的片段是一个使用支持分布式 CPU 训练的基础镜像的 Dockerfile 示例,然后将 Transformers 发布提取到/workspace目录中,以便示例脚本包含在镜像中:

代码语言:javascript
复制
FROM intel/ai-workflows:torch-2.0.1-huggingface-multinode-py3.9

WORKDIR /workspace

# Download and extract the transformers code
ARG HF_TRANSFORMERS_VER="4.35.2"
RUN mkdir transformers && \
    curl -sSL --retry 5 https://github.com/huggingface/transformers/archive/refs/tags/v${HF_TRANSFORMERS_VER}.tar.gz | tar -C transformers --strip-components=1 -xzf -

在将 PyTorchJob 部署到集群之前,需要构建并将镜像复制到集群的节点或推送到容器注册表。

PyTorchJob 规范文件

Kubeflow PyTorchJob用于在集群上运行分布式训练作业。PyTorchJob 的 yaml 文件定义了参数,例如:

  • PyTorchJob 的名称
  • 副本数(workers)的数量
  • 将用于运行训练作业的 Python 脚本及其参数
  • 每个 worker 所需的资源类型(节点选择器、内存和 CPU)
  • Docker 容器使用的图像/标签
  • 环境变量
  • PVC 的卷挂载

卷挂载定义了 PVC 将在每个 worker pod 的容器中挂载的路径。此位置可用于数据集、检查点文件以及训练完成后保存的模型。

下面的片段是一个 PyTorchJob 的 yaml 文件示例,其中有 4 个 worker 运行问答示例

代码语言:javascript
复制
apiVersion: "kubeflow.org/v1"
kind: PyTorchJob
metadata:
  name: transformers-pytorchjob
  namespace: kubeflow
spec:
  elasticPolicy:
    rdzvBackend: c10d
    minReplicas: 1
    maxReplicas: 4
    maxRestarts: 10
  pytorchReplicaSpecs:
    Worker:
      replicas: 4  # The number of worker pods
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: <image name>:<tag>  # Specify the docker image to use for the worker pods
              imagePullPolicy: IfNotPresent
              command:
                - torchrun
                - /workspace/transformers/examples/pytorch/question-answering/run_qa.py
                - --model_name_or_path
                - "bert-large-uncased"
                - --dataset_name
                - "squad"
                - --do_train
                - --do_eval
                - --per_device_train_batch_size
                - "12"
                - --learning_rate
                - "3e-5"
                - --num_train_epochs
                - "2"
                - --max_seq_length
                - "384"
                - --doc_stride
                - "128"
                - --output_dir
                - "/tmp/pvc-mount/output"
                - --no_cuda
                - --ddp_backend
                - "ccl"
                - --use_ipex
                - --bf16  # Specify --bf16 if your hardware supports bfloat16
              env:
              - name: LD_PRELOAD
                value: "/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4.5.9:/usr/local/lib/libiomp5.so"
              - name: TRANSFORMERS_CACHE
                value: "/tmp/pvc-mount/transformers_cache"
              - name: HF_DATASETS_CACHE
                value: "/tmp/pvc-mount/hf_datasets_cache"
              - name: LOGLEVEL
                value: "INFO"
              - name: CCL_WORKER_COUNT
                value: "1"
              - name: OMP_NUM_THREADS  # Can be tuned for optimal performance
-                value: "56"
              resources:
                limits:
                  cpu: 200  # Update the CPU and memory limit values based on your nodes
                  memory: 128Gi
                requests:
                  cpu: 200  # Update the CPU and memory request values based on your nodes
                  memory: 128Gi
              volumeMounts:
              - name: pvc-volume
                mountPath: /tmp/pvc-mount
              - mountPath: /dev/shm
                name: dshm
          restartPolicy: Never
          nodeSelector:  #  Optionally use the node selector to specify what types of nodes to use for the workers
            node-type: spr
          volumes:
          - name: pvc-volume
            persistentVolumeClaim:
              claimName: transformers-pvc
          - name: dshm
            emptyDir:
              medium: Memory

要运行此示例,请根据您的训练脚本和集群中的节点更新 yaml。

yaml 中的 CPU 资源限制/请求是以CPU 单位定义的,其中 1 个 CPU 单位等同于 1 个物理 CPU 核心或 1 个虚拟核心(取决于节点是物理主机还是虚拟机)。在 yaml 中定义的 CPU 和内存限制/请求量应小于单台机器上可用 CPU/内存容量的量。通常最好不要使用整个机器的容量,以便为 kubelet 和操作系统留下一些资源。为了为 worker pods 获得“guaranteed”服务质量,请为资源限制和请求设置相同的 CPU 和内存量。

部署

在为您的集群和训练作业更新了适当的值后,可以使用以下命令将 PyTorchJob 部署到集群中:

代码语言:javascript
复制
kubectl create -f pytorchjob.yaml

然后可以使用kubectl get pods -n kubeflow命令来列出kubeflow命名空间中的 pod。您应该看到刚刚部署的 PyTorchJob 的 worker pods。起初,它们可能会显示“Pending”状态,因为容器正在被拉取和创建,然后状态应该会变为“Running”。

代码语言:javascript
复制
NAME                                                     READY   STATUS                  RESTARTS          AGE
...
transformers-pytorchjob-worker-0                         1/1     Running                 0                 7m37s
transformers-pytorchjob-worker-1                         1/1     Running                 0                 7m37s
transformers-pytorchjob-worker-2                         1/1     Running                 0                 7m37s
transformers-pytorchjob-worker-3                         1/1     Running                 0                 7m37s
...

可以使用 kubectl logs -n kubeflow <pod name> 查看工作节点的日志。添加 -f 来实时查看日志,例如:

代码语言:javascript
复制
kubectl logs -n kubeflow transformers-pytorchjob-worker-0 -f

训练作业完成后,训练好的模型可以从 PVC 或存储位置复制。当作业完成后,可以使用 kubectl delete -f pytorchjob.yaml 命令从集群中删除 PyTorchJob 资源。

摘要

本指南涵盖了在裸金属和 Kubernetes 集群上使用多个 CPU 运行分布式 PyTorch 训练作业。这两种情况都利用了 Intel Extension for PyTorch 和 Intel oneCCL Bindings for PyTorch 来实现最佳的训练性能,并可以作为在多个节点上运行自己工作负载的模板。

使用 TensorFlow 在 TPU 上训练

原始文本:huggingface.co/docs/transformers/v4.37.2/en/perf_train_tpu_tf

如果您不需要长篇解释,只想要 TPU 代码示例来开始使用,请查看我们的 TPU 示例笔记本!

什么是 TPU?

TPU 是张量处理单元。它们是由 Google 设计的硬件,用于大大加速神经网络中的张量计算,类似于 GPU。它们可用于网络训练和推断。通常通过 Google 的云服务访问,但也可以通过 Google Colab 和 Kaggle Kernels 直接免费访问小型 TPU。

因为🤗 Transformers 中的所有 TensorFlow 模型都是 Keras 模型,因此本文档中的大多数方法通常适用于任何 Keras 模型的 TPU 训练!但是,有一些点是特定于 HuggingFace 生态系统(hug-o-system?)的 Transformers 和 Datasets,当我们到达这些点时,我们将确保标记它们。

有哪些类型的 TPU 可用?

新用户经常对各种 TPU 和访问方式感到困惑。要理解的第一个关键区别是TPU 节点TPU VM之间的区别。

当您使用TPU 节点时,实际上是间接访问远程 TPU。您将需要一个单独的 VM,该 VM 将初始化您的网络和数据管道,然后将它们转发到远程节点。当您在 Google Colab 上使用 TPU 时,您是以TPU 节点样式访问它。

对于不习惯使用 TPU 的人来说,使用 TPU 节点可能会产生一些意想不到的行为!特别是,因为 TPU 位于与运行 Python 代码的机器物理上不同的系统上,您的数据不能是本地的 - 从您机器的内部存储加载的任何数据管道将完全失败!相反,数据必须存储在 Google Cloud Storage 中,您的数据管道仍然可以访问它,即使管道在远程 TPU 节点上运行。

如果您可以将所有数据存储在内存中作为np.ndarraytf.Tensor,那么即使在使用 Colab 或 TPU 节点时,也可以在该数据上进行fit(),而无需将其上传到 Google Cloud Storage。

🤗具体的 Hugging Face 提示🤗:Dataset.to_tf_dataset()方法及其更高级别的包装器model.prepare_tf_dataset(),您将在我们的 TF 代码示例中看到,都会在 TPU 节点上失败。原因是即使它们创建了一个tf.data.Dataset,它也不是“纯粹”的tf.data管道,并且使用tf.numpy_functionDataset.from_generator()从底层 HuggingFaceDataset中流式传输数据。这个 HuggingFaceDataset由存储在本地磁盘上的数据支持,远程 TPU 节点将无法读取。

第二种访问 TPU 的方式是通过TPU VM。在使用 TPU VM 时,您直接连接到 TPU 连接的机器,就像在 GPU VM 上进行训练一样。TPU VM 通常更容易使用,特别是在处理数据管道时。所有上述警告不适用于 TPU VM!

这是一份主观的文件,所以这是我们的意见:**尽量避免使用 TPU Node。**它比 TPU VM 更令人困惑,更难以调试。未来也可能不受支持 - 谷歌最新的 TPU,TPUv4,只能作为 TPU VM 访问,这表明 TPU Node 越来越可能成为“传统”访问方法。但是,我们了解到唯一免费的 TPU 访问是在 Colab 和 Kaggle Kernels 上,它们使用 TPU Node - 因此,如果必须使用,我们将尝试解释如何处理!查看TPU 示例笔记本以获取更详细的代码示例。

可用的 TPU 尺寸是多少?

单个 TPU(v2-8/v3-8/v4-8)运行 8 个副本。TPU 存在于可以同时运行数百或数千个副本的pod中。当您使用多个 TPU 但少于整个 pod 时(例如 v3-32),您的 TPU 群被称为pod slice

当您通过 Colab 访问免费的 TPU 时,通常会获得一个 v2-8 TPU。

我一直听说这个 XLA。XLA 是什么,它与 TPU 有什么关系?

XLA 是一个优化编译器,被 TensorFlow 和 JAX 同时使用。在 JAX 中,它是唯一的编译器,而在 TensorFlow 中是可选的(但在 TPU 上是强制的!)。在训练 Keras 模型时启用它的最简单方法是将参数jit_compile=True传递给model.compile()。如果没有出现任何错误且性能良好,那么这是一个很好的迹象,表明您已准备好转移到 TPU!

在 TPU 上进行调试通常比在 CPU/GPU 上更困难,因此我们建议在尝试在 TPU 上运行之前,先在 CPU/GPU 上使用 XLA 使您的代码能够运行。当然,您不必训练很长时间 - 只需进行几个步骤,以确保您的模型和数据流水线按照您的预期工作。

XLA 编译的代码通常更快 - 因此,即使您不打算在 TPU 上运行,添加jit_compile=True也可以提高性能。但是,请注意下面关于 XLA 兼容性的注意事项!

**基于痛苦经验的提示:**虽然使用jit_compile=True是获得速度提升并测试您的 CPU/GPU 代码是否与 XLA 兼容的好方法,但如果在实际在 TPU 上训练时保留它,可能会导致许多问题。XLA 编译将在 TPU 上隐式发生,因此在实际在 TPU 上运行代码之前,请记得删除那行!

如何使我的模型与 XLA 兼容?

在许多情况下,您的代码可能已经与 XLA 兼容!但是,有一些在普通 TensorFlow 中有效但在 XLA 中无效的事情。我们将它们概括为以下三条核心规则:

**🤗具体的 HuggingFace 提示🤗:**我们已经付出了很多努力,将我们的 TensorFlow 模型和损失函数重写为 XLA 兼容。我们的模型和损失函数通常默认遵守规则#1 和#2,因此如果您使用transformers模型,则可以跳过它们。但是,在编写自己的模型和损失函数时,请不要忘记这些规则!

XLA 规则#1:您的代码不能具有“数据相关条件”

这意味着任何if语句都不能依赖于tf.Tensor内部的值。例如,此代码块无法使用 XLA 编译!

代码语言:javascript
复制
if tf.reduce_sum(tensor) > 10:
    tensor = tensor / 2.0

这一开始可能看起来非常受限制,但大多数神经网络代码不需要这样做。您通常可以通过使用tf.cond(请参阅此处的文档)或通过删除条件并找到一个巧妙的数学技巧来绕过此限制,例如:

代码语言:javascript
复制
sum_over_10 = tf.cast(tf.reduce_sum(tensor) > 10, tf.float32)
tensor = tensor / (1.0 + sum_over_10)

这段代码与上面的代码具有完全相同的效果,但通过避免条件语句,我们确保它将在 XLA 中编译而无问题!

XLA 规则#2:您的代码不能具有“数据相关形状”

这意味着代码中所有的tf.Tensor对象的形状不能依赖于它们的值。例如,函数tf.unique不能与 XLA 一起编译,因为它返回一个包含输入中每个唯一值的tensor。这个输出的形状显然会根据输入Tensor的重复程度而不同,因此 XLA 拒绝处理它!

一般来说,大多数神经网络代码默认遵守规则#2。但是,在一些常见情况下,这可能会成为一个问题。一个非常常见的情况是当您使用标签屏蔽时,将标签设置为负值以指示在计算损失时应忽略这些位置。如果您查看支持标签屏蔽的 NumPy 或 PyTorch 损失函数,您经常会看到类似于使用布尔索引的代码:

代码语言:javascript
复制
label_mask = labels >= 0
masked_outputs = outputs[label_mask]
masked_labels = labels[label_mask]
loss = compute_loss(masked_outputs, masked_labels)
mean_loss = torch.mean(loss)

这段代码在 NumPy 或 PyTorch 中完全正常,但在 XLA 中会出错!为什么?因为masked_outputsmasked_labels的形状取决于有多少位置被屏蔽 - 这使其成为**数据相关形状。**然而,就像规则#1 一样,我们通常可以重写这段代码,以产生完全相同的输出,而不涉及任何数据相关形状。

代码语言:javascript
复制
label_mask = tf.cast(labels >= 0, tf.float32)
loss = compute_loss(outputs, labels)
loss = loss * label_mask  # Set negative label positions to 0
mean_loss = tf.reduce_sum(loss) / tf.reduce_sum(label_mask)

在这里,我们通过为每个位置计算损失,但在计算均值时将被屏蔽的位置在分子和分母中归零,从而获得与第一个块完全相同的结果,同时保持 XLA 兼容性。请注意,我们使用与规则#1 相同的技巧 - 将tf.bool转换为tf.float32并将其用作指示变量。这是一个非常有用的技巧,所以如果您需要将自己的代码转换为 XLA,请记住它!

XLA 规则#3:XLA 将需要为每个不同的输入形状重新编译您的模型

这是一个重要的规则。这意味着如果您的输入形状非常不同,XLA 将不得不一遍又一遍地重新编译您的模型,这将导致巨大的性能问题。这在 NLP 模型中经常出现,因为输入文本在标记化后长度不同。在其他模态中,静态形状更常见,这个规则就不是那么大的问题了。

如何避开规则#3?关键是填充 - 如果您将所有输入填充到相同的长度,然后使用attention_mask,您可以获得与可变形状相同的结果,但没有任何 XLA 问题。然而,过度填充也会导致严重的减速 - 如果您将所有样本填充到整个数据集中的最大长度,您可能会得到由无尽填充标记组成的批次,这将浪费大量计算和内存!

解决这个问题并没有完美的方法。但是,你可以尝试一些技巧。一个非常有用的技巧是**将样本批次填充到 32 或 64 个标记的倍数。**这通常只会增加少量标记的数量,但会大大减少唯一输入形状的数量,因为现在每个输入形状都必须是 32 或 64 的倍数。更少的唯一输入形状意味着更少的 XLA 编译!

**🤗HuggingFace 专属提示🤗:**我们的分词器和数据整理器有助于解决这个问题的方法。在调用分词器时,您可以使用padding="max_length"padding="longest"来获取填充数据。我们的分词器和数据整理器还有一个pad_to_multiple_of参数,可以减少您看到的唯一输入形状的数量!

我如何在 TPU 上实际训练我的模型?

一旦您的训练是 XLA 兼容的,并且(如果您正在使用 TPU 节点/Colab)您的数据集已经准备就绪,那么在 TPU 上运行实际上非常容易!您真正需要在代码中做的改变只是添加几行代码来初始化您的 TPU,并确保您的模型和数据集都在TPUStrategy范围内创建。查看我们的 TPU 示例笔记本以查看实际操作!

总结

这里有很多内容,让我们用一个快速的清单来总结,当您想要准备好您的模型进行 TPU 训练时可以遵循:

  • 确保您的代码遵循 XLA 的三条规则
  • 在 CPU/GPU 上使用jit_compile=True编译您的模型,并确认您可以使用 XLA 进行训练
  • 要么将数据集加载到内存中,要么使用兼容 TPU 的数据集加载方法(请参阅notebook
  • 将您的代码迁移到 Colab(加速器设置为“TPU”)或 Google Cloud 上的 TPU VM
  • 添加 TPU 初始化代码(请参阅notebook
  • 创建您的TPUStrategy,并确保数据集加载和模型创建在strategy.scope()内(请参阅notebook
  • 当您转移到 TPU 时,不要忘记再次将jit_compile=True去掉!
  • 🙏🙏🙏🥺🥺🥺
  • 调用 model.fit()
  • 你做到了!

在 Apple 硅上进行 PyTorch 训练

原始文本:huggingface.co/docs/transformers/v4.37.2/en/perf_train_special

以前,在 Mac 上训练模型仅限于 CPU。随着 PyTorch v1.12 的发布,您可以利用使用 Apple 的硅 GPU 训练模型,以获得更快的性能和训练速度。在 PyTorch 中,这是通过将 Apple 的 Metal Performance Shaders(MPS)集成为后端来实现的。MPS 后端将 PyTorch 操作实现为自定义的 Metal 着色器,并将这些模块放置在mps设备上。

一些 PyTorch 操作尚未在 MPS 中实现,将会引发错误。为了避免这种情况,您应该设置环境变量PYTORCH_ENABLE_MPS_FALLBACK=1来使用 CPU 内核(仍会看到UserWarning)。

如果遇到其他错误,请在PyTorch存储库中打开问题,因为 Trainer 仅集成了 MPS 后端。

设置mps设备后,您可以:

  • 在本地训练更大的网络或批量大小
  • 减少数据检索延迟,因为 GPU 的统一内存架构允许直接访问完整的内存存储
  • 减少成本,因为您不需要在基于云的 GPU 上进行训练或添加额外的本地 GPU

首先确保您已安装 PyTorch。MPS 加速支持 macOS 12.3+。

代码语言:javascript
复制
pip install torch torchvision torchaudio

TrainingArguments 默认使用mps设备,如果可用的话,这意味着您不需要显式设置设备。例如,您可以在不进行任何更改的情况下自动启用 MPS 后端运行run_glue.py脚本。

代码语言:javascript
复制
export TASK_NAME=mrpc

python examples/pytorch/text-classification/run_glue.py \
  --model_name_or_path bert-base-cased \
  --task_name $TASK_NAME \
- --use_mps_device \
  --do_train \
  --do_eval \
  --max_seq_length 128 \
  --per_device_train_batch_size 32 \
  --learning_rate 2e-5 \
  --num_train_epochs 3 \
  --output_dir /tmp/$TASK_NAME/ \
  --overwrite_output_dir

gloonccl这样的分布式设置的后端不受mps设备支持,这意味着您只能在具有 MPS 后端的单个 GPU 上进行训练。

您可以在在 Mac 上介绍加速 PyTorch 训练博客文章中了解更多关于 MPS 后端的信息。

用于训练的定制硬件

原始文本:huggingface.co/docs/transformers/v4.37.2/en/perf_hardware

您用于运行模型训练和推理的硬件可能会对性能产生重大影响。要深入了解 GPU,请务必查看 Tim Dettmer 的优秀博客文章

让我们看看一些关于 GPU 设置的实用建议。

GPU

当您训练更大的模型时,您基本上有三个选择:

  • 更大的 GPU
  • 更多的 GPU
  • 更多的 CPU 和 NVMe(由 DeepSpeed-Infinity 卸载)

让我们从只有一个 GPU 的情况开始。

电源和冷却

如果您购买了昂贵的高端 GPU,请确保为其提供正确的电源和足够的冷却。

电源

一些高端消费级 GPU 卡有 2 个甚至 3 个 PCI-E 8 针电源插座。确保您有与插座数量相同的独立 12V PCI-E 8 针电缆插入卡中。不要使用同一电缆末端的 2 个分裂(也称为猪尾电缆)。也就是说,如果 GPU 上有 2 个插座,您希望从 PSU 到卡的有 2 个 PCI-E 8 针电缆,而不是一个末端有 2 个 PCI-E 8 针连接器的电缆!否则,您将无法充分发挥卡的性能。

每个 PCI-E 8 针电源电缆需要插入 PSU 侧的 12V 电源线,可以提供高达 150W 的功率。

一些其他卡可能使用 PCI-E 12 针连接器,这些连接器可以提供高达 500-600W 的功率。

低端卡可能使用 6 针连接器,提供高达 75W 的功率。

此外,您需要具有稳定电压的高端 PSU。一些质量较低的 PSU 可能无法为卡提供所需的稳定电压以使其在峰值状态下运行。

当然,PSU 需要有足够的未使用瓦特数来为卡供电。

冷却

当 GPU 过热时,它将开始降频,并且不会提供完整的性能,甚至在温度过高时可能会关闭。

很难确定 GPU 在负载严重时应该追求的确切最佳温度,但可能在+80C 以下是好的,但更低更好 - 也许 70-75C 是一个很好的范围。降频很可能会在 84-90C 左右开始。但除了降低性能外,长时间处于非常高的温度下可能会缩短 GPU 的寿命。

接下来让我们看看拥有多个 GPU 时最重要的一个方面:连接性。

多 GPU 连接

如果您使用多个 GPU,卡的互连方式对总训练时间有很大影响。如果 GPU 在同一物理节点上,您可以运行:

代码语言:javascript
复制
nvidia-smi topo -m

它将告诉您 GPU 是如何互连的。在具有双 GPU 且通过 NVLink 连接的机器上,您很可能会看到类似以下的内容:

代码语言:javascript
复制
        GPU0    GPU1    CPU Affinity    NUMA Affinity
GPU0     X      NV2     0-23            N/A
GPU1    NV2      X      0-23            N/A

在没有 NVLink 的不同机器上,我们可能会看到:

代码语言:javascript
复制
        GPU0    GPU1    CPU Affinity    NUMA Affinity
GPU0     X      PHB     0-11            N/A
GPU1    PHB      X      0-11            N/A

该报告包括此图例:

代码语言:javascript
复制
  X    = Self
  SYS  = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)
  NODE = Connection traversing PCIe as well as the interconnect between PCIe Host Bridges within a NUMA node
  PHB  = Connection traversing PCIe as well as a PCIe Host Bridge (typically the CPU)
  PXB  = Connection traversing multiple PCIe bridges (without traversing the PCIe Host Bridge)
  PIX  = Connection traversing at most a single PCIe bridge
  NV#  = Connection traversing a bonded set of # NVLinks

因此,第一个报告NV2告诉我们 GPU 是通过 2 个 NVLink 互连的,而第二个报告PHB则是典型的消费级 PCIe+Bridge 设置。

检查您的设置上有什么类型的连接性。其中一些将使卡之间的通信更快(例如 NVLink),而其他一些则更慢(例如 PHB)。

根据所使用的可扩展性解决方案的类型,连接速度可能会产生重大或轻微影响。如果 GPU 需要很少同步,如 DDP,较慢连接的影响将不那么显著。如果 GPU 需要经常互发消息,如 ZeRO-DP,则更快的连接变得非常重要以实现更快的训练。

NVlink

NVLink是由 Nvidia 开发的基于线的串行多通道近距离通信链路。

每一代新产品都提供更快的带宽,例如,这里是来自Nvidia Ampere GA102 GPU Architecture的一句话:

第三代 NVLink® GA102 GPU 利用了 NVIDIA 的第三代 NVLink 接口,其中包括四个 x4 链接,每个链接在两个 GPU 之间的每个方向提供 14.0625 GB/sec 的带宽。四个链接在每个方向提供 56.25 GB/sec 的带宽,两个 GPU 之间总共提供 112.5 GB/sec 的带宽。两个 RTX 3090 GPU 可以使用 NVLink 连接在一起进行 SLI。(请注意,不支持 3 路和 4 路 SLI 配置。)

所以在nvidia-smi topo -m的输出中,NVX报告中的X值越高越好。这取决于您的 GPU 架构。

让我们比较在 wikitext 的小样本上训练 gpt2 语言模型的执行。

结果是:

NVlink

时间

Y

101 秒

N

131 秒

您可以看到,NVLink 完成训练速度比快约 23%。在第二个基准测试中,我们使用NCCL_P2P_DISABLE=1告诉 GPU 不要使用 NVLink。

以下是完整的基准测试代码和输出:

代码语言:javascript
复制
# DDP w/ NVLink

rm -r /tmp/test-clm; CUDA_VISIBLE_DEVICES=0,1 torchrun \
--nproc_per_node 2 examples/pytorch/language-modeling/run_clm.py --model_name_or_path gpt2 \
--dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 --do_train \
--output_dir /tmp/test-clm --per_device_train_batch_size 4 --max_steps 200

{'train_runtime': 101.9003, 'train_samples_per_second': 1.963, 'epoch': 0.69}

# DDP w/o NVLink

rm -r /tmp/test-clm; CUDA_VISIBLE_DEVICES=0,1 NCCL_P2P_DISABLE=1 torchrun \
--nproc_per_node 2 examples/pytorch/language-modeling/run_clm.py --model_name_or_path gpt2 \
--dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 --do_train
--output_dir /tmp/test-clm --per_device_train_batch_size 4 --max_steps 200

{'train_runtime': 131.4367, 'train_samples_per_second': 1.522, 'epoch': 0.69}

硬件:每个 2x TITAN RTX 24GB + 2 个 NVLink 的 NVlink(在nvidia-smi topo -m中为NV2) 软件:pytorch-1.8-to-be + cuda-11.0 / transformers==4.3.0.dev0

使用 Trainer API 进行超参数搜索

原始文本:huggingface.co/docs/transformers/v4.37.2/en/hpo_train

🤗 Transformers 提供了一个专为训练🤗 Transformers 模型优化的 Trainer 类,使得更容易开始训练而无需手动编写自己的训练循环。Trainer 提供了用于超参数搜索的 API。本文档展示了如何在示例中启用它。

超参数搜索后端

Trainer 目前支持四种超参数搜索后端:optunasigoptraytunewandb

在使用超参数搜索后端之前,您应该先安装它们

代码语言:javascript
复制
pip install optuna/sigopt/wandb/ray[tune] 

如何在示例中启用超参数搜索

定义超参数搜索空间,不同的后端需要不同的格式。

对于 sigopt,请参阅 sigopt object_parameter,就像下面这样:

代码语言:javascript
复制
>>> def sigopt_hp_space(trial):
...     return [
...         {"bounds": {"min": 1e-6, "max": 1e-4}, "name": "learning_rate", "type": "double"},
...         {
...             "categorical_values": ["16", "32", "64", "128"],
...             "name": "per_device_train_batch_size",
...             "type": "categorical",
...         },
...     ]

对于 optuna,请参阅 optuna object_parameter,就像下面这样:

代码语言:javascript
复制
>>> def optuna_hp_space(trial):
...     return {
...         "learning_rate": trial.suggest_float("learning_rate", 1e-6, 1e-4, log=True),
...         "per_device_train_batch_size": trial.suggest_categorical("per_device_train_batch_size", [16, 32, 64, 128]),
...     }

Optuna 提供多目标 HPO。您可以在hyperparameter_search中传递direction并定义自己的compute_objective来返回多个目标值。 Pareto 前沿(List[BestRun])将在hyperparameter_search中返回,您应该参考test_trainer中的测试用例TrainerHyperParameterMultiObjectOptunaIntegrationTest。就像下面这样

代码语言:javascript
复制
>>> best_trials = trainer.hyperparameter_search(
...     direction=["minimize", "maximize"],
...     backend="optuna",
...     hp_space=optuna_hp_space,
...     n_trials=20,
...     compute_objective=compute_objective,
... )

对于 raytune,请参阅 raytune object_parameter,就像下面这样:

代码语言:javascript
复制
>>> def ray_hp_space(trial):
...     return {
...         "learning_rate": tune.loguniform(1e-6, 1e-4),
...         "per_device_train_batch_size": tune.choice([16, 32, 64, 128]),
...     }

对于 wandb,请参阅 wandb object_parameter,就像下面这样:

代码语言:javascript
复制
>>> def wandb_hp_space(trial):
...     return {
...         "method": "random",
...         "metric": {"name": "objective", "goal": "minimize"},
...         "parameters": {
...             "learning_rate": {"distribution": "uniform", "min": 1e-6, "max": 1e-4},
...             "per_device_train_batch_size": {"values": [16, 32, 64, 128]},
...         },
...     }

定义一个model_init函数并将其传递给 Trainer,例如:

代码语言:javascript
复制
>>> def model_init(trial):
...     return AutoModelForSequenceClassification.from_pretrained(
...         model_args.model_name_or_path,
...         from_tf=bool(".ckpt" in model_args.model_name_or_path),
...         config=config,
...         cache_dir=model_args.cache_dir,
...         revision=model_args.model_revision,
...         token=True if model_args.use_auth_token else None,
...     )

使用您的model_init函数、训练参数、训练和测试数据集以及评估函数创建一个 Trainer:

代码语言:javascript
复制
>>> trainer = Trainer(
...     model=None,
...     args=training_args,
...     train_dataset=small_train_dataset,
...     eval_dataset=small_eval_dataset,
...     compute_metrics=compute_metrics,
...     tokenizer=tokenizer,
...     model_init=model_init,
...     data_collator=data_collator,
... )

调用超参数搜索,获取最佳试验参数,后端可以是"optuna"/"sigopt"/"wandb"/"ray"。方向可以是"minimize""maximize",表示是优化更大还是更小的目标。

您可以定义自己的compute_objective函数,如果未定义,将调用默认的compute_objective,并将类似 f1 的评估指标的总和作为目标值返回。

代码语言:javascript
复制
>>> best_trial = trainer.hyperparameter_search(
...     direction="maximize",
...     backend="optuna",
...     hp_space=optuna_hp_space,
...     n_trials=20,
...     compute_objective=compute_objective,
... )

DDP 微调的超参数搜索

目前,optuna 和 sigopt 已启用 DDP 的超参数搜索。只有排名为零的进程才会生成搜索试验并将参数传递给其他排名。

优化推理

CPU 推理

原文链接:huggingface.co/docs/transformers/v4.37.2/en/perf_infer_cpu

通过一些优化,可以在 CPU 上高效运行大型模型推理。其中一种优化技术涉及将 PyTorch 代码编译成高性能环境(如 C++)的中间格式。另一种技术是将多个操作融合成一个内核,以减少单独运行每个操作的开销。

您将学习如何使用BetterTransformer进行更快的推理,以及如何将您的 PyTorch 代码转换为TorchScript。如果您使用 Intel CPU,还可以使用Intel Extension for PyTorch中的图优化来进一步提高推理速度。最后,学习如何使用🤗 Optimum 来加速使用 ONNX Runtime 或 OpenVINO 进行推理(如果您使用 Intel CPU)。

BetterTransformer

BetterTransformer 通过其快速路径(Transformer 函数的本机 PyTorch 专用实现)执行加速推理。快速路径执行中的两个优化是:

  1. 融合,将多个顺序操作组合成一个“内核”,以减少计算步骤的数量
  2. 跳过填充令牌的固有稀疏性,以避免与嵌套张量一起进行不必要的计算

BetterTransformer 还将所有注意力操作转换为更节省内存的缩放点积注意力

并非所有模型都支持 BetterTransformer。查看此列表以查看模型是否支持 BetterTransformer。

在开始之前,请确保您已经安装了🤗 Optimum installed

使用 PreTrainedModel.to_bettertransformer()方法启用 BetterTransformer:

代码语言:javascript
复制
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("bigcode/starcoder")
model.to_bettertransformer()

TorchScript

TorchScript 是一个中间 PyTorch 模型表示,可以在性能重要的生产环境中运行。您可以在 PyTorch 中训练模型,然后将其导出到 TorchScript 中,以解放模型免受 Python 性能约束。PyTorch跟踪一个模型以返回一个经过即时编译(JIT)优化的ScriptFunction。与默认的急切模式相比,PyTorch 中的 JIT 模式通常通过操作融合等优化技术为推理提供更好的性能。

有关 TorchScript 的简要介绍,请参阅PyTorch TorchScript 简介教程。

使用 Trainer 类,您可以通过设置--jit_mode_eval标志为 CPU 推理启用 JIT 模式:

代码语言:javascript
复制
python run_qa.py \
--model_name_or_path csarron/bert-base-uncased-squad-v1 \
--dataset_name squad \
--do_eval \
--max_seq_length 384 \
--doc_stride 128 \
--output_dir /tmp/ \
--no_cuda \
--jit_mode_eval

对于 PyTorch >= 1.14.0,JIT 模式可以使任何模型受益于预测和评估,因为jit.trace支持字典输入。

对于 PyTorch < 1.14.0,如果模型的前向参数顺序与jit.trace中的元组输入顺序匹配,例如问答模型,JIT 模式可以使模型受益。如果前向参数顺序与jit.trace中的元组输入顺序不匹配,例如文本分类模型,jit.trace将失败,我们在此处捕获此异常以使其回退。使用日志记录通知用户。

IPEX 图优化

Intel® Extension for PyTorch (IPEX)为 Intel CPU 的 JIT 模式提供进一步优化,并建议将其与 TorchScript 结合使用以获得更快的性能。IPEX 图优化融合了多头注意力、Concat Linear、Linear + Add、Linear + Gelu、Add + LayerNorm 等操作。

要利用这些图优化,请确保已安装 IPEX installed

代码语言:javascript
复制
pip install intel_extension_for_pytorch

在 Trainer 类中设置--use_ipex--jit_mode_eval标志以启用带有图优化的 JIT 模式:

代码语言:javascript
复制
python run_qa.py \
--model_name_or_path csarron/bert-base-uncased-squad-v1 \
--dataset_name squad \
--do_eval \
--max_seq_length 384 \
--doc_stride 128 \
--output_dir /tmp/ \
--no_cuda \
--use_ipex \
--jit_mode_eval

🤗 Optimum

Optimum Inference with ONNX Runtime指南中了解有关使用 ORT 与🤗 Optimum 的更多详细信息。本节仅提供了一个简短且简单的示例。

ONNX Runtime (ORT)是一个模型加速器,默认情况下在 CPU 上运行推理。ORT 受🤗 Optimum 支持,可以在🤗 Transformers 中使用,而无需对您的代码进行太多更改。您只需要将🤗 Transformers 的AutoClass替换为其等效的ORTModel以解决您正在解决的任务,并加载一个 ONNX 格式的检查点。

例如,如果您正在运行问题回答任务的推理,加载包含model.onnx文件的optimum/roberta-base-squad2检查点:

代码语言:javascript
复制
from transformers import AutoTokenizer, pipeline
from optimum.onnxruntime import ORTModelForQuestionAnswering

model = ORTModelForQuestionAnswering.from_pretrained("optimum/roberta-base-squad2")
tokenizer = AutoTokenizer.from_pretrained("deepset/roberta-base-squad2")

onnx_qa = pipeline("question-answering", model=model, tokenizer=tokenizer)

question = "What's my name?"
context = "My name is Philipp and I live in Nuremberg."
pred = onnx_qa(question, context)

如果您有 Intel CPU,请查看🤗 Optimum Intel,该支持各种压缩技术(量化、剪枝、知识蒸馏)和将模型转换为OpenVINO格式以获得更高性能推理的工具。

GPU 推理

原文:huggingface.co/docs/transformers/v4.37.2/en/perf_infer_gpu_one

与 CPU 不同,GPU 是机器学习的标准硬件选择,因为它们针对内存带宽和并行性进行了优化。为了跟上现代模型的更大尺寸或在现有和较旧的硬件上运行这些大型模型,您可以使用几种优化方法来加速 GPU 推理。在本指南中,您将学习如何使用 FlashAttention-2(一种更节省内存的注意力机制)、BetterTransformer(PyTorch 本地快速执行路径)和 bitsandbytes 将模型量化为较低精度。最后,学习如何使用🤗 Optimum 在 Nvidia 和 AMD GPU 上加速推理。

这里描述的大多数优化也适用于多 GPU 设置!

FlashAttention-2

FlashAttention-2 是实验性的,未来版本可能会发生较大变化。

FlashAttention-2是标准注意力机制的更快、更高效的实现,可以通过以下方式显著加速推理:

  1. 此外,可以通过在序列长度上并行化注意力计算来优化
  2. 将工作分区在 GPU 线程之间,以减少它们之间的通信和共享内存读/写

目前支持以下架构的 FlashAttention-2:

您可以通过打开 GitHub Issue 或 Pull Request 来请求为另一个模型添加 FlashAttention-2 支持。

在开始之前,请确保已安装 FlashAttention-2。

NVIDIAAMD

代码语言:javascript
复制
pip install flash-attn --no-build-isolation

我们强烈建议参考详细的安装说明以了解更多支持的硬件和数据类型!

要启用 FlashAttention-2,请将参数attn_implementation="flash_attention_2"传递给 from_pretrained():

代码语言:javascript
复制
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, LlamaForCausalLM

model_id = "tiiuae/falcon-7b"
tokenizer = AutoTokenizer.from_pretrained(model_id)

model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    torch_dtype=torch.bfloat16, 
    attn_implementation="flash_attention_2",
)

只有当模型的 dtype 为fp16bf16时,才能使用 FlashAttention-2。在使用 FlashAttention-2 之前,请确保将模型转换为适当的 dtype 并加载到支持的设备上。

您还可以设置use_flash_attention_2=True来启用 FlashAttention-2,但已被弃用,推荐使用attn_implementation="flash_attention_2"

FlashAttention-2 可以与其他优化技术(如量化)结合,以进一步加速推理。例如,您可以将 FlashAttention-2 与 8 位或 4 位量化结合使用:

代码语言:javascript
复制
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, LlamaForCausalLM

model_id = "tiiuae/falcon-7b"
tokenizer = AutoTokenizer.from_pretrained(model_id)

# load in 8bit
model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    load_in_8bit=True,
    attn_implementation="flash_attention_2",
)

# load in 4bit
model = AutoModelForCausalLM.from_pretrained(
    model_id, 
    load_in_4bit=True,
    attn_implementation="flash_attention_2",
)
预期的加速

您可以从推理中获得相当大的加速,特别是对于具有长序列的输入。但是,由于 FlashAttention-2 不支持使用填充令牌计算注意力分数,因此在序列包含填充令牌时,您必须手动填充/取消填充注意力分数以进行批量推理。这会导致使用填充令牌进行批量生成时出现显着减速。

为了克服这一点,在训练期间应该使用不带填充令牌的 FlashAttention-2(通过打包数据集或连接序列直到达到最大序列长度)。

对于在tiiuae/falcon-7b上进行单次前向传递,序列长度为 4096,各种批量大小且没有填充令牌,预期的加速是:

对于在meta-llama/Llama-7b-hf上进行单次前向传递,序列长度为 4096,各种批量大小且没有填充令牌,预期的加速是:

对于具有填充令牌的序列(使用填充令牌生成),您需要取消填充/填充输入序列以正确计算注意力分数。对于相对较小的序列长度,单次前向传递会产生额外开销,导致轻微加速(在下面的示例中,输入的 30%填充有填充令牌):

但是对于更大的序列长度,您可以期望获得更多的加速效益:

FlashAttention 更具内存效率,这意味着您可以在更大的序列长度上进行训练,而不会遇到内存不足的问题。对于更大的序列长度,您可以将内存使用量降低多达 20 倍。查看flash-attention存储库以获取更多详细信息。

PyTorch 缩放点积注意力

PyTorch 的torch.nn.functional.scaled_dot_product_attention(SDPA)也可以在底层调用 FlashAttention 和内存高效的注意力核。当可用实现时,SDPA 支持目前正在 Transformers 中本地添加,并且在torch>=2.1.1时默认用于torch

目前,Transformers 支持以下架构的 SDPA 推理和训练:

FlashAttention 只能用于具有fp16bf16 torch 类型的模型,因此请确保首先将您的模型转换为适当的类型。

默认情况下,SDPA 选择最高效的可用内核,但您可以使用torch.backends.cuda.sdp_kernel作为上下文管理器来检查在给定设置(硬件、问题大小)中是否有可用的后端:

代码语言:javascript
复制
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m", torch_dtype=torch.float16).to("cuda")
# convert the model to BetterTransformer
model.to_bettertransformer()

input_text = "Hello my dog is cute and"
inputs = tokenizer(input_text, return_tensors="pt").to("cuda")

+ with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
    outputs = model.generate(**inputs)

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

如果您看到下面的回溯中有错误,请尝试使用 PyTorch 的夜间版本,这可能对 FlashAttention 有更广泛的覆盖范围:

代码语言:javascript
复制
RuntimeError: No available kernel. Aborting execution.

# install PyTorch nightly
pip3 install -U --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu118

BetterTransformer

一些 BetterTransformer 功能正在被上游到 Transformers,支持本机torch.nn.scaled_dot_product_attention。BetterTransformer 仍然比 Transformers SDPA 集成具有更广泛的覆盖范围,但您可以期望越来越多的架构在 Transformers 中本地支持 SDPA。

查看我们在PyTorch 2.0 中使用 BetterTransformer 和缩放点积注意力的开箱即用加速和内存节省中的基准测试,并在BetterTransformer博客文章中了解更多关于快速执行的信息。

BetterTransformer 通过其快速路径(Transformer 函数的本机 PyTorch 专用实现)执行加速推断。快速路径执行中的两个优化是:

  1. 融合,将多个连续操作组合成一个单一的“内核”,以减少计算步骤的数量
  2. 跳过填充令牌的固有稀疏性,以避免使用嵌套张量进行不必要的计算

BetterTransformer 还将所有注意力操作转换为更节省内存的scaled dot product attention (SDPA),并在底层调用优化的内核,如FlashAttention

在开始之前,请确保您已安装🤗 Optimum (已安装)

然后,您可以使用 PreTrainedModel.to_bettertransformer()方法启用 BetterTransformer:

代码语言:javascript
复制
model = model.to_bettertransformer()

您可以使用 reverse_bettertransformer()方法返回原始的 Transformers 模型。在保存模型之前,应该使用这个方法来使用规范的 Transformers 建模:

代码语言:javascript
复制
model = model.reverse_bettertransformer()
model.save_pretrained("saved_model")

bitsandbytes

bitsandbytes 是一个包含对 4 位和 8 位量化支持的量化库。与其原生全精度版本相比,量化可以减小模型大小,使其更容易适应内存有限的 GPU。

确保您已安装 bitsandbytes 和🤗 Accelerate:

代码语言:javascript
复制
# these versions support 8-bit and 4-bit
pip install bitsandbytes>=0.39.0 accelerate>=0.20.0

# install Transformers
pip install transformers
4 位

要在 4 位模型中进行推断,使用load_in_4bit参数。device_map参数是可选的,但我们建议将其设置为"auto",以便🤗 Accelerate 根据环境中的可用资源自动高效地分配模型。

代码语言:javascript
复制
from transformers import AutoModelForCausalLM

model_name = "bigscience/bloom-2b5"
model_4bit = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_4bit=True)

要在多个 GPU 上加载 4 位模型进行推断,您可以控制要为每个 GPU 分配多少 GPU RAM。例如,将 600MB 的内存分配给第一个 GPU,将 1GB 的内存分配给第二个 GPU:

代码语言:javascript
复制
max_memory_mapping = {0: "600MB", 1: "1GB"}
model_name = "bigscience/bloom-3b"
model_4bit = AutoModelForCausalLM.from_pretrained(
    model_name, device_map="auto", load_in_4bit=True, max_memory=max_memory_mapping
)
8 位

如果您对 8 位量化的概念感兴趣并想了解更多信息,请阅读Hugging Face Transformers、Accelerate 和 bitsandbytes 使用规模化变压器进行 8 位矩阵乘法的初步介绍博客文章。

要在 8 位模型中进行推断,使用load_in_8bit参数。device_map参数是可选的,但我们建议将其设置为"auto",以便🤗 Accelerate 根据环境中的可用资源自动高效地分配模型:

代码语言:javascript
复制
from transformers import AutoModelForCausalLM

model_name = "bigscience/bloom-2b5"
model_8bit = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)

如果您要加载 8 位模型进行文本生成,应该使用 generate()方法,而不是未经优化的 Pipeline 函数,后者对 8 位模型不适用且速度较慢。一些采样策略,如核采样,也不受 Pipeline 支持。您还应该将所有输入放在与模型相同的设备上:

代码语言:javascript
复制
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "bigscience/bloom-2b5"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model_8bit = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)

prompt = "Hello, my llama is cute"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
generated_ids = model.generate(**inputs)
outputs = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)

要在多个 GPU 上加载 4 位模型进行推断,您可以控制要为每个 GPU 分配多少 GPU RAM。例如,要将 1GB 内存分配给第一个 GPU,将 2GB 内存分配给第二个 GPU:

代码语言:javascript
复制
max_memory_mapping = {0: "1GB", 1: "2GB"}
model_name = "bigscience/bloom-3b"
model_8bit = AutoModelForCausalLM.from_pretrained(
    model_name, device_map="auto", load_in_8bit=True, max_memory=max_memory_mapping
)

随意尝试在 Google Colab 的免费 GPU 上运行一个拥有 110 亿参数的T5 模型或 30 亿参数的BLOOM 模型进行推断!

🤗 Optimum

了解有关在NVIDIA GPU 上进行加速推断AMD GPU 上进行加速推断的指南中使用 ORT 的更多详细信息。本节仅提供简要且简单的示例。

ONNX Runtime(ORT)是一个模型加速器,支持在 Nvidia GPU 和使用ROCm堆栈的 AMD GPU 上进行加速推断。ORT 使用优化技术,如将常见操作融合为单个节点和常量折叠,以减少执行的计算量并加快推断速度。ORT 还将计算密集型操作放在 GPU 上,其余操作放在 CPU 上,智能地在两个设备之间分配工作负载。

ORT 受🤗 Optimum 支持,可以在🤗 Transformers 中使用。您需要使用一个ORTModel来解决您的任务,并指定provider参数,可以设置为CUDAExecutionProviderROCMExecutionProviderTensorrtExecutionProvider。如果要加载尚未导出为 ONNX 的模型,可以设置export=True将您的模型即时转换为 ONNX 格式:

代码语言:javascript
复制
from optimum.onnxruntime import ORTModelForSequenceClassification

ort_model = ORTModelForSequenceClassification.from_pretrained(
  "distilbert-base-uncased-finetuned-sst-2-english",
  export=True,
  provider="CUDAExecutionProvider",
)

现在您可以自由地使用模型进行推断:

代码语言:javascript
复制
from optimum.pipelines import pipeline
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")

pipeline = pipeline(task="text-classification", model=ort_model, tokenizer=tokenizer, device="cuda:0")
result = pipeline("Both the music and visual were astounding, not to mention the actors performance.")

结合优化

通常可以结合上述描述的多种优化技术,以获得最佳的推断性能。例如,您可以加载一个 4 位模型,然后启用带有 FlashAttention 的 BetterTransformer:

代码语言:javascript
复制
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

# load model in 4-bit
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16
)

tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m", quantization_config=quantization_config)

# enable BetterTransformer
model = model.to_bettertransformer()

input_text = "Hello my dog is cute and"
inputs = tokenizer(input_text, return_tensors="pt").to("cuda")

# enable FlashAttention
with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
    outputs = model.generate(**inputs)

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

实例化一个大模型

huggingface.co/docs/transformers/v4.37.2/en/big_models

当您想要使用非常大的预训练模型时,一个挑战是尽量减少 RAM 的使用。来自 PyTorch 的通常工作流程是:

  1. 用随机权重创建您的模型。
  2. 加载您的预训练权重。
  3. 将这些预训练权重放入您的随机模型中。

步骤 1 和 2 都需要内存中的完整模型版本,在大多数情况下这不是问题,但是如果您的模型开始占用数千兆字节,这两个副本可能会使您的 RAM 不足。更糟糕的是,如果您使用torch.distributed启动分布式训练,每个进程都会加载预训练模型并将这两个副本存储在 RAM 中。

请注意,随机创建的模型是用“空”张量初始化的,这些张量占用内存空间而不填充它(因此随机值是在给定时间内内存块中的内容)。适合模型/参数实例化的适当分布(例如正态分布)的随机初始化仅在第 3 步对未初始化的权重执行,以尽可能快地完成!

在本指南中,我们探讨了 Transformers 提供的解决此问题的解决方案。请注意,这是一个正在积极发展的领域,因此这里解释的 API 可能在未来略有变化。

分片检查点

自版本 4.18.0 以来,占用超过 10GB 空间的模型检查点会自动分片成较小的部分。在执行model.save_pretrained(save_dir)时,您将得到几个部分检查点(每个大小均小于 10GB)和一个将参数名称映射到存储文件的索引。

您可以使用max_shard_size参数控制分片之前的最大大小,因此为了举例,我们将使用一个具有小分片大小的正常大小模型:让我们使用传统的 BERT 模型。

代码语言:javascript
复制
from transformers import AutoModel

model = AutoModel.from_pretrained("bert-base-cased")

如果您使用 save_pretrained()保存它,您将获得一个新文件夹,其中包含模型的配置和权重:

代码语言:javascript
复制
>>> import os
>>> import tempfile

>>> with tempfile.TemporaryDirectory() as tmp_dir:
...     model.save_pretrained(tmp_dir)
...     print(sorted(os.listdir(tmp_dir)))
['config.json', 'pytorch_model.bin']

现在让我们使用最大分片大小为 200MB:

代码语言:javascript
复制
>>> with tempfile.TemporaryDirectory() as tmp_dir:
...     model.save_pretrained(tmp_dir, max_shard_size="200MB")
...     print(sorted(os.listdir(tmp_dir)))
['config.json', 'pytorch_model-00001-of-00003.bin', 'pytorch_model-00002-of-00003.bin', 'pytorch_model-00003-of-00003.bin', 'pytorch_model.bin.index.json']

除了模型的配置之外,我们看到三个不同的权重文件,以及一个index.json文件,这是我们的索引。可以使用 from_pretrained()方法完全重新加载这样的检查点:

代码语言:javascript
复制
>>> with tempfile.TemporaryDirectory() as tmp_dir:
...     model.save_pretrained(tmp_dir, max_shard_size="200MB")
...     new_model = AutoModel.from_pretrained(tmp_dir)

这样做大模型的主要优势在于,在上述工作流程的第 2 步中,检查点的每个分片在前一个分片之后加载,将 RAM 中的内存使用限制在模型大小加上最大分片大小的大小。

在幕后,索引文件用于确定检查点中的键以及相应权重存储的位置。我们可以像加载任何 json 一样加载该索引并获得一个字典:

代码语言:javascript
复制
>>> import json

>>> with tempfile.TemporaryDirectory() as tmp_dir:
...     model.save_pretrained(tmp_dir, max_shard_size="200MB")
...     with open(os.path.join(tmp_dir, "pytorch_model.bin.index.json"), "r") as f:
...         index = json.load(f)

>>> print(index.keys())
dict_keys(['metadata', 'weight_map'])

元数据目前只包含模型的总大小。我们计划在未来添加其他信息:

代码语言:javascript
复制
>>> index["metadata"]
{'total_size': 433245184}

权重映射是此索引的主要部分,它将每个参数名称(通常在 PyTorch 模型state_dict中找到)映射到其存储的文件:

代码语言:javascript
复制
>>> index["weight_map"]
{'embeddings.LayerNorm.bias': 'pytorch_model-00001-of-00003.bin',
 'embeddings.LayerNorm.weight': 'pytorch_model-00001-of-00003.bin',
 ...

如果您想在不使用 from_pretrained()(就像您为完整检查点执行model.load_state_dict()一样)的情况下直接加载这样的分片检查点,您应该使用 load_sharded_checkpoint():

代码语言:javascript
复制
>>> from transformers.modeling_utils import load_sharded_checkpoint

>>> with tempfile.TemporaryDirectory() as tmp_dir:
...     model.save_pretrained(tmp_dir, max_shard_size="200MB")
...     load_sharded_checkpoint(model, tmp_dir)

低内存加载

分片检查点减少了上述工作流程第 2 步中的内存使用,但为了在低内存环境中使用该模型,我们建议利用基于 Accelerate 库的工具。

请阅读以下指南以获取更多信息:使用 Accelerate 进行大型模型加载 model.save_pretrained(tmp_dir) … print(sorted(os.listdir(tmp_dir))) [‘config.json’, ‘pytorch_model.bin’]

代码语言:javascript
复制
现在让我们使用最大分片大小为 200MB:

```py
>>> with tempfile.TemporaryDirectory() as tmp_dir:
...     model.save_pretrained(tmp_dir, max_shard_size="200MB")
...     print(sorted(os.listdir(tmp_dir)))
['config.json', 'pytorch_model-00001-of-00003.bin', 'pytorch_model-00002-of-00003.bin', 'pytorch_model-00003-of-00003.bin', 'pytorch_model.bin.index.json']

除了模型的配置之外,我们看到三个不同的权重文件,以及一个index.json文件,这是我们的索引。可以使用 from_pretrained()方法完全重新加载这样的检查点:

代码语言:javascript
复制
>>> with tempfile.TemporaryDirectory() as tmp_dir:
...     model.save_pretrained(tmp_dir, max_shard_size="200MB")
...     new_model = AutoModel.from_pretrained(tmp_dir)

这样做大模型的主要优势在于,在上述工作流程的第 2 步中,检查点的每个分片在前一个分片之后加载,将 RAM 中的内存使用限制在模型大小加上最大分片大小的大小。

在幕后,索引文件用于确定检查点中的键以及相应权重存储的位置。我们可以像加载任何 json 一样加载该索引并获得一个字典:

代码语言:javascript
复制
>>> import json

>>> with tempfile.TemporaryDirectory() as tmp_dir:
...     model.save_pretrained(tmp_dir, max_shard_size="200MB")
...     with open(os.path.join(tmp_dir, "pytorch_model.bin.index.json"), "r") as f:
...         index = json.load(f)

>>> print(index.keys())
dict_keys(['metadata', 'weight_map'])

元数据目前只包含模型的总大小。我们计划在未来添加其他信息:

代码语言:javascript
复制
>>> index["metadata"]
{'total_size': 433245184}

权重映射是此索引的主要部分,它将每个参数名称(通常在 PyTorch 模型state_dict中找到)映射到其存储的文件:

代码语言:javascript
复制
>>> index["weight_map"]
{'embeddings.LayerNorm.bias': 'pytorch_model-00001-of-00003.bin',
 'embeddings.LayerNorm.weight': 'pytorch_model-00001-of-00003.bin',
 ...

如果您想在不使用 from_pretrained()(就像您为完整检查点执行model.load_state_dict()一样)的情况下直接加载这样的分片检查点,您应该使用 load_sharded_checkpoint():

代码语言:javascript
复制
>>> from transformers.modeling_utils import load_sharded_checkpoint

>>> with tempfile.TemporaryDirectory() as tmp_dir:
...     model.save_pretrained(tmp_dir, max_shard_size="200MB")
...     load_sharded_checkpoint(model, tmp_dir)

低内存加载

分片检查点减少了上述工作流程第 2 步中的内存使用,但为了在低内存环境中使用该模型,我们建议利用基于 Accelerate 库的工具。

请阅读以下指南以获取更多信息:使用 Accelerate 进行大型模型加载

本文参与 腾讯云自媒体同步曝光计划,分享自作者个人站点/博客。
原始发表:2024-06-26,如有侵权请联系 cloudcommunity@tencent.com 删除

本文分享自 作者个人站点/博客 前往查看

如有侵权,请联系 cloudcommunity@tencent.com 删除。

本文参与 腾讯云自媒体同步曝光计划  ,欢迎热爱写作的你一起参与!

评论
登录后参与评论
0 条评论
热度
最新
推荐阅读
目录
  • 多 GPU 上的高效训练
    • 可扩展性策略
      • 数据并行
        • DataParallel vs DistributedDataParallel
      • ZeRO 数据并行
        • 从天真的模型并行性到管道并行性
          • 张量并行
            • 数据并行 + 管道并行
              • 数据并行 + 流水线并行 + 张量并行
                • ZeRO 数据并行 + 流水线并行 + 张量并行
                  • FlexFlow
                    • GPU 选择
                      • GPU 的数量
                      • GPU 的顺序
                  • 完全分片数据并行
                    • FSDP 配置
                      • 分片策略
                      • CPU 卸载
                      • 包装策略
                      • 检查点
                      • TPU
                    • 启动训练
                      • 下一步
                      • 在 CPU 上高效训练
                        • 使用 IPEX 的混合精度
                          • IPEX 安装:
                          • Trainer 中的用法
                          • 实践示例
                      • 多 CPU 高效训练
                        • Intel® oneCCL 绑定的 PyTorch
                          • Intel® oneCCL 绑定的 PyTorch 安装
                        • Intel® MPI 库
                          • 在 Trainer 中的用法
                            • 在 Kubernetes 中的用法
                              • 设置
                              • PyTorchJob 规范文件
                              • 部署
                            • 摘要
                              • 什么是 TPU?
                              • 有哪些类型的 TPU 可用?
                              • 可用的 TPU 尺寸是多少?
                              • 我一直听说这个 XLA。XLA 是什么,它与 TPU 有什么关系?
                              • 如何使我的模型与 XLA 兼容?
                              • 我如何在 TPU 上实际训练我的模型?
                              • 总结
                          • 使用 TensorFlow 在 TPU 上训练
                          • 在 Apple 硅上进行 PyTorch 训练
                          • 用于训练的定制硬件
                            • GPU
                              • 电源和冷却
                              • 多 GPU 连接
                          • 使用 Trainer API 进行超参数搜索
                            • 超参数搜索后端
                              • 如何在示例中启用超参数搜索
                                • DDP 微调的超参数搜索
                                • 优化推理
                                • CPU 推理
                                  • BetterTransformer
                                    • TorchScript
                                      • IPEX 图优化
                                        • 🤗 Optimum
                                        • GPU 推理
                                          • FlashAttention-2
                                            • 预期的加速
                                          • PyTorch 缩放点积注意力
                                            • BetterTransformer
                                              • bitsandbytes
                                                • 4 位
                                                • 8 位
                                              • 🤗 Optimum
                                                • 结合优化
                                                • 实例化一个大模型
                                                  • 分片检查点
                                                    • 低内存加载
                                                      • 低内存加载
                                                      相关产品与服务
                                                      容器服务
                                                      腾讯云容器服务(Tencent Kubernetes Engine, TKE)基于原生 kubernetes 提供以容器为核心的、高度可扩展的高性能容器管理服务,覆盖 Serverless、边缘计算、分布式云等多种业务部署场景,业内首创单个集群兼容多种计算节点的容器资源管理模式。同时产品作为云原生 Finops 领先布道者,主导开源项目Crane,全面助力客户实现资源优化、成本控制。
                                                      领券
                                                      问题归档专栏文章快讯文章归档关键词归档开发者手册归档开发者手册 Section 归档