Pipeline stages and partition costs

A stage owns operations, parameters, buffers, and a contract for tensors crossing its boundaries. Module containment alone does not describe execution: a forward function can reuse a module, include functional operations, or carry a residual value across an apparent cut. The partition must route every required tensor with the correct identity, shape, dtype, and layout. PyTorch 2.8 distinguishes manual stage construction from graph-based splitting; a particular stage interface can impose a simpler sequential topology than an arbitrary model graph.

For a simple chain, stage s consumes the activation sent by stage s minus one, computes its output, and sends that output onward. Backward waits for the corresponding downstream activation gradient and uses saved or recomputed forward state to produce upstream gradients and parameter gradients. Attach a microbatch identifier to both directions. Receiving the correct-sized tensor for the wrong microbatch is still incorrect. Tied parameters, nonlocal residuals, and stateful buffers require additional explicit coordination rather than an assumption that every boundary is one interchangeable tensor.

Profile forward and backward separately at the intended microbatch shape and precision. Include activation-checkpointing choices, kernel efficiency, and any tensor-parallel communication inside a stage. A partition with balanced forward work can have an imbalanced backward or a large boundary tensor that crosses a slower link. Memory feasibility includes weights, optimizer state, communication buffers, temporary workspaces, and the maximum simultaneously live activation state. Parameter count alone misses several of those terms and does not guarantee a stage fits.

For a fixed partition, the slowest serial stage resource constrains steady-state admission, but dependencies, fill, drain, and contention can make actual execution slower. A bandwidth estimate for a boundary is a lower bound under stated transfer assumptions, not necessarily fully exposed time. Placement also matters: the same graph cut can use an intra-node link in one mapping and an inter-node path in another. Record the physical rank placement alongside the logical partition so a reproducible cost estimate names the communication path it assumes.

Choose whether the objective is iteration time, tokens per second, peak memory, or a constrained combination. Moving one layer can improve a compute balance while creating an infeasible activation boundary or duplicating a tied state. Evaluate candidate cuts against representative sequence lengths and avoid using one favorable batch as a universal partition. Once a small valid configuration runs, compare numerical losses and gradients with an unsplit reference before scaling it to additional ranks. This separates an execution mistake from an optimization hypothesis.

The queue lab isolates only a bottleneck-rate consequence. Reset, set arrival rate to 4 and service rate to 3, predict the backlog after one second, then Step twice. Compare a fresh Reset with both arrival rate at 2 and service rate at 3 before two Steps. The outcomes are five and three abstract units from an initial four. This fluid model does not partition layers, execute forward or backward passes, or account for communication and stage memory.