Задача: Нейросетевая архитектура с оптимизацией вывода на основе дистилляции

Puzzle: Distillation-Based NAS for Inference-Optimized LLMs

November 28, 2024

Авторы: Akhiad Bercovich, Tomer Ronen, Talor Abramovich, Nir Ailon, Nave Assaf, Mohammad Dabbah, Ido Galil, Amnon Geifman, Yonatan Geifman, Izhak Golan, Netanel Haber, Ehud Karpas, Itay Levy, Shahar Mor, Zach Moshe, Najeeb Nabwani, Omri Puny, Ran Rubin, Itamar Schen, Ido Shahaf, Oren Tropp, Omer Ullman Argov, Ran Zilberstein, Ran El-Yaniv

cs.AI

Аннотация

Большие языковые модели (LLM) продемонстрировали выдающиеся возможности, но их принятие ограничено высокими вычислительными затратами во время вывода. Увеличение количества параметров повышает точность, но также увеличивает разрыв между возможностями передовых технологий и практической применимостью. Мы представляем Puzzle, фреймворк для ускорения вывода LLM на конкретном оборудовании с сохранением их возможностей. Через инновационное применение поиска нейронной архитектуры (NAS) в масштабах, ранее не встречавшихся, Puzzle систематически оптимизирует модели с десятками миллиардов параметров в рамках ограничений оборудования. Наш подход использует блочное локальное дистиллирование знаний (BLD) для параллельного исследования архитектуры и применяет смешанное целочисленное программирование для точной оптимизации ограничений. Мы демонстрируем реальное воздействие нашего фреймворка через Llama-3.1-Nemotron-51B-Instruct (Nemotron-51B), публично доступную модель, производную от Llama-3.1-70B-Instruct. Nemotron-51B достигает ускорения скорости вывода в 2,17 раза, помещаясь на одном графическом процессоре NVIDIA H100, сохраняя 98,4% возможностей оригинальной модели. В настоящее время Nemotron-51B является наиболее точной языковой моделью, способной к выводу на одном графическом процессоре с большими размерами пакетов. Замечательно, что для этого преобразования потребовалось всего 45 миллиардов обучающих токенов, в сравнении с более чем 15 триллионами токенов, использованными для модели на 70 миллиардов, от которой она была производной. Это устанавливает новый парадигму, где мощные модели могут быть оптимизированы для эффективного развертывания с незначительным ущербом их возможностям, демонстрируя, что производительность вывода, а не только количество параметров, должна направлять выбор модели. С выпуском Nemotron-51B и представлением фреймворка Puzzle мы предоставляем практикующим немедленный доступ к передовым возможностям языкового моделирования с значительно сниженными вычислительными затратами.

English

Large language models (LLMs) have demonstrated remarkable capabilities, but their adoption is limited by high computational costs during inference. While increasing parameter counts enhances accuracy, it also widens the gap between state-of-the-art capabilities and practical deployability. We present Puzzle, a framework to accelerate LLM inference on specific hardware while preserving their capabilities. Through an innovative application of neural architecture search (NAS) at an unprecedented scale, Puzzle systematically optimizes models with tens of billions of parameters under hardware constraints. Our approach utilizes blockwise local knowledge distillation (BLD) for parallel architecture exploration and employs mixed-integer programming for precise constraint optimization. We demonstrate the real-world impact of our framework through Llama-3.1-Nemotron-51B-Instruct (Nemotron-51B), a publicly available model derived from Llama-3.1-70B-Instruct. Nemotron-51B achieves a 2.17x inference throughput speedup, fitting on a single NVIDIA H100 GPU while preserving 98.4% of the original model's capabilities. Nemotron-51B currently stands as the most accurate language model capable of inference on a single GPU with large batch sizes. Remarkably, this transformation required just 45B training tokens, compared to over 15T tokens used for the 70B model it was derived from. This establishes a new paradigm where powerful models can be optimized for efficient deployment with only negligible compromise of their capabilities, demonstrating that inference performance, not parameter count alone, should guide model selection. With the release of Nemotron-51B and the presentation of the Puzzle framework, we provide practitioners immediate access to state-of-the-art language modeling capabilities at significantly reduced computational costs.