Optimized Forward-Backward Rematerialization for Memory-Efficient Pipeline Parallel Training
Résumé
Pipeline parallelism is a key technique for scaling deep network training across multiple devices. Recent works have significantly reduced pipeline idle time by improving scheduling efficiency. Decoupling the computation of gradients with respect to weights and activations led to the development of schedules with almost no idle time. However, these methods still require substantial memory, limiting their applicability on resource-constrained hardware. Our first contribution is to introduce recomputation to the backward pass, extending rematerialization beyond the forward pass. This enables executing schedules with decoupled gradient computations under much tighter memory constraints. Our second contribution is a unified optimization approach that, given a model and hardware memory constraints, formulates and solves an Integer Linear Programming (ILP) problem to determine the optimal per-microbatch, per-GPU rematerialization strategy for a given schedule, applicable to both one-wave and multi-wave pipeline schedules. Our third contribution shows that, as device memory constraints vary, the relative advantages of different pipeline schedules also change in the presence of rematerialization. We provide corresponding insights and a PyTorch framework that enables finding and executing the optimal combination of pipeline scheduling and rematerialization strategies. Experiments demonstrate the effectiveness of all three contributions, showing that our approach enables efficient training of larger models under tight memory budgets, adapts optimally to varying memory capacities, and reduces recomputation overhead compared to existing recomputation solutions.
| Origine | Fichiers produits par l'(les) auteur(s) |
|---|---|
| Licence |