{"type":"ScholarlyArticle","title":"LLMem: Estimating GPU Memory Usage for Fine-Tuning Pre-Trained LLMs","authors":[{"name":"Taeho Kim","slug":"taeho-kim","given":"Taeho","family":"Kim","affiliation":"University of Colorado Boulder"},{"name":"Yanming Wang","slug":"yanming-wang","given":"Yanming","family":"Wang","affiliation":"Amazon Web Services"},{"name":"Vatshank Chaturvedi","slug":"vatshank-chaturvedi","given":"Vatshank","family":"Chaturvedi","affiliation":"Amazon Web Services"},{"name":"Lokesh Gupta","slug":"lokesh-gupta","given":"Lokesh","family":"Gupta","affiliation":"Amazon Web Services"},{"name":"Seyeon Kim","slug":"seyeon-kim","given":"Seyeon","family":"Kim","affiliation":"University of Colorado Boulder"},{"name":"Yongin Kwon","slug":"yongin-kwon","given":"Yongin","family":"Kwon","affiliation":"Electronics and Telecommunications Research Institute"},{"name":"Sangtae Ha","slug":"sangtae-ha","given":"Sangtae","family":"Ha","affiliation":"University of Colorado Boulder"}],"abstract":"Fine-tuning pre-trained large language models (LLMs) with limited hardware presents challenges due to GPU memory constraints. Various distributed fine-tuning methods have been proposed to alleviate memory constraints on GPU. However, determining the most effective method for achieving rapid fine-tuning while preventing GPU out-of-memory issues in a given environment remains unclear. To address this challenge, we introduce LLMem, a solution that estimates the GPU memory consumption when applying distributed fine-tuning methods across multiple GPUs and identifies the optimal method. We conduct GPU memory usage estimation prior to fine-tuning, leveraging the fundamental structure of transformer-based decoder models and the memory usage distribution of each method. Experimental results show that LLMem accurately estimates peak GPU memory usage on a single GPU, with an error rate of up to 1.6%. Additionally, it shows an average error rate of 3.0% when applying distributed fine-tuning methods to LLMs with more than a billion parameters on multi-GPU setups.","doi":"10.24963/ijcai.2024/699","url":"https://ijcai.org/proceedings/2024/699","pdf":"https://ijcai.org/proceedings/2024/0699.pdf","markdown":"https://ijcai.org/proceedings/2024/699.md","bibtex":"https://ijcai.org/proceedings/2024/bibtex/699","volume":{"slug":"2024","title":"Proceedings of the Thirty-Third International Joint Conference on Artificial Intelligence","short_title":"IJCAI 2024","url":"https://ijcai.org/proceedings/2024/","year":2024,"location":"Jeju, Korea","dates":"3-9 August 2024","editors":"Kate Larson","isbn":"978-1-956792-04-1","publisher":"International Joint Conferences on Artificial Intelligence Organization","doi_prefix":"10.24963/ijcai.2024"},"section":"Natural Language Processing","track":"Main Track","number":699,"pages":{"first":6324,"last":6332},"keywords":["Machine Learning: ML: Deep learning architectures","Natural Language Processing: NLP: Language models"],"videos":[],"references":[{"position":1,"raw":"[Albahar et al., 2022] Hadeel Albahar, Shruti Dongare, Yan- lin Du, Nannan Zhao, Arnab K Paul, and Ali R Butt. Schedtune: A heterogeneity-aware gpu scheduler for deep learning. In 2022 22nd IEEE International Symposium on Cluster, Cloud and Internet Computing (CCGrid), pages 695-705. IEEE, 2022.","title":"Schedtune: A heterogeneity-aware gpu scheduler for deep learning","authors":"Shruti Hadeel Albahar, Yanlin Dongare, Nannan Du, Arnab K Zhao, Ali R Paul, Butt","year":"2022","venue":"22nd IEEE International Symposium on Cluster, Cloud and Internet Computing (CCGrid)","doi":null,"url":null,"internal_url":null},{"position":2,"raw":"[Allal et al., 2023] Loubna Ben Allal, Raymond Li, De- nis Kocetkov, Chenghao Mou, Christopher Akiki, Car- los Munoz Ferrandis, Niklas Muennighoff, Mayank Mishra, Alex Gu, Manan Dey, et al. Santacoder: don't reach for the stars! arXiv preprint arXiv:2301.03988, 2023.","title":"Santacoder: don't reach for the stars!","authors":"Loubna Ben Allal, Raymond Li, Denis Kocetkov, Chenghao Mou, Christopher Akiki, Carlos Munoz Ferrandis, Niklas Muennighoff, Mayank Mishra, Alex Gu, Manan Dey","year":"2023","venue":null,"doi":null,"url":null,"internal_url":null},{"position":3,"raw":"[Bian et al., 2021] Zhengda Bian, Qifan Xu, Boxiang Wang, and Yang You. Maximizing parallelism in distributed training for huge neural networks. arXiv preprint arXiv:2105.14450, 2021.","title":"Maximizing parallelism in distributed training for huge neural networks","authors":"Zhengda Bian, Qifan Xu, Boxiang Wang, Yang You","year":"2021","venue":null,"doi":null,"url":null,"internal_url":null},{"position":4,"raw":"[Black et al., 2021] Sid Black, Gao Leo, Phil Wang, Con- nor Leahy, and Stella Biderman. Gpt-neo: Large scale autoregressive language modeling with mesh-tensorflow. https://doi.org/10.5281/zenodo.5297715, 2021. Accessed: 2024-01-15.","title":"Gpt-neo: Large scale autoregressive language modeling with mesh-tensorflow","authors":"Sid Black, Gao Leo, Phil Wang, Connor Leahy, Stella Biderman","year":"2024","venue":null,"doi":"10.5281/zenodo.5297715","url":"https://doi.org/10.5281/zenodo.5297715,2021","internal_url":null},{"position":5,"raw":"[CloudLab, 2024] CloudLab. Cloudlab. https://www. cloudlab.us/, 2024. Accessed: 2024-01-15.","title":"CloudLab. Cloudlab","authors":"Cloudlab","year":"2024","venue":null,"doi":null,"url":"https://www.cloudlab.us/,2024","internal_url":null},{"position":6,"raw":"[Devlin et al., 2018] Jacob Devlin, Ming-Wei Chang, Ken- ton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understand- ing. arXiv preprint arXiv:1810.04805, 2018.","title":"Bert: Pre-training of deep bidirectional transformers for language understanding","authors":"Jacob Devlin, Ming-Wei Chang, Kenton Lee, Kristina Toutanova","year":"2018","venue":null,"doi":null,"url":null,"internal_url":null},{"position":7,"raw":"[Fang et al., 2022] Jiarui Fang, Zilin Zhu, Shenggui Li, Hui Su, Yang Yu, Jie Zhou, and Yang You. Parallel training of pre-trained models via chunk-based dynamic memory management. IEEE Transactions on Parallel and Dis- tributed Systems, 34(1):304-315, 2022.","title":"Parallel training of pre-trained models via chunk-based dynamic memory management","authors":"Jiarui Fang, Zilin Zhu, Shenggui Li, Hui Su, Yang Yu, Jie Zhou, Yang You","year":"2022","venue":"IEEE Transactions on Parallel and Distributed Systems","doi":null,"url":null,"internal_url":null},{"position":8,"raw":"[Gao et al., 2020] Yanjie Gao, Yu Liu, Hongyu Zhang, Zhengxian Li, Yonghao Zhu, Haoxiang Lin, and Mao Yang. Estimating gpu memory consumption of deep learn- ing models. In Proceedings of the 28th ACM Joint Meeting on European Software Engineering Conference and Sym- posium on the Foundations of Software Engineering, pages 1342-1352, 2020.","title":"Estimating gpu memory consumption of deep learning models","authors":"Yanjie Gao, Yu Liu, Hongyu Zhang, Zhengxian Li, Yonghao Zhu, Haoxiang Lin, Mao Yang","year":"2020","venue":"Proceedings of the 28th ACM Joint Meeting on European Software Engineering Conference and Symposium on the Foundations of Software Engineering","doi":null,"url":null,"internal_url":null},{"position":9,"raw":"[Li et al., 2023] Shenggui Li, Hongxin Liu, Zhengda Bian, Jiarui Fang, Haichen Huang, Yuliang Liu, Boxiang Wang, and Yang You. Colossal-ai: A unified deep learning sys- tem for large-scale parallel training. In Proceedings of the 52nd International Conference on Parallel Processing, pages 766-775, 2023.","title":"Colossal-ai: A unified deep learning system for large-scale parallel training","authors":"Shenggui Li, Hongxin Liu, Zhengda Bian, Jiarui Fang, Haichen Huang, Yuliang Liu, Boxiang Wang, Yang You","year":"2023","venue":"Proceedings of the 52nd International Conference on Parallel Processing","doi":null,"url":null,"internal_url":null},{"position":10,"raw":"[Luo et al., 2022] Renqian Luo, Liai Sun, Yingce Xia, Tao Qin, Sheng Zhang, Hoifung Poon, and Tie-Yan Liu. Biogpt: generative pre-trained transformer for biomedical text generation and mining. Briefings in Bioinformatics, 23(6):bbac409, 2022.","title":"Biogpt: generative pre-trained transformer for biomedical text generation and mining","authors":"Luo","year":"2022","venue":"Briefings in Bioinformatics","doi":null,"url":null,"internal_url":null},{"position":11,"raw":"[Nie et al., 2022] Xiaonan Nie, Xupeng Miao, Zhi Yang, and Bin Cui. Tsplit: Fine-grained gpu memory management for efficient dnn training via tensor splitting. In 2022 IEEE 38th International Conference on Data Engineering (ICDE), pages 2615-2628. IEEE, 2022.","title":"Tsplit: Fine-grained gpu memory management for efficient dnn training via tensor splitting","authors":"Xiaonan Nie, Xupeng Miao, Zhi Yang, Bin Cui","year":"2022","venue":"IEEE 38th International Conference on Data Engineering (ICDE)","doi":null,"url":null,"internal_url":null},{"position":12,"raw":"[Nijkamp et al., 2022] Erik Nijkamp, Bo Pang, Hiroaki Hayashi, Lifu Tu, Huan Wang, Yingbo Zhou, Silvio Savarese, and Caiming Xiong. Codegen: An open large language model for code with multi-turn program synthe- sis. arXiv preprint arXiv:2203.13474, 2022.","title":"Codegen: An open large language model for code with multi-turn program synthesis","authors":"Erik Nijkamp, Bo Pang, Hiroaki Hayashi, Lifu Tu, Huan Wang, Yingbo Zhou, Silvio Savarese, Caiming Xiong","year":"2022","venue":null,"doi":null,"url":null,"internal_url":null},{"position":13,"raw":"[Rajbhandari et al., 2020] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory op- timizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1- 16. IEEE, 2020.","title":"Zero: Memory optimizations toward training trillion parameter models","authors":"Rajbhandari","year":"2020","venue":"SC20: International Conference for High Performance Computing, Networking, Storage and Analysis","doi":null,"url":null,"internal_url":null},{"position":14,"raw":"[Ren et al., 2021] Jie Ren, Samyam Rajbhandari, Reza Yaz- dani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Min- jia Zhang, Dong Li, and Yuxiong He. {ZeRO-Offload}: Democratizing {Billion-Scale} model training. In 2021 USENIX Annual Technical Conference (USENIX ATC 21), pages 551-564, 2021.","title":"{ZeRO-Offload}: Democratizing {Billion-Scale} model training","authors":"Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, Yuxiong He","year":"2021","venue":"USENIX Annual Technical Conference (USENIX ATC 21)","doi":null,"url":null,"internal_url":null},{"position":15,"raw":"[Shoeybi et al., 2019] Mohammad Shoeybi, Mostofa Pat- wary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-lm: Training multi-billion pa- rameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019.","title":"Megatron-lm: Training multi-billion parameter language models using model parallelism","authors":"Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick Legresley, Jared Casper, Bryan Catanzaro","year":"2019","venue":null,"doi":null,"url":null,"internal_url":null},{"position":16,"raw":"[Taori et al., 2023] Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford alpaca: An instruction-following llama model. https://github.com/ tatsu-lab/stanford alpaca, 2023.","title":"Stanford alpaca: An instruction-following llama model","authors":"Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, Tatsunori B Hashimoto","year":"2023","venue":null,"doi":null,"url":"https://github.com/tatsu-lab/stanfordalpaca","internal_url":null},{"position":17,"raw":"[Touvron et al., 2023] Hugo Touvron, Thibaut Lavril, Gau- tier Izacard, Xavier Martinet, Marie-Anne Lachaux, Tim- othée Lacroix, Baptiste Rozière, Naman Goyal, Eric Ham- bro, Faisal Azhar, et al. Llama: Open and efficient founda- tion language models. arXiv preprint arXiv:2302.13971, 2023.","title":"Hugo Touvron, Thibaut Lavril, Gautier Izacard","authors":"Touvron","year":"2023","venue":"Llama: Open and efficient foundation language models","doi":null,"url":null,"internal_url":null},{"position":18,"raw":"[Vaswani et al., 2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.","title":"Attention is all you need","authors":"Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, Illia Polosukhin","year":"2017","venue":"Advances in neural information processing systems","doi":null,"url":null,"internal_url":null},{"position":19,"raw":"[Wang et al., 2018] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bow- man. Glue: A multi-task benchmark and analysis plat- form for natural language understanding. arXiv preprint arXiv:1804.07461, 2018.","title":"Glue: A multi-task benchmark and analysis platform for natural language understanding","authors":"Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, Samuel R Bowman","year":"2018","venue":null,"doi":null,"url":null,"internal_url":null},{"position":20,"raw":"[Wang et al., 2021] Boxiang Wang, Qifan Xu, Zhengda Bian, and Yang You. 2.5-dimensional distributed model training. arXiv e-prints, pages arXiv-2105, 2021.","title":"2.5-dimensional distributed model training","authors":"Wang","year":"2021","venue":null,"doi":null,"url":null,"internal_url":null},{"position":21,"raw":"[Workshop et al., 2022] BigScience Workshop, Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Ilić, Daniel Hesslow, Roman Castagné, Alexan- dra Sasha Luccioni, Franc ¸ois Yvon, et al. Bloom: A 176b-parameter open-access multilingual language model. arXiv preprint arXiv:2211.05100, 2022.","title":"BigScience Workshop","authors":"Workshop","year":"2022","venue":"Bloom: A 176b-parameter open-access multilingual language model","doi":null,"url":null,"internal_url":null},{"position":22,"raw":"[Xu et al., 2021] Qifan Xu, Shenggui Li, Chaoyu Gong, and Yang You. An efficient 2d method for training super-large deep learning models. arXiv e-prints, pages arXiv-2104, 2021.","title":"An efficient 2d method for training super-large deep learning models","authors":"Xu","year":"2021","venue":null,"doi":null,"url":null,"internal_url":null},{"position":23,"raw":"[Zhang et al., 2022] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068, 2022.","title":"Opt: Open pre-trained transformer language models","authors":"Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin","year":"2022","venue":null,"doi":null,"url":null,"internal_url":null}],"license_note":"Copyright International Joint Conferences on Artificial Intelligence."}