H-Rockmate: Hierarchical Approach for Efficient Re-materialization of Large Neural Networks - Archive ouverte HAL Accéder directement au contenu
Pré-Publication, Document De Travail Année : 2023

H-Rockmate: Hierarchical Approach for Efficient Re-materialization of Large Neural Networks

Résumé

Training modern neural networks poses a significant memory challenge, as storing intermediate results during the forward and backward passes demands substantial memory resources. To address this issue while maintaining model accuracy, re-materialization techniques have been introduced to recompute selected intermediate results rather than storing them, thereby adhering to peak memory constraints. The main algorithmic problem is to compute a re-materialization schedule that minimizes the computational overhead within a given memory budget. Our H-Rockmate framework builds upon existing Rockmate solution and overcomes its limitation to work with sequential block structures by proposing a hierarchical approach. The framework performs an automatic decomposition of the data-flow graph into a hierarchy of small-scale subgraphs, and finds a re-materialization schedule for the whole graph by recursively solving optimization problems for each subgraph. H-Rockmate allows users to transform their PyTorch models into nn.Modules that execute forward and backward passes efficiently within the specified memory budget. This framework can handle neural networks with diverse data-flow graph structures, including U-Nets and encoder-decoder Transformers. H-Rockmate consistently outperforms existing re-materialization approaches both in terms of average training iteration time and peak memory trade-offs, demonstrating superior memory efficiency in training modern neural networks.
Fichier principal
Vignette du fichier
main.pdf (932.55 Ko) Télécharger le fichier
Origine : Fichiers produits par l'(les) auteur(s)

Dates et versions

hal-04403844 , version 1 (18-01-2024)

Identifiants

  • HAL Id : hal-04403844 , version 1

Citer

Julia Gusak, Xunyi Zhao, Théotime Le Hellard, Zhe Li, Lionel Eyraud-Dubois, et al.. H-Rockmate: Hierarchical Approach for Efficient Re-materialization of Large Neural Networks. 2023. ⟨hal-04403844⟩
28 Consultations
21 Téléchargements

Partager

Gmail Facebook X LinkedIn More