tensorplay.distributed.algorithms.ddp_comm_hooks.quantization_hooks.quantization_pertensor_hook

tensorplay.distributed.algorithms.ddp_comm_hooks.quantization_hooks.quantization_pertensor_hook(process_group, bucket: GradBucket)[source]

Apply quantize_per_tensor logic to DDP using allgather protocol.

Workers first allgather the scale and zero point of their own GradBucket prior to the quantization. After all workers have that information, the first then callback called quantize_and_allgather quantizes worker’s own gradient tensor, and uses allgather to communicate these across all workers. The final then callback called dequantize_and_aggregate, dequantizes and aggregates each quantized gradient tensor locally and returns the mean.

Warning

This is experimental, and uses allgather protocol which is considerably slower than allreduce protocol. It works only with flattened grads.

Example::

>>> # xdoctest: +SKIP
>>> ddp_model.register_comm_hook(process_group, quantization_pertensor_hook)