Skip to content

Gradient Synchronization

About

The d9d.internals.grad_sync package synchronizes gradients of DTensor parameters in distributed training. PyTorch DistributedDataParallel applies one communication strategy to the whole model. This package instead supports the mixed layouts of ND parallelism, which combines data, tensor, context and pipeline parallelism. It reads the DTensor placements of each parameter to find the mesh dimensions that need an all-reduce. Then it groups the parameters into communication buckets.

Internal API

If you use the standard d9d training loop, you do not need to call this package. d9d synchronizes the gradients itself. This page is for users who extend d9d.

Bucketing and Flattening

When many small tensors are reduced, latency dominates the communication cost. So GradientSynchronizer groups parameters into buckets.

A SyncGradientBucket stores the gradients of all its parameters in one contiguous buffer. Its reduction is a single all_reduce per process group on this buffer, instead of hundreds of small operations.

Parameters share a bucket only if they have the same:

  1. Optimizer parameter group
  2. Device
  3. Gradient dtype
  4. Reduce mesh: The mesh dimensions where the parameter has a Replicate placement.

The bucket_size_mb argument limits the size of a bucket in MiB. The limit counts the local shard of each gradient in the gradient dtype. Parameters are bucketed in reverse order, because the backward pass produces gradients roughly in that order. Parameters with no Replicate placement need no reduction. They go to local buckets that do not communicate.

Asynchronous Reduction

Training often accumulates gradients over several microbatches before an optimizer step. This package manages the DTensor gradients during the accumulation. You do not need a no_sync context manager.

  1. Local accumulation: During the backward pass of the first \(N-1\) microbatches, local gradients accumulate in the bucket buffer. The parameter DTensor is Replicate, so the gradient also has a Replicate placement across the data parallel mesh. Its data still differs between ranks at this point.

  2. Automatic trigger: Each bucket counts the gradient accumulations of its parameters. The all_reduce starts only when all parameters of the bucket reach the require_accumulations count. The trigger runs inside the backward hook of the last microbatch. So the communication overlaps with the backward pass of the remaining layers. The communication runs on a separate CUDA stream. You must wait for it before you use the gradients on your default stream.

  3. Synchronization: When the asynchronous reduction completes, the flat buffer holds the globally summed gradient. The gradients have the same Replicate placement as their parameters. So the optimizer can use them without any further synchronization.

API Reference

d9d.internals.grad_sync

Gradient bucketing and asynchronous gradient reduction.

It works like DistributedDataParallel, but on DTensor parameters and for internal use.

GradientSynchronizer

Gradient synchronizer for replicated parameters. It reduces the gradients during the backward pass.

It splits the parameters into buckets, allocates flat gradient buffers and runs asynchronous all-reduce operations.

__init__(param_groups, bucket_size_mb)

Constructs the GradientSynchronizer object.

Parameters:

Name Type Description Default
param_groups list[list[Parameter]]

The parameter groups.

 required
bucket_size_mb int

The maximum size of one gradient bucket in MiB.

 required
bind()

Builds the buckets, allocates their gradient buffers and registers the backward hooks.

It must be called before the backward pass.

scale_grads(factor)

Multiplies the gradients of all managed parameters in place.

Parameters:

Name Type Description Default
factor Tensor

The scalar to multiply the gradients by.

 required

Raises:

Type Description
ValueError

If the synchronizer is not bound.
set_required_accumulations(require_accumulations)

Sets the accumulation count of the current step for all buckets.

The count can change from step to step, because the pack length varies. Later binds, for example after offload and onload, reuse the latest count.

Parameters:

Name Type Description Default
require_accumulations int

The number of accumulations required before the gradients are reduced.

 required
unbind()

Destroys the buckets, frees their buffers and removes the hooks.

wait()

Makes the current stream wait for all asynchronous reductions and marks the buckets as synchronized.

Raises:

Type Description
ValueError

If the synchronizer is not bound.
zero_grad()

Resets gradients and accumulation counters for all managed parameters.