AI Infra Interviews logo

A training step follows a critical path through compute, exchange and update

Separate model work, communication, optimizer work and input preparation, then trace their dependencies. Calculate exposed time and training-duration sensitivity without assuming every stage overlaps or that low MFU identifies one cause.

18 MIN

TL;DR: Compute, communication, optimizer work and input preparation contribute to training time through a dependency graph. Some work overlaps; some must wait. Step time is generally neither the maximum of four isolated timings nor their simple sum. Measure the exposed critical path, then use a clearly defined throughput assumption for planning.

Where you are. First lesson of the training course. Suppose the matrix kernels look fast but the job processes fewer tokens than expected. Before changing hardware or parallelism, trace one representative optimizer step and determine where useful progress waits.

Four categories, with dependencies

A forward pass computes predictions from inputs and learned parameters. A loss measures the training objective. Backward computes gradients, which describe how parameter changes affect that loss; the optimizer then uses them to update the parameters. An activation is an intermediate result from the current input. These names describe different work and storage, even when a framework combines them in one training loop.

Work categoryWhat to inspect
Model executionForward, backward and recomputation. Check kernel shape, memory traffic and launch gaps.
CommunicationGradient synchronization, parameter gathers and model-parallel exchanges. Check readiness, link contention and late ranks.
OptimizerGradient processing and state updates. Check memory traffic, offload and update dependencies.
InputReading, decoding, transforming, packing and transfer. Check prefetch, host workers and storage.

For conventional synchronous training, an optimizer update needs the relevant synchronized gradients, and the next forward pass needs updated parameters. Input preparation can often run ahead. Gradient exchange can often begin during backward as buckets become ready. These dependencies limit overlap even when separate hardware resources are available.

An ideal pipeline's steady-state rate may be constrained by its slowest stage, but that does not make the latency of an entire training step equal to max(compute, communication, optimizer, input). Likewise, summing every profiler duration double-counts overlap. Use timestamps and dependencies on the slowest relevant rank.

A timed example

Consider an illustrative step with input ready at time zero. Forward takes 200 ms, backward computation takes 400 ms, gradient exchange runs from 300 to 650 ms, and the optimizer then takes 100 ms. The exchange overlaps backward for 300 ms and exposes a 50 ms tail.

forward + backward = 200 + 400 = 600 ms
exposed exchange tail = 650 - 600 = 50 ms
optimizer = 100 ms
step completes = 600 + 50 + 100 = 750 ms

The end of backward computation is a timeline marker; a synchronized backward call can still be waiting for communication. The communication duration is 350 ms, but adding all of it would predict 1,050 ms and double-count the overlap. Taking only the maximum category, 600 ms, would miss both the tail and update. If input becomes ready 80 ms late and nothing else hides that delay, completion moves to 830 ms.

The update waits for both compute and exchange

Illustrative timestamps. Input is ready at time zero.

Milliseconds from step start

0250500750
Forward0–200
Backward compute200–600
Gradient exchange300–650

Runs beside backward until 600 ms; the final 50 ms are exposed.

Optimizer update650–750
Complete step time is 750 ms. Summing every interval double-counts overlap; taking only their maximum misses dependencies.

Overlap has a budget and a cost

Smaller gradient buckets may launch earlier but increase collective overhead. Larger buckets can improve transfer efficiency while delaying readiness. Concurrent communication can contend with compute for device resources, so apparent overlap is useful only if step time improves. Prefetch depth also consumes host memory and cannot fix a storage system whose sustained rate is insufficient.

A conventional synchronized step often has some final reduction or update work exposed after backward, but its exact duration and overlap depend on the algorithm. Do not assume every training system has the same mandatory tail. Gradient accumulation, sharding and pipeline schedules alter the dependency graph.

Regular device-idle gaps suggest periodic waits, but not a unique cause. Inspect input readiness, host launches, rank synchronization, optimizer phases and checkpoint activity. A synthetic-input experiment can isolate part of the data path; a communication microbenchmark can test the fabric without proving the full job uses it efficiently.

A shorter tail can accompany a slower step

Return to the 750 ms example. Suppose a more aggressive overlap policy makes gradient exchange finish at 630 ms, but concurrent resource use stretches backward from 400 to 480 ms. Forward still ends at 200 ms, so backward now finishes at 680 ms. The optimizer starts there and ends at 780 ms.

old: max(600, 650) + 100 = 750 ms
new: max(680, 630) + 100 = 780 ms
communication tail: 50 ms → 0 ms
complete update: 30 ms slower, or 4% longer
sanity: removing the tail did not remove the new compute delay.

These are hypothetical timestamps, not a claim that overlap always causes contention. They show why “percentage of NCCL time overlapped” cannot be the acceptance metric. Inspect overlapping GEMM durations and the slowest rank's update completion. Keep the policy only when the same useful workload finishes sooner within the memory and numerical constraints.

There is another boundary to define: one optimizer step may contain several microbatches. With DP=4, two sequences per microbatch, eight accumulation passes and 2,048 tokens per sequence, one update consumes 4 × 2 × 8 × 2,048 = 131,072 tokens under the full-token assumption. Timing a single microbatch while dividing by the entire update's tokens exaggerates throughput eightfold. For useful-target throughput, count valid loss-bearing targets. For executed-work or MFU estimates, use the processed-token convention required by the FLOP model; report both when they differ.

The overlap interview exercise provides a small queue simulator: identical communication service times produce different tails when bucket readiness changes. Use it to predict the trace before changing a real reducer's bucket sizes.

Connect useful work to run duration

For a dense transformer approximation, 6 × N × D estimates dominant forward/backward parameter-matrix work for N parameters and D processed training tokens. Attention's sequence-dependent operations, embeddings, architecture details and recomputation require a consistent counting convention. For mixture-of-experts models, total parameters are not simply the active compute count.

Assume a hypothetical 7B dense model, 100B training tokens, 64 devices and 300 TFLOP/s relevant dense peak per device. With an explicitly assumed model FLOPs utilization of 40%:

model work ≈ 6 × 7e9 × 100e9 = 4.2e21 FLOPs
useful rate = 64 × 300e12 × 0.40 = 7.68e15 FLOP/s
active training time ≈ 4.2e21 / 7.68e15 = 546,875 s = 6.33 days

At 20% MFU the estimate doubles to 12.66 days; at 50% it becomes 5.06 days. These are sensitivity assumptions, not typical values guaranteed for large runs. Add downtime, evaluation, checkpointing and restart loss unless the measured rate already includes them. Keep the token and FLOP conventions consistent between numerator and measured MFU.

Low MFU can reflect inefficient compute kernels, communication, input, optimizer overhead, recomputation under the chosen convention, or a bad theoretical denominator. The gap between a compute floor and observed time is not automatically time spent in the other three categories.

Name the token count before comparing throughput

A completion-only training job can process prompt positions while scoring only the completion targets. Pretraining, Fine-Tuning and Distillation shows the shifted targets and mask. For this lesson, the consequence is a measurement choice: model execution sees a different number of positions from the loss's denominator.

Take a hypothetical fixture that processes eight positions in 4 ms, with two scored completion targets. Both rates below are correct if their numerator is named:

One 4 ms training fixture 8 processed positions 2 scored completion targets 1 2 3 4 5 6 7 8 Only positions 6 and 7 scored Processed: 8 / 0.004 = 2,000/s Scored: 2 / 0.004 = 500/s Same step, different numerator Removing loss terms does not remove prompt computation.

Suppose packing removes two padding positions and the measured fixture falls to 3 ms while still scoring those same two targets. The processed-position rate becomes 6 / 0.003 = 2,000/s, unchanged. The scored-target rate rises to 2 / 0.003 ≈ 667/s. That is useful evidence about this fixture, not a guaranteed speedup from packing. Check the actual mask and kernel path: some implementations already avoid padding work, and packed examples need the intended attention boundaries.

Report processed positions, scored targets, completed updates and wall time together when data composition changes. Use the processed-token convention required by the FLOP model when estimating MFU; replacing it with scored targets can make unchanged execution look less efficient. Scored-target throughput helps explain supervision density, but does not by itself establish better learning or time to a quality target.

Choose a change from evidence

Optimize the largest exposed contribution that is practical to reduce. Faster compute can shorten absolute time even in a communication-heavy job, but it may expose more communication and produce a modest speedup. Prefer a bucket or prefetch change when traces support that cause. Reverse it if contention or memory pressure increases total step time.

Do this before moving on

Remove 25 ms from the example's exposed communication tail: completion becomes 725 ms, a 750/725 ≈ 1.034× improvement. Compare that with halving the 100 ms optimizer phase, which gives 700 ms or 1.071×. Then identify which proposed change is actually feasible on the target system before selecting it.

Go deeper

Key takeaways

  • Training time follows dependencies and exposed work, not a universal max or sum rule.
  • The example completes in 750 ms: 600 ms compute, 50 ms exposed exchange and 100 ms update.
  • Overlap can trade latency against collective overhead and resource contention.
  • The hypothetical 7B/100B-token run takes 6.33 active days at the stated 40% MFU assumption.
  • Validate FLOP conventions and trace the cause before interpreting low utilization.

Check yourself

Answer before you look. Recalling it is what makes it stick; recognising it does not.

  1. 1Forward/backward finish at 600 ms, exchange finishes at 650 ms and the optimizer takes 100 ms afterward. When does the step finish?

  2. 2Does low MFU prove communication or input failed to overlap?

  3. 3Which measurement best checks whether smaller gradient buckets help?

  4. 4Packing changes a fixture from 8 processed positions in 4 ms to 6 in 3 ms, with 2 scored targets in both. Did throughput improve?

Sign in to track which lessons you have finished.