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。非經特殊聲明，原始代碼版權歸原作者所有，本譯文未經允許或授權，請勿轉載或複製。