tensorplay.distributed.algorithms

tensorplay.distributed.algorithms collects the pieces that plug into a running data-parallel job: communication hooks that change how gradients are exchanged between ranks, the join mechanism that lets uneven workloads participate in the same collectives, and the model averagers that keep replicas in sync outside the all-reduce path. The two subpackages are ddp_comm_hooks (hooks registered on DistributedDataParallel) and model_averaging (averaging strategies).

A DDP communication hook is a callable (state, bucket) -> future that replaces the default gradient all-reduce: when a backward pass fills a gradient bucket, DDP hands it to your hook instead of reducing it directly. That single seam is enough to implement gradient compression, fused optimizer steps, and local-SGD variants without touching the training loop.

import tensorplay as tp
import tensorplay.distributed as dist
from tensorplay.distributed.algorithms.ddp_comm_hooks import default_hooks

dist.init_process_group("nccl")
model = tp.nn.Linear(64, 64).cuda()
model = tp.nn.parallel.DistributedDataParallel(model, device_ids=[dist.get_rank()])

# halve the gradient traffic: compress to fp16 before the all-reduce
model.register_comm_hook(None, default_hooks.fp16_compress_hook)

Default hooks

`tensorplay.distributed.algorithms.ddp_comm_hooks.default_hooks.allreduce_hook`	Call `allreduce` using `GradBucket` tensors.
`tensorplay.distributed.algorithms.ddp_comm_hooks.default_hooks.fp16_compress_hook`	Compress by casting `GradBucket` to float16 divided by process group size.
`tensorplay.distributed.algorithms.ddp_comm_hooks.default_hooks.bf16_compress_hook`	Warning: This API is experimental, and it requires NCCL version later than 2.9.6.
`tensorplay.distributed.algorithms.ddp_comm_hooks.default_hooks.fp16_compress_wrapper`	Cast input tensor to `float16`, cast result of hook back to input dtype.
`tensorplay.distributed.algorithms.ddp_comm_hooks.default_hooks.bf16_compress_wrapper`	Warning: This API is experimental, and it requires NCCL version later than 2.9.6.

allreduce_hook() is the plain gradient all-reduce — the same communication DDP performs by default, useful as the baseline or the inner step of a composed hook.
fp16_compress_hook / bf16_compress_hook cast each bucket’s gradients to half precision before the all-reduce and cast the result back, halving the bytes on the wire. The *_compress_wrapper variants wrap another hook, so the compression applies before it runs — for example fp16_compress_wrapper(powerSGD_hook) compresses the residual that PowerSGD feeds to the collective.

PowerSGD

`tensorplay.distributed.algorithms.ddp_comm_hooks.powerSGD_hook.PowerSGDState`	Store both the algorithm's hyperparameters and internal state for all gradients during training.
`tensorplay.distributed.algorithms.ddp_comm_hooks.powerSGD_hook.powerSGD_hook`	Implement PowerSGD algorithm.
`tensorplay.distributed.algorithms.ddp_comm_hooks.powerSGD_hook.batched_powerSGD_hook`	Implement simplified PowerSGD algorithm.

PowerSGD compresses a gradient matrix to the product of two low-rank factors (rank matrix_approximation_rank), all-reduces the factors instead of the full matrix, and reconstructs on the other side; the residual (what the factors fail to represent) is carried into the next iteration when error feedback is on. PowerSGDState carries the per-parameter factors and the tuning knobs — the approximation rank, the warm-up iteration before compression starts (start_powerSGD_iter), and the minimum compression rate below which a tensor is left dense. batched_powerSGD_hook additionally batches same-shape tensors into one factorization for better throughput.

Quantization hooks

`tensorplay.distributed.algorithms.ddp_comm_hooks.quantization_hooks.quantization_pertensor_hook`	Apply `quantize_per_tensor` logic to DDP using `allgather` protocol.
`tensorplay.distributed.algorithms.ddp_comm_hooks.quantization_hooks.quantization_perchannel_hook`	Apply `quantize_per_channel` logic to DDP using `allgather` protocol.

The quantization hooks reduce gradient bytes by quantizing each bucket before the collective — per-tensor (one scale for the whole bucket) or per-channel (a scale per group of bucket_size elements), dequantizing after. Per-channel keeps more precision at the cost of bookkeeping.

Optimizer-in-hook (ZeRO-style)

`tensorplay.distributed.algorithms.ddp_comm_hooks.ddp_zero_hook.hook_with_zero_step`	Modify `hook` to overlap ZeRO's optimizer step with the DDP backward pass.
`tensorplay.distributed.algorithms.ddp_comm_hooks.ddp_zero_hook.hook_with_zero_step_interleaved`	Modify `hook` to overlap ZeRO's optimizer step with the DDP backward pass.

hook_with_zero_step() fuses the ZeroRedundancyOptimizer step into the gradient bucket’s collective: as each bucket’s all-reduce completes, the hook immediately applies the optimizer update for the parameters in that bucket, so the optimizer overlaps the backward pass instead of waiting for it. The shard_buckets option partitions buckets across ranks the way ZeRO partitions parameters. The _interleaved variant schedules the collectives and the parameter shard updates bucket-by-bucket. Both take the DistributedDataParallel module and the ZeroRedundancyOptimizer to drive.

Post-local SGD

`tensorplay.distributed.algorithms.ddp_comm_hooks.post_localSGD_hook.PostLocalSGDState`	Store state for all-reducing gradients globally until given step, then locally after.
`tensorplay.distributed.algorithms.ddp_comm_hooks.post_localSGD_hook.post_localSGD_hook`	Run post-localSGD algorithm.

post_localSGD_hook() switches the gradient reduction domain over time: before start_localSGD_iter the gradients are all-reduced across the full process_group; afterwards they are only all-reduced within the local subgroup (one node, typically). The per-replica parameters then drift, which is exactly what the post-local-SGD optimizer compensates for by averaging them periodically. PostLocalSGDState carries the two groups, the switch iteration, and whether to keep a global gradient all-reduce after the switch.

Model averaging

`tensorplay.distributed.algorithms.model_averaging.ModelAverager`	Base class for all model averagers.
`tensorplay.distributed.algorithms.model_averaging.PeriodicModelAverager`	Averages parameters periodically after the warm-up stage.
`tensorplay.distributed.algorithms.model_averaging.HierarchicalModelAverager`	Runs hierarchical model averaging (hierarchical SGD).

ModelAverager is the abstract interface the post-local-SGD optimizer consumes: it knows how to average parameters across a process group. PeriodicModelAverager averages every period steps after warmup_steps have elapsed. HierarchicalModelAverager averages in a tree of subgroups (given as period_group_size_dict), which reduces the latency of large-scale averaging at the cost of eventual (rather than immediate) consistency.

Join

`tensorplay.distributed.algorithms.Join`	This class defines the generic join context manager, which allows custom hooks to be called after a process joins.
`tensorplay.distributed.algorithms.JoinHook`	This defines a join hook, which provides two entry points in the join context manager.
`tensorplay.distributed.algorithms.Joinable`	This defines an abstract base class for joinable classes.

Join is a context manager for uneven inputs in data-parallel training: ranks that run out of work early “join” — their collective participation is emulated so the ranks still working do not hang on a barrier. A Joinable class (DDP is one) provides a JoinHook whose main_hook shadows the collectives while some ranks have joined, and whose post_hook runs once everyone has joined, letting the last ranks propagate any final state.

Where to go next

the distributed package — DDP and the collectives these hooks ride on.
the distributed optimizer — the ZeroRedundancyOptimizer the zero-step hooks fuse with, and the post-local-SGD optimizer the model averagers feed.
FSDP — the sharded alternative when optimizer state, not gradient traffic, is the bottleneck.