One-forward-one-backward schedules

Consider a two-stage chain and label microbatches zero through three. Stage zero can run forward zero, then forward one while stage one runs forward zero. Stage zero cannot yet backward zero until stage one computes that microbatch’s backward. The local steady-state order then alternates another forward with an earlier ready backward. Stage one, having no downstream stage, can start its backward soon after its loss-producing forward. Different stage positions therefore have different warmup and cooldown lengths.

A valid schedule tracks activation ownership separately from whether a device is busy. The state produced by a local forward remains needed until its corresponding backward consumes or recomputes it. It cannot be recycled because another microbatch’s backward has finished. Include send buffers and recomputation inputs in the actual memory accounting; a nominal saved-activation count omits them. A trace with explicit microbatch identifiers makes these lifetime errors visible before a large implementation turns them into intermittent numeric corruption.

In a standard noninterleaved 1F1B flush schedule, early stages maintain a bounded number of in-flight microbatches determined by their position and pipeline depth. By contrast, an all-forward schedule can retain M microbatches before backward begins. For sufficiently large M, this changes saved-activation memory substantially. The benefit comes from consuming old forward state earlier, not from removing parameters or optimizer state. Stage-specific peaks matter because the first and last stages need not retain the same number of activations.

Under the balanced compute-only assumptions used for the ordinary schedule comparison, the noninterleaved flush 1F1B schedule has the same modeled fill/drain bubble as GPipe. Its principal improvement in that comparison is lower activation residency. Real timing can still change through memory pressure, overlap, or kernel effects, but those require evidence. Keeping weights fixed until the flush also matters: the schedule’s memory behavior should not be confused with the asynchronous weight-stashing algorithm merely because both use the phrase one-forward-one-backward.

Warmup and cooldown are not optional corner cases. When M is smaller than the pipeline depth, the steady alternating portion can be short or absent. The schedule must still execute every required backward exactly once and delay the optimizer until the intended accumulation boundary. Test the last microbatch as carefully as the first; a missing cooldown can leave correct-looking outputs but incomplete gradients. Use a deterministic small model to compare all parameter gradients against the same synchronous unsplit update.

The queue lab supplies only an outstanding-work analogy. Reset, set arrival rate to 2 and service rate to 3, predict the backlog after one second, then Step twice: three units remain. Compare a fresh Reset with both arrival rate at 0 and service rate at 4 before two Steps: the initial four drain. This fluid instrument does not know a microbatch identifier, stage index, saved tensor, warmup, or backward readiness, so derive the actual activation bound from a hand trace.