SageMaker 分布式模型并行配置提示和陷阱 - Amazon SageMaker
Amazon Web Services 文档中描述的 Amazon Web Services 服务或功能可能因区域而异。要查看适用于中国区域的差异,请参阅中国的 Amazon Web Services 服务入门

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

SageMaker 分布式模型并行配置提示和陷阱

在使用 Amazon SageMaker 的分布式模型并行库之前,请查看以下提示和陷阱。此列表包括适用于各个框架的提示。有关 TensorFlow 和 PyTorch 的具体提示,请参阅修改 TensorFlow 训练脚本修改 PyTorch 训练脚本,分别。

Batch 大小和微批次数

  • 当批量大小增加时,库效率最高。对于模型适合于单个设备但只能用小批量进行训练的用例,可以而且应该在集成库后增加批量大小。模型并行性可为大型模型节省内存,使您能够使用以前不适合内存的批量大小进行训练。

  • 选择许多太小或太大的微粒可以降低性能。库在每台设备中按顺序执行每个微批处理,因此微批量大小(批量大小除以微批次数)必须足够大,以充分利用每个 GPU。与此同时,管道效率随着微批次的数量而提高,因此实现适当的平衡非常重要。通常,一个很好的起点是尝试 2 或 4 个微批,将批量大小增加到内存限制,然后尝试更大的批量和微批次数。随着微批次数量的增加,如果使用交错管道,可能会变得更大的批量大小可能变得可行。

  • 您的批量大小必须始终被微批量的数量整除。请注意,根据数据集的大小,有时每个时代的最后一批大小可能比其他时代小,而且这个较小的批次也需要被微批量的数量整除。如果不是,您可以设置drop_remainder=True中的tf.Dataset.batch()调用(在 TensorFlow 中),或者设置drop_last=TrueDataLoader(在 PyTorch 中),以便不使用最后一个小批量。如果您为数据管道使用不同的 API,则可能需要手动跳过最后一批次,只要不能被微批次数整除,则可能需要手动跳过该批次。

手动分区

  • 如果使用手动分区,请注意模型中的多个操作和模块所占用的参数,例如变压器体系结构中的嵌入表。为了正确起见,共享相同参数的模块必须放在同一设备中。使用自动分区时,库会自动强制执行此约束。

数据准备

  • 如果模型需要多个输入,请确保在数据管道中进行随机操作(例如,洗牌)smp.dp_rank(). 如果数据集在数据并行设备之间进行确定性分片,请确保分片的索引是smp.dp_rank(). 这是为了确保构成模型分区的所有级别上看到的数据顺序一致。

从返回张量smp.DistributedModel

  • 从中返回的任何张量smp.DistributedModel.call(对于 TensorFlow)或smp.DistributedModel.forward(对于 PyTorch)函数被广播到所有其他级别,从计算特定张量的排名中。因此,不应返回在调用和转发方法之外不需要的任何张量(例如中间激活),因为这会导致不必要的通信和内存开销并损害性能。

这些区域有:@smp.step装饰器

  • 如果smp.step-装饰函数有一个不具有批处理维度的 tensor 参数,参数名称必须在non_split_inputs打电话时列出smp.step. 这可以防止库尝试将张量分割成微批量。有关更多信息,请参阅smp.step在 API 文档中。

延迟参数初始化

对于超过 1000 亿个参数的大型型号,通过 CPU 内存进行权重初始化可能会导致内存不足错误。为了解决这个问题,图书馆提供smp.delay_param_initialization上下文管理器。这会延迟参数的物理分配,直到它们在首次执行smp.step-装饰功能。这避免了训练初始化期间 CPU 不必要的内存使用。创建模型对象时,请使用上下文管理器,如以下代码中所示。

with smp.delay_param_initialization(enabled=True): model = MyModel()

PyTorch 的张量并行性

  • 如果你使用种子获得确定性结果,请根据smp.dp_rank()(例如,torch.manual_seed(42 + smp.dp_rank()))。如果你不这样做,那么不同的分区nn.Parameter以同样的方式初始化,影响趋同。

  • SageMaker 的模型并行度库使用 NCCL 来实现分发模块所需的集体内容。特别是对于较小型号,如果同时在 GPU 上安排太多的 NCCL 调用,则由于 NCCL 使用的额外空间,内存使用量可能会增加。为了抵消这种情况,smp限制 NCCL 呼叫,以便任何给定时间正在进行的 NCCL 操作数量小于或等于给定的限制。默认限制为 8,但可以使用环境变量调整SMP_NCCL_THROTTLE_LIMIT. 如果在使用 tensor 并行性时观察到的内存使用量超出预期,则可以尝试降低此限制。但是,选择太小的限制可能会导致吞吐量损失。要完全禁用限制,你可以设置SMP_NCCL_THROTTLE_LIMIT=-1.

  • 当张量并行度为 1 时,以下标识在张量并行度大于 1 时将保持,当张量并行度大于 1 时不会保持:smp.mp_size() * smp.dp_size() == smp.size(). 这是因为张量并行组既是模型并行度组又是数据并行度组的一部分。如果你的代码现有引用mp_rankmp_sizeMP_GROUP等等,如果你想只使用管道并行组,则可能需要将引用替换为smp.pp_size(). 以下身份永远是真的:

    • smp.mp_size() * smp.rdp_size() == smp.size()

    • smp.pp_size() * smp.dp_size() == smp.size()

    • smp.pp_size() * smp.tp_size() * smp.rdp_size() == smp.size()

  • 由于smp.DistributedModel封装器在启用了张量并行度时修改模型参数,应在调用之后创建优化器smp.DistributedModel,带有分布式参数。例如,下面的操作不起作用:

    ## WRONG model = MyModel() optimizer = SomeOptimizer(model.parameters()) model = smp.DistributedModel(model)  # optimizer now has outdated parameters! 

    相反,应该使用smp.DistributedModel如下所示:

    ## CORRECT model = smp.DistributedModel(MyModel()) optimizer = SomeOptimizer(model.optimizers())
  • 当一个模块通过张量并行性替换为分布式对应模块时,分布式模块不会继承原始模块的权重,而是初始化新的权重。这意味着,例如,如果需要在特定调用中初始化权重(例如,通过load_state_dict打电话),这需要在smp.DistributedModel在模块分发发生后调用。

  • 直接访问分布式模块的参数时,请注意权重与原始模块的形状不同。例如, 

    with smp.tensor_parallelism():     linear = nn.Linear(60, 60) # will pass assert tuple(linear.weight.shape) == (60, 60) distributed_linear = smp.DistributedModel(linear) # will fail. the number of input channels will have been divided by smp.tp_size() assert tuple(distributed_linear.module.weight.shape) == (60, 60)
  • 使用torch.utils.data.distributed.DistributedSampler强烈建议用于张量并行性。这可确保每个数据并行排名接收相同数量的数据样本,从而防止因不同而可能导致的挂起dp_ranks 采取了不同数量的步骤。

  • 如果您将joinPyTorch 的 APIDistributedDataParallel类来处理不同数据并行排名具有不同批次数的情况,您仍然需要确保排名处于相同TP_GROUP批次数相同;否则模块的分布式执行中使用的通信集体可能会挂起。不同的等级TP_GROUPs 可以有不同数量的批次,只要join使用 API。

  • 如果要检查模型并使用张量并行性,请考虑以下几点:

    • 为避免在使用张量并行时保存和加载模型时停滞和竞争条件,请确保在减少的数据并行度排名中调用以下模型和优化程序状态中的适当函数。

    • 如果要转换现有的管道并行脚本并为脚本启用张量并行,请确保修改任何if smp.dp_rank() == 0用于保存和加载的块if smp.rdp_rank() == 0块。否则,这可能会导致你的训练工作停顿。

    • 张量并行度目前不支持优化程序状态的完整检查点。如果你想从检查点恢复训练,请改用部分检查点。

    有关检查点具有张量并行性的模型的更多信息,请参阅使用 Tensor 并行度进行检查点的说明.