Communication-Efficient Optimization

"They asked me to send a billion numbers every step. I sent ten that mattered and a sticky note of what I owed them. Nobody noticed the difference except the network bill."

A Gradient That Learned to Pack Light

The previous section traced what happens when gradients arrive late: stale and delayed updates still converge, but only if the staleness stays bounded. That discussion took the size of each message for granted. Here we attack the size itself. The communication tax was named in Chapter 4, where the cost of a single all-reduce was shown to grow with the length of the vector being reduced; Chapter 3 turned that into the $\alpha$-$\beta$ model, in which the per-step communication time is $\alpha + \beta P$ for a $P$-parameter gradient. The arithmetic of a training step is fixed by the model, but the bandwidth term $\beta P$ is something we can engineer down, and for large models it dominates. Every method in this section is a way of shrinking either $P$ (bytes per round) or the number of rounds, while keeping the optimizer converging to the same place.

1. Quantization: Fewer Bits per Coordinate Intermediate

The most direct way to shrink a gradient is to spend fewer bits on each number. A standard gradient stores every coordinate as a 32-bit float; if a coarse approximation of each coordinate suffices for the optimizer, we can transmit one, two, or a handful of bits instead and cut the bandwidth term by an order of magnitude. The earliest production instance, 1-bit SGD, sent only the sign of each coordinate (scaled by a per-bucket magnitude) and folded the rounding loss back through error feedback; it shipped in speech-recognition training where the bandwidth between servers was the binding constraint. The principle generalizes: replace the true gradient $g$ with a quantized $Q(g)$ that is cheap to send.

What separates a good quantizer from a lossy one is unbiasedness. QSGD makes the quantization stochastic, so that each coordinate rounds up or down at random with probabilities chosen to satisfy $\mathbb{E}[Q(g)] = g$. An unbiased compressor preserves the expected gradient, so SGD still descends in the right direction on average; the price is added variance, which the convergence analysis absorbs as a slightly smaller effective step. For a coordinate $g_i$ with magnitude bounded by $\lVert g \rVert$, a stochastic quantizer to one of $s$ levels rounds to the nearer level $\ell$ with probability proportional to its closeness:

TernGrad pushes the same idea to three levels $\{-1, 0, +1\}$ times a shared scale, an extreme that still trains deep networks because the stochastic rounding keeps the estimate unbiased. The common thread is that quantization touches every coordinate but spends almost nothing on each, so the index bookkeeping is trivial; what you give up is precision, recovered in expectation by randomization and in practice by error feedback.

2. Sparsification: Send Only What Matters Intermediate

Quantization keeps every coordinate and cheapens each. Sparsification does the opposite: it keeps every retained coordinate at full precision but transmits only a small subset. The observation that motivates it is empirical and striking. In deep-network training, the gradient is dominated by a few large entries; the vast majority of coordinates are near zero at any given step. Top-$k$ sparsification exploits this by sending only the $k$ largest-magnitude coordinates (with their indices) and zeroing the rest, where $k$ might be one percent of $P$ or less.

The famous result is that this is far more aggressive than it sounds. Deep Gradient Compression reported that transmitting roughly $0.1\%$ of gradient coordinates per step, with the right corrections, matched the accuracy of dense training across image and language models, a hundredfold to thousandfold reduction in the bytes per round. The corrections matter because top-$k$ is biased: a coordinate that is consistently medium-sized would be perpetually dropped, and its information would vanish. The fix is to never truly discard a dropped coordinate but to bank it. Each worker keeps an error buffer $e_t$, the running sum of everything compression has thrown away, and adds it back before the next compression. Formally, with a compressor $C$ and step $\eta$,

The buffer $e_{t+1}$ holds exactly the residual the compressor refused to send; on a later step, once those banked coordinates have grown large enough to enter the top-$k$, they are transmitted in full. No gradient information is lost, only delayed, which is why error feedback turns an aggressively biased compressor into a convergent optimizer. We make this mechanism concrete in Section 4 and watch it work in Code 10.7.1.

3. Low-Rank Compression: Send a Sketch Advanced

A third family treats the gradient of a weight matrix as a matrix, not a flat vector, and exploits the fact that a matrix can be approximated by a low-rank product. For a layer whose gradient is a matrix $G \in \mathbb{R}^{m \times n}$, a rank-$r$ approximation $G \approx P Q^\top$ with $P \in \mathbb{R}^{m \times r}$ and $Q \in \mathbb{R}^{n \times r}$ replaces $mn$ numbers with $r(m+n)$ numbers, a large saving whenever $r \ll \min(m, n)$. PowerSGD computes such a factorization cheaply with a single step of power iteration per round, reusing the previous step's factors as a warm start so the iteration converges essentially for free as training proceeds.

The decisive systems property of PowerSGD is that its compressed form is itself all-reducible. Top-$k$ sparsification produces different index sets on different workers, so combining the messages needs a sparse, irregular collective that is awkward on the GPU interconnects that Chapter 4 describes. PowerSGD's factors $P$ and $Q$ are dense and have the same shape on every worker, so they go straight through a standard all-reduce, the same primitive a dense gradient uses. That compatibility is why PowerSGD became a default option in production data-parallel training: it shrinks the message and keeps the fast collective. Like top-$k$, it is biased, and it relies on the same error-feedback buffer to stay convergent.

import torch
from torch.distributed.algorithms.ddp_comm_hooks import powerSGD_hook as psgd

model = torch.nn.parallel.DistributedDataParallel(model)   # standard DDP wrapper
state = psgd.PowerSGDState(                                 # rank and warm-start config
    process_group=None, matrix_approximation_rank=4,
    start_powerSGD_iter=1000,                               # warm up dense, then compress
)
model.register_comm_hook(state, psgd.powerSGD_hook)        # every all-reduce now low-rank

4. Error Feedback, in Code Intermediate

The claims above are quantitative, so we test them. The program below trains a linear-regression model with eight data-parallel workers two ways. The first sends each worker's full gradient every step. The second applies top-$k$ sparsification, keeping only the ten largest of one hundred coordinates per worker, with a per-worker error-feedback buffer following the recurrence from Section 2. We compare the final loss and count the coordinates each method puts on the wire.

import numpy as np

rng = np.random.default_rng(0)
N, d, K = 20_000, 100, 8                 # examples, features, workers
X = rng.standard_normal((N, d))
w_true = rng.standard_normal(d)
y = X @ w_true + 0.5 * rng.standard_normal(N)

shards = np.array_split(np.arange(N), K)        # one data shard per worker
lr, steps = 0.05, 60
keep = 10                                        # top-k: send only 10 of 100 coords

def grad_shard(w, s):
    """Mean-squared-error gradient on worker s's shard only."""
    Xs, ys = X[s], y[s]
    return (2.0 / len(s)) * (Xs.T @ (Xs @ w - ys))

def loss(w):
    return float(np.mean((X @ w - y) ** 2))

def topk(v, k):
    """Keep the k largest-magnitude coordinates of v, zero the rest."""
    out = np.zeros_like(v)
    idx = np.argpartition(np.abs(v), -k)[-k:]
    out[idx] = v[idx]
    return out

def run(compress):
    w = np.zeros(d)
    err = [np.zeros(d) for _ in shards]          # per-worker error-feedback memory
    sent = 0
    for _ in range(steps):
        agg = np.zeros(d)
        for j, s in enumerate(shards):
            g = grad_shard(w, s)
            if compress:
                corrected = g + err[j]           # add back what was dropped before
                msg = topk(corrected, keep)      # transmit only k coordinates
                err[j] = corrected - msg         # remember the residual locally
                sent += keep
            else:
                msg = g
                sent += d
            agg += msg
        w = w - lr * (agg / K)                   # average the received messages
    return loss(w), sent

L_full, sent_full = run(compress=False)
L_comp, sent_comp = run(compress=True)

print(f"workers K                 : {K}")
print(f"coordinates per gradient  : {d}   (top-k keeps {keep})")
print(f"uncompressed final loss    : {L_full:.6f}")
print(f"top-k + error-fb loss      : {L_comp:.6f}")
print(f"loss gap                   : {abs(L_comp - L_full):.2e}")
print(f"coords sent  uncompressed  : {sent_full}")
print(f"coords sent  compressed    : {sent_comp}")
print(f"communication reduction    : {100 * (1 - sent_comp / sent_full):.1f}%")

workers K                 : 8
coordinates per gradient  : 100   (top-k keeps 10)
uncompressed final loss    : 0.247179
top-k + error-fb loss      : 0.246605
loss gap                   : 5.74e-04
coords sent  uncompressed  : 48000
coords sent  compressed    : 4800
communication reduction    : 90.0%

The result is the section's thesis in three numbers. Ten coordinates per worker per step instead of a hundred is a tenfold cut in the bandwidth term $\beta P$, and the model reached the same loss. The discarded ninety coordinates were not lost, only deferred through the error buffer, which is why the gap is rounding-level rather than a visible degradation. Scale this from a hundred parameters to a billion and the saved bytes are the difference between a network-bound training step and a compute-bound one.

5. Local Updates: Fewer Rounds, Not Smaller Messages Intermediate

Quantization, sparsification, and low-rank compression all shrink the message while keeping one synchronization per step. Local-update methods attack the orthogonal axis: they cut the number of synchronizations. In local SGD, each worker takes $H$ ordinary optimizer steps on its own data before any communication, and only then are the worker models averaged. Setting $H = 1$ recovers synchronous distributed SGD; setting $H = 10$ communicates a tenth as often. Where compression divides the bytes per round, local SGD divides the rounds, and the two compose: a compressed message sent every tenth step is cheap on both axes.

The cost is statistical. Between synchronizations the workers drift apart, each descending the slightly different landscape of its own shard, so the averaged model makes a little less progress per gradient evaluated than fully synchronous SGD would. The trade is governed by $H$: larger $H$ means fewer rounds but more drift, and the sweet spot depends on how much faster communication is than computation. This is the same bytes-versus-rounds-versus-progress trade as compression, viewed from the rounds axis. Local SGD is the algorithmic core of federated averaging, where the workers are phones or hospitals that can communicate only rarely, and we develop that setting in full in Chapter 14, where infrequent communication is not an optimization but a hard constraint of the deployment.

6. The Unifying View and Its Limits Advanced

Lined up together, the families form a single design space with two axes and one safeguard. Quantization and sparsification and low-rank compression all reduce bytes-per-round; local SGD reduces rounds; error feedback is the correction that lets the biased, aggressive members of either axis stay convergent. Table 10.7.1 places each method on this map, naming what it sends, what it costs, and whether it needs feedback to converge.

The limits are as important as the savings. Compression helps only when communication, not computation, is the bottleneck; on a fast intra-node interconnect with a small model, compressing the gradient can cost more in encoding and decoding than it saves on the wire, and dense all-reduce wins. Sparsification's index overhead erodes its savings when $k$ is not tiny, and irregular sparse collectives are slower per byte than the dense all-reduce that PowerSGD preserves. Local SGD's drift grows with data heterogeneity, so the more unlike the workers' shards, the smaller the $H$ you can afford. None of these methods is free, and each is the right tool only when its specific bottleneck binds, the same discipline of matching the remedy to the ceiling that opened the book. We quantify exactly when each pays off with the communication-complexity bounds of the next sections.