H-Rockmate: Hierarchical Approach for Efficient Re-materialization of Large Neural Networks - Inria - Institut national de recherche en sciences et technologies du numérique Access content directly
Preprints, Working Papers, ... Year : 2023

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

Abstract

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
Origin Files produced by the author(s)

Dates and versions

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

Identifiers

  • HAL Id : hal-04403844 , version 1

Cite

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⟩
51 View
39 Download

Share

Gmail Mastodon Facebook X LinkedIn More