Quebra-cabeça: NAS Baseado em Destilação para LLMs Otimizados para Inferência

Puzzle: Distillation-Based NAS for Inference-Optimized LLMs

November 28, 2024

Autores: 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

Resumo

Os modelos de linguagem de grande escala (LLMs) têm demonstrado capacidades notáveis, mas sua adoção é limitada pelos altos custos computacionais durante a inferência. Aumentar o número de parâmetros melhora a precisão, porém amplia a lacuna entre as capacidades de ponta e a viabilidade prática. Apresentamos o Puzzle, um framework para acelerar a inferência de LLM em hardware específico, preservando suas capacidades. Através de uma aplicação inovadora de busca de arquitetura neural (NAS) em uma escala sem precedentes, o Puzzle otimiza sistematicamente modelos com dezenas de bilhões de parâmetros sob restrições de hardware. Nosso método utiliza a destilação de conhecimento local em blocos (BLD) para exploração de arquitetura paralela e emprega programação inteira mista para otimização precisa das restrições. Demonstramos o impacto do nosso framework no mundo real através do Nemotron-51B do Llama-3.1-Nemotron-51B-Instruct (Nemotron-51B), um modelo disponível publicamente derivado do Llama-3.1-70B-Instruct. O Nemotron-51B alcança um aumento de 2,17 vezes na velocidade de inferência, sendo executado em uma única GPU NVIDIA H100, preservando 98,4% das capacidades do modelo original. O Nemotron-51B atualmente é o modelo de linguagem mais preciso capaz de inferência em uma única GPU com tamanhos de lote grandes. Notavelmente, essa transformação exigiu apenas 45B de tokens de treinamento, em comparação com mais de 15T de tokens usados para o modelo de 70B do qual foi derivado. Isso estabelece um novo paradigma onde modelos poderosos podem ser otimizados para implantação eficiente com apenas comprometimento negligível de suas capacidades, demonstrando que o desempenho de inferência, e não apenas o número de parâmetros, deve guiar a seleção do modelo. Com o lançamento do Nemotron-51B e a apresentação do framework Puzzle, fornecemos aos profissionais acesso imediato às capacidades de modelagem de linguagem de ponta a custos computacionais significativamente reduzidos.

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.