Source-linked AI summary

PowerSGD: Practical Low-Rank Gradient Compression for Distributed Optimization

Thijs Vogels, Sai Praneeth Karimireddy, Martin Jaggi

arXiv:1905.13727v3cs.LGcs.DCmath.OCstat.ML

TL;DR

Large gradients make communication a scalability bottleneck, and existing compression methods may lose speedups or test accuracy. POWERSGD uses warm-started power iteration for linear low-rank compression with all-reduce aggregation, achieving SGD-level test performance and wall-clock speedups in the reported settings.

  • Problem

    Large gradients make communication limit synchronous data-parallel training, while existing compressors may lose speedups on optimized fast networks or reduce test accuracy.

  • Method

    POWERSGD computes warm-started low-rank gradient approximations with power iteration and uses linear compression to aggregate them through all-reduce.

  • Results

    POWERSGD is the only evaluated method with consistent wall-clock speedups over regular SGD with optimized communication, while matching SGD test performance.

  • Takeaways & Limitations

    POWERSGD provides a practical compression method whose speedup increases for larger models and supports training with more workers than full-precision SGD.

  • Takeaways & Limitations

    The paper identifies the generalization gap from increasing batch sizes as an orthogonal issue requiring further work.

Abstract

from arXiv · show

We study gradient compression methods to alleviate the communication bottleneck in data-parallel distributed optimization. Despite the significant attention received, current compression schemes either do not scale well or fail to achieve the target test accuracy. We propose a new low-rank gradient compressor based on power iteration that can i) compress gradients rapidly, ii) efficiently aggregate the compressed gradients using all-reduce, and iii) achieve test performance on par with SGD. The proposed algorithm is the only method evaluated that achieves consistent wall-clock speedups when benchmarked against regular SGD with an optimized communication backend. We demonstrate reduced training times for convolutional networks as well as LSTMs on common datasets. Our code is available at https://github.com/epfml/powersgd.

1 Introduction

Distributed training is limited by the cost of sharing large gradients, while existing compressors can lose speedups on fast networks or reduce test accuracy. POWERSGD combines linear low-rank compression, error feedback, and power iteration to address these constraints.

  • Gradient communication limits the scalability of synchronous data-parallel deep learning because model gradients can be large.
  • Existing sign, quantization, and low-rank compressors can lose speedups with fast networks and optimized communication, and some reduce test accuracy.
  • Nonlinear compressed messages cannot be hierarchically added, forcing gather operations instead of efficient all-reduce aggregation.
  • POWERSGD uses generalized power iteration with a warm start to compute a lightweight low-rank gradient approximation without Singular Value Decomposition.
  • More than 120× compression reduces communication time by 54% for RESNET18 on CIFAR10 and 90% for an LSTM on WIKITEXT-2.End-to-end training time falls by 24% for RESNET18 and 55% for the LSTM.

2 Related work

Related work explores quantization, signs, sparsity, low-rank approximation, linear aggregation, and error feedback. POWERSGD is positioned against methods that are computationally expensive, difficult to aggregate efficiently, or dependent on error feedback for biased compression.

  • Prior gradient compressors quantize coordinates, use signs or largest coordinates, or apply low-rank approximations.
  • Spectral Atomo samples singular vectors without bias, but requires a full Singular Value Decomposition every iteration.The paper characterizes this computation as impractical.
  • Commutativity between compression and addition enables efficient ring all-reduce, a property most compressors lack.POWERSGD has a similar property called linearity.
  • Error feedback stores compression errors and adds them to later gradients; in non-convex settings, it is reported as crucial for convergence and generalization with biased compressors.
  • Low-rank methods are motivated by evidence that over-parameterized networks and their gradients can exhibit low-ranked structure.

3 Method

POWERSGD approximates each parameter gradient matrix with a rank-r factorization computed by subspace iteration. Its linear operations support all-reduce aggregation, while error feedback and post-compression momentum preserve compatibility with momentum SGD.

  • Each worker computes stochastic gradients, aggregates them across workers, and updates shared model parameters in data-parallel optimization.
  • POWERSGD compression: For each gradient matrix M, rank-r POWERSGD computes factors P and Q through matrix multiplication, orthogonalization, and all-reduce mean operations.
  • POWERSGD compression: Decompression reconstructs the approximate gradient matrix as P̂Qᵀ.
  • POWERSGD compression: A single subspace-iteration step uses right multiplication, left multiplication, and orthogonalization, with previous approximations reused through warm-starting.Gram-Schmidt is the most expensive part because the factor matrices have few columns.
  • Efficient aggregation between workers: Linearity lets compressed gradients use all-reduce rather than gather; all-reduce computation and communication scale as O(log W), versus O(W) for all-gather.
  • Error-feedback SGD: Error-feedback momentum compensates for POWERSGD’s biased compression and reuses learning rates and hyperparameters tuned for SGD with momentum.

4 Analysis of POWERSGD

The analysis examines error feedback, warm-starting, and rank selection. It finds that error feedback improves biased low-rank compression, warm-starting closes an approximation-quality gap, and sufficient rank reaches SGD-level test quality.

  • Effect of error feedback: Error-feedback POWERSGD outperforms an unbiased linear rank-r compressor on image-classification test accuracy.
  • Effect of warm-start: POWERSGD uses a cheaper low-fidelity power-iteration approximation rather than the best rank-r approximation.
  • Effect of warm-start: Warm-starting improves POWERSGD test accuracy and can match the performance of the much more expensive best rank-2 approximation.
  • Effect of varying the rank: Increasing rank improves POWERSGD test quality in both image classification and language modeling, reaching quality as good as or slightly better than regular SGD.

5 Results

The experiments evaluate POWERSGD across compression regimes, worker counts, communication backends, and model types. POWERSGD combines strong accuracy with favorable scaling and substantial training-time reductions, especially for larger models.

  • Results overview: Across the experiments, POWERSGD is described as fast, accurate, scalable with increasing workers, and effective at reducing training time for larger models.The paper’s results section summarizes these as the three practical properties demonstrated by the compression scheme.
  • Experimental setting: POWERSGD is evaluated on CIFAR10 with RESNET18 and on WIKITEXT-2 language modeling, using optimized implementations and repeated measurements.The default CIFAR10 setting uses 16 workers with the NCCL backend; the LSTM evaluation uses 16 GPUs across 8 machines.
  • Comparison with other compressors: At 128× compression, POWERSGD is the only evaluated scheme that achieves the target test accuracy.At around 32× compression, most schemes also approach full-precision SGD accuracy, except Random Block.
  • Comparison with other compressors: POWERSGD and Random Block are the only schemes faster than full-precision SGD in both medium- and high-compression settings.POWERSGD relies on matrix multiplication and all-reduce, whereas Random K incurs substantial random-memory-access overhead.
  • Scalability of POWERSGD: Using all-reduce, POWERSGD’s gradient communication scales gracefully with worker count, while Signum’s all-gather becomes more expensive at 16 workers.Signum has comparable time to POWERSGD for 4 workers but becomes more expensive for 16 workers because communication and decompression scale differently.
  • Scalability of POWERSGD: With the slower GLOO backend, POWERSGD is the only evaluated method that retains excellent scaling because of its high compression rate.All three methods scale reasonably well with optimized NCCL, although Signum shows sub-linear scaling in the log-log plot.
  • Other tasks and methods: On the larger LSTM language-modeling task, rank-4 POWERSGD reduces communication by 90% and overall running time by 55% while matching full-precision SGD’s test score.Signum becomes slower than full-precision SGD and obtains a worse test score in this setting.

6 Conclusion

POWERSGD is presented as a practical compression method that combines efficient low-rank approximation and all-reduce aggregation while matching full-precision SGD test performance. Its benefits extend to larger models, though large-batch generalization remains an orthogonal limitation.

  • The method is presented as the first practical gradient compression approach because existing schemes either run slower than optimized SGD or miss its test performance.
  • POWERSGD combines cheap linear compression, all-reduce aggregation, and test performance matching full-precision SGD.These properties are identified as central to its practicality.
  • Its speedup increases for larger models, including models commonly used in natural language processing.
  • Larger batch sizes can still suffer a generalization gap, which the paper identifies as an orthogonal issue for future large-scale training work.

A.2 Subspace iteration

PowerSGD uses subspace iteration to approximate gradient matrices with low rank, reusing prior approximations to improve a single-step method. Under a fixed matrix and an eigenvalue gap, the iteration recovers the best rank-r approximation, while linearity preserves equivalence across workers and batch sizes.

  • Method: PowerSGD uses one subspace-iteration step to compute a fast low-rank approximation of each stochastic gradient.The method performs matrix multiplications and orthogonalization, then reuses the previous approximation as the next starting point.
  • Method: Warm-starting reuses prior low-rank factors because gradients change slowly as parameters receive small updates.This reuse can improve approximation quality even though the target matrices vary over time.
  • Convergence: Subspace iteration recovers the best rank-r approximation when the leading eigenvalues are separated from the remainder.The theorem assumes a fixed matrix and an eigen-decomposition ordered by decreasing eigenvalues.
  • Convergence: For a fixed matrix, an eigenvalue gap σr > σr+1 makes subspace iteration converge to the leading rank-r subspace.The initialization must have non-zero support on the first r eigenvectors.
  • Distributed aggregation: PowerSGD’s linearity makes compressed aggregation equivalent across multiple workers and a single worker using a larger batch size.The error computation and compression can be applied to the aggregated gradient rather than separately to each worker’s gradient.

B Cluster specifications

The experiments use an eight-node, GPU-based cluster with 10Gbit/s Ethernet and compare communication backends. Convergence-curve timings include testing, checkpointing, and other bookkeeping, with shaded ranges spanning three replications.

  • Hardware: The cluster contains 8 nodes, each with 2 Nvidia GeForce GTX Titan X GPUs and 48 CPU cores.GPU communication traverses PCIe and the SMP interconnect between NUMA nodes.
  • Network and software: The system uses 10Gbit/s SFI/SFP+ Ethernet with a fat-tree network topology.Experiments run with PyTorch 1.1 on Anaconda Python 3.7.
  • Communication backends: NCCL is the default communication backend, while GLOO is also measured; NCCL did not support gather in PyTorch at the time.This backend limitation is relevant to comparisons involving aggregation operations.
  • Convergence and timing: PowerSGD with sufficiently high rank can match full-precision SGD test quality with lower wall-clock duration in two tasks.Figure 4 timings include testing overhead, checkpointing, and other bookkeeping.
  • Comparative evaluation: Signum is the most competitive method among those compared in Figure 5.Its learning rate was tuned, and the plotted timings include end-of-epoch testing, checkpointing, and bookkeeping.

D Language Modeling with Transformers

In a Transformer language-modeling case study, PowerSGD requires a higher rank than in earlier experiments to approach uncompressed SGD’s validation loss. At rank 32, it achieves 14× compression and reduces training time from 20h to 12h.

  • Experimental setting: PowerSGD is evaluated on Transformer language modeling with 32 workers, new communication infrastructure, hardware, and model architecture.The implementation uses Facebook AI Research’s fairseq library and Google’s public cloud.
  • Efficiency: A speedup of around 1.5× in time-to-accuracy at loss = 5 is reported.The authors note that learning-rate-related hyperparameters could be retuned to improve the result.
  • Accuracy: Rank 32 achieves validation loss similar to uncompressed SGD in the same number of update steps.The experiment requires a higher rank than previous tasks, where rank 4 was used.
  • Efficiency: 14× compression and 12h training time are achieved versus 20h for the uncompressed-SGD baseline.These figures correspond to rank 32 in the Transformer language-modeling experiment.

E The need for error feedback

The section’s supplied passages provide figure and table references but do not state a substantive result about the need for error feedback.

  • Error feedback: Figure 7 compares PowerSGD with and without error feedback, but the supplied passage does not state the comparison’s outcome.The passage identifies the figure without providing its substantive conclusion.
  • Experimental references: The supplied materials include implementation-parameter references for ResNet18 and LSTM experiments without reporting their numerical findings.Tables 10 and 11 are described as showing parameter shapes and per-tensor compression ratios.

G Compressor implementation details

The paper compares several compressor implementations, including randomized, sign-based, top-k, Spectral Atomo, and PowerSGD variants, with differing aggregation and computational strategies.

  • Randomized compressors: Random Block samples one contiguous update slice, while Random K samples uniformly selected entries to match rank-r PowerSGD's communicated size.Both methods use the same seed across workers and reconstruct an n×m update matrix after aggregation.
  • Randomized compressors: Random K is significantly slower than Random Block because CPU index sampling and numerous random lookups are relatively expensive.The implementation samples indices with NumPy on the CPU.
  • Sign-based compressors: Sign+Norm communicates matrix signs with an L1 norm, reconstructing each worker's contribution through NCCL all-gather.Single-bit transmission requires a C++ extension because PyTorch does not natively support scalar types below 8 bits.
  • Sparse compressors: Top K selects the b largest-magnitude entries and communicates their values and indices, then reconstructs worker contributions using NCCL all-gather.The sample count b is set to (m + n)r to match rank-r PowerSGD.
  • Spectral Atomo: Spectral Atomo computes an SVD on the CPU, samples singular components, and was modified to repeat sampling until exactly r components are selected.The modification targets faster communication and does not significantly affect runtime; Atomo also uses a tensor-to-matrix treatment that yields lower compression at the same rank.
  • PowerSGD variants: A PowerSGD variant using four subspace-iteration steps, or eight matrix multiplications, is sufficient to converge to the best low-rank gradient approximation for final test accuracy.The standard variant uses one step and reuses previous iterates, whereas this variant does not reuse them.
  • Benchmark implementation: All compared optimizers were optimized for timing by packing gradient tensors into one flat buffer, with specialized sign packing and CPU/GPU benchmarking for Atomo.These implementation choices were intended to place methods at a similar optimization level.

I Learning rate tuning

Learning rates were tuned separately for methods without error feedback, while error-feedback methods with momentum reused SGD's tuned learning rate. Other optimization settings were not tuned.

  • Learning-rate procedure: Methods without error feedback received separate learning-rate tuning for each task and algorithm.For error-feedback methods with momentum, the learning rate tuned for SGD was reused.
  • Learning-rate procedure: Learning rates were tuned in the 16-worker setting, scaled linearly from the single-worker rate with a five-epoch warmup.The rates were defined for one worker before scaling to the default worker count.
  • Learning-rate procedure: The selected learning rate came from one full-training replication among rates initially spaced by factors of two, with the search range extended when an endpoint performed best.Selection was based on final test accuracy after the full number of epochs.
  • Untuned settings: Momentum, L2 or weight-decay parameters, and learning-rate schedules were not tuned in any experiment.Thus, the tuning procedure varied learning rates but left these other optimization choices fixed.
Loading 1905.13727v3…