Hugging Face在PyTorch / XLA TPUs上的应用

Hugging Face在PyTorch/XLA TPUs上的应用

使用PyTorch / XLA在云TPUs上训练您喜爱的Transformers模型

PyTorch-TPU项目是Facebook PyTorch团队和Google TPU团队的合作项目，于2019年PyTorch开发者大会正式推出。此后，我们与Hugging Face团队合作，为使用PyTorch / XLA在云TPUs上训练提供了一流的支持。这种新的集成使PyTorch用户能够在云TPUs上运行和扩展其模型，同时保持与Hugging Face训练器接口完全一致。

本博文提供了Hugging Face库中所做更改的概述，PyTorch / XLA库的功能，一个用于在云TPUs上训练您喜爱的Transformers模型的示例，以及一些性能基准。如果您迫不及待地想要开始使用TPUs，请直接跳到“在云TPUs上训练您的Transformer”部分 – 我们在Trainer模块中为您处理所有的PyTorch / XLA机制！

XLA:TPU设备类型

PyTorch / XLA为PyTorch添加了一个新的xla设备类型。这个设备类型的使用方式与其他PyTorch设备类型相同。例如，下面是如何创建和打印一个XLA张量：

import torch
import torch_xla
import torch_xla.core.xla_model as xm

t = torch.randn(2, 2, device=xm.xla_device())
print(t.device)
print(t)

这段代码应该很熟悉。PyTorch / XLA使用与普通PyTorch相同的接口，只是增加了一些功能。导入torch_xla会初始化PyTorch / XLA，xm.xla_device()会返回当前的XLA设备。这可能是CPU、GPU或TPU，取决于您的环境，但在本博文中，我们主要关注TPU。

Trainer模块利用一个TrainingArguments数据类来定义训练的具体参数。它处理多个参数，包括批量大小、学习率、梯度累积等等，以及所使用的设备。基于上述情况，在使用XLA:TPU设备时，在TrainingArguments._setup_devices()中，我们只需返回要由Trainer使用的TPU设备：

@dataclass
class TrainingArguments:
    ...
    @cached_property
    @torch_required
    def _setup_devices(self) -> Tuple["torch.device", int]:
        ...
        elif is_torch_tpu_available():
            device = xm.xla_device()
            n_gpu = 0
        ...

        return device, n_gpu

XLA设备步骤计算

在典型的XLA:TPU训练场景中，我们在多个TPU核心上并行训练（一个云TPU设备包含8个TPU核心）。因此，我们需要确保所有梯度在数据并行副本之间进行交换，通过合并梯度并执行优化器步骤。为此，我们提供了xm.optimizer_step(optimizer)，它执行梯度合并和步骤执行。在Hugging Face训练器中，我们相应地更新训练步骤，使用PyTorch / XLA的API：

class Trainer:
…
   def train(self, *args, **kwargs):
       ...
                    if is_torch_tpu_available():
                        xm.optimizer_step(self.optimizer)

PyTorch / XLA输入管道

运行PyTorch / XLA模型有两个主要部分：（1）追踪和延迟执行模型的图（请参阅下面的“PyTorch / XLA库”部分，了解更深入的解释）和（2）提供输入给模型。如果没有任何优化，模型的追踪/执行和输入提供将被串行执行，导致主机CPU和TPU加速器在执行过程中都会处于空闲状态。为了避免这种情况，我们提供了一个API，将这两个过程进行流水线处理，从而能够在执行步骤n时重叠步骤n+1的追踪。

import torch_xla.distributed.parallel_loader as pl
...
  dataloader = pl.MpDeviceLoader(dataloader, device)

检查点的写入和加载

当从 XLA 设备中对张量进行检查点操作，并从检查点加载时，它将被加载回原始设备。在对模型的张量进行检查点操作之前，你需要确保所有的张量都在 CPU 设备上，而不是 XLA 设备上。这样，当你加载张量时，你将通过 CPU 设备加载它们，然后有机会将它们放置在你想要的任何 XLA 设备上。我们提供了 xm.save() API 来实现这一点，它已经能够确保只有一个进程（在每个主机上）或一个全局进程（如果使用跨主机共享文件系统）将数据写入存储位置。

class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin):
…
    def save_pretrained(self, save_directory):
        ...
        if getattr(self.config, "xla_device", False):
            import torch_xla.core.xla_model as xm

            if xm.is_master_ordinal():
                # 保存配置文件
                model_to_save.config.save_pretrained(save_directory)
            # xm.save 负责只从主进程保存
            xm.save(state_dict, output_model_file)

class Trainer:
…
   def train(self, *args, **kwargs):
       ...
       if is_torch_tpu_available():
           xm.rendezvous("saving_optimizer_states")
           xm.save(self.optimizer.state_dict(),
                   os.path.join(output_dir, "optimizer.pt"))
           xm.save(self.lr_scheduler.state_dict(),
                   os.path.join(output_dir, "scheduler.pt"))

PyTorch / XLA 库

PyTorch / XLA 是一个使用 XLA 线性代数编译器将 PyTorch 深度学习框架与 XLA 设备（包括 CPU、GPU 和云 TPU）连接起来的 Python 包。以下部分的内容也可以在我们的 API_GUIDE.md 中找到。

PyTorch / XLA 张量是惰性的

使用 XLA 张量和设备只需要改变几行代码。但是，尽管 XLA 张量的行为很像 CPU 和 CUDA 张量，但它们的内部实现是不同的。CPU 和 CUDA 张量会立即或急切地执行操作。而 XLA 张量则是惰性的。它们会将操作记录在图中，直到需要结果为止。这种延迟执行的方式让 XLA 优化操作。多个独立操作的图可能会被融合成一个优化的操作。

惰性执行对调用者来说通常是透明的。PyTorch / XLA 会自动构建图，并将它们发送到 XLA 设备，并在 XLA 设备和 CPU 之间复制数据时进行同步。在进行优化步骤时插入一个屏障可以显式地同步 CPU 和 XLA 设备。

这意味着当你调用 model(input) 进行前向传播、计算损失 loss.backward() ，以及进行优化步骤 xm.optimizer_step(optimizer) 时，所有操作的图都在后台构建。只有当你明确地评估张量（例如打印张量或将其移动到 CPU 设备）或标记一个步骤（每次迭代通过 MpDeviceLoader 时都会自动完成）时，才会执行完整的步骤。

追踪、编译、执行和重复

从用户的角度来看，运行在 PyTorch / XLA 上的模型的典型训练过程包括前向传播、反向传播和优化步骤。从 PyTorch / XLA 库的角度来看，情况有所不同。

当用户运行前向和后向传播时，会实时跟踪中间表示（IR）图。可以如下所示检查导致每个根/输出张量的 IR 图：

>>> import torch
>>> import torch_xla
>>> import torch_xla.core.xla_model as xm
>>> t = torch.tensor(1, device=xm.xla_device())
>>> s = t*t
>>> print(torch_xla._XLAC._get_xla_tensors_text([s]))
IR {
  %0 = s64[] prim::Constant(), value=1
  %1 = s64[] prim::Constant(), value=0
  %2 = s64[] xla::as_strided_view_update(%1, %0), size=(), stride=(), storage_offset=0
  %3 = s64[] aten::as_strided(%2), size=(), stride=(), storage_offset=0
  %4 = s64[] aten::mul(%3, %3), ROOT=0
}

在用户程序上运行前向和后向传递时，该实时图会累积，并且一旦调用xm.mark_step()（由pl.MpDeviceLoader间接调用），实时张量的图就会被截断。这个截断标志着一个步骤的完成，随后我们将IR图降低到XLA Higher Level Operations（HLO），这是XLA的IR语言。

然后，将此HLO图编译成TPU二进制文件，并在TPU设备上执行。然而，这个编译步骤可能代价高昂，通常比单个步骤花费的时间长，因此如果我们每个步骤都编译用户的程序，开销将很高。为了避免这种情况，我们有一个缓存，通过其HLO图的唯一哈希标识符来存储编译的TPU二进制文件。因此，一旦在第一步骤中填充了此TPU二进制文件缓存，随后的步骤通常不需要重新编译新的TPU二进制文件；而是可以直接从缓存中查找所需的二进制文件。

由于TPU编译通常比步骤执行时间慢得多，这意味着如果图形的形状保持变化，我们将会有缓存未命中并且编译过于频繁。为了最小化编译成本，我们建议尽可能保持张量形状静态。Hugging Face库的形状在很大程度上已经是静态的，输入令牌会被适当地填充，因此在整个训练过程中，缓存应该会一直被命中。可以使用PyTorch / XLA提供的调试工具来检查这一点。在下面的示例中，您可以看到编译仅发生了5次（CompileTime），而执行发生在1220个步骤中的每一个（ExecuteTime）：

>>> import torch_xla.debug.metrics as met
>>> print(met.metrics_report())
Metric: CompileTime
  TotalSamples: 5
  Accumulator: 28s920ms153.731us
  ValueRate: 092ms152.037us / second
  Rate: 0.0165028 / second
  Percentiles: 1%=428ms053.505us; 5%=428ms053.505us; 10%=428ms053.505us; 20%=03s640ms888.060us; 50%=03s650ms126.150us; 80%=11s110ms545.595us; 90%=11s110ms545.595us; 95%=11s110ms545.595us; 99%=11s110ms545.595us
Metric: DeviceLockWait
  TotalSamples: 1281
  Accumulator: 38s195ms476.007us
  ValueRate: 151ms051.277us / second
  Rate: 4.54374 / second
  Percentiles: 1%=002.895us; 5%=002.989us; 10%=003.094us; 20%=003.243us; 50%=003.654us; 80%=038ms978.659us; 90%=192ms495.718us; 95%=208ms893.403us; 99%=221ms394.520us
Metric: ExecuteTime
  TotalSamples: 1220
  Accumulator: 04m22s555ms668.071us
  ValueRate: 923ms872.877us / second
  Rate: 4.33049 / second
  Percentiles: 1%=045ms041.018us; 5%=213ms379.757us; 10%=215ms434.912us; 20%=217ms036.764us; 50%=219ms206.894us; 80%=222ms335.146us; 90%=227ms592.924us; 95%=231ms814.500us; 99%=239ms691.472us
Counter: CachedCompile
  Value: 1215
Counter: CreateCompileHandles
  Value: 5
...

在云TPU上训练您的Transformer

要配置您的虚拟机和云TPU，请按照“设置计算引擎实例”和“启动云TPU资源”（写作时的pytorch-1.7版本）部分。一旦您创建了虚拟机和云TPU，使用它们就像SSH到您的GCE虚拟机并运行以下命令来启动bert-large-uncased训练（批量大小是为v3-8设备，可能会在v2-8上OOM）：

conda activate torch-xla-1.7
export TPU_IP_ADDRESS="输入你的TPU IP地址"  # 例如 10.0.0.2
export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
git clone -b v4.2.2 https://github.com/huggingface/transformers.git
cd transformers && pip install .
pip install datasets==1.2.1
python examples/xla_spawn.py \
  --num_cores 8 \
  examples/language-modeling/run_mlm.py \
  --dataset_name wikitext \
  --dataset_config_name wikitext-103-raw-v1 \
  --max_seq_length 512 \
  --pad_to_max_length \
  --logging_dir ./tensorboard-metrics \
  --cache_dir ./cache_dir \
  --do_train \
  --do_eval \
  --overwrite_output_dir \
  --output_dir language-modeling \
  --overwrite_cache \
  --tpu_metrics_debug \
  --model_name_or_path bert-large-uncased \
  --num_train_epochs 3 \
  --per_device_train_batch_size 8 \
  --per_device_eval_batch_size 8 \
  --save_steps 500000

以上代码将在大约不到200分钟的时间内完成训练，并得到一个评估困惑度约为3.25的结果。

性能基准测试

以下表格显示在一个包含4个TPU v3芯片的v3-8云TPU系统上运行PyTorch / XLA训练bert-large-uncased的性能。用于所有基准测试测量的数据集是WikiText103数据集，我们使用Hugging Face examples中提供的run_mlm.py脚本。为了确保工作负载不受主机CPU限制，我们在这些测试中使用n1-standard-96 CPU配置，但您也可以使用较小的配置而不影响性能。

开始使用PyTorch / XLA和TPU

请参阅Hugging Face examples中的“在TPU上运行”部分以开始。要获取有关我们的API的更详细描述，请查看我们的API指南，并且要获取性能最佳实践，请查看我们的故障排除指南。对于通用的PyTorch / XLA示例，请运行我们提供的带有免费云TPU访问权限的以下Colab笔记本。要直接在GCP上运行，请参阅我们文档网站上标记为“PyTorch”的教程。

有其他问题或疑问吗？请在https://github.com/huggingface/transformers/issues 或 https://github.com/pytorch/xla/issues 上提出问题或疑问。

我的旅程：在Google Cloud上构建无服务器的tra...
合作伙伴关系：Amazon SageMaker 和 Hugging Face
了解BigBird的块稀疏注意力
使用和混合 Hugging Face 模型与 Gradio 2.0
实践中的少样本学习：GPT-Neo 和 🤗 加速推理 API
深度学习在互联网上：协作训练语言模型
人工智能

人工智能

探索人工智能的世界及AI的未来潜力