PyTorch - 亚马逊 SageMaker AI
PyTorch 带有 Hugging Face 变形金刚的模型

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

PyTorch

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

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

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

提示

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

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

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

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)函数调用从 SageMaker AI 分布式运行时 for PyTorch (pytorchxla) 重定向到训练脚本的main()函数。

def _mp_fn(index):
    main()

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

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

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

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

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

PyTorch 直接使用大型语言模型（不带 Hugging Face Transformers 训练器API）

如果您有 PyTorch 直接使用的训练脚本，则需要对 PyTorch 训练脚本进行其他更改才能实现 PyTorch/XLA。按照说明修改脚本以正确设置 PyTorch/p XLA rimatives。

适用于单次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. 如果你使用的是分布式数据加载器，请将你的数据加载器封装在/的类中 PyTorch：XLAParallelLoader

    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训练节中列出的更改外，还要添加以下更改以正确分配工作负载GPUs。

  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 的情况下使用任何 PyTorch 模型启动分布式训练API。请注意，这些指令可用于单节点多节点GPU和多节点多节点。GPU

  5. 适用于 PyTorch v1.11.0 及更高版本

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

    def _mp_fn(index):
    main()

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

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

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

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

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

如果你想在原生 SageMaker 训练脚本上使用 PyTorch 训练编译器，你可能需要先PyTorch 在XLA设备上熟悉一下。以下各节列出了要启用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 APIs使用张量和运算符构建计算时， 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优化器

在自动混合精度 (AMP) 模式下训练可利用的 Tensor 内核显著加快训练速度。NVIDIA GPUs SageMaker Training Com syncfree piler 提供经过优化的优化器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 开始，T SageMaker raining Compiler 通过XLA让 PyTorch /自动覆盖优化器（例如SGD，、）中的优化torch.optim器 transformers.optimizationtorch_xla.amp.syncfree（例如 Adam、AdamW）来提高性能。torch_xla.amp.syncfree.SGD torch_xla.amp.syncfree.Adam torch_xla.amp.syncfree.AdamW您无需更改训练脚本中定义优化器的代码行。