Puzzle: Distillation-Based NAS for Inference-Optimized LLMs

• URL: http://arxiv.org/abs/2411.19146v4
• Date: Thu, 20 Mar 2025 14:50:04 GMT
• Title: Puzzle: Distillation-Based NAS for Inference-Optimized LLMs
• Authors: Akhiad Bercovich, Tomer Ronen, Talor Abramovich, Nir Ailon, Nave Assaf, Mohammad Dabbah, Ido Galil, Amnon Geifman, Yonatan Geifman, Izhak Golan, Netanel Haber, Ehud Karpas, Roi Koren, Itay Levy, Pavlo Molchanov, Shahar Mor, Zach Moshe, Najeeb Nabwani, Omri Puny, Ran Rubin, Itamar Schen, Ido Shahaf, Oren Tropp, Omer Ullman Argov, Ran Zilberstein, Ran El-Yaniv,
• Abstract summary: Large language models (LLMs) offer remarkable capabilities, yet their high inference costs restrict wider adoption.<n>We present Puzzle, a hardware-aware framework that accelerates the inference of LLMs while preserving their capabilities.<n>We showcase our framework's impact via Llama-3.1-Nemotron-51B-Instruct (Nemotron-51B), a publicly available model derived from Llama-3.1-70B-Instruct.
• Score: 17.72841008597783
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Large language models (LLMs) offer remarkable capabilities, yet their high inference costs restrict wider adoption. While increasing parameter counts improves accuracy, it also broadens the gap between state-of-the-art capabilities and practical deployability. We present Puzzle, a hardware-aware framework that accelerates the inference of LLMs while preserving their capabilities. Using neural architecture search (NAS) at a large-scale, Puzzle optimizes models with tens of billions of parameters. Our approach utilizes blockwise local knowledge distillation (BLD) for parallel architecture exploration and employs mixed-integer programming for precise constraint optimization. We showcase our framework's impact via 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 retaining 98.4% of the original model's benchmark accuracies. Notably, it is the most accurate model supporting single H100 GPU inference with large batch sizes, despite training on only 45B tokens, far fewer than the 15T used to train Llama-70B. Lastly, we derive Llama-3.3-Nemotron-49B-Super-Base to demonstrate Puzzle can retain long-context and that lightweight alignment on these derived models allows them to surpass the parent model in specific capabilities. Our work establishes that powerful LLM models can be optimized for efficient deployment with only negligible loss in quality, underscoring that inference performance, not parameter count alone, should guide model selection.

Related papers

• Pangu Ultra: Pushing the Limits of Dense Large Language Models on Ascend NPUs [123.25404278506585]
  We present Pangu Ultra, a Large Language Model (LLM) with 135 billion parameters and dense Transformer modules trained on Ascend Neural Processing Units (NPUs) To perform such large-scale training efficiently, we utilize 8,192 Ascend NPUs with a series of system optimizations. Our exploration demonstrates that Ascend NPUs are capable of efficiently and effectively training dense models with more than 100 billion parameters.
  arXiv  Detail & Related papers  (2025-04-10T15:41:51Z)
• Every FLOP Counts: Scaling a 300B Mixture-of-Experts LING LLM without Premium GPUs [96.68469559192846]
  We present two differently sized MoE large language models (LLMs) Ling-Lite contains 16.8 billion parameters with 2.75 billion activated parameters, while Ling-Plus boasts 290 billion parameters with 28.8 billion activated parameters. We propose innovative methods for (1) optimization of model architecture and training processes, (2) refinement of training anomaly handling, and (3) enhancement of model evaluation efficiency.
  arXiv  Detail & Related papers  (2025-03-07T04:43:39Z)
• Democratizing AI: Open-source Scalable LLM Training on GPU-based Supercomputers [65.35142508909892]
  We present a novel four-dimensional hybrid parallel algorithm implemented in a highly scalable, portable, open-source framework called AxoNN. We demonstrate fine-tuning of a 405-billion parameter LLM using AxoNN on Frontier.
  arXiv  Detail & Related papers  (2025-02-12T06:05:52Z)
• CMoE: Fast Carving of Mixture-of-Experts for Efficient LLM Inference [33.871080938643566]
  Large language models (LLMs) achieve impressive performance by scaling model parameters, but this comes with significant inference overhead. We propose CMoE, a novel framework to efficiently carve MoE models from dense models. CMoE achieves remarkable performance through efficient expert grouping and lightweight adaptation.
  arXiv  Detail & Related papers  (2025-02-06T14:05:30Z)
• MOFHEI: Model Optimizing Framework for Fast and Efficient Homomorphically Encrypted Neural Network Inference [0.8388591755871735]
  Homomorphic Encryption (HE) enables us to perform machine learning tasks over encrypted data. We propose MOFHEI, a framework that optimize the model to make HE-based neural network inference, fast and efficient. Our framework achieves up to 98% pruning ratio on LeNet, eliminating up to 93% of the required HE operations for performing PI.
  arXiv  Detail & Related papers  (2024-12-10T22:44:54Z)
• Pruning Large Language Models with Semi-Structural Adaptive Sparse Training [17.381160429641316]
  We propose a pruning pipeline for semi-structured sparse models via retraining, termed Adaptive Sparse Trainer (AST) AST transforms dense models into sparse ones by applying decay to masked weights while allowing the model to adaptively select masks throughout the training process. Our work demonstrates the feasibility of deploying semi-structured sparse large language models and introduces a novel method for achieving highly compressed models.
  arXiv  Detail & Related papers  (2024-07-30T06:33:44Z)
• Enabling High-Sparsity Foundational Llama Models with Efficient Pretraining and Deployment [56.44025052765861]
  Large language models (LLMs) have revolutionized Natural Language Processing (NLP), but their size creates computational bottlenecks. We introduce a novel approach to create accurate, sparse foundational versions of performant LLMs. We show a total speedup on CPUs for sparse-quantized LLaMA models of up to 8.6x.
  arXiv  Detail & Related papers  (2024-05-06T16:03:32Z)
• XMoE: Sparse Models with Fine-grained and Adaptive Expert Selection [30.687511115573038]
  tool is a novel MoE designed to enhance both the efficacy and efficiency of sparse MoE models. tool can enhance model performance while decreasing the computation load at MoE layers by over 50% without sacrificing performance.
  arXiv  Detail & Related papers  (2024-02-27T08:18:02Z)
• MoE-LLaVA: Mixture of Experts for Large Vision-Language Models [27.930351465266515]
  We propose a simple yet effective training strategy MoE-Tuning for LVLMs. MoE-LLaVA, a MoE-based sparse LVLM architecture, uniquely activates only the top-k experts through routers. Experiments show the significant performance of MoE-LLaVA in a variety of visual understanding and object hallucination benchmarks.
  arXiv  Detail & Related papers  (2024-01-29T08:13:40Z)
• Distributed Inference and Fine-tuning of Large Language Models Over The Internet [91.00270820533272]
  Large language models (LLMs) are useful in many NLP tasks and become more capable with size. These models require high-end hardware, making them inaccessible to most researchers. We develop fault-tolerant inference algorithms and load-balancing protocols that automatically assign devices to maximize the total system throughput.
  arXiv  Detail & Related papers  (2023-12-13T18:52:49Z)
• SqueezeLLM: Dense-and-Sparse Quantization [80.32162537942138]
  Main bottleneck for generative inference with LLMs is memory bandwidth, rather than compute, for single batch inference. We introduce SqueezeLLM, a post-training quantization framework that enables lossless compression to ultra-low precisions of up to 3-bit. Our framework incorporates two novel ideas: (i) sensitivity-based non-uniform quantization, which searches for the optimal bit precision assignment based on second-order information; and (ii) the Dense-and-Sparse decomposition that stores outliers and sensitive weight values in an efficient sparse format.
  arXiv  Detail & Related papers  (2023-06-13T08:57:54Z)
• Fine-Tuning Language Models with Just Forward Passes [92.04219196752007]
  Fine-tuning language models (LMs) has yielded success on diverse downstream tasks, but as LMs grow in size, backpropagation requires a large amount of memory. We propose a memory-efficient zerothorder (MeZO) to operate in-place, thereby fine-tuning LMs with the same memory footprint as inference.
  arXiv  Detail & Related papers  (2023-05-27T02:28:10Z)
• Energy-efficient Task Adaptation for NLP Edge Inference Leveraging Heterogeneous Memory Architectures [68.91874045918112]
  adapter-ALBERT is an efficient model optimization for maximal data reuse across different tasks. We demonstrate the advantage of mapping the model to a heterogeneous on-chip memory architecture by performing simulations on a validated NLP edge accelerator.
  arXiv  Detail & Related papers  (2023-03-25T14:40:59Z)
• Efficiently Scaling Transformer Inference [8.196193683641582]
  We study the problem of efficient generative inference for Transformer models, in one of its most challenging settings. We develop a simple analytical model for inference efficiency to select the best multi-dimensional partitioning techniques optimized for TPU v4 slices. We achieve a low-batch-size latency of 29ms per token during generation (using int8 weight quantization) and a 76% MFU during large-batch-size processing of input tokens.
  arXiv  Detail & Related papers  (2022-11-09T18:50:38Z)
• MoEfication: Conditional Computation of Transformer Models for Efficient Inference [66.56994436947441]
  Transformer-based pre-trained language models can achieve superior performance on most NLP tasks due to large parameter capacity, but also lead to huge computation cost. We explore to accelerate large-model inference by conditional computation based on the sparse activation phenomenon. We propose to transform a large model into its mixture-of-experts (MoE) version with equal model size, namely MoEfication.
  arXiv  Detail & Related papers  (2021-10-05T02:14:38Z)
• Scalable and Efficient MoE Training for Multitask Multilingual Models [55.987536562357086]
  We develop a system capable of scaling MoE models efficiently to trillions of parameters. We also present new training methods to improve MoE sample efficiency and leverage expert pruning strategy to improve time efficiency. A model trained with 10 billion parameters on 50 languages can achieve state-of-the-art performance in Machine Translation (MT) and multilingual natural language generation tasks.
  arXiv  Detail & Related papers  (2021-09-22T00:57:46Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.