Stage imbalance and repartitioning

A stage operation waits for its own input and for the stage’s device resource. For a forward-only chain with one serial operation per stage, these are the previous stage’s completion for the same microbatch and the same stage’s completion for the previous microbatch. The maximum of those readiness times determines its start. Backward, communication, recomputation, shared links, and multiple chunks add more dependencies; they cannot be handled by replacing every stage with an average measured number.

Profile representative inputs because attention costs can change with sequence length, and output projections or embeddings can give edge stages different work. Separate genuine compute from time waiting for communication or another rank. A stage that appears slow in a synchronized measurement may be a victim of an upstream bottleneck rather than its cause. Build a trace that identifies operation readiness, device occupancy, and transfer completion before deciding which layer to move. End-to-end iteration time is the objective, not equality of one selected profiler column.

Moving operations changes more than arithmetic balance. A new boundary can increase payload size, shift a transfer to a slower physical path, or extend the lifetime of a residual activation. Extra retained memory can force a different microbatch size or recomputation policy. Tied parameters and optimizer shards also constrain legal movement. Reprofile a candidate after its graph and memory contracts are made valid. A compute-only improvement is useful evidence for a hypothesis, not a guarantee that the whole distributed job becomes faster.

Use critical-path and resource-work lower bounds as checks on a predicted makespan. A dependency path cannot complete faster than the sum of its required durations, and a device cannot perform more serial work than fits in elapsed time under the model. These bounds can expose an impossible speedup claim, but matching a bound does not certify a real schedule. Contention, incomplete overlap, or changed kernels can still add time. State which operations and resources are represented in each comparison.

A practical repartition often drains the intended update, saves a consistent checkpoint, rebuilds the partition and rank mapping, and restores parameters, optimizer moments, scheduler state, buffers, and data progress. Preserve random-state and update-count semantics as required by the experiment. Moving only weights can reset optimizer history or duplicate data even if outputs look plausible. Revalidate tensor contracts and compare a small subsequent update with the unrepartitioned reference. The time spent saving, moving, rebuilding, and reloading belongs in the cost decision.

The queue lab illustrates sustained imbalance only. Reset, set arrival rate to 4 and service rate to 3, predict backlog at one second, then Step twice: five units. Compare a fresh Reset with both arrival rate at 4 and service rate at 5 before two Steps: three units. The initial backlog is four. The instrument does not identify a critical path, move layers, rebuild checkpoints, or simulate an unequal-stage timeline; use the recurrence and an actual trace for those questions.