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

当批次大小增加时，库的效率最高。对于模型可以放在单个设备中、但只能用小批次训练的使用场景，在集成库后，可以而且应该增加批次大小。模型并行性可以为大型模型节省内存，使您能够使用以前无法放入内存的批次大小进行训练。

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

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

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

如果模型接受多个输入，请确保使用 smp.dp_rank()，在数据管道中将随机操作（例如随机排序）设置为种子。如果要在数据并行设备上确定性地对数据集进行分片，请确保分片按照 smp.dp_rank() 编制索引。这是为了确保在构成模型分区的所有排序上看到的数据顺序是一致的。

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

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

如果您使用种子来获得确定性结果，请根据 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()。这是因为张量并行组既是模型并行性组的一部分，也是数据并行性组的一部分。如果您的代码已有对 mp_rank、mp_size、MP_GROUP 等的引用，并且您只想使用管道并行组，则可能需要将这些引用替换为 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())

当通过张量并行性将模块替换为分布式的对应模块时，分布式模块不会从原始模块继承其权重，而是初始化新的权重。举例而言，这意味着如果需要在特定调用中初始化权重（例如，通过 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_rank 采取的不同步骤数而可能导致的挂起。

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

如果要对使用张量并行性的模型执行检查点操作，请考虑以下几点：

有关对使用张量并行度的模型执行检查点操作的更多信息，请参阅对分布式模型执行检查点操作。