Transformer Training Explained
← /training

Chapter 13 of 16

Training over slow links

Top-k with error feedback, DeMo and DiLoCo when workers share a WAN.

Slides for this chapter: Transformer Training 12, Training over Slow Links

Inside a data centre, chapter 11's gradient sync rides links of hundreds of gigabytes a second. Between data centres, or over ordinary networks, the link can be a thousand times slower, and sending the whole gradient every step stops being affordable. Below, four workers each compress their momentum DeMo's way, from the trained model, in your browser: each sends 4 numbers of every 64 and keeps the rest for later.

Loading the animation…

Concept

This chapter is a toy, and says how far it is from the real thing. The Decoder's model has 5,376 parameters; here four simulated workers in one process each train it on four of the default run's sixteen windows per step, for 2,000 steps, on one book, on a CPU, three seeds per setting. The Compute in Transit results site measured the same ideas on a 15.8-million-parameter GPT trained by eight simulated workers for 3,000 steps on TinyStories, on GPUs, and then simulated data-centre-scale links with those measurements. Its numbers, not these, say what works at scale; this chapter explains the mechanisms its pages rely on and shows them running on a model small enough to follow.

WAN and inter-data-centre bandwidth. Inside a server, NVLink moves hundreds of gigabytes a second per GPU; across a data centre's network, tens (chapter 11). A wide-area network (WAN) link between data centres is shared by a whole site and can offer each training job far less: a 100 GbE link is 12.5 GB/s for everything that crosses it. At those rates chapter 11's sum, step time equals compute plus communication, is dominated by the communication: LLaMA 7B's dense BF16 gradient (13.5 GB) takes over a second to send at 12.5 GB/s, every step, before any other job shares the link. Training over slow links means sending less, sending less often, or both.

Top-k sparsification. Send only the kk largest values of a vector (and where they are), here 1 in 16 per tensor. Most of a gradient's coordinates are small at any one step, so most of the step survives. But each kept value needs its position too: for the Decoder's model 2 bytes of BF16 value and 2 bytes of position, so 1 value in 16 costs 4/32 of the dense bytes, not 1/16. Deep Gradient Compression (Lin et al. 2018) does this to an accumulated momentum: each worker adds its gradient into a local momentum and an accumulator, sends the top 1/16 of the accumulator (that is 1,344 bytes a step here, against 10,752 for the dense gradient), and plain SGD steps with the average of what was sent.

Error feedback (the residual). What a compressor drops is not thrown away: each worker keeps it, the residual, and adds it to the next step's vector before compressing again. A coordinate too small to be sent this step keeps growing until it is large enough to be sent, so every update arrives eventually, only late. Without this, the steady small coordinates of the gradient would never be applied at all; with it, compressed SGD provably converges at the uncompressed rate in the long run (Stich et al. 2018, "sparsified SGD with memory"; Karimireddy et al. 2019, "error feedback fixes signSGD"). 1-bit SGD (Seide et al. 2014) had used the same trick to send one bit per gradient value. DeMo's momentum is its residual: the animation's second sync starts from a momentum that still holds what the first did not send.

Why compressed vectors are all-gathered, not all-reduced. A ring all-reduce (chapter 11) adds the workers' chunks as they pass, which needs every worker's vector to have the same layout. Top-k vectors don't: each worker picks its own positions (in the animation's chunk, worker 1 sends positions that worker 2 may not). So the workers all-gather instead: every worker receives every other worker's values and positions and adds them itself. A ring all-gather has each worker send its own payload and forward the other (p−2)(p-2), so a worker's bytes on the wire grow with the number of workers: with four workers DeMo's 1,008 bytes of payload become 3,024 on the wire, against the dense all-reduce's 16,128. With eight workers the same payload costs 7 times itself, which is why the results site quotes both a payload ratio and a smaller wire ratio.

DeMo (decoupled momentum). Peng et al. (2024) let each worker keep its own momentum instead of syncing gradients (Algorithm 1): fold the gradient into it, cut each tensor into chunks, take each chunk's discrete cosine transform (DCT), send the kk largest coefficients per chunk and subtract them from the momentum, so the rest stays behind as error feedback. The DCT rearranges a chunk into frequency components without losing anything (it is orthonormal: what is sent and what is kept add up to the whole, checked in the frame tests), and the momentum's large, slowly changing components tend to collect in a few coefficients, the "fast components" that get sent. Everyone then averages the coefficients gathered, transforms back, and steps by the sign of the result. Here chunks are up to 64 values, k=k = 4 of them, values in BF16 and positions in one byte: 1,008 bytes a step per worker.

DiLoCo: syncing every H steps. The other way to save bytes is to sync less often. In DiLoCo (Douillard et al. 2023, Algorithm 1) each worker trains alone with AdamW for HH steps; then the workers average how far they moved (the "pseudo-gradient") and an outer optimiser, Nesterov momentum (learning rate 0.7, momentum 0.9, the paper's setting), moves the shared weights; then everyone starts the next HH steps from them. It sends the dense vector, but only once per HH steps: 1613 bytes a step on average at H=10H = 10, 323 at H=50H = 50. Local SGD, its ancestor, is analysed by Stich (2018).

Measured on the toy. Each method at the best of three learning rates (chapter 6: a sign step, a plain SGD step and Adam's normalised step want very different rates; the grid was placed by a one-seed pilot), three seeds, mean ± standard deviation of the final validation loss: dense AdamW 2.277 ± 0.090 (lr 0.01 ); top-k with error feedback 2.225 ± 0.039 (lr 1; at 3.0 every seed diverged); DeMo 2.073 ± 0.047 (lr 0.03, the top of its grid, so it might do better higher still); DiLoCo 2.578 ± 0.078 at H=10H = 10 and 2.706 ± 0.094 at H=50H = 50. Steps to a validation loss of 2.6: dense 1150, top-k 533, DeMo 317 (means over seeds); DiLoCo did not get there in every seed. On this toy the compressors leave chapter 1's plateau sooner than dense AdamW and DeMo ends lowest, while DiLoCo, whose workers each see only 4 windows a step between syncs, lags. Do not read that order as general: on the results site's GPT, dense AdamW ended lowest and DeMo's top-k some way behind it. A 5,376-parameter character model on one book is a different optimisation problem, and the compressors' big early steps (DeMo's are a full learning rate on every weight) suit its plateau.

validation loss (mean of 3 seeds)22.533.540500100015002000step
dense AdamW (lr 0.01) · top-k + error feedback (lr 1) · DeMo (lr 0.03) · DiLoCo, H = 10 (lr 0.03) · DiLoCo, H = 50 (lr 0.1).Each method at the best of three learning rates; final validation loss (mean ± sd): dense AdamW 2.277 ± 0.090; top-k + error feedback 2.225 ± 0.039; DeMo 2.073 ± 0.047; DiLoCo, H = 10 2.578 ± 0.078; DiLoCo, H = 50 2.706 ± 0.094. results.md §13.
validation loss22.533.5410 kB100 kB1 MB10 MBbytes sent per worker so far (log scale)
dense AdamW (lr 0.01) · top-k + error feedback (lr 1) · DeMo (lr 0.03) · DiLoCo, H = 10 (lr 0.03) · DiLoCo, H = 50 (lr 0.1).The same runs against what each worker has put on the wire (BF16 values, ring collectives among 4 workers): dense AdamW 16,128 bytes per step; top-k + error feedback 4,032 bytes per step; DeMo 3,024 bytes per step; DiLoCo, H = 10 1,613 bytes per step; DiLoCo, H = 50 323 bytes per step.

The step-time model. Put the measured steps and bytes together with an assumed compute time and a link bandwidth and you get the time to reach a target. The widget below does that for the toy and, as an explicitly labelled extrapolation, for larger models (scaling the bytes with the parameter count and keeping the toy's steps, an assumption nothing here tests). Slide the bandwidth down and the dense all-reduce's time grows fastest, because it sends the most; the results site does the same arithmetic with its own measurements and a discrete-event simulation of the links.

Loading the animation…

Queueing on a shared link, and stragglers. The step-time model assumes the link is the job's alone. A shared WAN carries other traffic, and a sync that arrives while the link is busy waits its turn. If other jobs' transfers arrive at random and the link is busy a fraction ρ\rho of the time, a newcomer waits on average ρ/(2(1−ρ))\rho/(2(1 - \rho)) transfer-times for an M/D/1 queue (the Pollaczek–Khinchine formula): 0.21 at 30% load, 4.5 at 90% (a seeded simulation gives 0.21 and 4.5; results.md §11d). The wait grows without bound as the link fills, and it hurts most the methods that send most. Synchronous training has a second waiting problem: every step waits for the slowest worker, a straggler. With eight workers whose step times vary by a random 10% on average, the mean worker takes 1.099 time units a step but synchronised steps take 1.270; syncing every 10 steps, as DiLoCo does, waits for the slowest sum of 10 steps instead: 1.148 per step (and 1.121 every 50). The results site's transit page shows the queueing effect on its recorded simulation runs.

The hardware underneath. Inside a data centre the slower links are the scale-out network's: Racks, chapter 5 builds its leaf, spine and core tiers and their oversubscription, and chapter 9 shows which of a job's traffic crosses them (the Racks site, in the Hardware track).

Maths

Top-k with error feedback (DGC's form, per worker ii, momentum μ\mu, compressor CC = keep the top kk of each tensor):

ui←μui+gi,vi←vi+ui,si=C(vi),vi←vi−si,θ←θ−η 1W∑isiu_i \leftarrow \mu u_i + g_i, \quad v_i \leftarrow v_i + u_i, \quad s_i = C(v_i), \quad v_i \leftarrow v_i - s_i, \quad \theta \leftarrow \theta - \eta \,\tfrac{1}{W}\textstyle\sum_i s_i

(and uiu_i is zeroed where sis_i was sent).

DiLoCo, every HH inner steps, with the outer Nesterov momentum bb:

Δ=1W∑i(θ−θi),b←μb+Δ,θ←θ−ηout(Δ+μb)\Delta = \tfrac{1}{W}\textstyle\sum_i \big(\theta - \theta_i\big), \quad b \leftarrow \mu b + \Delta, \quad \theta \leftarrow \theta - \eta_{\text{out}}(\Delta + \mu b)

Bytes on the wire per worker per step, WW workers, PP parameters, 2-byte values, payload cc bytes for a compressor:

dense: 2 W−1W 2P,compressed (all-gather): (W−1) c,DiLoCo: 2 W−1W 2P / H\text{dense: } 2\,\tfrac{W-1}{W}\,2P, \qquad \text{compressed (all-gather): } (W - 1)\,c, \qquad \text{DiLoCo: } 2\,\tfrac{W-1}{W}\,2P \,/\, H

M/D/1 waiting time for a load ρ\rho of transfers each taking ss: Wq=ρ s/(2(1−ρ))W_q = \rho\,s / (2(1 - \rho)).

Code

DeMo's compression of one chunk on one worker, from demoSync in src/lib/training/slowlinks.ts (a port of scripts/training/slowlinks.py, matched to its fixtures):

const ch = mom.slice(c * m, (c + 1) * m);
const coef = dctForward(ch, D);
const idx = topkIdx(coef, k);
const q = new Array<number>(m).fill(0);
const sent = idx.map((i) => coef[i]!);
idx.forEach((i, j) => {
  q[i] = sent[j]!;
  total[c]![i] = total[c]![i]! + sent[j]!;
  count[c]![i] = count[c]![i]! + 1;
});
const back = dctInverse(q, D);
for (let i = 0; i < m; i++) resid[c * m + i] = ch[i]! - back[i]!;

Comments

Be the first to leave a comment on this section.

Sign in (top-right) to leave a comment.