GPT-J 是由 Eleuther AI 发布的开源 60 亿参数模型。该模型在 Pile 上进行训练，可执行各种语言处理任务。它可以支持广泛的用例，包括文本分类、标记分类、文本生成、问答、实体提取、摘要、情感分析等等。GPT-J 是使用 Ben Wang 的 Mesh Transformer JAX 训练的变压器模型。
在本文中，我们介绍了使用 Amazon SageMaker 分布式模型并行库训练大型语言模型（LLMs）的指南和最佳实践，以减少训练时间和成本。您将学习如何轻松地在 SageMaker 上训练 60 亿参数的 GPT-J 模型。最后，我们分享了 SageMaker 分布式模型并行性的主要特点，以帮助加快训练时间。
变压器神经网络
变压器神经网络是一种流行的深度学习架构，用于解决序列到序列任务。它使用注意力作为学习机制，以实现接近人类水平的性能。与自然语言处理（NLP）模型的先前几代相比，该架构的一些其他有用属性包括能够分布、扩展和预训练。基于变压器的模型可以应用于处理文本数据的不同用例，例如搜索、聊天机器人等等。变压器使用预训练的概念从大型数据集中获取智能。预训练的变压器可以直接使用，也可以在您的数据集上进行微调，这些数据集可能比较小且特定于您的业务。
Hugging Face 在 SageMaker 上
Hugging Face 是一家开发一些最受欢迎的开源库的公司，这些库提供基于变压器架构的最先进的 NLP 技术。Hugging Face 的 transformers、tokenizers 和 datasets 库提供了多种语言的 API 和工具，可下载和预测使用预训练模型。SageMaker 使您能够直接使用 Hugging Face estimator 在 SageMaker SDK 中从其 Hugging Face Model Hub 训练、微调和运行推断使用 Hugging Face 模型。这种集成使得更容易针对特定领域的用例定制 Hugging Face 模型。在幕后，SageMaker SDK 使用 AWS Deep Learning Containers（DLC），这是 SageMaker 提供的一组预构建的 Docker 镜像，用于训练和服务模型。DLC 是 AWS 和 Hugging Face 之间的协作开发。该集成还提供了 Hugging Face transformers SDK 和 SageMaker 分布式训练库之间的集成，使您能够在 GPU 簇上扩展训练作业。
SageMaker 分布式模型并行库概述
模型并行是一种分布式训练策略，它将深度学习模型分区到多个设备中，设备可以在实例内或跨实例间。具有更多层和参数的深度学习（DL）模型在复杂任务（如计算机视觉和 NLP）中表现更好。然而，单个 GPU 可存储的最大模型大小受到限制。 GPU 内存约束会在以下方面影响训练 DL 模型：
- 它们限制可以训练的模型的大小，因为模型的内存占用量与参数数量成比例增加
- 它们通过限制训练期间每个 GPU 批处理大小来降低 GPU 利用率和训练效率
SageMaker 包括分布式模型并行库，以帮助有效地分发和训练 DL 模型，克服在单个 GPU 上训练模型时出现的限制。此外，该库使您能够使用支持 EFA 的设备获得最优的分布式训练，从而提高了低延迟、高吞吐量和操作系统旁路的节点间通信性能。
因为像 GPT-J 这样具有数十亿参数的大型模型的 GPU 内存占用量超过了单个芯片，所以将它们分区到多个 GPU 上变得至关重要。SageMaker 模型并行（SMP）库使得模型可以自动分区到多个 GPU 上。使用 SageMaker 模型并行，SageMaker 会代表您运行初始的分析作业，以分析模型的计算和内存要求。然后使用此信息来决定如何在 GPU 上分区模型，以最大化目标，例如最大化速度或最小化内存占用量。
它还支持可选的管道运行调度，以最大化可用 GPU 的整体利用率。在正向传递期间传播激活和在反向传递期间传播梯度需要顺序计算，这限制了 GPU 利用率的数量。SageMaker 使用管道运行计划通过将小批处理拆分为不同 GPU 上并行处理的微批来克服顺序计算约束。SageMaker 模型并行支持两种管道运行模式：
- 简单管道 – 此模式在开始向后传递之前完成每个微批次的前向传递。
- 交错管道 – 在此模式中，尽可能优先处理微批次的后向运行。这可以更快地释放用于激活的内存，从而更有效地使用内存。
张量并行
个别层或 nn.Modules 通过张量并行划分到设备上，以便可以同时运行。 以下是该库将具有四个层的模型分成两个张量并行（ "tensor_parallel_degree": 2 ）的最简单示例，如下图所示。每个模型副本的层都被划分为两半并分配到两个GPU之间。此示例中数据并行度为8，因为模型并行配置还包括 "pipeline_parallel_degree": 1 和 "ddp": True。该库管理张量分布模型的副本之间的通信。
这个功能的好处是，您可以选择应用张量并行的层或层的子集。要深入了解PyTorch的张量并行和其他节省内存的功能，并了解如何设置管道和张量并行的组合，请参见PyTorch的SageMaker模型并行库的扩展功能。
SageMaker分片数据并行
分片数据并行是一种节省内存的分布式训练技术，将模型的训练状态（模型参数、梯度和优化器状态）分割到数据并行组中的GPU上。
当将培训作业扩展到大型GPU集群时，您可以通过将训练状态分片到多个GPU上来减少模型的每个GPU内存占用量。这有两个好处：您可以适应更大的模型，否则这些模型将使用标准数据并行法运行内存不足，或者您可以使用释放的GPU内存增加批处理大小。
标准数据并行技术将训练状态复制到数据并行组中的GPU上，并根据AllReduce操作执行梯度聚合。实际上，分片数据并行介绍了通信开销和GPU内存效率之间的权衡。使用分片数据并行会增加通信成本，但是每个GPU的内存占用量（不包括由激活引起的内存使用）将被分片数据并行度除以，因此更大的模型可以适合于GPU集群。
SageMaker通过MiCS实现分片数据并行。有关更多信息，请参见在AWS上巨型模型训练的近线性缩放。
有关如何将分片数据并行应用于您的培训作业的详细信息，请参见分片数据并行。
使用SageMaker模型并行库
SageMaker模型并行库附带SageMaker Python SDK。您需要安装SageMaker Python SDK才能使用该库，并且它已经安装在SageMaker笔记本内核上。要使PyTorch培训脚本利用SMP库的功能，您需要进行以下更改：
- 通过调用
smp.init()来导入和初始化smp库。 - 初始化后，您可以使用
smp.DistributedModel包装器包装模型，并使用返回的DistributedModel对象而不是用户模型。 - 对于您的优化器状态，请使用
smp.DistributedOptimizer包装器包装您的模型优化器，使smp能够保存和加载优化器状态。前向和后向传递逻辑可以抽象为单独的函数，并在函数中添加smp.step装饰器。本质上，需要在smp.step装饰器上运行前向传递和反向传播。这使得smp能够将函数的张量输入分成指定启动训练作业时的多个微批次。 - 接下来，我们可以使用
torch.cuda.set_deviceAPI将输入张量移动到当前进程使用的GPU，然后使用.to()API调用。 - 最后，对于反向传播，我们替换了
torch.Tensor.backward和torch.autograd.backward。
看一下下面的代码：
@smp.step
def train_step(model, data, target):
output = model(data)
loss = F.nll_loss(output, target, reduction="mean")
model.backward(Loss)
return output, loss
with smp.tensor_parallelism():
model = AutoModelForCausalLM.from_config(model_config)
model = smp.DistributedModel (model)
optimizer = smp. DistributedOptimizer(optimizer)SageMaker模型并行库的张量并行提供了以下Hugging Face Transformer模型的开箱即用的支持：
- GPT-2、BERT和RoBERTa（在SMP库v1.7.0及更高版本中可用）
- GPT-J（在SMP库v1.8.0及更高版本中可用）
- GPT-Neo（在SMP库v1.10.0及更高版本中可用）
使用SMP库进行性能调优的最佳实践
当训练大型模型时，请考虑以下步骤，以便您的模型以合理的批量大小适合GPU内存：
- 建议使用具有更高GPU内存和高带宽互连的实例以提高性能，例如p4d和p4de实例。
- 在大多数情况下可以启用优化器状态分片，并在您拥有多个模型副本（启用数据并行）时将有所帮助。您可以通过在
modelparallel配置中设置"shard_optimizer_state": True来开启优化器状态分片。 - 使用激活检查点，一种通过清除某些层的激活并在模型中选择的模块的后向传递期间重新计算它们来减少内存使用的技术。
- 使用激活卸载，这是一项额外的功能，可以进一步减少内存使用。要使用激活卸载，请在
modelparallel配置中设置"offload_activations": True。在启用激活检查点和管道并行和微批次数大于1时使用。 - 启用张量并行，并增加并行度，其中并行度是2的幂。通常出于性能原因，张量并行受限于节点内部。
我们进行了许多实验，以优化在SageMaker上使用SMP库训练和调整GPT-J。我们已经成功将GPT-J在SageMaker上的一个时代的训练时间从58分钟缩短到不到10分钟，每个时代的训练时间快了6倍。它只需要从Amazon Simple Storage Service（Amazon S3）下载初始化、模型和数据集不到一分钟，使用GPU作为跟踪设备进行跟踪和自动分区不到1分钟，使用一个ml.p4d.24xlarge实例、FP16精度和SageMaker Hugging Face估算器进行张量并行训练一个时代8分钟。
为了将训练时间缩短为最佳实践，在SageMaker上训练GPT-J时，我们建议执行以下操作：
- 将预训练模型存储在Amazon S3上
- 使用FP16精度
- 使用GPU作为跟踪设备
- 使用自动分区、激活检查点和优化器状态分片：
auto_partition: Trueshard_optimizer_state: True
- 使用张量并行
- 使用具有多个GPU的SageMaker训练实例，例如ml.p3.16xlarge、ml.p3dn.24xlarge、ml.g5.48xlarge、ml.p4d.24xlarge或ml.p4de.24xlarge。
在SageMaker上使用SMP库训练和调整GPT-J模型
Amazon SageMaker Examples公共存储库上提供了一个工作的逐步代码示例。导航到training/distributed_training/pytorch/model_parallel/gpt-j文件夹。选择gpt-j文件夹并打开用于张量并行示例的train_gptj_smp_tensor_parallel_notebook.jpynb Jupyter笔记本，以及用于管道并行示例的train_gptj_smp_notebook.ipynb。您可以在我们的Amazon SageMaker生成AI工作坊中找到代码演练。
本文将向您介绍如何使用SageMaker模型并行库提供的张量并行功能。您将学习如何在GLUE sst2数据集上使用张量并行和管道并行运行GPT-J模型的FP16训练。
摘要
SageMaker模型并行库提供了多种功能。您可以在SageMaker上减少成本并加快LLM的训练速度。您还可以在Amazon SageMaker示例公共存储库中学习和运行BERT、GPT-2和GPT-J的示例代码。如需了解有关使用SMP库进行LLMS训练的AWS最佳实践，请参阅以下资源：
- SageMaker分布式模型并行最佳实践
- 在Amazon SageMaker上训练大型语言模型：最佳实践
如需了解我们的客户如何在SageMaker上实现低延迟GPT-J推理，请参阅How Mantium achieves low-latency GPT-J inference with DeepSpeed on Amazon SageMaker。
如果您希望加快LLM的上市时间并降低成本，请使用SageMaker。让我们知道您构建了什么！