GPU Memory Prediction for Multimodal Model Training

Jinwoo Jeong; Minchul Kang; Younghun Go; Changyong Shin; Hyunho Lee; Junho Yoon; Gyeongsik Yang; Chuck Yoo

Paper summary

Count memory according to what is being trained

A multimodal model’s memory needs depend on module size and which modules are updated.

Vision-language models combine modules with different roles and training behavior. Frozen layers do not need the same gradient and optimizer storage as trainable layers. A memory predictor must therefore distinguish the model’s components and the current training configuration.

Decompose modules into memory factors

The framework parses vision, projection and language modules into layers. For each layer it estimates parameter, gradient, optimizer-state and activation memory, accounting for whether that layer is being trained.

Per-factor estimates are aggregated into a peak GPU-memory prediction for the supplied hyperparameters. The estimate accounts for the storage required by frozen and updated modules.

Module and layer descriptions feed factor-specific memory estimates, which combine into a peak-memory prediction.

Evaluate two LLaVA training configurations

The evaluation uses LLaVA-1.5 7B with DeepSpeed ZeRO-2 on eight H100 80 GB GPUs. Data parallelism varies from 1 to 8. For sequence length 1,024 and micro-batch size 16, average prediction error is 13% MAPE; for sequence length 2,048 and micro-batch size 8, it is 8.7%.

Peak-memory prediction error for LLaVA-1.5 7B with ZeRO-2; data parallelism varies from 1 to 8.
Sequence / micro-batch	MAPE
1,024 / 16	13%
2,048 / 8	8.7%

Key takeaways

Multimodal memory use follows the model’s internal structure. Treating all modules alike obscures their different memory contributions.
Trainability changes what must be stored. Frozen and trainable layers require different gradient and optimizer state, making training configuration central to prediction.