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

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

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

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

Batch 大小和微批次数

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

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

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

手动分区

  • 如果您使用手动分区,请注意模型中多个操作和模块消耗的参数,例如变压器架构中的嵌入表。为确保正确性,必须将共享相同参数的模块放置在同一设备中。使用自动分区时,库会自动强制执行此限制。

数据准备

  • 如果模型需要多个输入,请确保将数据流水线中的随机运算(例如 shuffling)种子为smp.dp_rank()。如果数据集是在数据parallel 设备之间进行确定性分片的,请确保使用该分区进行索引smp.dp_rank()。这是为了确保在构成模型分区的所有等级上看到的数据顺序是一致的。

从中返回张量smp.DistributedModel

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

@smp.step装饰者

  • 如果smp.step-drofecteed 函数的张量参数没有批处理维度,则在调用时必须在non_split_inputs列表中提供参数名称smp.step。这可以防止库尝试将张量拆分成微批次。有关更多信息,请参见 smp.stepAPI 文档。

延迟参数初始化

对于超过 1000 亿个参数的超大型模型,通过 CPU 内存进行权重初始化可能会导致 out-of-memory 错误。为了解决这个问题,库提供了smp.delay_param_initialization上下文管理器。这会延迟参数的物理分配,直到它们在第一次执行smp.step-drofection 函数期间移至 GPU。这样可以避免在训练初始化期间不必要地使用 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。如果您在使用张量并行时观察到的内存使用量超出预期,则可以尝试降低此限制。但是,选择太小的限制可能会导致吞吐量损失。要完全禁用限制,可以设置SMP_NCCL_THROTTLE_LIMIT=-1

  • 当张量并行度为 1 时,以下恒等式成立,但当张量并行度大于 1 时不成立:smp.mp_size() * smp.dp_size() == smp.size()。这是因为张量parallel 组是模型并行度组和数据并行度组的一部分。如果您的代码已存在对mp_rankmp_sizeMP_GROUP、等的引用,并且您只想使用管道parallel 组,则可能需要将引用替换为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())
  • 当通过张量并行将一个模块替换为其分布式对应模块时,分布式模块不会继承其原始模块的权重,而是初始化新的权重。这意味着,例如,如果需要在特定的调用中(例如,通过调用)初始化权重,则需要在smp.DistributedModel调用之后,即模块分配发生后。load_state_dict

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

    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建议使用张量并行性。这样可以确保每个数据parallel 等级都收到相同数量的数据样本,从而防止因不同的dp_rank s 采取不同步数而可能导致的挂起。

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

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

    • 在使用张量并行度时,为避免在保存和加载模型时出现停顿和争用情况,请确保在减少的数据并行度等级内从以下模型和优化器状态中调用相应的函数。

    • 如果您要过渡现有的流水线并行脚本并为该脚本启用 tensor parallel,请确保修改任何if smp.dp_rank() == 0用于保存和加载if smp.rdp_rank() == 0块的块。否则,它可能会导致你的训练工作停滞不前。

    • 张量并行性目前不支持对优化器状态进行全面检查。如果您想从检查点恢复训练,请改用部分检查点。

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