Data Movement in LLM Training

Topics We'll Cover

• The Traffic Each Kind of Parallelism Creates
• How Big Is a Gradient Sync?
• Ring All-Reduce Arithmetic
• Interactive: Where Links Bind
• In-Network Aggregation
• Low-Communication Training: the DiLoCo Family
• Frequency-Domain Gradient Compression
• DeMo: Decoupled Momentum and the DCT
• Mapping DeMo onto a Photonic Link Module
• The Fixed-Band Experiment (Planned)
• From Convergence to Wall-Clock
• Other Training Traffic
• What to Take Away

The Traffic Each Kind of Parallelism Creates

Large-model training splits work along several axes at once, and each axis puts a different kind of data on a different kind of link. What a compute-in-transit module could act on depends on which traffic it sees.

• The bulk transfer is the data-parallel sync: once per step, one model-sized buffer per worker (slide 02). It is also the only traffic that can cross a data centre boundary without stalling every layer, which is where low-communication methods act (slide 06).
• Per-layer traffic is latency-critical: tensor-parallel and ring-attention exchanges sit on the critical path of each layer, so anything inserted on those links must add little delay.
• Sources: Megatron-LM for tensor parallelism (Shoeybi et al., arXiv:1909.08053); ZeRO (Rajbhandari et al., arXiv:1910.02054) and PyTorch FSDP (Zhao et al., arXiv:2304.11277); GPipe (Huang et al., arXiv:1811.06965); ring attention (Liu, Zaharia and Abbeel, arXiv:2310.01889); GShard for MoE all-to-all (Lepikhin et al., arXiv:2006.16668).

How Big Is a Gradient Sync?

A BF16 gradient is 2 bytes per parameter. A ring all-reduce (slide 03) makes each worker send a little under twice that:

Source: results.md, section “3. Training”

• Training syncs are the largest movements in LLM work. A Llama-3-70B BF16 update is 141.10 GB, against 268.4 MB for a 2,048-token Llama-3-8B KV hand-off in inference (results.md section 6; deck 03 covers inference). Bar widths are linear; the two smallest bars are widened so their labels fit.
• The gradient is the same size as the weights, so the per-step volume does not shrink with batch size: larger batches only make the sync rarer per token.
• FSDP moves weights too. Gathering BF16 weights forward and backward, then reduce-scattering gradients, costs about 1.5 times the data-parallel volume by ZeRO's own count (Rajbhandari et al., arXiv:1910.02054), in exchange for memory.

Ring All-Reduce Arithmetic

N workers each hold a buffer of B bytes and all need the sum. Split each buffer into N chunks and arrange the workers in a ring.

• Bandwidth-optimal: the bytes each worker sends approach 2B as N grows, independent of N, so adding workers does not add traffic per link. That is why results.md's 8- and 64-worker columns differ by only about an eighth (28.10 against 31.62 GB for Llama-3-8B).
• The cost is latency: 2(N−1) sequential steps, each paying the link's latency. Collective libraries switch to tree or hierarchical schemes for small messages and large N; the byte count of the bulk sync stays near 2B.
• The arithmetic is tiny. The reduce-scatter adds each value once per hop: one add per BF16 value, 0.5 operations per byte (results.md section 2). The work is moving the bytes, not summing them.
• Where a link module sits matters: in a ring, each link carries a partial sum from one neighbour, so a module on that link sees one stream. It can transform or compress that stream; it cannot sum workers whose data never pass through it (slide 05).

Interactive: Where Links Bind

Time to move one model-sized BF16 update once over one link: latency + bytes / bandwidth, with the link presets of Disaggregated_Inference_Sim (milliseconds):

Source: results.md, section “3. Training”. 2x = BF16 to FP8; 8x = BF16 to 2-bit (MuLoCo's setting); 85x = DeMo's reported reduction. The ratios are the papers' figures, used only to size transfers.

• Inside a node the sync is short: 35.7 ms for Llama-3-8B and 313.6 ms for Llama-3-70B on NVLink 4. Between nodes on InfiniBand NDR a 70B update takes 2,822.1 ms, which training frameworks hide by overlapping it with the backward pass.
• Across data centres it binds. Over a 100 GbE link a Llama-3-70B update takes 11.3 s uncompressed. This is the traffic DiLoCo-family methods shrink (slide 06), and the place compression in transit would act.

In-Network Aggregation

If the network itself sums the workers' updates, each worker sends its gradient once and receives the sum once, instead of the 2(N−1)/N · B of a ring, and the switch does the adds.

Low-Communication Training: the DiLoCo Family

Instead of syncing every step, each worker (or island of accelerators) takes many local optimiser steps, then the workers exchange the pseudo-gradient, the change in their parameters, and apply one outer optimiser step to it.

for each worker i, in parallel:
    θi ← θ                        # start from the shared parameters
    repeat H times: θi ← AdamW(θi, local data)   # inner steps, no communication
    Δi ← θ − θi                    # pseudo-gradient: one model-sized buffer
Δ ← all_reduce_mean(Δi)               # the only sync: once per H steps
θ ← Nesterov(θ, Δ)                  # outer step

• The sync is still one model-sized buffer, just rarer and smaller. The 2x, 8x and 85x columns of slide 04 size it: a 2-bit Llama-3-70B update over 100 GbE is 1,411.1 ms rather than 11,288.4 ms.
• Every member of the family compresses at the endpoint, on the accelerator that holds the optimiser state. Error feedback, in particular, needs memory of what was not sent: state that a link module would have to hold per parameter if it did the compressing.

Frequency-Domain Gradient Compression

Sparsification usually keeps the largest gradient entries in their own domain. Wang et al. asked whether a transform first would keep more of the information at the same compression ratio (“FFT-based Gradient Sparsification for the Distributed Training of Deep Neural Networks”, HPDC 2020, doi:10.1145/3369583.3392681; preprint arXiv:1811.08596, titled “SuperNeurons: FFT-based Gradient Sparsification…”).

• The premise: the gradient's energy concentrates in fewer frequency coefficients than spatial entries, so the same number of kept values carries more of the signal.
• Quantisation of what is kept uses a range-based floating-point representation, fitted to the span of the surviving coefficients, to cut the bits per value further.
• Convergence: the authors prove convergence when the fraction of coefficients dropped is reduced during training, so the compression error shrinks as training proceeds.
• Why it matters here: the expensive step is a transform. A photonic module's published operation set includes the FFT; top-k selection does not (it needs comparisons), which is the same split slide 09 meets in DeMo.

DeMo: Decoupled Momentum and the DCT

DeMo (Peng, Quesnelle and Kingma; the later version adds Chen, Su and Liu; ICLR 2026; arXiv:2411.19870) drops the per-step all-reduce of full gradients. Each worker's momentum evolves locally, and only its fast-moving components are shared.

mi ← β mi + gi                    # local momentum, never all-reduced
C ← DCT(chunks(mi))                  # transform each fixed-size chunk
qi ← top_k(C)                       # keep the k largest per chunk: values + indices
mi ← mi − IDCT(qi)                  # remove what was sent: momentum is the error feedback
Q ← all_gather(qi)                    # the only communication: small
θ ← θ − η · IDCT(Q)                   # same update on every worker (see the paper for its exact form)

• The momentum buffer acts as error feedback: what was not sent stays in m_i and accumulates until it is large enough to be picked by top-k.
• Reported: up to 85× less data per GPU than AdamW with DDP, at comparable loss and accuracy, for 300M and 1B-parameter models.
• A DiLoCo-style variant (Nedelkoski et al., “Distributed Low-Communication Training with Decoupled Momentum Optimization”, NeurIPS 2025 DynaFront workshop, arXiv:2510.03371) decomposes Nesterov momentum into high- and low-frequency components with the DCT and synchronises only the high-frequency components every H steps, reporting up to 16× less communication than DiLoCo.

Mapping DeMo onto a Photonic Link Module

The module's published operation set is FFT, NTT, MAC, and modular and floating-point arithmetic (Kundu et al., arXiv:2608.21536). Step by step:

Source: results.md, section “2. Work per byte”. Top-k is costed at about 2 comparisons per value (illustrative).

The DCT costs 10 to 15 FLOPs per BF16 byte if charged against a budget of 0.40 to 3.96, so it fits only if the optics does the transform without charge. If it does, the fixed band adds no arithmetic and fits; top-k does not, being outside the set.

Source: results.md, section “5. A fixed frequency band against top-k”: share of each 64 × 64 chunk's energy kept after a 2-D DCT, keeping 1/16 of the coefficients. Synthetic signals, illustrative.

• A fixed band works only when energy sits at low frequencies (the smooth case, 0.957 against 0.995). On noise or spikes it keeps about its own share of the energy, 0.062, far below top-k's 0.325 and 0.320.
• Whether real momentum looks like the smooth case or the noisy one is not known from synthetic signals. The next slide plans the measurement.

The Fixed-Band Experiment (Planned)

Train a small GPT, tens of millions of parameters (nanoGPT-style, on a small text corpus), four ways, at the same keep fraction in arms (b) to (d):

• Metric: validation loss against bytes communicated, so arms with different compression are compared at equal traffic, not equal steps.
• Sweeps: chunk size (for example 16 and 64, the two sizes costed in results.md section 2); keep fraction (for example 1/16, as in section 5); ENOB in arm (d), with noise added by the rule of the sister series (the noise rule for floating-point workloads).
• Diagnostics: the share of each chunk's momentum energy that falls in the fixed band, logged during training, which ties the outcome back to section 5's synthetic cases.

Which band?

Low frequencies within a chunk are the obvious choice, and section 5 shows they work for smooth signals. But DeMo's own premise is that the components worth synchronising are the fast-moving ones, and the NeurIPS 2025 workshop variant (Nedelkoski et al.) synchronises the high-frequency components of momentum (slide 08). Is momentum energy within a chunk concentrated at low frequencies at all? And if what matters is high-frequency content, does a fixed high band do as well as top-k? The experiment should test a low band, a high band and top-k side by side rather than assume the answer.

From Convergence to Wall-Clock

The experiment gives loss against bytes. Whether that saves time depends on the system: link rates, latencies, how much of each sync overlaps with compute, and whether the module adds delay or energy. That question belongs to a discrete-event simulation, which is the part of this work that falls to a simulation engineer.

• Calibration: compression ratios and loss curves from the experiment of slide 10; link presets from Disaggregated_Inference_Sim; the module's budget from results.md section 1.
• Implementation: a SimPy model that extends Disaggregated_Inference_Sim with a training mode, so training and inference (deck 03) share one hardware description and cannot drift apart. Planned, not yet built.

Other Training Traffic

• Tensor-parallel and MoE traffic stay on NVLink inside a node or rack and are latency-bound; this series does not treat them as targets.
• Aggregation of any of these needs a switch, not a link (slide 05).

What to Take Away

• The data-parallel sync is the bulk transfer of LLM training: 16.06 GB of BF16 gradient for Llama-3-8B and 141.10 GB for Llama-3-70B, with a ring all-reduce sending 2(N−1)/N times that per worker.
• Inside a cluster it is overlapped; across data centres it binds: 11.3 s per uncompressed Llama-3-70B update over 100 GbE (the interactive, slide 04).
• Summing belongs in a switch. SwitchML and SHARP aggregate in the network; a link module sees one stream, so it can compress but not sum.
• Low-communication training already cuts the traffic by orders of magnitude at the endpoint: DiLoCo 500×, Streaming DiLoCo two orders of magnitude in bandwidth, MuLoCo 2-bit deltas, DeMo up to 85× less data per GPU.
• DeMo's compressor is half in a photonic module's operation set: the DCT is, top-k is not. A fixed frequency band is the comparison-free substitute, and it works only if momentum energy sits in the chosen band, which a planned experiment will measure.
• Next: deck 03 turns to inference: KV hand-off between prefill and decode, weight streaming, and what a link module could do to them.