Elamur Research⌕
☰
home / papers / 2024 / LLM Pruning and Distillation in Practice: The Minitron Approach

LLM Pruning and Distillation in Practice: The Minitron Approach

Paper · 2024 · Sharath Turuvekere Sreenivas, Saurav Muralidharan, Raviraj Joshi et al.

Summary

Minitron-style compression trains only the largest model in a family and derives smaller ones by structured pruning followed by knowledge distillation. It normally assumes access to the original pretraining corpus, which many frontier LLMs, including open ones, do not release. This paper adapts the recipe to that setting. Before pruning, a "teacher correction" step briefly fine-tunes the unpruned model (about 100B tokens) on the authors' own distillation data, so the teacher fits the new token distribution. Saliency is then estimated from activations on a small calibration set. For depth pruning, a new criterion uses downstream accuracy (Winogrande) to choose blocks of contiguous layers. The pruned student is retrained with logit-only forward KL distillation. Applied to Mistral NeMo 12B and Llama 3.1 8B, the method yields MN-Minitron-8B and width- and depth-pruned Llama-3.1-Minitron-4B models, and the instruct variants were aligned with NeMo Aligner. MN-Minitron-8B leads similarly sized open models on common LM Evaluation Harness benchmarks while using far less training data than training from scratch. The base weights are released openly on Hugging Face.

Key points

  • Released models: Mistral-NeMo-Minitron-8B-Base, Llama-3.1-Minitron-4B-Width-Base and Llama-3.1-Minitron-4B-Depth-Base, published on Hugging Face under a permissive license.
  • Training efficiency: distillation used under 400B tokens, about 380B for MN-Minitron-8B and 94B for the 4B models. The latter figure is roughly 150× fewer than the 15T tokens behind Llama 3.1 8B.
  • Effect of teacher correction: it lowers LM validation loss by over 6%, and it works equally well when run in parallel with distillation.
  • MN-Minitron-8B results: 69.5 on MMLU (5-shot). It beats its 12B teacher on GSM8k (55.7 to 58.5) and HumanEval (23.8 to 36.2).
  • Width vs. depth pruning for Llama 3.1:
    • The width-pruned 4B model is more accurate: MMLU 60.5 vs. 58.7, and GSM8k 41.2 vs. 16.8 for the base models.
    • The depth-pruned model is faster with TensorRT-LLM (FP8, H100): about 2.7× throughput over Llama 3.1 8B, versus about 1.8× for the width variant.
    • For width pruning, keeping all attention heads and pruning MLP and embedding channels worked best.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Fields, topics and concepts

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 worksComputer Vision and Pattern Recognition · 2,869 works

🧭 Topics 2Research subjects and languages the work is about.Language models & pretraining · 3,320 worksEfficient models & compression · 242 works

Cite

Cite as Sharath Turuvekere Sreenivas, Saurav Muralidharan, Raviraj Joshi, Marcin Chochowski, Ameya Sunil Mahabaleshwarkar, Gerald Shen, Jiaqi Zeng, Zijia Chen, Yoshi Suhara, Shizhe Diao, Chenhan Yu, Wei-Chun Chen, Hayley Ross, Oluwatobi Olabiyi, Ashwath Aithal, Oleksii Kuchaiev, Daniel Korzekwa, Pavlo Molchanov, Mostofa Patwary, Mohammad Shoeybi, Jan Kautz and Bryan Catanzaro. 2024. LLM Pruning and Distillation in Practice: The Minitron Approach. arXiv:2408.11796.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2408.11796 · arXiv 2408.11796

@misc{sreenivas2024llm,
  title = {{LLM Pruning and Distillation in Practice: The Minitron Approach}},
  author = {Sharath Turuvekere Sreenivas, Saurav Muralidharan, Raviraj Joshi, Marcin Chochowski, Ameya Sunil Mahabaleshwarkar, Gerald Shen, Jiaqi Zeng, Zijia Chen, Yoshi Suhara, Shizhe Diao, Chenhan Yu, Wei-Chun Chen, Hayley Ross, Oluwatobi Olabiyi, Ashwath Aithal, Oleksii Kuchaiev, Daniel Korzekwa, Pavlo Molchanov, Mostofa Patwary, Mohammad Shoeybi, Jan Kautz and Bryan Catanzaro},
  year = {2024},
  doi = {10.48550/arxiv.2408.11796},
  eprint = {2408.11796},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2408.11796}
}

References

34 references, 18 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Abhimanyu Dubey and Abhinav Jauhri et al. The Llama 3 Herd of Models. arXiv 2407.21783, 2024. Figure 10 \| TensorRT-LLM FP8 throughput comparison for the L l a m a 3 . 1 - M i n i t ro n - 4 B models with the Llama 3.1 8B model w.r.t. increasing input and output sequence lengths. \#\# Acknowledgments T…
  2. Saurav Muralidharan, Sharath Turuvekere Sreenivas, Raviraj Joshi, Marcin Chochowski, Mostofa Patwary, Mohammad Shoeybi, Bryan Catanzaro, Jan Kautz, and Pavlo Molchanov. Compact language models via pruning and knowledge distillation. arXiv preprint arXiv:2407.14679, 2024. find
  3. Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the Knowledge in a Neural Network. arXiv preprint arXiv:1503.02531, 2015.
  4. Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, Pouya Tafti, Léonard Hussenot, Pier Giuseppe Sessa, Aakanksha Chowdhery, Adam Roberts, Aditya Barua, Alex Botev, Alex Castro-Ros, Ambrose Sl… find
  5. Mistral https://mistral.ai/news/mistral-nemo, Accessed: 2024. nemo. 2024. ↗
  6. Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. Sheared llama: Accelerating language model pre-training via structured pruning. In The Twelfth International Conference on Learning Representations, 2023. find
  7. Saleh Ashkboos, Maximilian L Croci, Marcelo Gennari do Nascimento, Torsten Hoefler, and James Hensman. Slicegpt: Compress large language models by deleting rows and columns. In The Twelfth International Conference on Learning Representations, 2023. find
  8. Xin Men, Mingyu Xu, Qingyu Zhang, Bingning Wang, Hongyu Lin, Yaojie Lu, Xianpei Han, and Weipeng Chen. ShortGPT: Layers in Large Language Models are More Redundant Than You Expect, 2024.
  9. Bo-Kyeong Kim, Geonmin Kim, Tae-Ho Kim, Thibault Castells, Shinkook Choi, Junho Shin, and Hyoung-Kyu Song. Shortened LLaMA: A simple depth pruning for large language models. In ICLR 2024 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2024. find
  10. Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. WinoGrande: An adversarial winograd schema challenge at scale. Commun. ACM, 64(9), 2021. find
  11. Andrey Gromov, Kushal Tirumala, Hassan Shapourian, Paolo Glorioso, and Daniel A. Roberts. The unreasonable ineffectiveness of the deeper layers. 2024.
  12. Solomon Kullback and Richard A. Leibler. On information and sufficiency. Annals of Mathematical Statistics, 22(1):79–86, 1951. find
  13. Jupinder Parmar, Shrimai Prabhumoye, Joseph Jennings, Mostofa Patwary, Sandeep Subramanian, Dan Su, Chen Zhu, Deepak Narayanan, Aastha Jhunjhunwala, Ayush Dattagupta, Vibhu Jawa, Jiwei Liu, Ameya Mahabaleshwarkar, Osvald Nitski, Annika Brundyn, James Maki, Miguel Martinez, Jiaxuan You, John Kamalu,… — Nemotron-4 15B Technical Report
  14. Jupinder Parmar, Sanjev Satheesh, Mostofa Patwary, Mohammad Shoeybi, and Bryan Catanzaro. Reuse, don’t retrain: A recipe for continued pretraining of language models, 2024. — Reuse, Don't Retrain: A Recipe for Continued Pretraining of Language Models
  15. Edward J Hu, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2021.
  16. Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016.
  17. Gerald Shen, Zhilin Wang, Olivier Delalleau, Jiaqi Zeng, Yi Dong, Daniel Egert, Shengyang Sun, Jimmy Zhang, Sahil Jain, Ali Taghibakhshi, Markel Sanz Ausin, Ashwath Aithal, and Oleksii Kuchaiev. Nemo-aligner: Scalable toolkit for efficient model alignment, 2024. find
  18. Nvidia, :, Bo Adler, Niket Agarwal, Ashwath Aithal, Dong H. Anh, Pallab Bhattacharya, Annika Brundyn, Jared Casper, Bryan Catanzaro, Sharon Clay, Jonathan Cohen, Sirshak Das, Ayush Dattagupta, Olivier Delalleau, Leon Derczynski, Yi Dong, Daniel Egert, Ellie Evans, Aleksander Ficek, Denys Fridman, Sh… find
  19. Shoaib Ahmed Siddiqui, Xin Dong, Greg Heinrich, Thomas Breuel, Jan Kautz, David Krueger, and Pavlo Molchanov. A deeper look at depth pruning of llms. arXiv preprint arXiv:2407.16286, 2024. [24] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, J…
  20. Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller,… — Llama 2: Open Foundation and Fine-Tuned Chat Models
  21. Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. In International Conference on Learning Representations, 2021.
  22. Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde, Jared Kaplan, Harrison Edwards, Yura Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Aleth… — Evaluating Large Language Models Trained on Code
  23. Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, find
  24. Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try ARC, the AI2 reasoning challenge. ArXiv, abs/1803.05457, 2018.
  25. Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. HellaSwag: Can a machine really finish your sentence? In Anna Korhonen, David Traum, and Lluís Màrquez, editors, Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, Florence, Italy, July 2019… find
  26. Stephanie Lin, Jacob Hilton, and Owain Evans. Truthfulqa: Measuring how models mimic human falsehoods, 2022.
  27. Tahmid Hasan, Abhik Bhattacharjee, Md Saiful Islam, Kazi Samin, Yuan-Fang Li, Yong-Bin Kang, M. Sohel Rahman, and Rifat Shahriyar. Xl-sum: Large-scale multilingual abstractive summarization for 44 languages, 2021. find
  28. David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang, Julien Dirani, Julian Michael, and Samuel R. Bowman. Gpqa: A graduate-level google-proof q&a benchmark, 2023.
  29. Fanjia Yan, Huanzhi Mao, Charlie Cheng-Jie Ji, Tianjun Zhang, Shishir G. Patil, Ion Stoica, and Joseph E. Gonzalez. Berkeley function calling leaderboard. 2024. find
  30. Jeffrey Zhou, Tianjian Lu, Swaroop Mishra, Siddhartha Brahma, Sujoy Basu, Yi Luan, Denny Zhou, and Le Hou. Instruction-following evaluation for large language models. arXiv preprint arXiv:2311.07911, 2023.
  31. Zhilin Wang, Yi Dong, Olivier Delalleau, Jiaqi Zeng, Gerald Shen, Daniel Egert, Jimmy J. Zhang, Makesh Narsimhan Sreedhar, and Oleksii Kuchaiev. Helpsteer2: Open-source dataset for training top-performing reward models, 2024. find
  32. Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, Hao Zhang, Joseph E Gonzalez, and Ion Stoica. Judging llm-asa-judge with mt-bench and chatbot arena. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levin… find
  33. Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. ArXiv, abs/1904.09751, 2019.
  34. Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhai… — Evaluating Large Language Models Trained on Code

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/sreenivas-2024-llm-pruning-distillation-practice-minitron/