tensorplay.distributed.algorithms.ddp_comm_hooks.quantization_hooks API

Apply quantize_per_channel logic to DDP using allgather protocol.

Compared to per-tensor, the main motivation of per-channel is for considerably large tensors such as a tensor that contains 6 million elements quantizing per a bucket size of 512 (or 128) elements may significantly increase the resolution.

It first splits GradBucket tensor into multiple chunks (channels) of bucket_size elements. Then, workers allgather the scales and zero points 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, flattens, and aggregates each quantized gradient tensor locally and returns the mean.

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.

Example::

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