Python PyTorch powerSGD_hook用法及代码示例

本文简要介绍python语言中 torch.distributed.algorithms.ddp_comm_hooks.powerSGD_hook.powerSGD_hook 的用法。

用法:

torch.distributed.algorithms.ddp_comm_hooks.powerSGD_hook.powerSGD_hook(state, bucket)

参数：

• state(PowerSGDState) -用于配置压缩率和支持错误反馈、热启动等的状态信息。要调整压缩配置，主要需要调整 matrix_approximation_rank 、 start_powerSGD_iter 和 min_compression_rate 。

• bucket(dist.GradBucket) -存储一维扁平梯度张量的桶，该张量批量处理多个每个变量的张量。请注意，由于 DDP 通信钩子仅支持单进程单设备模式，因此该存储桶中仅存储了一个张量。

返回：

通信的未来处理程序，它更新到位的梯度。

此 DDP 通信钩子实现 paper 中说明的 PowerSGD 梯度压缩算法。一旦梯度张量在所有工作人员中聚合，此钩子将按如下方式应用压缩：

1. 将输入的扁平化一维梯度张量视为每个参数张量的列表，并将所有张量分为两组：

   1.1 allreduce之前应该压缩的张量，因为压缩可以节省足够的带宽。

   1.2 其余的张量将直接全部归约而不压缩，包括所有的向量张量(用于偏差)。

2. 处理未压缩的张量：

   2.1。为那些未压缩的张量分配连续的内存，并将所有未压缩的张量作为一个批次全部减少，不进行压缩；

   2.2.将单个未压缩张量从连续内存复制回输入张量。

3. 处理应该通过PowerSGD 压缩来压缩的张量：

   3.1.对于每个张量 M，创建两个低秩张量 P 和 Q 来分解 M，使得 M = PQ^T，其中 Q 从标准正态分布初始化并正交化；

   3.2.计算 Ps 中的每个 P，等于 MQ；

   3.3. Allreduces Ps 作为一个批次；

   3.4.正交化 Ps 中的每个 P；

   3.5.计算 Qs 中的每个 Q，大约等于 M^TP；

   3.6. Allreduce Qs 作为一个批次；

   3.7.计算所有压缩张量中的每个 M，大约等于 PQ^T。

请注意，此通信挂钩对第一个 state.start_powerSGD_iter 迭代强制执行 vanilla allreduce。这不仅让用户可以更好地控制加速和准确性之间的权衡，还有助于为未来的通信钩子开发人员抽象出 DDP 内部优化的一些复杂性。

例子：

>>> state = PowerSGDState(process_group=process_group, matrix_approximation_rank=1,
                          start_powerSGD_iter=10, min_compression_rate=0.5)
>>> ddp_model.register_comm_hook(state, powerSGD_hook)