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.
Domaines
Intelligence artificielle [cs.AI]
Origine : Fichiers produits par l'(les) auteur(s)