分布式包

此包包含各种实用程序，用于在优化器步骤之前完成每个 rank 上的模型权重梯度。这包括一个分布式数据并行包装器，用于在数据并行副本之间进行 all-reduce 或 reduce-scatter 梯度，以及一个 finalize_model_grads 方法，用于同步不同并行模式下的梯度（例如，不同流水线阶段的“tied”层，或由于专家并行性而导致 MoE 中不同 rank 上的专家梯度）。

分布式数据并行模型包装器。将梯度存储在连续缓冲区中，并支持通过将完整模型的梯度分解为更小的 bucket 并在每个 bucket 上异步运行 all-reduce / reduce-scatter，来实现通信（all-reduce 或 reduce-scatter）与反向传播计算的重叠。

class core.distributed.distributed_data_parallel.DistributedDataParallel(*args: Any, **kwargs: Any)

基类: core.distributed.data_parallel_base._BaseDataParallel

DDP 包装器，将梯度存储在连续缓冲区中。还具有通过将完整模型的梯度分解为更小的 bucket 并在每个 bucket 上异步运行 all-reduce / reduce-scatter，来实现通信与反向传播计算重叠的选项。此类还提供了以参数类型以外的类型（例如，bf16 模型的 fp32）执行梯度累积的选项。

参数

• config – Transformer 配置对象。

• ddp_config – DistributedDataParallel 配置对象。

• module – 底层模型。

• disable_bucketing – 如果为 true，则强制将所有参数分配到单个 bucket。如果为 false，则使用标准分桶策略：将参数分配到较小的 bucket，并且 _如果_ overlap_grad_reduce 为 True 且 pp_rank 为 0，则每个 bucket 执行 all-reduce。

broadcast_params()

在所有 DP rank 之间同步参数。

disable_forward_pre_hook(param_sync: bool = True)

禁用与前向计算重叠的参数 all-gather 所需的前向预钩子。如果 param_sync 为 False，则跳过同步参数 all-gather。

enable_forward_pre_hook()

启用与前向计算重叠的参数 all-gather 所需的前向预钩子。

finish_grad_sync()

完成所有模型梯度的梯度同步（all-reduce 或 reduce-scatter）通信操作。

当 overlap_grad_reduce 设置为 True 时，等待异步通信调用完成。当 overlap_grad_reduce 设置为 False 时，调用同步通信操作。

no_sync()

关闭梯度同步的上下文管理器。

scale_gradients(scaling_factor: float)

按 scaling_factor 缩放缓冲区内的所有梯度。

start_grad_sync(*unused)

启动所有模型梯度的梯度同步（all-reduce 或 reduce-scatter）通信操作。

当 overlap_grad_reduce 设置为 True 时，调度异步通信调用。当 overlap_grad_reduce 设置为 False 时，调用同步通信操作。

start_param_sync(*unused, force_sync: bool = False, force_dispatch: bool = False)

启动所有模型参数的参数同步（all-gather）通信操作。

默认情况下，当 overlap_param_gather 设置为 True 时，调度异步通信调用；当 overlap_param_gather 设置为 False 时，调用同步通信操作。可以使用以下标志覆盖此默认行为。

参数

• force_sync (bool, optional) – 强制同步集体通信，无论其他设置如何。

• force_dispatch (bool, optional) – 强制调度，无论其他设置如何。

zero_grad_buffer()

将所有梯度缓冲区归零。需要在每个训练迭代开始时调用。

在所有使用的并行模式中，为优化器步骤完成模型梯度。同步跨 DP 副本的模型梯度的 all-reduce / reduce-scatter，all-reduce 序列并行性的 layernorm 梯度，跨第一和最后流水线阶段的嵌入梯度（如果未绑定），以及专家并行性的专家梯度。

core.distributed.finalize_model_grads.finalize_model_grads(model: List[torch.nn.Module], num_tokens: Optional[torch.Tensor] = None)

在 DP 副本之间 all-reduce 所有模型梯度，序列并行性的 layernorm 梯度，跨第一和最后流水线阶段的嵌入梯度（如果未绑定），按 num_tokens 缩放梯度。

包含在优化器步骤之前跨不同 rank 同步梯度的功能。