GPU Memory Prediction for Multimodal Model Training

TL;DR

This paper introduces a framework for accurately predicting GPU memory usage in multimodal model training, addressing the limitations of previous unimodal-focused methods and reducing out-of-memory errors.

Contribution

We propose a novel decomposition-based framework that estimates peak GPU memory for multimodal models, improving prediction accuracy over prior approaches.

Findings

Achieves ~8.7% average MAPE in memory prediction

Addresses the gap in GPU memory prediction for multimodal architectures

Reduces out-of-memory errors during training

Abstract

As deep learning models in agentic AI systems grow in scale and complexity, GPU memory requirements increase and often exceed the available GPU memory capacity, so that out-of-memory (OoM) errors occur. It is well known that OoM interrupts the whole training itself and wastes substantial computational resources. Therefore, to prevent OoM, accurate prediction of GPU memory usage is essential. However, previous studies focus only on unimodal architectures and fail to generalize to multimodal models, even though the multimodal models are a common choice in agentic AI systems. To address this limitation, we propose a framework that predicts the peak GPU memory usage by analyzing the model architecture and training behavior of multimodal models. Specifically, the framework decomposes the multimodal model into its constituent layers and applies factorization to estimate the memory usage of each layer. Our evaluation shows that our framework achieves high prediction accuracy of ~8.7% average MAPE.