Python PyTorch bf16_compress_wrapper用法及代码示例

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

用法: torch.distributed.algorithms.ddp_comm_hooks.default_hooks.bf16_compress_wrapper(hook)

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

此包装器将给定 DDP 通信挂钩的输入梯度张量转换为 half-precisionBrain floating point format <https://en.wikipedia.org/wiki/Bfloat16_floating-point_format> `_ (``torch.bfloat16`)，并将给定钩子的结果张量转换回输入数据类型，例如float32.

因此，bf16_compress_hook 等价于 bf16_compress_wrapper(allreduce_hook) 。

例子：

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

相关用法

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