PyTorch - 亚马逊 SageMaker AI
Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异,请参阅 中国的 Amazon Web Services 服务入门 (PDF)

本文属于机器翻译版本。若本译文内容与英语原文存在差异,则一律以英文原文为准。

PyTorch

将自己的 PyTorch 模型引入 SageMaker AI,然后使用 SageMaker Training Compiler 运行训练作业。

带 Hugging Face Transformers 的 PyTorch 模型

Hugging Face Transformers 的 PyTorch 模型基于 PyTorch 的 torch.nn.Module API。Hugging Face Transformers 还为 PyTorch 提供 Trainer 和预训练的模型类,以便减少配置自然语言处理 (NLP) 模型所需的工作量。准备训练脚本后,在继续启用 SageMaker Training Compiler 上的下一个主题时,可以启动一个将 SageMaker AI PyTorchHuggingFace 估算器与 SageMaker Training Compiler 配置结合使用的训练作业。

提示

在训练脚本中使用 Transformers 为 NLP 模型创建分词器时,请务必通过指定 padding='max_length' 来使用静态输入张量形状。请勿使用 padding='longest',因为填充批处理中的最长序列可能会改变每个训练批处理的张量形状。动态输入形状可以触发对模型的重新编译,并可能会增加总训练时间。有关 Transformers 分词器的填充选项的更多信息,请参阅 Hugging Face Transformers 文档中的填充和截断

使用 Hugging Face Transformers Trainer 类的大型语言模型

如果您使用转换器库的 Trainer 类,则无需对训练脚本进行任何其他更改。如果您通过估算器类启用 Trainer 模型,则 SageMaker Training Compiler 会自动编译该模型。以下代码显示了带 Hugging Face Trainer API 的 PyTorch 训练脚本的基本形式。

from transformers import Trainer, TrainingArguments training_args=TrainingArguments(**kwargs) trainer=Trainer(args=training_args, **kwargs)

对于单个 GPU 训练

在使用 transformers.Trainer 类时,无需更改代码。

对于分布式训练

PyTorch v1.11.0 及更高版本

要使用 SageMaker Training Compiler 运行分布式训练,必须在训练脚本中添加以下 _mp_fn() 函数并包装 main() 函数。它将 _mp_fn(index) 函数调用从适用于 PyTorch 的 SageMaker AI 分布式运行时 (pytorchxla) 重定向到训练脚本的 main() 函数。

def _mp_fn(index): main()

此函数接受 index 参数以指示当前 GPU 在集群中的排名以进行分布式训练。要查找更多示例脚本,请参阅 Hugging Face Transformers 语言建模示例脚本

对于 Transformers v4.17 及早期版本与 PyTorch v1.10.2 及早期版本

SageMaker Training Compiler 使用另一种机制来启动分布式训练作业,并且您无需在训练脚本中进行任何修改。相反,SageMaker Training Compiler 要求您在 SageMaker AI Hugging Face 估算器中将 SageMaker AI 分布式训练启动器脚本传递给 entry_point 参数,并将训练脚本传递给 hyperparameters 参数。

将 SageMaker Training Compiler 与 Trainer 结合使用的最佳实践

  • 通过在设置 transformers.TrainingArgument 时将 optim 参数设置为 adamw_torch_xla 来确保使用 SyncFree 优化器。另请参阅 Hugging Face Transformers 文档中的优化器

  • 确保数据处理管道的吞吐量高于训练吞吐量。您可以调整 transformers.TrainingArgument 类的 dataloader_num_workerspreprocessing_num_workers 参数来实现这一点。通常,它们需要大于或等于 GPU 数,但小于 CPU 数。

调整完训练脚本后,继续到 使用 SageMaker Training Compiler 运行 PyTorch 训练作业

直接使用 PyTorch 的大型语言模型(没有 Hugging Face Transformers Trainer API)

如果您有直接使用 PyTorch 的训练脚本,则需要对 PyTorch 训练脚本进行其他更改以实施 PyTorch/XLA。按照说明进行操作,修改脚本以正确设置 PyTorch/XLA 基元。

对于单个 GPU 训练

  1. 导入优化库。

    import torch_xla import torch_xla.core.xla_model as xm
  2. 将目标设备更改为 XLA 而不是 torch.device("cuda")

    device=xm.xla_device()
  3. 如果您使用 PyTorch 的自动混合精度 (AMP),请执行以下操作:

    1. torch.cuda.amp 替换为以下项:

      import torch_xla.amp
    2. torch.optim.SGDtorch.optim.Adam 替换为以下项:

      import torch_xla.amp.syncfree.Adam as adam import torch_xla.amp.syncfree.SGD as SGD
    3. torch.cuda.amp.GradScaler 替换为以下项:

      import torch_xla.amp.GradScaler as grad_scaler
  4. 如果您未使用 AMP,请将 optimizer.step() 替换为以下项:

    xm.optimizer_step(optimizer)
  5. 如果您使用的是分布式数据加载器,请将数据加载器包装在 PyTorch/XLA 的 ParallelLoader 类中:

    import torch_xla.distributed.parallel_loader as pl parallel_loader=pl.ParallelLoader(dataloader, [device]).per_device_loader(device)
  6. 在不使用 parallel_loader 时,将 mark_step 添加到训练循环的结尾:

    xm.mark_step()
  7. 要为训练设置检查点,请使用 PyTorch/XLA 的模型检查点方法:

    xm.save(model.state_dict(), path_to_save)

调整完训练脚本后,继续到 使用 SageMaker Training Compiler 运行 PyTorch 训练作业

对于分布式训练

除了上一 对于单个 GPU 训练 部分中列出的更改外,还要添加以下更改以便在 GPU 之间正确分配工作负载。

  1. 如果您使用 AMP,请在 scaler.scale(loss).backward() 的后面添加 all_reduce

    gradients=xm._fetch_gradients(optimizer) xm.all_reduce('sum', gradients, scale=1.0/xm.xrt_world_size())
  2. 如果您需要设置为 local_ranksworld_size 设置变量,请使用类似于以下内容的代码:

    local_rank=xm.get_local_ordinal() world_size=xm.xrt_world_size()
  3. 对于任何大于 1world_size (num_gpus_per_node*num_nodes),您必须定义一个训练取样器,它应类似于以下内容:

    import torch_xla.core.xla_model as xm if xm.xrt_world_size() > 1: train_sampler=torch.utils.data.distributed.DistributedSampler( train_dataset, num_replicas=xm.xrt_world_size(), rank=xm.get_ordinal(), shuffle=True ) train_loader=torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, sampler=train_sampler, drop_last=args.drop_last, shuffle=False if train_sampler else True, num_workers=args.num_workers )
  4. 进行以下更改,确保使用由 torch_xla distributed 模块提供的 parallel_loader

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

    train_device_loader 函数类似于常规 PyTorch 加载器,如下所示:

    for step, (data, target) in enumerate(train_device_loader): optimizer.zero_grad() output=model(data) loss=torch.nn.NLLLoss(output, target) loss.backward()

    在进行所有这些更改后,您应能够在没有 Transformer Trainer API 的情况下使用任何 PyTorch 模型启动分布式训练。请注意,这些说明适用于单节点多 GPU 和多节点多 GPU。

  5. 对于 PyTorch v1.11.0 及更高版本

    要使用 SageMaker Training Compiler 运行分布式训练,必须在训练脚本中添加以下 _mp_fn() 函数并包装 main() 函数。它将 _mp_fn(index) 函数调用从适用于 PyTorch 的 SageMaker AI 分布式运行时 (pytorchxla) 重定向到训练脚本的 main() 函数。

    def _mp_fn(index): main()

    此函数接受 index 参数以指示当前 GPU 在集群中的排名以进行分布式训练。要查找更多示例脚本,请参阅 Hugging Face Transformers 语言建模示例脚本

    对于 Transformers v4.17 及早期版本与 PyTorch v1.10.2 及早期版本

    SageMaker Training Compiler 使用另一种机制来启动分布式训练作业,同时要求您在 SageMaker AI Hugging Face 估算器中将 SageMaker AI 分布式训练启动器脚本传递给 entry_point 参数,并将训练脚本传递给 hyperparameters 参数。

调整完训练脚本后,继续到 使用 SageMaker Training Compiler 运行 PyTorch 训练作业

将 SageMaker Training Compiler 与 PyTorch/XLA 结合使用的最佳实践

如果您想在原生 PyTorch 训练脚本上利用 SageMaker Training Compiler,您可能需要先熟悉 XLA 设备上的 PyTorch。以下部分列出了启用适用于 PyTorch 的 XLA 的一些最佳实践。

注意

此部分中的最佳实践假定您使用以下 PyTorch/XLA 模块:

import torch_xla.core.xla_model as xm import torch_xla.distributed.parallel_loader as pl
了解 PyTorch/XLA 中的延迟模式

PyTorch/XLA 和原生 PyTorch 之间的一个显著区别是 PyTorch/XLA 系统在延迟模式下运行,而原生 PyTorch 则在紧急模式下运行。延迟模式下的张量是用于构建计算图的占位符,直到编译和评估完成后才会被具体化。当您调用 PyTorch API 以使用张量和运算符构建计算时,PyTorch/XLA 系统会即时构建计算图。在以下情况下编译和执行计算图:pl.MpDeviceLoader/pl.ParallelLoader 显式或隐式调用 xm.mark_step() 时,或在通过调用 loss.item()print(loss) 显式请求张量值时。

使用 pl.MpDeviceLoader/pl.ParallelLoaderxm.step_closure 最大限度地减少 compilation-and-executions

为了获得最佳性能,您应记住启动 compilation-and-executions 的可能方法(如了解 PyTorch/XLA 中的延迟模式中所述),并最大限度地减少 compilation-and-executions 数。理想情况下,每次训练迭代只需一次 compilation-and-execution,并由 pl.MpDeviceLoader/pl.ParallelLoader 自动启动。MpDeviceLoader 针对 XLA 进行了优化,如果可能,应始终使用它以获得最佳性能。在训练期间,您可能需要检查一些中间结果,例如损失值。在这种情况下,延迟张量的打印应使用 xm.add_step_closure() 进行包装,以避免不必要的 compilation-and-executions。

使用 AMP 和 syncfree 优化器

通过利用 NVIDIA GPU 的 Tensor 核心,在自动混合精度 (AMP) 模式下进行训练可显著加快训练速度。SageMaker Training Compiler 提供了 syncfree 优化器,它们针对 XLA 进行了优化以提高 AMP 性能。目前,有以下三个 syncfree 优化器可用,应尽可能使用这些优化器以获得最佳性能。

torch_xla.amp.syncfree.SGD torch_xla.amp.syncfree.Adam torch_xla.amp.syncfree.AdamW

这些 syncfree 优化器应与 torch_xla.amp.GradScaler 配对以实施梯度扩展/取消扩展。

提示

从 PyTorch 1.13.1 开始,SageMaker Training Compiler 可让 PyTorch/XLA 自动将 torch.optimtransformers.optimization 中的优化器(例如 SGD、Adam、AdamW)替换为这些优化器在 torch_xla.amp.syncfree 中的 syncfree 版本(例如 torch_xla.amp.syncfree.SGDtorch_xla.amp.syncfree.Adamtorch_xla.amp.syncfree.AdamW),从而提高性能。您无需更改训练脚本中定义优化器的代码行。