Pré-Publication, Document De Travail Année : 2025

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.

Fichier principal
Vignette du fichier
preprint.pdf (1.41 Mo) Télécharger le fichier
Origine Fichiers produits par l'(les) auteur(s)
Licence

Dates et versions

hal-05151601 , version 1 (08-07-2025)
hal-05151601 , version 2 (12-11-2025)

Licence

Identifiants

  • HAL Id : hal-05151601 , version 2

Citer

Adrien Aguila--Multner, Olivier Beaumont, Lionel Eyraud-Dubois, Julia Gusak. Optimized Forward-Backward Rematerialization for Memory-Efficient Pipeline Parallel Training. 2025. ⟨hal-05151601v2⟩
540 Consultations
537 Téléchargements

Partager

  • More