Python PyTorch bf16_compress_hook用法及代码示例

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

用法: torch.distributed.algorithms.ddp_comm_hooks.default_hooks.bf16_compress_hook(process_group, bucket)

警告：此 API 是实验性的，需要 NCCL 版本高于 2.9.6。

此 DDP 通信钩子实现了一种简单的梯度压缩方法，将 GradBucket 张量转换为 half-precision Brain floating point format ( torch.bfloat16 )，然后将其除以进程组大小。它全部减少了那些 bfloat16 梯度张量。一旦压缩梯度张量全部归约，链式回调 decompress 会将其转换回输入数据类型(例如 float32 )。

例子：

>>> ddp_model.register_comm_hook(process_group, bf16_compress_hook)

相关用法

注：本文由纯净天空筛选整理自pytorch.org大神的英文原创作品 torch.distributed.algorithms.ddp_comm_hooks.default_hooks.bf16_compress_hook。非经特殊声明，原始代码版权归原作者所有，本译文未经允许或授权，请勿转载或复制。