PowerSGD State#

classtorch.distributed.algorithms.ddp_comm_hooks.powerSGD_hook.PowerSGDState(process_group,matrix_approximation_rank=1,start_powerSGD_iter=1000,min_compression_rate=2,use_error_feedback=True,warm_start=True,orthogonalization_epsilon=0,random_seed=0,compression_stats_logging_frequency=10000,batch_tensors_with_same_shape=False)[source]#

Store both the algorithm’s hyperparameters and internal state for all gradients during training.

Particularly,matrix_approximation_rank andstart_powerSGD_iter are the main hyperparameters that should be tuned by the user.For performance, we suggest to keep binary hyperparametersuse_error_feedback andwarm_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.

   1.1. Ifmatrix_approximation_rank is too low, the full model quality will need more training steps to reach or will never reach and yield loss in accuracy.

   1.2. The increase ofmatrix_approximation_rank can substantially increase the computation costs of the compression, and the accuracy may not be further improved beyond a certainmatrix_approximation_rank threshold.

To tunematrix_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. Typically only a small value 1-4 is used. For some NLP tasks (as shown in Appendix D of the original paper), this value has been increased to 32.

2. start_powerSGD_iter defers PowerSGD compression until stepstart_powerSGD_iter, and vanilla allreduce runs prior to stepstart_powerSGD_iter. This hybrid scheme ofvanilla allreduce + PowerSGD can effectively improve the accuracy, even a relatively smallmatrix_approximation_rank is used. This is because that, the beginning of training phase is usually very sensitive to inaccurate gradients, and compressing gradients too early may make the training quickly take a suboptimal trajectory, which can result in an irrecoverable impact on the accuracy.

To tunestart_powerSGD_iter, we suggest to start with 10% of total training steps, and increase it until a satisfactory accuracy is reached. If there is a warm-up stage in the training,start_powerSGD_iter typically should be no less than the number of warm-up steps.

3. min_compression_rate is the minimum compression rate required when a layer is compressed. Due to the computation overheads incurred by the compression, a tensor is worth compressing only if there can be sufficient saving in bandwidth, where(num_rows+num_cols)*matrix_approximation_rank*min_compression_rate<num_rows*num_cols. If the specified compression rate threshold cannot be satisfied, the tensor will be directly allreduced without compression.

Compression statistics are logged everycompression_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. If this can already be prevented (e.g., by batch normalization), an epsilon of 0 is recommended for accuracy.

5. batch_tensors_with_same_shape controls whether to compress and decompress tensors with same shape in a batched operation to achieve higher parallelism. Note that you should also increase the bucket size (i.e.,bucket_cap_mb arg in DDP constructor) to make more same-shaped tensors appear in the same bucket, however this may reduce the overlap between computation and communication, and increase the memory footprint due to stacking the tensors of the same shape. Set toTrue if the compression / decompression computation is a bottleneck.

Warning

If error feedback or warm-up is enabled, the minimum value ofstart_powerSGD_iter allowed in DDP is 2.This is because there is another internal optimization that rebuilds buckets at iteration 1 in DDP,and this can conflict with any tensor memorized before the rebuild process.

PowerSGD Hooks#

torch.distributed.algorithms.ddp_comm_hooks.powerSGD_hook.powerSGD_hook(state,bucket)[source]#

Implement PowerSGD algorithm.

This DDP communication hook implements PowerSGD gradient compressionalgorithm described in thepaper.Once gradient tensors are aggregated across all workers, this hook appliescompression 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 firststate.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.

Parameters:

• state (PowerSGDState) – State information to configure the compression rate and support error feedback, warm start, etc.To tune the compression configs, mainly need to tunematrix_approximation_rank,start_powerSGD_iterandmin_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.

Return type:

Future[Tensor]

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)

torch.distributed.algorithms.ddp_comm_hooks.powerSGD_hook.batched_powerSGD_hook(state,bucket)[source]#

Implement simplified PowerSGD algorithm.

This DDP communication hook implements a simplified PowerSGD gradient compressionalgorithm described in thepaper.This variant does not compress the gradients layer by layer,but instead compresses the flattened input tensor that batches all the gradients.Therefore, it isfaster thanpowerSGD_hook(),but usually results in amuch lower accuracy, unlessmatrix_approximation_rank is 1.

Warning

Increasingmatrix_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 considerpowerSGD_hook() first,and only consider this variant when a satisfactory accuracy can be achieved whenmatrix_approximation_rank is 1.

Once gradient tensors are aggregated across all workers, this hook appliescompression 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 firststate.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.

Parameters:

• state (PowerSGDState) – State information to configure the compression rate and support error feedback, warm start, etc.To tune the compression configs, mainly need to tunematrix_approximation_rank andstart_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.

Return type:

Future[Tensor]

Example::

>>>state=PowerSGDState(process_group=process_group,matrix_approximation_rank=1)>>>ddp_model.register_comm_hook(state,batched_powerSGD_hook)