tensorplay.distributed.algorithms.ddp_comm_hooks.ddp_zero_hook API

Modify hook to overlap ZeRO’s optimizer step with the DDP backward pass.

Once a bucket’s gradients have been computed, the optimizer computation using those gradients launches, yielding an interleaving of all-reduces and broadcasts in the communication stream. Preferred over hook_with_zero_step() when communication is relatively fast.

Modify hook to overlap ZeRO’s optimizer step with the DDP backward pass.

The optimizer computation follows the backward computation, overlapping with outstanding backward communication. May be preferred over hook_with_zero_step_interleaved() when communication is relatively slow compared to computation.

Parameters:

• hook – the hook to modify.

• ddp – the DDP instance to use.

• zero – the ZeRO instance to use.

• shard_buckets (bool) – if True, each DDP bucket assignment is partitioned across possibly multiple ranks.

Raises:

ValueError – if zero was constructed with overlap_with_ddp=False.