tensorplay.distributed.algorithms.ddp_comm_hooks.powerSGD_hook.powerSGD_hook

tensorplay.distributed.algorithms.ddp_comm_hooks.powerSGD_hook.powerSGD_hook(state: PowerSGDState, bucket: GradBucket)[source]

Implement PowerSGD algorithm.

This DDP communication hook implements PowerSGD gradient compression algorithm described in the paper.

Note that this communication hook enforces vanilla allreduce for the first state.start_powerSGD_iter iterations.

Parameters:

• state (PowerSGDState) – State information to configure the compression rate and support error feedback, warm start, etc.

• bucket (dist.GradBucket) – Bucket that stores a 1D flattened gradient tensor that batches multiple per-variable tensors.

Returns:

Future handler of the communication, which updates the gradients in place.

Example::

>>> # xdoctest: +SKIP
>>> state = PowerSGDState(process_group=process_group, matrix_approximation_rank=1,
                          start_powerSGD_iter=10, min_compression_rate=0.5)
>>> ddp_model.register_comm_hook(state, powerSGD_hook)