# Source code for tensorplay.distributed.algorithms.ddp_comm_hooks.powerSGD_hook Source: https://www.tensorplay.cn/docs/_modules/tensorplay/distributed/algorithms/ddp_comm_hooks/powerSGD_hook.html ``` # # linalg.qr(out=) (copy the Q factor back in place), and no resize_ (the # batched variant pads via a temporary buffer and returns it). import logging import math from collections import defaultdict import tensorplay as tp from tensorplay.linalg import qr as _qr, vector_norm as _vector_norm import tensorplay.distributed as dist from tensorplay.distributed import distributed_core from . import default_hooks as default __all__ = ["PowerSGDState", "powerSGD_hook", "batched_powerSGD_hook"] logger = logging.getLogger(__name__) def _not_none(x): if x is None: raise ValueError("Expected non-None value") return x def _bmm_out(a, b, out): batch = a.shape[0] for i in range(batch): out[i].copy_(a[i] @ b[i]) return out def _orthogonalize(matrices, epsilon=0): """ Decide between Gram-Schmidt or QR factorization to orthogonalize a batch of matrices. QR factorization doesn't work with half-precision, but it is usually faster with a rank > 2. """ if not (len(matrices.shape) == 3 and matrices.shape[2] <= matrices.shape[1]): raise AssertionError num_matrices = matrices.shape[0] rank = matrices.shape[2] dtype = matrices.dtype if rank <= 2 or dtype in [tp.float16, tp.bfloat16]: _orthogonalize_gram_schmidt(matrices, epsilon=epsilon) else: q, _ = _qr(matrices) matrices.copy_(q) def _orthogonalize_gram_schmidt(matrices, epsilon=0): """ Apply Gram-Schmidt procedure to orthogonalize a batch of matrices. If epsilon is 0, this is equivalent to `linalg.qr(matrices, out=(matrices, _))`, """ num_cols = matrices.shape[2] for i in range(num_cols): # Normalize the i'th column. col = matrices[:, :, i : i + 1] # If no epsilon is added here, division by zero may be caused by vanishing gradients. # This epsilon is not needed if the input batch of matrices covers the gradients of at least one entire layer # in the neural network. if epsilon == 0: # Note that col ** 2 can underflow/overflow if we use FP16. # May need to consider multiplying a scaling factor and dividing it later, or using bfloat16 instead. try: col.div_(_vector_norm(col, dim=1, keepdim=True)) except ZeroDivisionError: logger.error( "The matrices to be orthogonalized has at least a column of all 0s. Please set a small value such as 1e-8 " "as `orthogonalization_epsilon` in PowerSGD state." ) # Recover the values from NaNs to 0s. col.fill_(0.0) else: col.div_(_vector_norm(col, dim=1, keepdim=True) + epsilon) # Project it on the rest and remove it. if i + 1 < num_cols: rest = matrices[:, :, i + 1 :] rest.sub_(tp.sum(col * rest, dim=1, keepdim=True) * col) def _should_compress( num_rows, num_cols, matrix_approximation_rank, min_compression_rate ): """ Recommend if tensor given is worth compressing. Returns a recommendation as to whether the 2D tensor described by the arguments is worth compressing, including statistics describing the expected savings from compression. We consider a tensor worth compressing when ``min_compression_rate`` < uncompressed size / compressed size, where uncompressed size = ``num_rows`` * ``num_cols``, and compressed size = (``num_rows`` + ``num_cols``) * ``matrix_approximation_rank``. The result of this function is a tuple of the form (compression_recommendation, uncompressed_el_count, compressed_el_count), where: compression_recommendation is true if the tensor is worth compressing, and false otherwise (see above); uncompressed_el_count is the uncompressed element count, i.e. ``num_rows`` * ``num_cols``; and, compress_el_count is the element count after compression, i.e. (``num_rows`` + ``num_cols``) * ``matrix_approximation_rank``. """ uncompressed_size = num_rows * num_cols compressed_size = (num_rows + num_cols) * matrix_approximation_rank return ( compressed_size * min_compression_rate < uncompressed_size, uncompressed_size, compressed_size, ) def _report_compression_stats(bucket, state): """Report compression stats at frequency of ``compression_stats_logging_frequency`` specified in PowerSGD state.""" if bucket.is_last() and state.iter >= state.next_stats_report: stats = state.compression_stats() logger.info( "Compression stats: iter %s, total before compression %s, total after compression %s, " "rate %s", state.iter, stats[1], stats[2], stats[0], ) state.next_stats_report = state.iter + state.compression_stats_logging_frequency [docs] class PowerSGDState: r""" Store both the algorithm's hyperparameters and internal state for all gradients during training. Particularly, ``matrix_approximation_rank`` and ``start_powerSGD_iter`` are the main hyperparameters that should be tuned by the user. For performance, we suggest to keep binary hyperparameters ``use_error_feedback`` and ``warm_start`` on. 1. ``matrix_approximation_rank`` controls the size of compressed low-rank tensors, which determines the compression rate. The lower the rank, the stronger the compression. To tune ``matrix_approximation_rank``, we suggest to start from 1 and increase by factors of 2 (like an exponential grid search, 1, 2, 4, ...), until a satisfactory accuracy is reached. 2. ``start_powerSGD_iter`` defers PowerSGD compression until step ``start_powerSGD_iter``, and vanilla allreduce runs prior to step ``start_powerSGD_iter``. 3. ``min_compression_rate`` is the minimum compression rate required when a layer is compressed. Compression statistics are logged every ``compression_stats_logging_frequency`` iterations once PowerSGD compression starts. 4. ``orthogonalization_epsilon`` can be a very small value (e.g., 1e-8) added to every normalized matrix column in orthogonalization step, to prevent div-by-zero error if any column has all 0s. 5. ``batch_tensors_with_same_shape`` controls whether to compress and decompress tensors with same shape in a batched operation to achieve higher parallelism. .. warning :: If error feedback or warm-up is enabled, the minimum value of ``start_powerSGD_iter`` allowed in DDP is 2. """ __slots__ = [ "process_group", # The fields below are the hyperparameters that often need to be tuned by the user. "matrix_approximation_rank", "start_powerSGD_iter", # The fields below are the hyperparameters that seldom need be tuned by the user. "min_compression_rate", "orthogonalization_epsilon", # The fields below are the binary hyperparameters recommended to be turned on for performance and accuracy. "use_error_feedback", "warm_start", "batch_tensors_with_same_shape", # The fields below are internal state. "rng", "error_dict", "p_memory_dict", "q_memory_dict", "iter", # The fields below are for recording compression stats. "total_numel_before_compression", "total_numel_after_compression", "compression_stats_logging_frequency", "next_stats_report", ] def __init__( self, process_group, matrix_approximation_rank=1, start_powerSGD_iter=1_000, min_compression_rate=2, use_error_feedback=True, warm_start=True, orthogonalization_epsilon=0, random_seed=0, compression_stats_logging_frequency=10_000, batch_tensors_with_same_shape: bool = False, ): logger.info( "PowerSGD config: matrix_approximation_rank = %s; start_powerSGD_iter = %s; " "min_compression_rate = %s; orthogonalization_epsilon = %s; use_error_feedback = %s; warm_start = %s; " "random_seed = %s; compression_stats_logging_frequency = %s; batch_tensors_with_same_shape = %s", matrix_approximation_rank, start_powerSGD_iter, min_compression_rate, orthogonalization_epsilon, use_error_feedback, warm_start, random_seed, compression_stats_logging_frequency, batch_tensors_with_same_shape, ) self.process_group = process_group self.matrix_approximation_rank = matrix_approximation_rank if (use_error_feedback or warm_start) and start_powerSGD_iter <= 1: raise ValueError( "Expect `start_powerSGD_iter` > 1 if `use_error_feedback` or `warm_start` is enabled, " "because PowerSGD can only be applied after the first two iterations in DDP." ) self.start_powerSGD_iter = start_powerSGD_iter self.min_compression_rate = min_compression_rate self.use_error_feedback = use_error_feedback self.warm_start = warm_start self.orthogonalization_epsilon = orthogonalization_epsilon # The purpose of this RNG is to generate different random seeds for initializing Q across iterations, # but in the same order for all the DDP replicas. import numpy as np self.rng = np.random.RandomState(random_seed) self.error_dict: dict[int, tp.Tensor] = {} self.p_memory_dict: dict[int, tp.Tensor] = {} self.q_memory_dict: dict[int, tp.Tensor] = {} self.iter = 0 self.total_numel_before_compression = 0 self.total_numel_after_compression = 0 self.compression_stats_logging_frequency = max( 1, compression_stats_logging_frequency ) self.next_stats_report = 0 self.batch_tensors_with_same_shape = batch_tensors_with_same_shape def __getstate__(self): r"""Return picklable state; process group is excluded.""" logger.warning( "NOTE: Process group is not serializable and excluded from a saved state." ) return { slot: getattr(self, slot) for slot in self.__slots__ if slot != "process_group" } def __setstate__(self, state): r"""Restore state; process group is set to default.""" self.process_group = distributed_core._get_default_group() logger.warning( "NOTE: Process group will be set to a default group (i.e. the world size).\ If a different group is desired, please set `self.process_group` after PowerSGD state is loaded." ) for slot, value in state.items(): setattr(self, slot, value) [docs] def maybe_increase_iter(self, bucket): """Track iterations and trigger log message at start of local SGD.""" # Since bucket 0 is the last bucket to allreduce in an iteration. # Only increase `iter` when bucket 0 is processed. if bucket.is_last(): self.iter += 1 if self.iter == self.start_powerSGD_iter: logger.info("Start to apply PowerSGD after %s iterations.", self.iter) [docs] def compression_stats(self): r""" Return latest compression statistics as tuple. Returns tuple of form (compress_rate, numel_before_compression, numel_after_compression). """ compress_rate = ( self.total_numel_before_compression / self.total_numel_after_compression if self.total_numel_after_compression > 0 else 0 ) return ( compress_rate, self.total_numel_before_compression, self.total_numel_after_compression, ) [docs] def powerSGD_hook( state: PowerSGDState, bucket: dist.GradBucket ): r""" Implement PowerSGD algorithm. This DDP communication hook implements PowerSGD gradient compression algorithm described in the `paper `_. Note that this communication hook enforces vanilla allreduce for the first ``state.start_powerSGD_iter`` iterations. Args: 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. Returns: Future handler of the communication, which updates the gradients in place. Example:: >>> # xdoctest: +SKIP >>> 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) """ process_group = state.process_group group_to_use = ( process_group if process_group is not None else _not_none(dist.GroupMember.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: logger.info( "A zero tensor of length %s that represents local error is created.", total_length, ) state.error_dict[bucket_index] = tp.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 = input_tensor.detach().clone() # 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.reshape(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 = ( tp.cat([tensor.reshape(-1) for tensor in uncompressed_tensors]) if uncompressed_tensors else tp.tensor([], device=device, dtype=dtype) ) # Step III: Handle the tensors that should be compressed. need_randomize_qs = False if not state.warm_start or bucket_index not in state.p_memory_dict: need_randomize_qs = True if state.warm_start: logger.info( "Allocating contiguous memory of length %s for Ps, and of length %s for Qs, respectively.", total_Ps_size, total_Qs_size, ) state.p_memory_dict[bucket_index] = tp.empty( total_Ps_size, device=device, dtype=dtype ) state.q_memory_dict[bucket_index] = tp.empty( total_Qs_size, device=device, dtype=dtype ) # Batch tensors to compress by shape. shape_to_tensors = defaultdict(list) for tensor in tensors_to_compress: shape_to_tensors[tensor.shape].append(tensor) def maybe_batched_tensors_to_compress(): for tensors in shape_to_tensors.values(): if state.batch_tensors_with_same_shape: batch_size = len(tensors) if batch_size == 1: # Use the original tensor to avoid copy. yield tensors[0].unsqueeze(0) else: yield tp.stack(tensors) else: for tensor in tensors: yield tensor.unsqueeze(0) # Create Ps and Qs that point to the allocated memory. tensors_to_compress = [] ps = [] qs = [] p_idx = 0 q_idx = 0 for tensor in maybe_batched_tensors_to_compress(): batch_size, n, m = tensor.shape matrix_approximation_rank = min(n, m, state.matrix_approximation_rank) tensors_to_compress.append(tensor) ps.append( state.p_memory_dict[bucket_index][ p_idx : p_idx + batch_size * n * matrix_approximation_rank ].view((batch_size, n, matrix_approximation_rank)) ) qs.append( state.q_memory_dict[bucket_index][ q_idx : q_idx + batch_size * m * matrix_approximation_rank ].view((batch_size, m, matrix_approximation_rank)) ) p_idx += batch_size * n * matrix_approximation_rank q_idx += batch_size * m * matrix_approximation_rank # If warm-start is enabled, reuse Qs from the previous iteration if possible and skip filling random values. if not need_randomize_qs: for q in qs: _orthogonalize(q, state.orthogonalization_epsilon) else: with tp.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. tp.manual_seed(state.rng.randint(1_000_000_000)) for q in qs: q.copy_( tp.randn( list(q.shape), device="cpu", dtype=dtype, ).to(q.device) ) _orthogonalize(q, state.orthogonalization_epsilon) # Compute Ps. for tensor, q, p in zip(tensors_to_compress, qs, ps): _bmm_out(tensor, q, 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. 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_memory = fut.value()[0].div_(world_size) idx = 0 for tensor in uncompressed_tensors: tensor.copy_( uncompressed_memory[idx : idx + tensor.numel()].view(tensor.shape) ) 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): _bmm_out(tensor.transpose(1, 2), p, q) # 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): _bmm_out(p, q.transpose(1, 2), tensor) # Copy batched tensors back to original buffer. if state.batch_tensors_with_same_shape: for tensor in tensors_to_compress: if tensor.shape[0] == 1: # Skip tensor with batch_size == 1 since itself is the original tensor. continue original_tensors = shape_to_tensors[tensor.shape[1:]] for i, original_tensor in enumerate(original_tensors): original_tensor.copy_(tensor[i]) if tp.cuda.is_available(): tp.cuda.synchronize(device.index if device.index >= 0 else None) if state.use_error_feedback: # Memorize the local errors. if input_tensor_cp is None: raise AssertionError 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) ) [docs] def batched_powerSGD_hook( state: PowerSGDState, bucket: dist.GradBucket ): r""" Implement simplified PowerSGD algorithm. This DDP communication hook implements a simplified PowerSGD gradient compression algorithm described in the `paper `_. 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. Args: 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. Returns: Future handler of the communication, which updates the gradients in place. Example:: >>> # xdoctest: +SKIP >>> state = PowerSGDState(process_group=process_group, matrix_approximation_rank=1) >>> ddp_model.register_comm_hook(state, batched_powerSGD_hook) """ process_group = state.process_group group_to_use = ( process_group if process_group is not None else _not_none(dist.GroupMember.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. # tp has no resize_; pad through a temporary buffer instead. 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 padded_input = tp.zeros(padded_total_length, device=device, dtype=input_tensor.dtype) padded_input[:total_length].copy_(input_tensor) input_tensor = padded_input _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: logger.info( "A zero tensor of length %s that represents local error is created.", padded_total_length, ) state.error_dict[bucket_index] = tp.zeros( padded_total_length, device=device, dtype=input_tensor.dtype ) input_tensor_cp = input_tensor.detach().clone() matrix = input_tensor.view((square_side_length, square_side_length)) # Reuse P and Q from the previous iteration if possible. if not state.warm_start or bucket_index not in state.p_memory_dict: if state.warm_start: logger.info( "Initializing low-rank tensors P and Q, each of which has a shape of %s x %s.", square_side_length, state.matrix_approximation_rank, ) def create_low_rank_tensor(fill_random_values, rng): """Return a low-rank 2D tensor of square_side_length * matrix_approximation_rank.""" if fill_random_values: with tp.fork_rng(devices=[]): tp.manual_seed(rng.randint(1_000_000_000)) return tp.randn( [square_side_length, state.matrix_approximation_rank], device="cpu", dtype=input_tensor.dtype, ).to(device) else: return tp.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]) state.p_memory_dict[bucket_index].copy_( matrix @ state.q_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]) state.q_memory_dict[bucket_index].copy_( matrix.t() @ state.p_memory_dict[bucket_index] ) 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) approx_matrix = ( state.p_memory_dict[bucket_index] @ state.q_memory_dict[bucket_index].t() ) if state.use_error_feedback: # Memorize the local errors. if input_tensor_cp is None: raise AssertionError state.error_dict[bucket_index] = input_tensor_cp - input_tensor if tp.cuda.is_available(): tp.cuda.synchronize(device.index if device.index >= 0 else None) if not state.warm_start: state.p_memory_dict.clear() state.q_memory_dict.clear() # Truncate back to the original length; the reducer copies this # result buffer back into the parameter gradients. ret = approx_matrix.reshape(-1)[:total_length] state.maybe_increase_iter(bucket) return ret return allreduce_p_fut.then(compute_q).then(decompress) ```