在 Amazon SageMaker 上使用分片数据并行性训练具有近线性扩展的巨大模型

为了追求卓越的准确性，在过去几年中，自然语言处理和计算机视觉等领域的深度学习模型规模显着增长，经常有数百到数千亿的参数。 训练这些巨大的模型具有挑战性，并且需要复杂的分布策略。 数据科学家和机器学习工程师一直在寻找优化他们的训练计算的最佳方法，但同时也在努力应对可能随着整体集群大小而增加的通信开销。

这就是我们最近推出的原因 s硬数据并行 on 亚马逊SageMaker，一种新的节省内存的分布式训练技术 SageMaker 模型并行 (SMP) 库. 分片数据并行性专为超大规模模型构建，并在内部使用 Amazon 麦克风 幕后技术，一项科学努力，通过降低源于参数收集和梯度同步的昂贵通信开销来最小化通信规模。 使用序列长度为 30 的 2B 参数 GPT-2048 模型，这一新功能实现了 141 TFLOP，与 DeepSpeed ZeRO-39.7 相比，速度提高了 3%。 对于序列长度为 10 的 2B GPT-512 模型，这一新功能还实现了每秒 564 个样本，与 PyTorch 的全分片数据并行 (FSDP) 相比，速度提高了 13.9%。 请记住，在巨大的模型训练中，每一个百分比的加速都转化为节省的资金和提高团队的生产力。

在这篇博文中，我们将首先仔细研究分片数据并行性的关键区别以及何时使用它。 然后，您将学习如何使用此新功能在 SageMaker 上轻松训练 30B 参数 GPT-2 模型。 最后，我们将与其他开源选项的性能进行比较，特别是在 39.7 个 GPU 上的性能比 DeepSpeed ZeRO 高出 256%。

分片数据并行性如何工作以及何时使用它

在我们介绍分片数据并行之前，让我们看看它更广泛的技术系列。 最近针对大型模型的分布式训练方法已经转变为模型参数、梯度和优化器状态在数据并行节点之间共享的范式。 与流水线并行不同，它具有选择层以跨设备分区的先天复杂性，尤其是当您的框架不支持时 自动模型拆分，这种范式优雅地保留了数据并行性的简单性，同时消除了模型必须适合单个 GPU 的数据并行性约束。

在属于这种范式的现有框架中，尤其是从 FairScale 上游的 DeepSpeed ZeRO-3 和 PyTorch 的 FSDP，模型状态被分片 所有 GPU 是一种降低每个 GPU 上的内存消耗的策略，其代价是会产生大量通信开销，而通信开销会随着集群大小的增加而增加，因此会导致可扩展性在规模上显着下降。 相比之下，SMP 库分区模型中的分片数据并行性处于 规模感知 通过仅在内部划分模型状态的每个副本的方式 一个子集 GPU 的数量。

让我们仔细看看 尺度感知模型分区 在 MiCS 中，分片数据并行背后的核心技术。 这种设计背后的直觉是，可能不需要在整个数据并行组中划分训练状态来训练具有数百亿参数的模型。 例如，8 个 V100 GPU（每个 32GB）足以容纳 10B 参数模型的模型状态副本，当使用混合精度使用 Adam 优化器进行训练时，该模型需要大约 200GB 内存。 通过限制模型状态的完整副本 最少 作为 GPU 的子集，与 DeepSpeed 和 PyTorch FSDP 相比，我们可以有效地减少通信开销的规模。 分片数据并行还利用了 MiCS 中的其他技术，例如分层通信和 2 跳梯度同步。 欲了解更多信息，请查看 AWS 上的巨型模型训练的近线性扩展 or MiCS：在公共云上训练巨型模型的近线性缩放.

现在，您如何知道何时选择并行分片数据而不是其他分布式训练技术？ 一般规则是，如果您的模型的参数少于 1 亿个并且可以放入 GPU 内存中， SageMaker 数据并行库 or SageMaker 训练编译器 对你来说就足够了。 如果您有更大的语言或计算机视觉模型，我们的建议是使用分片数据并行技术结合 激活检查点 和 激活卸载 首先在 SageMaker 模型并行库中，然后是其他技术，例如 张量并行 或管道并行性。

使用分片数据并行性在 Amazon SageMaker 上训练 GPT-2

现在让我们学习如何使用分片数据并行训练 GPT-2 模型，使用 SMP 为您封装复杂性。 这个 完整的教程笔记本 引导您完成从数据处理、定义和提交训练作业到监控训练日志的整个过程。 以下是简要概述，重点介绍使用此功能的关键步骤。

1。 开始吧

分片数据并行性在 PyTorch v1.12.0+ 中可用，并且适用于 FP16 和 BF16。 使用 SMP 库的最简单方法是通过用于 PyTorch 的预构建 AWS 深度学习容器。 不过如果想自带Docker容器，可以参考 使用 SageMaker 分布式模型并行库创建您自己的 Docker 容器。 要开始，请按照 修改 PyTorch 训练脚本 在您的训练脚本中调整 SMP 的 API。 在本节中，我们仅使用现成的训练脚本中的代码片段调用几个主要步骤 train_gpt_simple.py. 您可以按照脚本中的注释和 API文件 详细了解 SMP API 的使用位置。

首先，通过调用导入和初始化库 smdistributed.modelparallel.torch.init() 在训练脚本的开头：

import smdistributed.modelparallel.torch as smp

smp.init(smp_config)

二、包装要分区的模型 smdistributed.modelparallel.torch.DistributedModel 并使用返回的 DistributedModel 前进的对象：

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_config(model_config)
model = smp.DistributedModel(model, trace_device="gpu", backward_passes_per_step=args.gradient_accumulation)

包装优化器 smdistributed.modelparallel.torch.DistributedOptimizer 用于保存和加载优化器状态。

from torch import optim

optimizer = optim.Adam(
    param_groups, betas=(args.beta1, args.beta2), lr=args.lr, weight_decay=args.weight_decay
)

optimizer = smp.DistributedOptimizer(
        optimizer, 
        static_loss_scale=None, 
        dynamic_loss_scale=True,
        dynamic_loss_args={"scale_window": 1000, "min_scale": 1, "delayed_shift": 2},
        )

将前向和后向逻辑放在一个 step 函数中，并用 smdistributed.modelparallel.torch.step.  内部定义的任何计算 smp.step-decorated 功能以分布式方式执行。

@smp.step
def train_step(model, optimizer, input_ids, attention_mask, args):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    model.backward(loss)

    return loss

@smp.step
def test_step(model, input_ids, attention_mask):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    
    return loss

2.准备数据集

我们使用 打开网页文本 是我们在这个例子中使用的数据集。 笔记本使用脚本 data_prep_512.py 下载和预处理数据集。 您还可以通过修改使用其他数据集进行训练 data_pipeline.py. 在处理大型数据集和模型时，您可以通过使用存储在 适用于Lustre的Amazon FSx，它提供了一个与本机集成的高性能文件系统 亚马逊简单存储服务 (S3)。 请参阅说明 配置数据输入通道以使用 Amazon FSx for Lustre 有关将 FSx Lustre 文件系统设置为数据输入通道的指导。

3. 开始培训工作

此步骤假设您已经 修改了你的训练脚本 并准备好前面几节中提到的数据集。 至 启用分片数据并行，只需设置 sharded_data_parallel_degree ，在 PyTorch 估计器. 在本教程中，我们设置 sharded_data_parallel_degree=128 和 instace_count=32 对于 p4d.24xlarge 节点，这表示模型状态将在总共 128 个 GPU 中的 256 个 GPU 上分片。 基于这个选定的值，SMP 会自动将数据并行度设置为 2（因为 256/128=2），这意味着我们将有两个副本用于数据并行。 选择理想值的一般规则 sharded_data_parallel_degree 就是每3B模型参数多增加一个节点到共享组。 在本教程中，我们的模型大小为 30B，因此我们应该使用至少 10 个节点进行分片。 并且因为 16 个节点（128 个 GPU）是高于阈值的最小 2 次方，我们设置 sharded_data_parallel_degree=128.

对于检查点，我们还提供了一组检查点实用程序 sharded_data_parallel_checkpoint.py ，包括重建完整的实用程序 state_dict 用于高级用例。 最后，我们可以通过调用 Estimator 上的 fit() 来启动分布式训练作业。

smp_estimator = PyTorch(
    entry_point="train_gpt_simple.py",
    instance_type="ml.p4d.24xlarge",
    source_dir=os.getcwd(),
    volume_size=500,
    instance_count=32,
    distribution={
        "mpi": {
            "enabled": True,
            "processes_per_host": processes_per_host,
            "custom_mpi_options": mpioptions,
        },
        "smdistributed": {
            "modelparallel": {
                "enabled": True,
                "parameters": {
                    "ddp": True,
                    "skip_tracing": True,
                    "delayed_parameter_initialization": True,
                    "offload_activations": True,
                    "activation_loading_horizon": 4,
                    # To enable sharded data parallelism.
                    # Here we shard model states across 128 GPUs. 
                    "sharded_data_parallel_degree": 128, 
                    "fp16": False,
                    "bf16": True,
                    # This is to disable pipeline parallelism.
                    "partitions": 1,
                },
            }
        },
    },
    framework_version="1.12",
    py_version="py38",
    hyperparameters=hyperparameters,
    checkpoint_s3_uri=checkpoint_s3_uri if not use_fsx else None,
    checkpoint_local_path=hyperparameters["checkpoint-dir"] if use_fsx else None,
    ...
)

smp_estimator.fit(inputs=data_channels)

4. 监控培训工作

您可以访问训练日志并跟踪 GPU 和内存利用率 亚马逊CloudWatch. 确保查看“algo-1”的日志 因为这是主节点，其输出流包含来自所有实例的训练作业日志。

基准性能

我们分别在序列长度为 16 和 32 的 4 个和 24 个 p512d.2048xlarge 节点上对 SMP 库中的分片数据并行性进行了基准测试。 30B 参数的 GPT2 模型配置为使用 7168、48 层和 64 个头的隐藏宽度。 您可以通过设置采用完全相同的配置，其中序列长度为 2048 model_config = "gpt2-30b" 在教程笔记本中。 使用此设置，SMP 每秒可实现 73.52 个样本，与 DeepSpeed ZeRO-39.7 相比，速度提高了 3%。 如果您的代币大小为 500 亿，这种加速意味着在 p367d.4xlarge 节点上节省近 24 小时，相当于每次训练节省超过 12,000 美元的预算！ 下表总结了我们的基准测试结果。

1/ 对于每个模型配置，我们在 DeepSpeed ZeRO 中测试了不同的功能、阶段和配置，并选择了提供最佳吞吐量的一个作为 DeepSpeed 基线。 基准运行于 亚马逊弹性计算云 （Amazon EC2）。 2/ 这些结果依赖于为 AWS 优化的改进的通信集合，这些集合很快就会提供。 3/ 训练时间是根据处理的令牌数量从速度预测的。

总之，在一系列模型和配置中，与 DeepSpeed 相比，我们观察到 SMP 中的分片数据并行性始终具有更高的吞吐量。 与 DeepSpeed 相比，此新功能还展示了更好的内存效率，使 SMP 能够适应更大的批量大小并降低适应特定全局批量大小所需的梯度累积水平。

结论

在这篇文章中，我们介绍了一种新的分布式训练技术 - 分片数据并行性 - 以及它如何在 Amazon SageMaker 上以接近线性扩展的速度加速巨大的模型训练。 我们还介绍了如何使用以下新技术训练 GPT-2 模型 完整示例. 您可以按照 Amazon SageMaker 示例 GitHub 存储库 跟踪所有 SageMaker 模型并行示例或参加我们的下一个 分布式培训研讨会. 要了解有关分片数据并行性的更多信息，请参阅 文件.

关于作者

艾米丽·韦伯（Emily Webber） 在 SageMaker 推出后就加入了 AWS，从那时起就一直在努力向全世界宣传它！ 除了为客户打造新的机器学习体验之外，Emily 还喜欢冥想和研究藏传佛教。

坎卡拉库斯 是 AWS 的高级应用科学家，负责优化 AWS 上的大规模分布式深度学习。 他的研究兴趣包括深度学习、分布式优化、分布式系统和信息论。 工作之余，他喜欢骑自行车、旅行、阅读和学习。

拉胡尔·惠尔戈尔 是 AWS 的高级软件工程师。 他致力于分布式深度学习系统，旨在使在云中训练大型深度学习模型变得容易和高效。 在业余时间，他喜欢摄影、骑自行车和园艺。

苏希特·科古勒 是 AWS 人工智能团队的一名软件开发工程师，致力于深度学习框架。 业余时间，他喜欢远足、旅游和烹饪。

何艾琳 是 AWS 深度学习的产品经理。 她致力于开发使客户更容易在 AWS 上训练深度学习模型的产品。 为了享受工作之外的乐趣，她喜欢远足和滑雪。