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

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

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

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

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

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

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

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

如果您使用种子来获得确定性结果，请基于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_rank、mp_sizeMP_GROUP、等的引用，并且您只想使用管道parallel 组，则可能需要将引用替换为smp.pp_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 的排名就可以有不同的批次数。

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

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