def fp16_compress_hook(process_group: object, bucket: dist._GradBucket):
"""
This DDP communication hook implements a simple gradient compression
approach that converts ``GradBucket`` tensors whose type is assumed to be
``torch.float32`` to half-precision floating point format (``torch.float16``).
It allreduces those ``float16`` gradient tensors. Once compressed gradient
tensors are allreduced, its then callback called ``decompress`` converts the
aggregated result back to ``float32`` and takes the mean.
Example::
>>> ddp_model._register_comm_hook(process_group, fp16_compress_hook)
"""
group_to_use = process_group if process_group is not None else dist.group.WORLD
world_size = (process_group.size()
if process_group is not None else dist.get_world_size())
compressed_tensor = bucket.get_tensors()[0].to(torch.float16)
fut = dist.all_reduce(compressed_tensor, group=group_to_use,
async_op=True).get_future()
def decompress(fut):
return [fut.value()[0].to(torch.float32).div_(world_size)]
return fut.then(decompress)
def fp16_compress_hook(process_group: dist.ProcessGroup,
bucket: dist._GradBucket) -> torch.futures.Future:
"""
This DDP communication hook implements a simple gradient compression
approach that converts ``GradBucket`` tensors whose type is assumed to be
``torch.float32`` to half-precision floating point format (``torch.float16``).
It allreduces those ``float16`` gradient tensors. Once compressed gradient
tensors are allreduced, its then callback called ``decompress`` converts the
aggregated result back to ``float32`` and takes the mean.
Example::
>>> ddp_model.register_comm_hook(process_group, fp16_compress_hook)
"""
group_to_use = process_group if process_group is not None else dist.group.WORLD
world_size = group_to_use.size()
compressed_tensor = bucket.get_tensors()[0].to(torch.float16)
fut = dist.all_reduce(compressed_tensor, group=group_to_use,
async_op=True).get_future()
def decompress(fut):
decompressed_tensor = bucket.get_tensors()[0]
# Decompress in place to reduce the peak memory.
# See: https://github.com/pytorch/pytorch/issues/45968
decompressed_tensor.copy_(fut.value()[0].div_(world_size))
return [decompressed_tensor]
return fut.then(decompress)
def allreduce_hook(process_group: dist.ProcessGroup,
bucket: dist._GradBucket) -> torch.futures.Future:
"""
This DDP communication hook just calls ``allreduce`` using ``GradBucket``
tensors. Once gradient tensors are aggregated across all workers, its ``then``
callback takes the mean and returns the result. If user registers this hook,
DDP results is expected to be same as the case where no hook was registered.
Hence, this won't change behavior of DDP and user can use this as a reference
or modify this hook to log useful information or any other purposes while
unaffecting DDP behavior.
Example::
>>> ddp_model.register_comm_hook(process_group, allreduce_hook)
"""
group_to_use = process_group if process_group is not None else dist.group.WORLD
world_size = group_to_use.size()
tensor = bucket.get_tensors()[0]
fut = dist.all_reduce(tensor, group=group_to_use,
async_op=True).get_future()
def then_callback(fut):
return [fut.value()[0].div_(world_size)]
return fut.then(then_callback)
def fp16_compress_hook_gloo(state: object, bucket: dist._GradBucket):
group_to_use = dist.group.WORLD
world_size = dist.get_world_size()
compressed_tensor = bucket.get_tensors()[0].div_(world_size).half()
dist.all_reduce(compressed_tensor, async_op=False)
future = torch.futures.Future()
future.set_result([compressed_tensor])
return future
def allreduce_hook(process_group: dist.ProcessGroup,
bucket: dist._GradBucket) -> torch.futures.Future:
"""
This DDP communication hook just calls ``allreduce`` using ``GradBucket``
tensors. Once gradient tensors are aggregated across all workers, its ``then``
callback takes the mean and returns the result. If user registers this hook,
DDP results is expected to be same as the case where no hook was registered.
Hence, this won't change behavior of DDP and user can use this as a reference
or modify this hook to log useful information or any other purposes while
unaffecting DDP behavior.
Example::
>>> ddp_model.register_comm_hook(process_group, allreduce_hook)
"""
return allreduce_fut(process_group, bucket.get_tensors()[0])
def _allgather_then_aggregate_hook(
process_group: dist.ProcessGroup,
bucket: dist._GradBucket) -> torch.futures.Future:
"""
Similar to ``allreduce_hook``, this hook first gathers ``GradBucket`` tensors
and its ``then`` callback aggregates the gathered gradient tensors and takes
mean. Instead of ``allreduce`` this hook uses ``allgather``. Note that with
W workers, both the computation and communication time scale as O(W) for
allgather compared to O(logW) for allreduce. Therefore, this hook is expected
to be much slower than ``allreduce_hook`` although both essentially do the
same thing with the gradients.
.. warning ::
This is for test and experiments. User is suggested to use a faster
alternative called ``allreduce_hook`` that uses ``allreduce`` protocol
instead of ``allgather`` protocol.
Example::
>>> ddp_model.register_comm_hook(process_group, allreduce_hook)
"""
group_to_use = process_group if process_group is not None else dist.group.WORLD
rank = process_group.rank(
) if process_group is not None else dist.get_rank()
world_size = (process_group.size()
if process_group is not None else dist.get_world_size())
tensor = bucket.get_tensors()[0]
fut = dist.all_gather(
_get_allgather_out_list(tensor, world_size),
tensor,
group=group_to_use,
async_op=True,
).get_future()
def aggregate(fut):
all_ranks_tensor = fut.value()[0]
tensor = bucket.get_tensors()[0]
for r, gathered_tensor in enumerate(all_ranks_tensor):
if r != rank:
tensor += gathered_tensor
return [tensor.div_(world_size)]
return fut.then(aggregate)
def powerSGD_hook(
state: PowerSGDState,
bucket: dist._GradBucket,
) -> torch.futures.Future:
"""
This DDP communication hook implements a simplified PowerSGD gradient compression
algorithm described in 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 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;
2) Computes P, which is equal to MQ;
3) Allreduces P;
4) Orthogonizes P;
5) Computes Q, which is approximately equal to M^TP;
6) Allreduces Q;
7) Computes M, which is approximately equal to PQ^T.
8) Truncates the input tensor to the original length.
TODO(wayi@): 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.
Arguments:
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.
Note that since DDP comm hook only supports single process single device mode at this time,
only exactly one tensor is stored in this bucket.
matrix_approximation_rank (int): The low rank for matrix approximation.
Typically only 1 or 2 is used. See https://arxiv.org/pdf/1905.13727.pdf.
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, powerSGD_hook)
"""
process_group = state.process_group
group_to_use = process_group if process_group is not None else dist.group.WORLD
world_size = (
process_group.size() if process_group is not None else dist.get_world_size()
)
# The input tensor is a flattened 1D tensor.
input_tensor = bucket.get_tensors()[0]
device = input_tensor.device
total_length = input_tensor.shape[0]
# View the input tensor as a 2D square-shape tensor, and pad 0s if necessary.
square_side_length = math.ceil(math.sqrt(total_length))
padded_total_length = square_side_length ** 2
input_tensor.resize_(padded_total_length)
input_tensor[total_length:padded_total_length].fill_(0)
matrix = input_tensor.view(square_side_length, square_side_length)
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.
# Such 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"
).to(device)
else:
return torch.empty(
square_side_length, state.matrix_approximation_rank, device=device
)
p = create_low_rank_tensor(fill_random_values=False, rng=state.rng)
q = create_low_rank_tensor(fill_random_values=True, rng=state.rng)
_orthogonalize(q, 0)
torch.matmul(matrix, q, out=p)
allreduce_p_fut = dist.all_reduce(p, group=group_to_use, async_op=True).get_future()
def compute_q(fut):
p = fut.value()[0]
_orthogonalize(p, 0)
torch.matmul(matrix.t(), p, out=q)
return [
dist.all_reduce(q, group=group_to_use, async_op=True)
.get_future()
.value()[0]
]
def decompress(fut):
q = fut.value()[0].div_(world_size)
torch.matmul(p, q.t(), out=matrix)
ret = input_tensor.resize_(total_length)
return [ret]
return allreduce_p_fut.then(compute_q).then(decompress)
def quantization_pertensor_hook(
process_group: dist.ProcessGroup,
bucket: dist._GradBucket) -> torch.futures.Future:
"""
Applies the ``torch.quantize_per_tensor`` logic to DDP using ``allgather``
protocol. Workers first allgather the scale and zero point of their own
``GradBucket`` prior to the quantization. After all workers have that information,
the first ``then`` callback called ``quantize_and_allgather`` quantizes worker's
own gradient tensors, and uses ``allgather`` to communicate these accross all workers.
The final ``then`` callback called ``dequantize_and_aggregate``, dequantizes and
aggregates each quantized gradient tensors locally and returns the mean.
.. warning ::
This is experimental, and uses ``allgather`` protocol which is considerably slower than
``allreduce`` protocol. It works only with flattened grads.
Example::
>>> ddp_model.register_comm_hook(process_group, quantization_pertensor_hook)
"""
group_to_use = process_group if process_group is not None else dist.group.WORLD
rank = process_group.rank(
) if process_group is not None else dist.get_rank()
world_size = group_to_use.size()
tensor = bucket.get_tensors()[0]
myObserver = torch.quantization.MinMaxObserver().cuda(tensor.device)
myObserver(tensor)
s, z = myObserver.calculate_qparams()
s_and_z = torch.FloatTensor([s, z]).cuda(tensor.device)
all_ranks_s_and_z = _get_allgather_out_list(s_and_z, world_size)
# First, allgather scale and zeros.
fut = dist.all_gather(all_ranks_s_and_z,
s_and_z,
group=group_to_use,
async_op=True).get_future()
def quantize_and_allgather(fut):
# Store scale and zeros accross all workers.
all_ranks_s_and_z = fut.wait()[0]
# All workers quantize their own ``GradBucket`` tensors.
quantized_tensor = _quantize_per_tensor_cuda(
tensor, all_ranks_s_and_z[rank][0], all_ranks_s_and_z[rank][1])
# Allgather quantized tensors.
fut = dist.all_gather(
_get_allgather_out_list(quantized_tensor, world_size),
quantized_tensor,
group=group_to_use,
async_op=True,
).get_future()
return fut.wait()
def dequantize_and_aggregate(fut):
all_ranks_quantized_tensor = fut.wait()[0]
aggregated_dequantized_tensor = torch.zeros_like(
all_ranks_quantized_tensor[0],
device=tensor.device,
dtype=torch.float32)
# Using previously allgathered scales and zeros, dequantize gradient tensors
# locally and then aggregate them.
for r, quantized_tensor in enumerate(all_ranks_quantized_tensor):
aggregated_dequantized_tensor += _dequantize_per_tensor_cuda(
quantized_tensor, all_ranks_s_and_z[r][0],
all_ranks_s_and_z[r][1])
return [aggregated_dequantized_tensor / world_size]
return fut.then(quantize_and_allgather).then(dequantize_and_aggregate)
def quantization_perchannel_hook(process_group: dist.ProcessGroup,
bucket: dist._GradBucket,
bucket_size=512) -> torch.futures.Future:
"""
Applies the ``torch.quantize_per_channel`` logic to DDP using ``allgather``
protocol. Compared to pertensor, the main motivation of perchannel is
for considerably large tensors such as a tensor that contains 6 million
elements quantizing per a bucket size of 512 (or 128) elements may significantly
increase the resolution.
It first splits ``GradBucket`` tensors into multiple chunks (channels) of ``bucket_size``
elements. Then, workers allgather the scales and zero points of their own
``GradBucket`` prior to the quantization. After all workers have that information,
the first ``then`` callback called ``quantize_and_allgather`` quantizes worker's
own gradient tensors, and uses ``allgather`` to communicate these accross all workers.
The final ``then`` callback called ``dequantize_and_aggregate``, dequantizes, flattens, and
aggregates each quantized gradient tensors locally and returns the mean.
.. warning ::
This is experimental, and uses ``allgather`` protocol which is considerably slower than
``allreduce`` protocol. It works only with flattened grads.
Example::
>>> ddp_model.register_comm_hook(process_group, quantization_perchannel_hook)
"""
group_to_use = process_group if process_group is not None else dist.group.WORLD
rank = process_group.rank(
) if process_group is not None else dist.get_rank()
world_size = group_to_use.size()
tensor = bucket.get_tensors()[0]
tensor_in_channels = (nn.functional.pad(
input=tensor,
pad=(0, bucket_size - len(tensor) % bucket_size),
mode="constant",
value=0,
).view(-1, bucket_size).cuda(tensor.device))
myPerChannelObserver = torch.quantization.PerChannelMinMaxObserver().cuda(
tensor.device)
myPerChannelObserver(tensor_in_channels)
s_ch, z_ch = myPerChannelObserver.calculate_qparams()
s_and_z = torch.stack((s_ch, z_ch)).cuda(tensor.device)
all_ranks_s_and_z = _get_allgather_out_list(s_and_z, world_size)
# First, allgather scale and zeros.
fut = dist.all_gather(all_ranks_s_and_z,
s_and_z,
group=group_to_use,
async_op=True).get_future()
def quantize_and_allgather(fut):
# Store scale and zeros accross all workers.
all_ranks_s_and_z = fut.wait()[0]
# All workers quantize their corresponding ``GradBucket`` tensors.
quantized_tensor = _quantize_per_channel_cuda(
tensor_in_channels,
all_ranks_s_and_z[rank, 0, :],
all_ranks_s_and_z[rank, 1, :],
)
# Allgather quantized tensors.
fut = dist.all_gather(
_get_allgather_out_list(quantized_tensor, world_size),
quantized_tensor,
group=group_to_use,
async_op=True,
).get_future()
return fut.wait()
def dequantize_and_aggregate(fut):
all_ranks_quantized_tensor = fut.wait()[0]
aggregated_dequantized_tensor = torch.zeros_like(
all_ranks_quantized_tensor[0],
device=tensor.device,
dtype=torch.float32)
# Using previously allgathered scales and zeros, dequantize gradient tensors
# locally and then aggregate them.
for r, quantized_tensor in enumerate(all_ranks_quantized_tensor):
aggregated_dequantized_tensor += _dequantize_per_channel_cuda(
quantized_tensor, all_ranks_s_and_z[r][0],
all_ranks_s_and_z[r][1])
return [
torch.flatten(aggregated_dequantized_tensor).cuda(
tensor.device)[:tensor.size()[0]] / world_size
]
return fut.then(quantize_and_allgather).then(dequantize_and_aggregate)
def fp16_compress_hook_nccl(state: object, bucket: dist._GradBucket):
group_to_use = dist.group.WORLD
world_size = dist.get_world_size()
compressed_tensor = bucket.get_tensors()[0].div_(world_size).half()
future = dist.all_reduce(compressed_tensor, async_op=True).get_future()
return future
def powerSGD_hook(
process_group: dist.ProcessGroup,
bucket: dist._GradBucket,
matrix_approximation_rank: int = 1,
) -> torch.futures.Future:
"""
This DDP communication hook implements a simplified PowerSGD gradient compression
algorithm described in 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 square-shaped tensor M with 0 paddings;
2) Decomposes M into two low-rank tensors P and Q,
such that M = PQ^T, where Q is initialized from a standard normal distribution and orthogonalized;
2) Allreduces P;
3) Orthogonizes P;
4) Compute Q, which is approximately equal to M^TP;
5) Allreduces Q;
6) Computes M, which is approximately equal to PQ^T.
7) Truncates the input tensor to the original length.
TODO(wayi@): 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.
Arguments:
process_group (dist.ProcessGroup): Process group to communicate.
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 at this time,
only exactly one tensor is stored in this bucket.
matrix_approximation_rank (int): The low rank for matrix approximation.
Typically only 1 or 2 is used. See https://arxiv.org/pdf/1905.13727.pdf.
Returns:
Future handler of the communication, which updates the gradients in place.
Example::
PowerSGDState state(process_group, 1)
>>> ddp_model.register_comm_hook(state, powerSGD_hook)
"""
group_to_use = process_group if process_group is not None else dist.group.WORLD
world_size = (process_group.size()
if process_group is not None else dist.get_world_size())
# The input tensor is a flattened 1D tensor.
input_tensor = bucket.get_tensors()[0]
device = input_tensor.device
total_length = input_tensor.shape[0]
# View the input tensor as a 2D square-shape tensor, and pad 0s if necessary.
square_side_length = math.ceil(math.sqrt(total_length))
padded_total_length = square_side_length**2
input_tensor.resize_(padded_total_length)
input_tensor[total_length:padded_total_length].fill_(0)
matrix = input_tensor.view(square_side_length, square_side_length)
def create_low_rank_tensor(fill_random_values):
"Returns a low-rank 2D tensor of square_side_length * matrix_approximation_rank."
if fill_random_values:
with torch.random.fork_rng(devices=[device]):
# The seed makes sure that the initial random values are the same across all the DDP replicas.
# Such seed should differ at every step.
# Currently use the length of input tensor as the seed, which should be mostly different.
# TODO(wayi@): Should read the random seed from the state of this hook provided by the constructor.
torch.manual_seed(total_length)
return torch.randn(square_side_length,
matrix_approximation_rank,
device=device)
else:
return torch.empty(square_side_length,
matrix_approximation_rank,
device=device)
p = create_low_rank_tensor(fill_random_values=False)
q = create_low_rank_tensor(fill_random_values=True)
_orthogonalize(q, 0)
torch.matmul(matrix, q, out=p)
allreduce_p_fut = dist.all_reduce(p, group=group_to_use,
async_op=True).get_future()
def compute_q(fut):
p = fut.value()[0]
_orthogonalize(p, 0)
torch.matmul(matrix.t(), p, out=q)
return [
dist.all_reduce(q, group=group_to_use,
async_op=True).get_future().value()[0]
]
def decompress(fut):
q = fut.value()[0].div_(world_size)
torch.matmul(p, q.t(), out=matrix)
ret = input_tensor.resize_(total_length)
return [ret]
return allreduce_p_fut.then(compute_q).then(decompress)
def batched_powerSGD_hook(state: PowerSGDState,
bucket: dist._GradBucket) -> torch.futures.Future:
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 at this time,
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
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.get_tensors()[0]
# 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]
# View the input tensor as a 2D square-shape tensor, and pad 0s if necessary.
square_side_length = math.ceil(math.sqrt(total_length))
padded_total_length = square_side_length**2
input_tensor.resize_(padded_total_length)
input_tensor[total_length:padded_total_length].fill_(0)
# Incorporate the error from the previous state into the gradients.
bucket_index = bucket.get_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], 0)
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], 0)
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()[0].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
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:
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 two groups of per-parameter tensors: high-rank tensors and vector-like rank-1 tensors (for biases).
2. Divides all the tensors into two groups:
2.1 High-rank tensors that can have enough saving in bandwidth after the compression should be compressed before allreduce.
2.2 Rest of the tensors will be directly allreduced without compression (this group is referred to as rank-1 tensors below).
3. Handles rank-1 tensors by allreducing them without compression:
3.1. Allocate contiguous memory for those rank-1 tensors, and allreduces all the rank-1 tensors as a batch, without compression;
3.2. Copies the individual rank-1 tensors from the contiguous memory back to the input tensor.
4. Handles high-rank tensors by PowerSGD compression:
4.1. For each high-rank 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;
4.2. Computes each P in Ps, which is equal to MQ;
4.3. Allreduces Ps as a batch;
4.4. Orthogonalizes each P in Ps;
4.5. Computes each Q in Qs, which is approximately equal to M^TP;
4.6. Allreduces Qs as a batch;
4.7. Computes each M among all the high-rank 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 at this time,
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
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.get_tensors()[0]
# 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.get_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.get_per_parameter_tensors()
# Step I: Divide all the tensors into two groups,
# one will be compressed before allreduce and the other will be directly allreduced without compression.
rank1_tensors, high_rank_tensors, high_rank_tensors_to_compress = [], [], []
for tensor in tensors:
if tensor.ndimension() <= 1:
rank1_tensors.append(tensor)
else:
high_rank_tensors.append(tensor.view(tensor.shape[0], -1))
total_Ps_size = 0
total_Qs_size = 0
# Treat high-rank tensors that do not gain compression benefit as rank-1 tensors
while len(high_rank_tensors):
tensor = high_rank_tensors.pop()
n, m = tensor.shape
matrix_approximation_rank = min(n, m, state.matrix_approximation_rank)
if _should_compress(n, m, matrix_approximation_rank,
state.min_compression_rate):
high_rank_tensors_to_compress.append(tensor)
total_Ps_size += n * matrix_approximation_rank
total_Qs_size += m * matrix_approximation_rank
else:
rank1_tensors.append(tensor.view(-1))
# Step II: Handle rank-1 tensors (including the high-rank tensors that not worth compression).
# Allocate contiguous memory for rank-1 tensors to allreduce them without compression efficiently.
rank1_tensors_memory = (torch.cat([
tensor.view(-1) for tensor in rank1_tensors
]) if rank1_tensors else torch.tensor([], device=device, dtype=dtype))
# Step III: Handle high-rank tensors that should be compressed.
# Allocate contiguous memory for Ps and Qs to allreduce compressed high-rank tensors 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 high_rank_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)
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)
# Compute Ps.
for tensor, q, p in zip(high_rank_tensors_to_compress, qs, ps):
torch.matmul(tensor, q, out=p)
# This allreduce is only applied to rank-1 tensors,
# so it should have been kicked off before the above computation on the high-rank tensors to hide more communication costs.
# However, this somehow requires a separate future chain at this time.
allreduce_contiguous_rank1_tensors_fut = dist.all_reduce(
rank1_tensors_memory, group=group_to_use, async_op=True).get_future()
def unpack_rank1_tensors_and_allreduce_ps(fut):
rank1_tensors_memory = fut.value()[0].div_(world_size)
idx = 0
for tensor in rank1_tensors:
tensor.copy_(rank1_tensors_memory[idx:idx + tensor.shape[0]])
idx += tensor.shape[0]
# 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()[0]
for p in ps:
_orthogonalize(p)
# Compute Qs.
for tensor, p, q in zip(high_rank_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()[0].div_(world_size)
for p, q, tensor in zip(ps, qs, high_rank_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_rank1_tensors_fut.then(
unpack_rank1_tensors_and_allreduce_ps).then(compute_qs).then(
decompress))
