本文属于机器翻译版本。若本译文内容与英语原文存在差异,则一律以英文原文为准。
PyTorch
将自己的 PyTorch 模型引入 SageMaker AI,然后使用 SageMaker Training Compiler 运行训练作业。
带 Hugging Face Transformers 的 PyTorch 模型
带 Hugging Face TransformersPyTorch 或 HuggingFace 估算器与 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_workers和preprocessing_num_workers参数来实现这一点。通常,它们需要大于或等于 GPU 数,但小于 CPU 数。
调整完训练脚本后,继续到 使用 SageMaker Training Compiler 运行 PyTorch 训练作业。
直接使用 PyTorch 的大型语言模型(没有 Hugging Face Transformers Trainer API)
如果您有直接使用 PyTorch 的训练脚本,则需要对 PyTorch 训练脚本进行其他更改以实施 PyTorch/XLA。按照说明进行操作,修改脚本以正确设置 PyTorch/XLA 基元。
对于单个 GPU 训练
-
导入优化库。
import torch_xla import torch_xla.core.xla_model as xm -
将目标设备更改为 XLA 而不是
torch.device("cuda")device=xm.xla_device() -
如果您使用 PyTorch 的自动混合精度
(AMP),请执行以下操作: -
将
torch.cuda.amp替换为以下项:import torch_xla.amp -
将
torch.optim.SGD和torch.optim.Adam替换为以下项:import torch_xla.amp.syncfree.Adam as adam import torch_xla.amp.syncfree.SGD as SGD -
将
torch.cuda.amp.GradScaler替换为以下项:import torch_xla.amp.GradScaler as grad_scaler
-
-
如果您未使用 AMP,请将
optimizer.step()替换为以下项:xm.optimizer_step(optimizer) -
如果您使用的是分布式数据加载器,请将数据加载器包装在 PyTorch/XLA 的
ParallelLoader类中:import torch_xla.distributed.parallel_loader as pl parallel_loader=pl.ParallelLoader(dataloader, [device]).per_device_loader(device) -
在不使用
parallel_loader时,将mark_step添加到训练循环的结尾:xm.mark_step() -
要为训练设置检查点,请使用 PyTorch/XLA 的模型检查点方法:
xm.save(model.state_dict(), path_to_save)
调整完训练脚本后,继续到 使用 SageMaker Training Compiler 运行 PyTorch 训练作业。
对于分布式训练
除了上一 对于单个 GPU 训练 部分中列出的更改外,还要添加以下更改以便在 GPU 之间正确分配工作负载。
-
如果您使用 AMP,请在
scaler.scale(loss).backward()的后面添加all_reduce:gradients=xm._fetch_gradients(optimizer) xm.all_reduce('sum', gradients, scale=1.0/xm.xrt_world_size()) -
如果您需要设置为
local_ranks和world_size设置变量,请使用类似于以下内容的代码:local_rank=xm.get_local_ordinal() world_size=xm.xrt_world_size() -
对于任何大于
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 ) -
进行以下更改,确保使用由
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。
-
对于 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 模块:
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.ParallelLoader 和 xm.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.optim 或 transformers.optimization 中的优化器(例如 SGD、Adam、AdamW)替换为这些优化器在 torch_xla.amp.syncfree 中的 syncfree 版本(例如 torch_xla.amp.syncfree.SGD、torch_xla.amp.syncfree.Adam、torch_xla.amp.syncfree.AdamW),从而提高性能。您无需更改训练脚本中定义优化器的代码行。