LLaMaFlex: Many-in-one LLMs via Generalized Pruning and Weight Sharing

[W1: Extend to other Pre-trained Model Type] Thank you for the suggestion. We validated our method on Minitron-4B [1], a representative efficient large language model of relatively small scale. Using the same hyperparameter settings as in our original experiments on Llama-3.1-8B, we trained the model for 10k iterations and evaluated the perplexity (PPL) on four resultant sub-networks.

Table D1: Evaluation of our method using Minitron-4B as the pre-training model. We report the anchor model (the sub-models directly received gradient during training), as well as the interpolated model (the sub-models generated by router interpolation).

[W2: Clarify our approach, and Metric Ablation] Thank you for the suggestion. Using the proposed importance metric (cumulative activated values), we rank the channels and permute the weight matrix to reorder the channels/heads in decreasing order of importance. To ablate the metric, we compare this ranking with permutation to a setup without permutation. After training the model for 2k iterations, we observe a significant improvement in validation loss when incorporating the importance-based metric.

Table D2: Ablation experiments on our proposed importance metric.

It is worth noting that our importance metric builds upon previous work. Specifically, both the Minitron[1] and Flextron[2] papers include detailed ablations comparing various metrics for importance estimation (see Sections 4.2 and 5.5 in the Minitron and Flextron papers, respectively). We use the best metrics proposed in these works.

[W3: sampled network] Thank you for the question. As detailed in Equation 4, our overall training objective is to minimize the loss of the sampled sub-networks while ensuring they satisfy the parameter budget constraints. This is achieved by incorporating an additional loss term that is activated only when the actual number of activated parameters exceeds the budget. This approach is standard and has been adopted in existing works [2]. Further details are provided in Section 2.3. In terms of joint distribution of Q, we independently model the categorical distribution of each architectural choice by a MLP and sample one category with gumbel softmax. Specifically, $Q(D^j, N_A^j, H^j, N^j, \boldsymbol{\lambda}^j) = P(D^j) P(N_A^j) P(H^j) P(N^j) \prod_i P(\lambda^j_i)$.

[W4: Limitation discussion] Thanks for your suggestion. We have included a limitations section in the revised version (marked in blue) as below:

Our method is not training free, as it requires router tuning and necessitates modification to backbone weights to convert the pre-trained model into an elastic one. There is a potential application of training LlamaFlex from scratch; we leave this to future work.

[Q1: Router interpolation] Thank you for the interesting question. The router outputs cannot be interpolated because they are designed to model the categorical distribution of each architectural choice. The model samples one architecture choice from this distribution, and as training progresses, the sampling process becomes increasingly deterministic, consistently selecting the optimal choice. This results in the router outputs converging to a nearly one-hot distribution. Interpolating the router outputs would disrupt this one-hot distribution, undermining the intended selection mechanism.

[Q2: Configuration details] Thanks for the question. We provide the choices of \mathcal{D}, \mathcal{N_A}, \mathcal{N}, \mathcal{H} in Table 1 in Section 3.1. Specifically:

• \mathcal{N}_A = {25%, 50%, 75%, 100%}
• \mathcal{D} = {25%, 37.5%, 50%, 62.5%, 75%, 87.5% , 100%}
• \mathcal{H} = {50%, 62.5%, 75%, 87.5% , 100%}

We observe that the final results are largely independent of specific choices within these sets; in other words, the selections can be made flexibly, as long as there exist combinations of these sets that satisfy the parameter budget constraints and the elements within each set are reasonably spaced.

[1] Compact language models via pruning and knowledge distillation

[2] Flextron: Many-in-One Flexible Large Language Model

Dear Reviewer,

Could you kindly respond and indicate whether authors have addressed your concerns?

Thanks, AC

Dear authors, thank you for your thorough responses, which addressed my concerns, therefore I raise the score to 6.

[W3: Comparison with Distillation] Thank you for your insightful question. First, we would like to emphasize that elastic frameworks, including LlamaFlex, are orthogonal to traditional fixed-size compression methods. While our primary focus is on the elastic component, our framework is fully compatible with various fixed-size compression approaches and can potentially benefit from advanced methods in this domain. To validate this compatibility, we incorporated knowledge distillation into our framework by adding an additional knowledge distillation loss. We trained the model for 2k iterations and reported the validation loss in Table C2 below. This experiment demonstrates the flexibility of our approach in integrating with other compression techniques.

Table C2: Compatibility of the proposed framework with existing fix-sized compression methods.

Second, we compared our method with representative fixed size compression methods - Minitron [1]. To ensure fair comparison, we use the same pre-trained model and similar data sources as Minitron, and compare the downstream performance in Table C3.

Table C3: Downstream comparison with Minitron method (summarized from Table 2).

[W4: Discussion on optimal architectures] Thank you for the suggestion! In our framework, a lightweight sub-network can be derived by reducing the hidden dimension ($H$), MLP intermediate dimension ($D$), number of attention heads ($N_A$), and the number of remaining layers ($N$). This framework offers a wide range of possible combinations to achieve approximately the same parameter budget (e.g. 50% remaining parameters). However, different configurations would yield different performance. Thus we apply the learnable framework that aims to search for the sub-network with optimal configurations.

Table C4: The learnable framework yields optimal sub-network configurations. With a 50% remaining parameter budget, the sub-network can adopt diverse configurations. For example, we can prune the MLP by retaining only half of its intermediate dimension, prune the hidden dimension by keeping only half of the hidden feature size, or skip half of the layers to meet the 50% target. However, our learnable framework identifies the optimal configurations that combine multiple strategies. All models are fine-tuned for 2k iterations in these experiments.

[1] Compact language models via pruning and knowledge distillation

[W1: Comparison with Supernet-based NAS] Thanks for the insightful suggestion. Our work primarily focuses on LLMs, which sets it apart from previous studies that primarily target relatively small models. Consequently, our work offers unique contributions specifically designed to address the challenges associated with LLMs:

• LLMs demand substantial computational resources and extensive datasets for training. To mitigate these high costs, our framework focuses on transforming a pre-trained model into an elastic framework, enabling efficient training and adaptation. In contrast, supernet-based methods such as HAT [1], HELP [2], DARTS-EGS [3] and MetaNAS [4], as well as recent work like Matformer [5], primarily rely on training models from scratch. Additionally, our method enables zero-shot generation of sub-networks for any parameter budget through router interpolation, a capability not achieved by previous supernet-based methods.
• LLMs typically consist of billions of parameters, making GPU memory usage a critical bottleneck during training. To address this, we adopt weight sharing, ensuring the number of training parameters remains comparable to that of a dense pre-trained model. In contrast, methods like HELP [2], DARTS-EGS [3], and MetaNAS [4] train a supernet with multiple branches. While HAT [1] also employs weight sharing, its use of heterogeneous layers complicates system deployment.
• Matformer [5] relies on manual mix-and-match to extract sub-models, requiring manual selection of optimal configurations. In contrast, we present a fully learnable router that can do this sub-model extraction automatically.

We have added the following paragraph to the related work section in the revised PDF:

HAT[1] and HELP[2] have explored supernet-based approaches for generating sub-networks, primarily targeting relatively small-sized models, while Meta-NAS[4] and DARTS-EGS[3] have utilized Gumbel-Softmax-based supernet techniques. Our work focuses on addressing the unique challenges of LLMs, which differ significantly from prior studies targeting smaller models. To avoid costly training processes and huge GPU memory usage, we transform a pre-trained model into an elastic framework, and employ weight-sharing, unlike previous attempts. To avoid repeated computation, we enable zero-shot sub-network generation for any parameter budget through router interpolation - capabilities not achieved by supernet-based methods.

[W2: Training dynamics] Our framework leverages pre-trained LLMs and applies weight matrix permutation based on an importance metric, ensuring that the largest model variant (i.e., the full pre-trained model) is already converged while other sub-models are reasonably initialized. Furthermore, the algorithm inherently avoids selecting the converged large model variant, as it is penalized by the parameter loss for exceeding the budget requirement.

To mitigate the interference caused by weight sharing, we propose the modulation technique (detailed in Sec 2.4). Specifically, the nesting constraint can potentially limit the representational capacity of elastic networks, as the same set of weights must accommodate inputs across all possible sub-networks. To address this limitation, we propose the use of policy-aware modulation, a technique inspired by methods in the literature on diffusion models. We introduce lightweight non-linear modulation heads, which are applied after the elastic components (i.e., elastic MLP/elastic MHA). These heads modulate the outputs of elastic operations based on the elastic choice. We ablate the modulation technique in Section 4 and provide a detailed comparison in Table C1 below.

Table C1: Results of our ablation study on policy-aware modulation. We run LLAMAFLEX for 800 iterations and report the validation loss. All sub-networks show improved performance when modulation is enabled, reducing validation loss by 0.08 on average.

[1] HAT: Hardware-Aware Transformers for Efficient Natural Language Processing

[2] HELP: Hardware-Adaptive Efficient Latency Predictor for NAS via Meta-Learning

[3] Differentiable Architecture Search with Ensemble Gumbel-Softmax

[4] Meta-Learning of Neural Architectures for Few-Shot Learning

[5] MatFormer: Nested Transformer for Elastic Inference

Dear Reviewer,

Could you kindly respond and indicate whether authors have addressed your concerns?

Thanks, AC

[Q1: Learn routers while retaining pre-trained model weights] Thank you for your insightful suggestion. Our method remains effective when retaining model weights, and the routers are able to search the best configurations of the sub-networks. We conduct an experiment where we fix the backbone weights and enable only the learnable routers to search for the optimal sub-network configurations under two parameter budgets: 25% and 50%. The results, compared against heuristic methods, are presented in Table A1 below.

Table A1: Comparison of learnable searching framework (ours) with heuristic based methods, on elastic training while retaining the original weights.

[Q2: Extend to other Pre-trained Model Type] Thank you for the suggestion. We validated our method on Minitron-4B [1], a representative efficient large language model of relatively small scale. Using the same hyperparameter settings as in our original experiments on Llama-3.1-8B, we trained the model for 10k iterations and evaluated the perplexity (PPL) on four resultant sub-networks. Table A2 below provides the results for this experiment.

Table A2: Evaluation of our method using Minitron-4B as the pre-training model. We report the anchor model (the sub-models directly received gradient during training), as well as the interpolated model (the sub-models generated by router interpolation).

[1] Compact language models via pruning and knowledge distillation

Thanks for your detailed response. I maintain my score.

Many thanks for your thoughtful comments and suggestions. We value your support and will be incorporating all your feedback in our revised version.

[W1: Pipeline Figure] Thank you for the suggestion. We have updated Figure 3 to explicitly illustrate how the routers determine the sub-network configurations. Additionally, we have included a new Figure 4 to provide a clearer depiction of the modulation process.

[W2: Text and subtitles] Thank you for your valuable suggestions. We have updated the revised PDF and marked the changes in blue based on your suggestions. Please let us know if there are any further questions or points of confusion. Regarding data resources, we utilize in-house data but are unfortunately unable to disclose specific details at this stage to avoid revealing author identities, in compliance with the double-blind review policy. More detailed information will be provided upon paper acceptance.

[W3: Confusion on Equation 2] Thanks for pointing this out and our apologies for the confusion. We have updated Equations 1 and Equation 2 in the revised PDF, and have added the skip connection term in the second equation. In this case, the layer will be skipped and only the skip connection is enabled if $\lambda_i$ equals 0.

[W4: Unfair comparison, and the comparison with layer-pruning-only Minitron] Thank you for pointing this out. Our method primarily compares with Minitron, a representative compression approach that integrates pruning and distillation. To ensure a fair comparison, we used Llama-3.1-8B as the pre-trained base model and the similar data distribution (confirmed with the Minitron authors), aligning with the setup of Minitron-4B to the best of our knowledge.

As shown in Table B1 below, our method outperforms the Minitron models (both the layer-pruning and width-pruning versions). Notably, our approach requires only 60 billion tokens, compared to Minitron's 200 billion tokens (100 billion tokens per model version, with computations repeated). Despite using significantly fewer tokens, we notice that our approach achieves better performance, underscoring its effectiveness.

Table B1: Downstream comparison with Minitron method (summarized from Table 2).

We acknowledge that the claim "LlamaFlex can incubate SLMs similar to Minitron accuracy" may not fully capture the nuances of the comparison. As shown in Table B1, LlamaFlex-50% achieves higher average accuracy (70.6%) compared to Minitron-Depth (69.3%) and is on par with Minitron-Width (70.5%). Moreover, LlamaFlex requires significantly fewer tokens for fine-tuning.

[Q1: Heterogeneous Layer] Thank you for your insightful question. While our primary focus is on a system-friendly elastic framework, we observed the trade-off between performance and system efficiency. Loosening the homogeneous restrictions in LlamaFlex could enhance its potential to produce stronger models at the same scale. Specifically, we explored using layer-wise routers instead of a single router for all layers, and this approach demonstrated improved performance while maintaining the same model size.

[Q2: Comparison with Phi] Thank you for the comments. Our model cannot be directly compared with the Phi series or Llama 3.2 due to differences in data resources. To ensure a fair comparison, we focus on benchmarking against the Minitron method, where our model demonstrates exceptional performance while utilizing fewer training resources.

[Q3: Training details] Thank you for the question. We used a total of 60.4 billion training tokens and approximately 12,288 GPU hours. Specifically, we set the sequence length to 4096, the batch size to 128, and fine-tuned the model for 28,800 iterations.