Pufferfish: Communication-efficient Models At No Extra Cost

TL;DR

Pufferfish is a novel distributed training framework that reduces communication costs and computational overheads by integrating gradient compression into low-rank, pre-factorized models, maintaining accuracy and improving training speed.

Contribution

It introduces a method to incorporate gradient compression directly into model training via low-rank factorization, eliminating additional computational overheads and preserving accuracy.

Findings

Achieves up to 1.64x speedup in distributed training.

Maintains state-of-the-art accuracy with compressed models.

Produces smaller, more accurate models than existing pruning methods.

Abstract

To mitigate communication overheads in distributed model training, several studies propose the use of compressed stochastic gradients, usually achieved by sparsification or quantization. Such techniques achieve high compression ratios, but in many cases incur either significant computational overheads or some accuracy loss. In this work, we present Pufferfish, a communication and computation efficient distributed training framework that incorporates the gradient compression into the model training process via training low-rank, pre-factorized deep networks. Pufferfish not only reduces communication, but also completely bypasses any computation overheads related to compression, and achieves the same accuracy as state-of-the-art, off-the-shelf deep models. Pufferfish can be directly integrated into current deep learning frameworks with minimum implementation modification. Our extensive experiments over real distributed setups, across a variety of large-scale machine learning tasks, indicate that Pufferfish achieves up to 1.64x end-to-end speedup over the latest distributed training API in PyTorch without accuracy loss. Compared to the Lottery Ticket Hypothesis models, Pufferfish leads to equally accurate, small-parameter models while avoiding the burden of "winning the lottery". Pufferfish also leads to more accurate and smaller models than SOTA structured model pruning methods.