PyTorch - Amazon SageMaker
PyTorch 带有 Hugging Face 变形金刚的模型

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

PyTorch

将自己的 PyTorch 模型带到训练编译器中 SageMaker，然后使用训练编译器运行 SageMaker 训练作业。

PyTorch 带有 Hugging Face 变形金刚的模型

PyTorch 带有 Hugging Face Transf ormers 的模型基 PyTorch于 Torch.nn.Modu le AP I。Hugging Face Transformers 还提供训练器和预训练模型类 PyTorch ，以帮助减少配置自然语言处理 (NLP) 模型的工作量。准备好训练脚本后，你可以使用或估算器启动训练作业，并使用训练编译器配置中的 SageMaker PyTorchHuggingFace或估算器启动 SageMaker 训练作业，然后继续讨论下一个主题。启用 SageMaker 训练编译器

提示

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

使用 Hugging Face Transformers 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 及更高版本

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

def _mp_fn(index):
    main()

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

适用于《变形金刚》v4.17 及之前版本 1.10.2 及更 PyTorch 低版本

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

使用 SageMaker 训练编译器的最佳实践 Trainer

调整完训练脚本后，继续到 使用 PyTorch 训练编译器运行 SageMaker 训练作业。

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.SGD 和 torch.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. 如果你使用的是分布式数据加载器，请将你的数据加载器封装在 /XLA 的类中 PyTorch：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)

调整完训练脚本后，继续到 使用 PyTorch 训练编译器运行 SageMaker 训练作业。

对于分布式训练

除了上一 对于单个 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_ranks 和 world_size 设置变量，请使用类似于以下内容的代码：

    local_rank=xm.get_local_ordinal()
world_size=xm.xrt_world_size()

  3. 对于任何大于 1 的 world_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 及更高版本

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

    def _mp_fn(index):
    main()

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

    适用于《变形金刚》v4.17 及之前版本 1.10.2 及更 PyTorch低版本

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

调整完训练脚本后，继续到 使用 PyTorch 训练编译器运行 SageMaker 训练作业。

将 SageMaker 训练编译器与 PyTorch /XLA 配合使用的最佳实践

如果您想在原生 SageMaker 训练脚本上使用 PyTorch 训练编译器，则可能需要先熟悉 XLA 设备PyTorch 上的训练编译器。以下各节列出了为其启用 XLA 的一些最佳实践。 PyTorch

注意

本节最佳实践假设您使用以下 PyTorch /XLA 模块：

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

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

尽量减少compilation-and-executions使用次数pl.MpDeviceLoader/pl.ParallelLoader和 xm.step_closure

为了获得最佳性能，您应记住可能的启动方式，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 训练编译器提供针对 XLA 进行了优化的优化syncfree器，以提高 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 开始，T SageMaker raining Compiler 通过让 PyTorch /XLA 自动覆盖优化器（例如 SGD、Adam、AdamW）torch.optim或transformers.optimization其中的同步版本（例如、、），来提高性能。torch_xla.amp.syncfree torch_xla.amp.syncfree.SGD torch_xla.amp.syncfree.Adam torch_xla.amp.syncfree.AdamW您无需更改训练脚本中定义优化器的代码行。