Optical Engines in a Training Cluster

Topics We'll Cover

• The Question for a Cluster, and the Short Answer
• Anatomy of a Training Step
• Parallelism, and Where the Engine Sits
• Masks Change Every Step
• Interactive: Step Time Against Mask Time
• Layer-Major Scheduling Against the Pipeline
• Gradient Accumulation as the Compromise
• Data-by-Data Products Stay Digital
• The Static Amdahl Bound for a Step
• What 1.1x Means at Cluster Scale
• The Optimizer Step and DeMo
• Frequency Transforms in Transit: a Fixed Band
• Frequency Transforms in Transit: Where the Band Is Computed
• Training Across Data Centres: Time and Energy to a Loss
• What Does Not Map at Cluster Scale
• What to Take Away

The Question for a Cluster, and the Short Answer

Deck 01 counted the transform work in a training step; deck 02 looked at how optics trains itself. This deck asks the systems question: if every accelerator in a training cluster had a Fourier-optical transform engine beside it (FOptInf 01 for the physics), what would a training step look like, and would it be shorter?

Short answer: at today's mask rates the engine lengthens the step rather than shortening it, and the frequency transform that genuinely helps a training cluster (DeMo's DCT of the momentum) is worth link bytes, not FLOPs. All numbers come from results.md, written by analysis/train_flop_share.py; the step size (1,048,576 tokens), the cluster (64 GPUs at 400 TFLOP/s sustained) and the device (2-megapixel DMD at 1,031 Hz) are illustrative.

Anatomy of a Training Step

One optimizer step processes a global batch (here 1,048,576 tokens) and ends with a single weight update. Four phases, of which only the first two contain transforms:

The FLOP ledger is checked against the 6N rule: for the transformer variant it gives 46.64 GFLOP per token at 2,048 tokens and 70.80 at 32,768, a ratio of 1.0000 to 6N plus the attention term (results.md §1).

Parallelism, and Where the Engine Sits

Suppose each accelerator has a transform engine beside it that takes the FFTs and weight-masked products of the layers that accelerator owns. What each engine's mask must hold depends on how the model is split:

Megatron's own sequence parallelism splits only the norm and dropout regions along the sequence (Korthikanti et al., arXiv:2205.05198), which a long convolution would tolerate; splitting the convolution itself would not suit an engine that transforms whole rows at once.

Masks Change Every Step

A 2-megapixel mask holds a small fraction of the model's filter spectra (FOptInf 02, “Fourier-Plane Mask Capacity”), so an engine cycles through frames as it walks the layers. In inference that cycle serves each batch with fixed weights. In training the weights are new every step, and the backward pass needs them again:

Source: results.md, section “4. Holding the masks”. Illustrative: 2-megapixel DMD at 1,031 Hz, one complex value per pixel (optimistic); 64 GPUs at 400 TFLOP/s; layer-major order.

• The mask bill rivals the whole step. For Hyena-2 it is a quarter of the digital step at 2,048 tokens and about all of it at 8,192, to remove 0.13–0.15% of the FLOPs. For block-circulant weights the step is 0.15 s and the masks take 0.54–2.10 s.
• Longer sequences cost more frames, because a filter spectrum has a bin per padded sample; fewer sequences per step does not reduce the count, since every layer's spectra are written once per step regardless.
• The backward does not strictly need a separate conj(H) mask: for a real filter, correlating with h is convolving the time-reversed gradient with h and reversing the result. But by the time the backward reaches a layer the mask holds other layers, so its spectra are written again either way; the table counts both writes.
• Compare inference, where the same device writes 269 frames per forward pass of a 2,048-token Hyena-2 prompt (FOptInf 02) and the simulator found mask rewrites dominate the engine's time (FOptInf 03, “Where the Time Goes”). Training doubles the frames and removes any hope of writing them once.

Interactive: Step Time Against Mask Time

Per-token step FLOPs and static shares are embedded from results.json for the four sequence lengths the analysis covers; mask values use the same formula as its mask_budget(). At the defaults the explorer reproduces the slide 04 table.

The model, stated plainly. Digital step = per-token step FLOPs × tokens ÷ (GPUs × TFLOP/s), with GPU time proportional to FLOPs (FFTs at the matmul rate). Removable = static share × step: the engine does the static work in zero time. Each engine holds the masks of 1/k of the model (k = tensor × pipeline degree) and writes them twice per pass (forward and backward): 2 × ⌈values ÷ (k × pixels)⌉ frames. Layer-major makes one pass per step; sequence-major makes one per microbatch of m sequences on each of the GPUs ÷ k replicas. Not overlapped (conservative): step − removable + mask time. Fully overlapped: the larger of (step − removable) and mask time. Per-token FLOPs and shares are those at 1,048,576 tokens, held fixed for other step sizes (an approximation: the once-per-step filter transforms shift them slightly at long sequences). The data-by-data figure is what writing the weight gradient's conj(U) masks would cost, spread evenly over all engines.

Layer-Major Scheduling Against the Pipeline

Layer-major order passes every sequence of the step through one layer, with that layer's spectra on the mask, before moving on: one write per layer per direction per step. It is the order the slide 04 numbers assume, and the opposite of how large models are trained.

Pipeline parallelism does help in one way: a stage's engine holds only its stage's layers, so each pass has fewer frames (the model-parallel degree in the explorer). It does not help enough: for Hyena-2 at 2,048 tokens even one of 32 stages has more spectra than a 2-megapixel mask holds, so each microbatch still costs frames.

Gradient Accumulation as the Compromise

Gradient accumulation already splits an optimizer step into rounds whose gradients are summed before one update. Run each round layer-major over its sequences and the trade is explicit:

• In the explorer, sequence-major with microbatch m is exactly this, with a = sequences per replica ÷ m. At the defaults (512 sequences of 2,048 tokens over 64 GPUs, 8 per GPU) m = 4 doubles the mask time and m = 1 multiplies it by eight; the weights are the same throughout the step, so only capacity forces the rewrites.
• Activation checkpointing (Chen et al., arXiv:1604.06174) is the other memory lever. It recomputes the forward during the backward, which for this engine means repeating the forward FFTs and, unless the recomputed layer's spectra are still on the mask, another mask pass. The analysis does not count recomputation, so the static shares here are slightly optimistic for a checkpointed run.
• Where the compromise lands: whatever a, the engine needs the mask time to hide behind digital work it cannot itself remove. The explorer's “fully overlapped” case shows the best this can do.

Data-by-Data Products Stay Digital

The filter gradient dH = IFFT(Σ G·conj(U)) multiplies two spectra that both depend on the sequence (the gradient of a convolution is a correlation). With block-circulant weights the weight gradient conj(X)·G depends on every token. Either way, a mask would carry data, and would be new for every sequence or token:

Source: results.md, sections 2 and 4 (illustrative device as before).

• Mask writes that cost more than the work they remove. On one device the data-by-data writes take minutes per step. Spread over 64 engines (the explorer's last tile) the Hyena-2 figure is still about one whole digital step, to remove about 2% of an optical share that is itself 0.14% of the step; the block-circulant figure is many digital steps.
• So the static share is the ceiling, not the full optical share: for block-circulant weights 9.66% at 4,096 tokens rather than 12.19%.
• Speculation: a classical joint transform correlator computes the correlation of two inputs without a mask, by placing both in the input plane and transforming the detected intensity of their joint spectrum again. It moves the problem to the input modulator's frame rate and adds a square-law detection and a second pass. Not modelled here.

The Static Amdahl Bound for a Step

Static share = FFTs plus the weight-masked products, over all step FLOPs: what an engine with once-per-step masks could take (glossary). If that work became free, the step would be at most 1/(1−share) faster.

Source: results.md, section “2. A training step”

• Training does not rescue the share. The backward has as many FFTs as the forward but twice the matmuls, so the step share sits below the forward share in every row.
• The bound is a FLOP bound. It assumes the engine and its masks cost no time and that GPUs run FFTs at their matmul rate. If GPUs ran FFTs more slowly (GPU FFT efficiency), the time share would be larger; the analysis does not quantify that for training.

What 1.1x Means at Cluster Scale

A bound of 1.105–1.112x is not nothing at cluster scale: it is roughly a tenth of a training run's GPU-hours, which for a large run is a large bill. That is the case for writing it down. The case against is everything the bound leaves out:

The Optimizer Step and DeMo

AdamW's update is element-wise and memory-bound: no transforms. One recent optimizer does use them. DeMo (Peng, Quesnelle, Kingma, “DeMo: Decoupled Momentum Optimization”, arXiv:2411.19870, 2024) lets each worker keep its own momentum, takes a DCT of it in small chunks, sends only the largest coefficients (top-k) to the other workers, and removes what it sent from its local momentum. The slow components stay local, decoupled from the exchange.

Source: results.md, section “3. The optimizer step and DeMo”: a DCT and an inverse DCT per chunk of 64, at the real-FFT cost.

DeMo as measured. The Compute in Transit experiment (CiT 04) ran it on a 15.8M-parameter GPT (8 layers, width 384, context 256) on TinyStories, with 8 simulated workers of 8 sequences each per step for 3,000 steps. Its mechanics, step by step:

• Error feedback. What is compressed is the delta: the momentum plus the residual of everything not sent before, which decays by 0.999 per step. What is sent is removed from it, so nothing is dropped, only delayed.
• Chunks. The DCT runs over 1-D 64-point, 1-D 256-point or 2-D 64×64 chunks; each keeps 1/16 of its coefficients.
• Sign. The workers gather each other's values and apply the sign of the resulting update, so every weight moves by a full step.
• The cost of compressing. Final validation loss, 3 seeds each: AdamW 1.646 ± 0.005, DeMo top-k (2-D 64×64) 1.785 ± 0.007. AdamW passes 1.785 by step 2,000; top-k never reaches 1.646 in 3,000 steps.

Payload is not wire bytes. A ring all-reduce of B bytes sends 2(N−1)/N·B per worker, whatever N is; DeMo's workers pick different coefficients, so they all-gather instead, and a ring all-gather of P bytes per worker sends (N−1)·P, which grows with N. On 8 workers, per worker per step (BF16 values):

Source: the experiment's results.md, section 1 (results site, “What each worker sends”). Top-k's payload counts each index at the bytes a position within a chunk needs; the reference code sends int64 indices, 9,895,920 bytes per step, 2.5x the compact count.

• As compute, it is nothing. Five millionths of the step, in 64-point transforms: far too small and too short for an optical engine, whose conversions and passes cost more per value than a GPU spends on a 64-point DCT.
• As communication, it is the point, but less than the payload suggests. An 8x smaller payload is 2x fewer wire bytes on 8 workers, and the gap widens as workers are added. What DeMo removes is gradient bytes on the network between replicas, the phase of the step (slide 02) that no transform engine touches. The next three slides follow it onto the link.

Frequency Transforms in Transit: a Fixed Band

Gradient compression is an established line of work: sparsification (Deep Gradient Compression, Lin et al., arXiv:1712.01887), low-rank updates (PowerSGD, Vogels et al., arXiv:1905.13727) and DeMo's DCT. All of them compress at the worker. Could the transform run in the link's data path instead, on bytes that have to move anyway?

Why a fixed band. A photonic link module (Kundu et al., arXiv:2608.21536) can run the transform but not the comparisons top-k needs to find the largest coefficients. A fixed band, the same 1/16 of positions in every chunk every step, needs no comparisons, sends no indices (16x smaller payload, against top-k's 8x), and puts every worker's values at the same positions. The Compute in Transit experiment measured what that substitution costs, on the same 15.8M-parameter model:

Source: the experiment's results.md, section 1 and “Steps to reach a loss” (loss curves, steps ladder). Steps are the first evaluation (every 200 steps) at or below the target, the mean over seeds; 3.158 is the 2-D low band's own final loss.

Why it loses: captured energy. The share of a tensor's energy (its squared DCT coefficients) a selector keeps at the same k. On worker 0's momentum, averaged after step 200:

Source: the experiment's results.md, section 2, runs C, tensor “momentum”, group “all” (results site, “Why a fixed band loses”). A random set keeps about the keep fraction, 1/16.

• The momentum's energy is not where a fixed band looks. At 2-D 64×64 the low band keeps 0.013 of it, less than a random set; top-k keeps 0.643.
• Error feedback makes it worse. The band is removed from the delta every step, so on the delta the low band keeps 0.000 (to three places) and the residual piles up out of band, where it is never sent. By step 3,000 the band holds 1.6–1.8e-5 of the delta's energy in the 1-D 256-point runs (optics investigation, section 2).
• Every fixed band loses. The best, a random band at 2.700, is still 0.915 above top-k. A fixed band sends the same positions forever.

Frequency Transforms in Transit: Where the Band Is Computed

A DCT is linear, so the sum of the workers' bands is the band of the summed delta: unlike top-k, a fixed band can be summed rather than gathered. That gives three places to compute it, beside the measured C (band on the GPU, all-gathered):

Source: the simulation's results.md, section 1. C′ and E′ reuse C's and E's measured steps, because summing the band gives the same update as gathering it.

The emulated optics (E). The first design, with the input full scale fixed after calibration and error feedback that removed the exact band, diverged in all six runs (final losses 11.153 to 558.834): the optics' errors were never corrected and the sign update turned a repeated bias into full steps the wrong way. The fix was to feed back the band actually sent, which each sender sees in the all-gather. Repaired, 8 planes: ENOB 6 ends at 3.423, ENOB 8 at 3.017, ENOB 10 at 2.878; perfect optics at 3.141. The best is still 1.09 above top-k (results site, “ENOB and bit planes”). That the repaired optics needs fewer steps than the exact band (1,800 to 3.158 at ENOB 8, against 2,667) is a hypothesis from single seeds.

When it pays: break-even WAN bandwidth (2 data centres of 4 workers; time to 3.158; transform not charged to the module):

Source: the simulation's results.md, sections 7a and 11 (results site, break-even; CiT 05).

• The pluggable never pays. Its all-gather still sends 13,854,288 bytes per worker, a quarter of AdamW's 55,417,152, while needing 1,800 steps to AdamW's 400. With the transform charged at 10 GFLOP/s, its step is 33.348 ms against 1.443 ms for the band on the GPU, 95.0% of it module compute; the budget would need 140.0 GFLOP/s before the optics set the rate, above the series' bracket of 10–99 GFLOP/s.
• Summing in the network is where a fixed band could gain, and only on a slow WAN. The same summation on the GPUs (C′) does better without a module: only E's single-seed step advantage lets E′ beat C′ at all.

Training Across Data Centres: Time and Energy to a Loss

Saved bytes matter most where the link is slowest: data-parallel workers split across data centres, syncing over a WAN. The training-transit simulation takes the experiment's bytes and measured steps to a loss and asks how long, and how much energy, each way of syncing needs to reach 3.158. Background on where training's bytes go: CiT 02.

Source: the simulation's results.md, sections 2, 5 and 9 (results site, explorer; shared WAN). [high]: module budget 99 GFLOP/s, transform charged; [free]: transform not charged. E's steps are the single-seed ENOB 8 run; the queueing column is the discrete-event engine's mean over 200 syncs, not run for E.

• At this scale AdamW's step is mostly communication. On 2 DC × 4 over 100 GbE, compute is 12.1% of it and the inter-DC link 63.1%; in one data centre its step is 1.740 ms, on 8 DC × 1 over 25 GbE 18.505 ms.
• Steps decide it. Top-k reaches 3.158 first because it needs 400 steps, like AdamW, with half the wire bytes. The band's cheaper steps cannot make up 2,667 (or 1,800) of them. Further down the ladder only AdamW reaches 1.785, in 2,000 steps (8.15 s), and nothing reaches 1.646 on every seed.
• Energy is the devices'. Compute (100 W static per H100 for the whole step plus 1 pJ per FLOP, illustrative) is the largest column in every row; the modules (9.7 W static from Kundu et al., plus an illustrative 2 W laser) add 7.2–7.8% in the rows above (10.0% with the transform charged at 10 GFLOP/s).
• A shared WAN hurts whoever sends most across it. With other jobs at 0.9 of its capacity, AdamW's step grows 4.075 to 26.291 ms, E′'s 3.726 to 5.442 ms. The closed forms, and the break-evens on slide 13, leave this out.

Source: the simulation's results.md, section 10 (results site, extrapolations). Syncing every H steps assumes the steps to a loss are unchanged, which DiLoCo (Douillard et al., arXiv:2311.08455) does not promise, since it changes the optimizer too. The model-size row scales bytes and compute with the parameter count. Summary: “What the simulation says”.

Rarer syncs shrink the communication the band saves, and with it the band's only advantage: E′ is behind AdamW at every H. At the Llama-3-8B shape the order is unchanged: top-k first, the in-transit variants last, bound by the module.

What Does Not Map at Cluster Scale

And the costs that apply even where it maps: a mask rewritten every step, conversions per value, static power on every accelerator for the whole run, and one engine per data-parallel replica because the masks are replicated.

What to Take Away

For this deck and for the series:

• Transforms are a small share of training, as of inference. A Llama-3-8B-shaped Hyena-2 model spends 0.13–0.18% of a step on work a once-per-step mask could hold (bound 1.002x at most); only speculative block-circulant weights reach 9.53–10.07% (1.105–1.112x).
• Training adds two costs inference does not have. Masks must be rewritten every optimizer step, 0.52–2.10 s on the illustrative device against digital steps of 0.15–2.05 s, and the weight gradient's data-by-data products would need a mask per sequence, minutes per step, so they stay digital.
• In a cluster, scheduling decides it. Only layer-major order amortises the writes, and it fights pipelines and activation memory; data parallelism replicates the masks on every engine. Even fully hidden, a free engine buys at most about a tenth of the GPU time, against static power and an optical subsystem per accelerator.
• Where frequency transforms do help training is the link. DeMo's DCTs are 4.76e-06 of a step's FLOPs and save gradient bytes: a communication result, not a compute one. Its all-gather turns an 8x smaller payload into 2x fewer wire bytes on 8 workers.
• Putting the transform in the link, measured, does not pay. On a 15.8M-parameter GPT, a fixed band in place of top-k ends at 3.124–3.158 against 1.785; the band in a pluggable module never breaks even; summing it at the switch pays only below a few GB/s of WAN, and a GPU all-reduce of the band does better. Across data centres, steps to a loss decide the race, and top-k wins it.
• Next: this is the final deck of Fourier Optics for Training. The measurements behind slides 11–14 are on the Compute in Transit results site and in CiT 04, “The Fixed-Band Experiment” and CiT 05, “Training Transit, Simulated”; the series hub has the analysis, the glossary and all three decks.