def _collect_ddp_bucket_info(
bucket: dist.GradBucket,
zero: ZeroRedundancyOptimizer,
rank: int,
assigned_rank: int,
):
r"""
Collects :class:`DistributedDataParallel` gradient bucket information for
the :class:`ZeroRedundancyOptimizer` instance ``zero`` to use when
overlapping.
Arguments:
bucket (dist.GradBucket): the current gradient bucket.
zero (ZeroRedundancyOptimizer): the calling process's
:class:`ZeroRedundancyOptimizer` instance.
rank (int): the calling process's rank.
assigned_rank (int): the rank assigned to update the parameters
corresponding to ``bucket``.
"""
overlap_info = zero._overlap_info
bucket_index = bucket.index()
bucket_params = bucket.parameters()
assert len(bucket_params) > 0, "Bucket {bucket_index} is empty"
params_per_rank = overlap_info.params_per_rank
params_per_bucket = overlap_info.params_per_bucket
# Collect relevant information
if assigned_rank == rank:
overlap_info.offsets[bucket_index] = len(
params_per_rank[assigned_rank])
params_per_rank[assigned_rank].extend(bucket_params)
params_per_bucket.append(bucket_params)
def _perform_local_step(
bucket: dist.GradBucket,
zero: ZeroRedundancyOptimizer,
rank: int,
):
r"""
Performs a local optimizer step using the gradients provided by ``bucket``.
Arguments:
bucket (dist.GradBucket): the bucket providing the gradients.
zero (ZeroRedundancyOptimizer): the :class:`ZeroRedundancyOptimizer`
instance to perform the :meth:`_local_step`.
rank (int): the calling process's rank.
.. warning::
This function assumes that appropriate synchronization has taken place
so that the bucket's gradients can be used.
"""
overlap_info = zero._overlap_info
bucket_index = bucket.index()
assert len(zero.optim.param_groups) == 1, \
"Overlapping DDP with ZeRO only supports a single parameter group"
# Construct the `gradients` input for the local optimizer step, which
# expects `None` in a list position to indicate that the corresponding
# parameter should not be updated
num_local_optim_params = len(zero.optim.param_groups[0]["params"])
gradients: List[Optional[torch.Tensor]] = \
[_NO_PARAM_UPDATE for _ in range(num_local_optim_params)]
assert bucket_index in overlap_info.offsets, \
f"Bucket index {bucket_index} was not assigned to rank {rank}"
gradients_offset = overlap_info.offsets[bucket_index]
bucket_assignment = zero._bucket_assignments_per_rank[rank][bucket_index]
bucket_offset = bucket_assignment.offset
length = len(bucket_assignment.parameters)
bucket_gradients = bucket.gradients()[bucket_offset:bucket_offset + length]
for i, grad in enumerate(bucket_gradients):
gradients[gradients_offset + i] = grad
zero._local_step(gradients)
def _hook_with_zero_step_setup(
ddp_ref: weakref.ReferenceType,
zero: ZeroRedundancyOptimizer,
bucket: dist.GradBucket,
):
r"""
Encapsulates the setup logic for :func:`hook_with_zero_step` and
:func:`hook_with_zero_step_interleaved`, meaning the logic to run in the
hook before the backward pass and optimizer step can actually be
overlapped. This is factored out since it is common to both
:func:`hook_with_zero_step` and :func:`hook_with_zero_step_interleaved`.
Arguments:
ddp_ref (weakref.ReferenceType): weak reference to the process's
:class:`DistributedDataParallel` instance.
zero (ZeroRedundancyOptimizer): the calling process's
:class:`ZeroRedundancyOptimizer` instance.
bucket (dist.GradBucket): the current gradient bucket.
"""
# Proceed as normal until the DDP buckets have been rebuilt
if not ddp_ref()._has_rebuilt_buckets: # type: ignore[union-attr]
assert zero._overlap_info.status == _OverlapStatus.UNINITIALIZED
return
bucket_index = bucket.index()
overlap_info = zero._overlap_info
if overlap_info.status == _OverlapStatus.UNINITIALIZED:
overlap_info.status = _OverlapStatus.DDP_HAS_REBUILT_BUCKETS
if overlap_info.status == _OverlapStatus.DDP_HAS_REBUILT_BUCKETS:
if bucket_index == 0 and len(overlap_info.params_per_bucket) > 0:
# This corresponds to the first bucket of the backward pass
# immediately after all information has been saved, so we
# can perform the delayed ZeRO initialization
zero._init_zero_for_overlap()
else:
# Once DDP buckets have been rebuilt but ZeRO has not been
# properly initialized yet, save the information needed
_save_ddp_bucket_info(bucket, zero)
def batched_powerSGD_hook(
state: PowerSGDState,
bucket: dist.GradBucket) -> torch.futures.Future[torch.Tensor]:
r"""
This DDP communication hook implements a simplified PowerSGD gradient compression
algorithm described in the `paper <https://arxiv.org/abs/1905.13727>`_.
This variant does not compress the gradients layer by layer,
but instead compresses the flattened input tensor that batches all the gradients.
Therefore, it is **faster** than :meth:`powerSGD_hook`,
but usually results in a **much lower accuracy**, unless ``matrix_approximation_rank`` is 1.
.. warning ::
Increasing ``matrix_approximation_rank`` here may not necessarily increase the accuracy,
because batching per-parameter tensors without column/row alignment can destroy low-rank structure.
Therefore, the user should always consider :meth:`powerSGD_hook` first,
and only consider this variant when a satisfactory accuracy can be achieved when ``matrix_approximation_rank`` is 1.
Once gradient tensors are aggregated across all workers, this hook applies
compression as follows:
1. Views the input flattened 1D gradient tensor as a square-shaped tensor M with 0 paddings;
2. Creates two low-rank tensors P and Q for decomposing M, such that M = PQ^T, where Q is initialized from a standard normal distribution and orthogonalized;
3. Computes P, which is equal to MQ;
4. Allreduces P;
5. Orthogonalizes P;
6. Computes Q, which is approximately equal to M^TP;
7. Allreduces Q;
8. Computes M, which is approximately equal to PQ^T.
9. Truncates the input tensor to the original length.
Note that this communication hook enforces vanilla allreduce for the first ``state.start_powerSGD_iter`` iterations.
This not only gives the user more control over the tradeoff between speedup and accuracy,
but also helps abstract away some complexity of the internal optimization of DDP for future communication hook developers.
Args:
state (PowerSGDState): State information to configure the compression rate and support error feedback, warm start, etc.
To tune the compression configs, mainly need to tune ``matrix_approximation_rank`` and ``start_powerSGD_iter``.
bucket (dist.GradBucket): Bucket that stores a 1D flattened gradient tensor that batches multiple per-variable tensors.
Note that since DDP comm hook only supports single process single device mode,
only exactly one tensor is stored in this bucket.
Returns:
Future handler of the communication, which updates the gradients in place.
Example::
>>> state = PowerSGDState(process_group=process_group, matrix_approximation_rank=1)
>>> ddp_model.register_comm_hook(state, batched_powerSGD_hook)
""" # noqa: B950
process_group = state.process_group
group_to_use = process_group if process_group is not None else dist.group.WORLD
world_size = group_to_use.size()
# The input tensor is a flattened 1D tensor.
input_tensor = bucket.buffer()
# Run vanilla allreduce in the first `start_powerSGD_iter` iterations.
if state.iter < state.start_powerSGD_iter:
state.maybe_increase_iter(bucket)
return default._allreduce_fut(group_to_use, input_tensor)
# Apply PowerSGD after `start_powerSGD_iter` iterations.
device = input_tensor.device
total_length = input_tensor.shape[0]
state.total_numel_before_compression += total_length
# View the input tensor as a 2D square-shape tensor, and pad 0s if necessary.
square_side_length = math.ceil(math.sqrt(total_length))
state.total_numel_after_compression += (square_side_length *
state.matrix_approximation_rank *
2)
padded_total_length = square_side_length**2
input_tensor.resize_(padded_total_length)
input_tensor[total_length:padded_total_length].fill_(0)
_report_compression_stats(bucket, state)
# Incorporate the error from the previous state into the gradients.
bucket_index = bucket.index()
input_tensor_cp = None
if state.use_error_feedback:
if bucket_index in state.error_dict:
input_tensor.add_(state.error_dict[bucket_index])
else:
logging.info(
"A zero tensor of length {} that represents local error is created."
.format(padded_total_length))
state.error_dict[bucket_index] = torch.zeros(
padded_total_length, device=device, dtype=input_tensor.dtype)
# Keep a copy of the input tensor,
# so that we can compute the local error caused by compression later,
# by comparing this copy and the input tensor updated after decompression.
input_tensor_cp = torch.clone(input_tensor).detach()
matrix = input_tensor.view(square_side_length, square_side_length)
# Reuse P and Q from the previous iteration if possible.
# The memory spaces of P and Q need to be allocated in the first iteration when PowerSGD is applied.
if not state.warm_start or bucket_index not in state.p_memory_dict:
# If warm-start is disabled, low-rank tensors will be initialized at every step.
# Only log this if warm-start to avoid spamming.
if state.warm_start:
logging.info(
"Initializing low-rank tensors P and Q, each of which has a shape of {} x {}."
.format(square_side_length, state.matrix_approximation_rank))
def create_low_rank_tensor(fill_random_values, rng):
"Returns a low-rank 2D tensor of square_side_length * matrix_approximation_rank."
if fill_random_values:
with torch.random.fork_rng(devices=[]):
# Fork this RNG to avoid changing the seed globally and affecting the random sampling
# anywhere else in the training.
# The seed makes sure that the initial random values are the same across all the DDP replicas.
# This seed should differ at every step.
# Since it is very slow to fork RNG state across all the CUDA devices,
# only fork on CPU and then move the generated tensor to the CUDA device.
torch.manual_seed(rng.randint(1_000_000_000))
return torch.randn(
square_side_length,
state.matrix_approximation_rank,
device="cpu",
dtype=input_tensor.dtype,
).to(device)
else:
return torch.empty(
square_side_length,
state.matrix_approximation_rank,
device=device,
dtype=input_tensor.dtype,
)
state.p_memory_dict[bucket_index] = create_low_rank_tensor(
fill_random_values=False, rng=state.rng)
state.q_memory_dict[bucket_index] = create_low_rank_tensor(
fill_random_values=True, rng=state.rng)
_orthogonalize(state.q_memory_dict[bucket_index])
torch.matmul(matrix,
state.q_memory_dict[bucket_index],
out=state.p_memory_dict[bucket_index])
allreduce_p_fut = dist.all_reduce(state.p_memory_dict[bucket_index],
group=group_to_use,
async_op=True).get_future()
def compute_q(fut):
state.p_memory_dict[bucket_index] = fut.value()[0]
_orthogonalize(state.p_memory_dict[bucket_index])
torch.matmul(
matrix.t(),
state.p_memory_dict[bucket_index],
out=state.q_memory_dict[bucket_index],
)
# TODO: The above procedure does two matmul+allreduce steps per iteration --
# one left multiplication and one right multiplication.
# For warm-start, can take one such step at a time, and alternate between them.
return (dist.all_reduce(state.q_memory_dict[bucket_index],
group=group_to_use,
async_op=True).get_future().wait()[0])
def decompress(fut):
state.q_memory_dict[bucket_index] = fut.value().div_(world_size)
torch.matmul(
state.p_memory_dict[bucket_index],
state.q_memory_dict[bucket_index].t(),
out=matrix,
)
if state.use_error_feedback:
# Memorize the local errors.
state.error_dict[bucket_index] = input_tensor_cp - input_tensor
# Removing this seemingly unnecessary sync somehow may cause faliures.
# See: https://github.com/pytorch/pytorch/pull/54838
if torch.cuda.is_available():
torch.cuda.synchronize(device)
if not state.warm_start:
state.p_memory_dict.clear()
state.q_memory_dict.clear()
ret = input_tensor.resize_(total_length)
state.maybe_increase_iter(bucket)
return ret
return allreduce_p_fut.then(compute_q).then(decompress)
def powerSGD_hook(
state: PowerSGDState,
bucket: dist.GradBucket) -> torch.futures.Future[torch.Tensor]:
r"""
This DDP communication hook implements PowerSGD gradient compression
algorithm described in the `paper <https://arxiv.org/abs/1905.13727>`_.
Once gradient tensors are aggregated across all workers, this hook applies
compression as follows:
1. Views the input flattened 1D gradient tensor as a list of per-parameter tensors, and divides all the tensors into two groups:
1.1 The tensors that should be compressed before allreduce, because the compression can give enough saving in bandwidth.
1.2 Rest of the tensors will be directly allreduced without compression, including all the vector tensors (for biases).
2. Handles uncompressed tensors:
2.1. Allocate contiguous memory for those uncompressed tensors, and allreduces all the uncompressed tensors as a batch, without compression;
2.2. Copies the individual uncompressed tensors from the contiguous memory back to the input tensor.
3. Handles the tensors that should be compressed by PowerSGD compression:
3.1. For each tensor M, creates two low-rank tensors P and Q for decomposing M,
such that M = PQ^T, where Q is initialized from a standard normal distribution and orthogonalized;
3.2. Computes each P in Ps, which is equal to MQ;
3.3. Allreduces Ps as a batch;
3.4. Orthogonalizes each P in Ps;
3.5. Computes each Q in Qs, which is approximately equal to M^TP;
3.6. Allreduces Qs as a batch;
3.7. Computes each M among all the compressed tensors, which is approximately equal to PQ^T.
Note that this communication hook enforces vanilla allreduce for the first ``state.start_powerSGD_iter`` iterations.
This not only gives the user more control over the tradeoff between speedup and accuracy,
but also helps abstract away some complexity of the internal optimization of DDP for future communication hook developers.
Args:
state (PowerSGDState): State information to configure the compression rate and support error feedback, warm start, etc.
To tune the compression configs, mainly need to tune ``matrix_approximation_rank``, ``start_powerSGD_iter``
and ``min_compression_rate``.
bucket (dist.GradBucket): Bucket that stores a 1D flattened gradient tensor that batches multiple per-variable tensors.
Note that since DDP comm hook only supports single process single device mode,
only exactly one tensor is stored in this bucket.
Returns:
Future handler of the communication, which updates the gradients in place.
Example::
>>> 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)
""" # noqa: B950
process_group = state.process_group
group_to_use = process_group if process_group is not None else dist.group.WORLD
world_size = group_to_use.size()
# The input tensor is a flattened 1D tensor.
input_tensor = bucket.buffer()
# Run vanilla allreduce in the first `start_powerSGD_iter` iterations.
if state.iter < state.start_powerSGD_iter:
state.maybe_increase_iter(bucket)
return default._allreduce_fut(group_to_use, input_tensor)
# Apply PowerSGD after `start_powerSGD_iter` iterations.
device = input_tensor.device
dtype = input_tensor.dtype
# Incorporate the error from the previous state into the gradients.
bucket_index = bucket.index()
input_tensor_cp = None
total_length = input_tensor.shape[0]
if state.use_error_feedback:
if bucket_index in state.error_dict:
input_tensor.add_(state.error_dict[bucket_index])
else:
logging.info(
"A zero tensor of length {} that represents local error is created."
.format(total_length))
state.error_dict[bucket_index] = torch.zeros(total_length,
device=device,
dtype=dtype)
# Keep a copy of the input tensor,
# so that we can compute the local error caused by compression later,
# by comparing this copy and the input tensor updated after decompression.
input_tensor_cp = torch.clone(input_tensor).detach()
# Unflatten the input tensor into per-parameter tensors, for layer-wise compression.
tensors = bucket.gradients()
# Step I: Divide all the tensors into two groups,
# one will be compressed before allreduce and the other will be directly allreduced without compression.
tensors_to_compress, uncompressed_tensors = [], []
total_Ps_size = 0
total_Qs_size = 0
for tensor in tensors:
matrix = tensor.view(tensor.shape[0], -1)
n, m = matrix.shape
matrix_approximation_rank = min(n, m, state.matrix_approximation_rank)
compress_test = _should_compress(n, m, matrix_approximation_rank,
state.min_compression_rate)
state.total_numel_before_compression += compress_test[1]
if compress_test[0]:
tensors_to_compress.append(matrix)
total_Ps_size += n * matrix_approximation_rank
total_Qs_size += m * matrix_approximation_rank
state.total_numel_after_compression += compress_test[2]
else:
uncompressed_tensors.append(tensor)
state.total_numel_after_compression += compress_test[1]
_report_compression_stats(bucket, state)
# Step II: Handle uncompressed tensors.
# Allocate contiguous memory for these tensors to allreduce efficiently.
uncompressed_tensors_memory = (
torch.cat([tensor.view(-1) for tensor in uncompressed_tensors]) if
uncompressed_tensors else torch.tensor([], device=device, dtype=dtype))
# Step III: Handle the tensors that should be compressed.
# Allocate contiguous memory for Ps and Qs to allreduce efficiently.
# If warm-start is enabled, reuse Ps and Qs from the previous iteration if possible.
# The memory spaces of Ps and Qs need to be allocated in the first iteration when PowerSGD is applied.
need_randomize_qs = False
if not state.warm_start or bucket_index not in state.p_memory_dict:
need_randomize_qs = True
# If warm-start is disabled, low-rank tensors will be initialized at every step.
# Only log this if warm-start to avoid spamming.
if state.warm_start:
logging.info(
"Allocating contiguous memory of length {} for Ps, and of length {} for Qs, respectively."
.format(total_Ps_size, total_Qs_size))
state.p_memory_dict[bucket_index] = torch.empty(total_Ps_size,
device=device,
dtype=dtype)
state.q_memory_dict[bucket_index] = torch.empty(total_Qs_size,
device=device,
dtype=dtype)
# Create Ps and Qs that point to the allocated memory.
ps = []
qs = []
p_idx = 0
q_idx = 0
for tensor in tensors_to_compress:
n, m = tensor.shape
matrix_approximation_rank = min(n, m, state.matrix_approximation_rank)
ps.append(
state.p_memory_dict[bucket_index][p_idx:p_idx + n *
matrix_approximation_rank].view(
n,
matrix_approximation_rank))
qs.append(
state.q_memory_dict[bucket_index][q_idx:q_idx + m *
matrix_approximation_rank].view(
m,
matrix_approximation_rank))
p_idx += n * matrix_approximation_rank
q_idx += m * matrix_approximation_rank
# If warm-start is enabled, reuse Qs from the previous iteration if possible and skip filling random values.
# The exception is the first iteration when PowerSGD is applied.
if not need_randomize_qs:
for q in qs:
_orthogonalize(q, state.orthogonalization_epsilon)
else:
with torch.random.fork_rng(devices=[]):
# Fork this RNG to avoid changing the seed globally and affecting the random sampling anywhere else in the training.
# The seed makes sure that the initial random values are the same across all the DDP replicas.
# This seed should differ at every step.
# Since it is very slow to fork RNG state across all the CUDA devices,
# only fork on CPU and then move the generated tensor to the CUDA device (by overwriting q).
torch.manual_seed(state.rng.randint(1_000_000_000))
for q in qs:
q.copy_(torch.randn(
*q.shape,
device="cpu",
dtype=dtype,
))
_orthogonalize(q, state.orthogonalization_epsilon)
# Compute Ps.
for tensor, q, p in zip(tensors_to_compress, qs, ps):
torch.matmul(tensor, q, out=p)
# This allreduce is only applied to uncompressed tensors,
# so it should have been kicked off before the above computation on the compressed tensors to hide more communication costs.
# However, this somehow requires a separate future chain at this time.
allreduce_contiguous_uncompressed_tensors_fut = dist.all_reduce(
uncompressed_tensors_memory, group=group_to_use,
async_op=True).get_future()
def unpack_uncompressed_tensors_and_allreduce_ps(fut):
uncompressed_tensors_memory = fut.value()[0].div_(world_size)
idx = 0
for tensor in uncompressed_tensors:
tensor.copy_(
uncompressed_tensors_memory[idx:idx +
tensor.numel()].view_as(tensor))
idx += tensor.numel()
# Since these Ps will be orthogonalized later, no need to divide them by world size.
return (dist.all_reduce(state.p_memory_dict[bucket_index],
group=group_to_use,
async_op=True).get_future().wait()[0])
def compute_qs(fut):
state.p_memory_dict[bucket_index] = fut.value()
for p in ps:
_orthogonalize(p, state.orthogonalization_epsilon)
# Compute Qs.
for tensor, p, q in zip(tensors_to_compress, ps, qs):
torch.matmul(tensor.t(), p, out=q)
# TODO: The above procedure does two matmul+allreduce steps per iteration --
# one left multiplication and one right multiplication.
# For warm-start, can take one such step at a time, and alternate between them.
# Allreduce Qs.
return (dist.all_reduce(state.q_memory_dict[bucket_index],
group=group_to_use,
async_op=True).get_future().wait()[0])
def decompress(fut):
state.q_memory_dict[bucket_index] = fut.value().div_(world_size)
for p, q, tensor in zip(ps, qs, tensors_to_compress):
torch.matmul(p, q.t(), out=tensor)
if torch.cuda.is_available():
torch.cuda.synchronize(device)
if state.use_error_feedback:
# Memorize the local errors.
state.error_dict[bucket_index] = input_tensor_cp - input_tensor
if not state.warm_start:
state.p_memory_dict.clear()
state.q_memory_dict.clear()
state.maybe_increase_iter(bucket)
return input_tensor
return (allreduce_contiguous_uncompressed_tensors_fut.then(
unpack_uncompressed_tensors_and_allreduce_ps).then(compute_qs).then(
decompress))
def hook_with_zero_fn(
state: Any,
bucket: dist.GradBucket,
) -> torch.futures.Future[torch.Tensor]:
r"""
Returns a :class:`Future` that gives a gradient bucket tensor and
performs the equivalent of a :class:`ZeroRedundancyOptimizer`
:meth:`step` if ``bucket`` is the last gradient bucket.
The function performs additional computation on the iteration that
the :class:`DistributedDataParallel` buckets are rebuilt to collect
information used to implement the modified hook.
Arguments:
state (Any): any state for the hook.
bucket (dist.GradBucket): the :class:`DistributedDataParallel`
gradient bucket.
"""
fut = hook(state, bucket)
_hook_with_zero_step_setup(ddp_ref, zero, bucket)
if zero._overlap_info.status != _OverlapStatus.INITIALIZED:
return fut
overlap_info = zero._overlap_info
bucket_index = bucket.index()
rank = zero.global_rank
assert overlap_info.status == _OverlapStatus.INITIALIZED
assert len(overlap_info.assigned_ranks_per_bucket) > bucket_index, \
"`assigned_ranks_per_bucket` is not fully constructed"
assigned_to_bucket = rank in overlap_info.assigned_ranks_per_bucket[bucket_index]
# Save the bucket reference and all-reduce future for the final bucket
if assigned_to_bucket:
overlap_info.bucket_index_to_bucket[bucket_index] = bucket
overlap_info.bucket_index_to_future[bucket_index] = fut
# Check that buckets are indexed incrementally starting from 0 in the
# order of their autograd hooks firing
if len(overlap_info.bucket_indices_seen) > 0:
assert overlap_info.bucket_indices_seen[-1] == bucket_index - 1, \
"Bucket indices are not in incremental order"
else:
assert bucket_index == 0, "Bucket indices do not start from 0"
overlap_info.bucket_indices_seen.append(bucket_index)
# Directly return the future without any optimizer computation if this
# is not the last bucket
num_buckets = len(overlap_info.params_per_bucket)
is_last_bucket = bucket_index == num_buckets - 1
if not is_last_bucket:
return fut
# Perform partial optimizer step on all buckets after the final
# bucket has been computed
# NOTE: This should not be chained as a callback to the last bucket's
# all-reduce future since that would add synchronization that delays
# all optimizer computation to wait for that last all-reduce
for bucket_index in range(num_buckets):
assigned_ranks = overlap_info.assigned_ranks_per_bucket[bucket_index]
if rank in assigned_ranks:
# Wait on the bucket's all-reduce future to ensure correct
# gradients
assert bucket_index in overlap_info.bucket_index_to_future, \
f"All-reduce future for bucket {bucket_index} not saved " \
f"on rank {rank}"
allreduce_future = overlap_info.bucket_index_to_future[bucket_index]
allreduce_future.wait()
# Perform the partial optimizer step
curr_bucket = overlap_info.bucket_index_to_bucket[bucket_index]
_perform_local_step(curr_bucket, zero, rank)
_broadcast_bucket(bucket_index, zero)
# Ensure that all parameter updates are finished before the
# next forward pass
overlap_info.wait_for_broadcasts()
overlap_info.clear_per_iter_info()
return fut
def hook_with_zero_fn(
state: Any,
bucket: dist.GradBucket,
) -> torch.futures.Future[torch.Tensor]:
r"""
Returns a :class:`Future` that gives a gradient bucket tensor and
performs the equivalent of a :class:`ZeroRedundancyOptimizer`
:meth:`step` if ``bucket`` is the last gradient bucket.
The function performs additional computation on the iteration that
the :class:`DistributedDataParallel` buckets are rebuilt to collect
information used to implement the modified hook.
Arguments:
state (Any): any state for the hook.
bucket (dist.GradBucket): the :class:`DistributedDataParallel`
gradient bucket.
"""
fut = hook(state, bucket)
overlap_info = zero._overlap_info
bucket_index = bucket.index()
# Proceed as normal until the DDP buckets have been rebuilt
if not ddp._has_rebuilt_buckets:
assert overlap_info.status == _OverlapStatus.UNINITIALIZED
return fut
if overlap_info.status == _OverlapStatus.UNINITIALIZED:
overlap_info.status = _OverlapStatus.DDP_HAS_REBUILT_BUCKETS
rank = zero.global_rank
rank_to_update = zero._ddp_bucket_index_to_rank(bucket_index)
# Once DDP buckets have been rebuilt but ZeRO has not been
# properly initialized yet, collect the information needed
if overlap_info.status == _OverlapStatus.DDP_HAS_REBUILT_BUCKETS:
bucket_params = bucket.parameters()
assert len(bucket_params) > 0, "Empty bucket"
params_per_rank = overlap_info.params_per_rank
params_per_bucket = overlap_info.params_per_bucket
if rank_to_update == rank:
overlap_info.offsets[bucket_index] = len(
params_per_rank[rank_to_update])
params_per_rank[rank_to_update].extend(bucket_params)
params_per_bucket.append(bucket_params)
return fut
assert overlap_info.status == _OverlapStatus.INITIALIZED
# Save the bucket reference and all-reduce future for the final bucket
if rank_to_update == rank:
overlap_info.bucket_index_to_bucket[bucket_index] = bucket
overlap_info.bucket_index_to_future[bucket_index] = fut
# NOTE: The implementation from this point forward assumes that the
# buckets are indexed incrementally starting from 0 in the order of
# their autograd hooks firing
num_buckets = len(overlap_info.params_per_bucket)
is_last_bucket = bucket_index == num_buckets - 1
if not is_last_bucket:
return fut
# Perform partial optimizer step on all buckets after the final
# bucket has been computed
# NOTE: This should not be chained as a callback to the last bucket's
# all-reduce future since that would add synchronization that delays
# all optimizer computation to wait for that last all-reduce
for bucket_index in range(num_buckets):
rank_to_update = zero._ddp_bucket_index_to_rank(bucket_index)
num_local_optim_params = len(zero.optim.param_groups[0]["params"])
if rank_to_update == rank:
gradients: List[Optional[torch.Tensor]] = \
[_NO_PARAM_UPDATE for _ in range(num_local_optim_params)]
assert bucket_index in overlap_info.offsets, \
f"Bucket index {bucket_index} was not assigned to rank {rank}"
offset = overlap_info.offsets[bucket_index]
# Ensure that the all-reduce completes before performing the
# the parameter update
assert bucket_index in overlap_info.bucket_index_to_future, \
f"All-reduce future for bucket {bucket_index} not saved " \
f"on rank {rank}"
allreduce_future = overlap_info.bucket_index_to_future[
bucket_index]
allreduce_future.wait()
bucket_gradients = overlap_info.bucket_index_to_bucket[
bucket_index].gradients()
for i, grad in enumerate(bucket_gradients):
gradients[offset + i] = grad
zero._local_step(gradients)
device = overlap_info.params_per_bucket[bucket_index][0].device
device_index = zero._device_to_device_index[device]
assert bucket_index in zero._buckets[device_index][rank_to_update]
overlap_info.broadcast_handles.append(
dist.broadcast(
zero._buckets[device_index][rank_to_update][bucket_index],
src=rank_to_update,
async_op=True))
# Ensure that all parameter updates are finished before the
# next forward pass
assert len(overlap_info.broadcast_handles) == num_buckets, \
f"Missing at least one broadcast handle on rank {rank}"
_ = list(map(lambda x: x.wait(), overlap_info.broadcast_handles))
overlap_info.broadcast_handles.clear()
# Reset per-iteration information
overlap_info.bucket_index_to_future.clear()
overlap_info.bucket_index_to_bucket.clear()
return fut
