PaperlayerView exact revision on arXiv

Source-linked AI summary

Flextron: Many-in-One Flexible Large Language Model

Ruisi Cai, Saurav Muralidharan, Greg Heinrich, Hongxu Yin, Zhangyang Wang, Jan Kautz, Pavlo Molchanov

arXiv:2406.10260v2cs.CLcs.LG

TL;DR

Large language models are costly to customize for deployments with different compute and memory limits. FLEXTRON converts an existing trained LLM into a nested, input-adaptive elastic network using post-training routing and sample-efficient optimization. It reports superior zero-shot performance over smaller trained variants and Matformer while using a single pretraining run.

  • Problem

    Large language models require substantial compute and memory, making repeated training of multiple deployment-specific variants impractical.

    Source: §1 · p. 1

  • Method

    FLEXTRON converts a trained LLM into a nested elastic network with static or input-adaptive routing, surrogate-model router training, and sample-efficient candidate sampling.

    Source: §1 · p. 2, §3.2 · p. 4

  • Results

    FLEXTRON demonstrates superior zero-shot performance over multiple smaller end-to-end trained variants and Matformer on GPT-3 and Llama-2 models.

    Source: §4.2 · p. 5, §7 · p. 11

  • Takeaways & Limitations

    One pretrained model can support deployment-specific latency and accuracy targets, including input-adaptive sub-network selection without additional fine-tuning.

    Source: §1 · p. 2

Abstract

from arXiv · show

Training modern LLMs is extremely resource intensive, and customizing them for various deployment scenarios characterized by limited compute and memory resources through repeated training is impractical. In this paper, we introduce Flextron, a network architecture and post-training model optimization framework supporting flexible model deployment. The Flextron architecture utilizes a nested elastic structure to rapidly adapt to specific user-defined latency and accuracy targets during inference with no additional fine-tuning required. It is also input-adaptive, and can automatically route tokens through its sub-networks for improved performance and efficiency. We present a sample-efficient training method and associated routing algorithms for systematically transforming an existing trained LLM into a Flextron model. We evaluate Flextron on the GPT-3 and LLama-2 family of LLMs, and demonstrate superior performance over multiple end-to-end trained variants and other state-of-the-art elastic networks, all with a single pretraining run that consumes a mere 7.63% tokens compared to original pretraining.

1. Introduction

FLEXTRON addresses the cost of deploying large language models across constrained hardware by converting one trained LLM into a flexible, input-adaptive elastic network. Its post-training framework combines nested sub-networks, routing algorithms, surrogate-model training, and efficient sampling to support varied deployment targets without additional fine-tuning.

1. Introduction

  • Existing deployment practice trains multiple model variants because billion-parameter LLMs impose substantial memory and compute constraints.Llama-2 provides 7B, 13B, and 70B variants, while Pythia provides eight sizes from 80M to 12B parameters.Source: §1 · p. 1
  • FLEXTRON converts a standard trained LLM into one model containing heterogeneous nested sub-networks selected according to input data and deployment constraints.The architecture extends mixture-of-experts ideas to attention and feed-forward layers while supporting multiple models in one deployment.Source: §1 · p. 2
  • FLEXTRON is designed to adapt to user-defined latency and accuracy targets during inference without additional fine-tuning.Its overview also describes fast, zero-shot generation of hardware- and input-adaptive sub-networks under accuracy, latency, and parameter constraints.Source: Figure 1 · p. 1, §1 · p. 2
  • FLEXTRON uses static and dynamic routers to select sub-networks for latency targets and, in the dynamic case, individual input tokens.A surrogate model predicts language loss from router choices, enabling router optimization despite limited gradient propagation.Source: §1 · p. 2
  • The framework includes sample-efficient elastic-network training that requires significantly less compute than existing methods.It randomly samples candidate networks during training instead of jointly evaluating the prohibitively large candidate pool.Source: §1 · p. 2

2. Background and Notation

FLEXTRON constructs elastic networks by replacing layer parameters with selectable candidates, yielding many sub-networks with different runtime and accuracy characteristics. Nested MLPs share hidden neurons, while elastic attention selects subsets of heads to expand the searchable design space.

2. Background and Notation

  • Replacing each layer’s parameters with one of K candidates produces K^N possible elastic sub-networks across N layers.These sub-networks can differ in runtime and accuracy characteristics.Source: §2 · p. 2
  • Nested elastic MLP candidates share hidden neurons through indexing, saving memory and improving efficiency.Candidate hidden dimensions satisfy d_1 < d_2 < ... < d_K, with the largest candidate using the full dimension D.Source: §2 · p. 2
  • Each elastic MLP candidate selects only its first d_j hidden neurons from shared weight matrices.The selection can be implemented with slicing rather than an explicit diagonal matrix.Source: §2 · p. 2
  • FLEXTRON makes multi-head attention elastic by selecting subsets of attention heads, expanding the candidate operation space beyond elastic MLPs.The paper states that this is the first work supporting both elastic MLP and elastic MHA layers.Source: §2 · pp. 2–3

3. FLEXTRON Framework

FLEXTRON transforms an existing LLM into a nested elastic network whose sub-networks can be selected for deployment constraints and input conditions. Its framework combines importance-ordered continued-training with static or input-adaptive routing trained through a surrogate model.

3.1. Elastic Network Continued-Training

  • Importance ranking orders neurons and attention heads so sub-networks can be formed by indexing the most important channels while sharing parameters.The ranking uses accumulated activation magnitudes from a small dataset, and the first channels or heads are treated as the most important.Source: §3.1 · p. 3
  • Elastic continued-training samples a tractable subset of candidate networks and trains them simultaneously with a combined loss instead of enumerating every combination.A 32-layer LLaMA2-7B model can have 464 possible combinations, motivating random sampling of k candidates.Source: §3.1 · p. 3

3.2. Automatic Network Selection

  • The routing objective combines language-model loss with a constraint loss so selected topologies meet requested latency or other deployment targets.The framework defines the selected topology from a latency constraint and notes that the loss can also extend to GPU-memory constraints.Source: §3.2 · p. 3
  • Static routing maps a target latency to layer-wise expert choices, determining the number of channels or heads used at each layer.Each router embeds the scalar latency target into logits over four predefined expert candidates.Source: §3.2 · p. 4
  • The surrogate model predicts the full LLM language loss from router outputs, providing a smaller proxy objective for training routers when direct gradients are insufficient.Routers receive latency-loss updates throughout training and receive surrogate-loss updates when the surrogate’s error is below a predefined threshold; the surrogate is then discarded before joint fine-tuning.Source: §3.2 · pp. 4–5, Figure 3 · p. 4

3. FLEXTRON Framework

  • FLEXTRON uses nested elastic experts across attention and feed-forward layers, allowing one model to support heterogeneous sub-networks selected by deployment constraints and input data.The architecture uses weight sharing and array-based indexing to reduce storage, memory-bandwidth, and communication overheads.Source: §3.2 · p. 4

3.2. Automatic Network Selection

  • Dynamic routing additionally conditions token-wise decisions on hidden states, selecting sub-networks according to both latency requirements and input features.Hidden and latency embeddings are combined before a linear layer generates routing decisions.Source: §3.2 · p. 4

4. Experiments

FLEXTRON is evaluated on GPT3 and Llama2 families against nested elastic, smaller pretrained, open-source, and post-hoc compression baselines. The results cover downstream quality, latency–capacity trade-offs, and training cost.

4.1. Experimental Settings

  • The study uses GPT3-2B, GPT3-8B, and Llama2-7B, with elastic continued-training, router tuning, and shared domain data.The models support contexts up to 4096, and the training procedure begins with importance sorting before elastic-model training.Source: framework (Kuchaiev et al., 2019) · p. 5
  • The experiments compare FLEXTRON with Matformer, smaller GPT3 models, open-source families, and post-hoc compression methods.Baselines include Pythia, OpenLLaMA, Sheared-LLaMA, Compresso, LLM-Pruner, SliceGPT, and LaCo.Source: framework (Kuchaiev et al., 2019) · p. 5

4.2. Results

  • FLEXTRON is evaluated across downstream tasks using static and input-adaptive variants on GPT3-8B and Llama2-7B backbones.The evaluation includes ARC-easy, LAMBADA, PIQA, WinoGrande, MMLU, and HellaSwag, with dynamic models reporting averaged activated non-embedding parameters.Source: §4.2 · p. 5, Table 1 · p. 6
  • FLEXTRON extends scaling analysis by relating sub-model capacity to both non-embedding parameters and input latency.The fitted trade-off is presented for static and input-adaptive routing, providing guidance for jointly meeting latency, parameter, and capacity constraints.Source: §4.2 · pp. 6–7, Table 3 · p. 7
  • 89.9 billion tokens are used for elastic continued-training and 1.049 billion for router tuning on both GPT3 and Llama2.The reported training-cost comparison includes additional joint tuning for Llama2-7B.Source: §4.2 · p. 7, Table 4 · p. 7

5. Analysis

The analysis explains how FLEXTRON allocates elastic computation across layers and inputs, trains routers, and constructs useful nested sub-networks. It also shows that learned routing and weight permutation improve the reliability of selected sub-models.

5.1. FLEXTRON Learnings & Insights

  • FLEXTRON assigns more computation to deeper MLP layers, while lightweight MHA layers preserve performance better under limited latency.Replacing full MLP layers and deep layers causes larger averaged language-model-loss degradation, and learned static architectures follow this pattern.Source: §5.1 · p. 7, Figure 6 · p. 8, Figure 7 · p. 8
  • Input-adaptive routing assigns smaller candidates to easier data and larger candidates to harder data across domains.Almost half of HTML tokens select the smallest candidate, whereas Books3 tokens tend to select the full layer; code data also needs fewer parameters than multilingual data.Source: §5.1 · p. 7, Figure 9 · p. 9

5.2. Training Trajectory of Elastic Continued-Training

  • Elastic continued-training preserves full-model performance while sub-networks converge synchronously, with larger models achieving lower validation loss overall.Middle-sized sub-models show smoother convergence, and uniform selection confirms that lightweight MHAs are advantageous in limited-resource regimes.Source: §5.2 · p. 8, Figure 5 · p. 8

5.3. Router Training Dynamics.

  • Router training uses a surrogate model to stabilize optimization through a three-stage progression from surrogate tuning to joint tuning and router tuning.The L2 loss quickly drops during surrogate tuning, both losses decrease during joint tuning, and LM loss continues decreasing once the surrogate error remains below threshold.Source: §5.3 · pp. 8–9, Figure 8 · p. 9

5.4. Effectiveness of Learned Routers

  • At 65% latency, randomly selected sub-models have averaged LM loss ranging from 2.32 to 3.06, whereas FLEXTRON routers identify optimal sub-models.The comparison uses box plots of random selections and router-selected models at different latency targets.Source: §5.4 · p. 9

5.5. Effectiveness and Necessity of Weight Permutation

  • Weight permutation reduces Wikipedia perplexity for first-half sub-networks from above 1000 to 193.6.The same experiment observes a similar enhancement for MHA modules, supporting importance-based ordering of neurons and heads.Source: §5.5 · p. 9, Table 5 · p. 9

6. Related Work

Prior work provides elastic, input-adaptive, and static acceleration approaches for neural networks and Transformers. FLEXTRON combines these directions with shared nested weights, elasticity across MLP and MHA layers, and latency-conditioned routing.

Elastic Inference

  • Elastic inference methods derive multiple deployable models from one network, but prior supernets can require costly non-standard architectures and prolonged training.Related approaches include slimmable networks, Once-for-all, Matformer, and sampling-based training.Source: §6 · p. 9
  • FLEXTRON extends nested elastic modeling to both MLP and MHA layers, with a larger operation pool, efficient training, and latency-based input-adaptive selection.These features distinguish it from Matformer and related slimmable Transformer methods.Source: §6 · pp. 9–10

Input Adaptivity

  • MoE and heterogeneous-expert systems motivate input adaptivity, but typically use fixed budgets or separately stored expert weights.FLEXTRON instead shares one weight matrix across experts and selects sub-networks through array indexing.Source: §6 · p. 10

Static Acceleration

  • Static acceleration methods such as quantization, structured sparsity, NAS, pruning, and checkpoint reuse are presented as complementary to dynamic inference.The paper positions these methods as additional opportunities for performance improvement rather than direct alternatives.Source: §6 · p. 10

7. Conclusion

FLEXTRON adapts to latency and accuracy targets during inference without additional fine-tuning while supporting input-adaptive routing. It converts trained LLMs into flexible models with strong zero-shot performance using only 7.63% of full-pretraining tokens.

7. Conclusion

  • FLEXTRON outperforms smaller end-to-end trained variants on GPT-3 and Llama-2-7B, and surpasses Matformer.These results are achieved through a single pretraining run consuming 7.63% of the tokens required for full pretraining.Source: §7 · p. 11
  • FLEXTRON adapts to different latency and accuracy targets during inference without additional fine-tuning.Source: §7 · p. 11
  • FLEXTRON supports input-adaptive routing to maximize performance during deployment.Source: §7 · p. 11
  • The post-training framework converts standard trained GPT-3 and Llama2 models into FLEXTRON models using a sample-efficient procedure.Source: §7 · p. 11

A. Architecture Visualization

The paper provides searched GPT-3-family architectures for 50% and 70% latency targets. These architectures are presented as validating the heuristic described in Section 5.1.

A. Architecture Visualization

Related papers

Loading 2406.10260v2…