AmoebaLLM: Constructing Any-Shape Large Language Models for Efficient and Instant Deployment

Motivated by the transformative capabilities of large language models (LLMs) across various natural language tasks, there has been a growing demand to deploy these models effectively across diverse real-world applications and platforms. However, the challenge of efficiently deploying LLMs has become increasingly pronounced due to the varying application-specific performance requirements and the rapid evolution of computational platforms, which feature diverse resource constraints and deployment flows. These varying requirements necessitate LLMs that can adapt their structures (depth and width) for optimal efficiency across different platforms and application specifications. To address this critical gap, we propose AmoebaLLM, a novel framework designed to enable the instant derivation of LLM subnets of arbitrary shapes, which achieve the accuracy-efficiency frontier and can be extracted immediately after a one-time fine-tuning. In this way, AmoebaLLM significantly facilitates rapid deployment tailored to various platforms and applications. Specifically, AmoebaLLM integrates three innovative components: (1) a knowledge-preserving subnet selection strategy that features a dynamic-programming approach for depth shrinking and an importance-driven method for width shrinking; (2) a shape-aware mixture of LoRAs to mitigate gradient conflicts among subnets during fine-tuning; and (3) an in-place distillation scheme with loss-magnitude balancing as the fine-tuning objective. Extensive experiments validate that AmoebaLLM not only sets new standards in LLM adaptability but also successfully delivers subnets that achieve state-of-the-art trade-offs between accuracy and efficiency.

Paper

Similar papers

Peer review

Reviewer U1ac5/10 · confidence 3/52024-07-11

Summary

The authors proposed a novel pipeline for once-for-all training multiple subnets in a supernet LLM under different resource constraints. The entire pipeline consists of a knowledge preserving subnet selection utilizing DP to sample depth and width and a new LoRA to resolve the gradient conflicts during training multiple subnets jointly. The experimental results on two widely used decoder-only architectures show the efficacy of the pipeline.

Strengths

- The paper is well-motivated and well-written. To deploy the LLM on multiple resource-constrained devices, fine-tuning multiple subnets jointly is a promising direction to save the computation and training efforts. - Resolving the complexity when jointly selecting the removal layers in layer-wise pruning is a remaining question for a long time. Although the authors did not entirely solve the problem, the proposed DP solution may potentially also benefit other layer-wise network architectures, e.g., CNN, encoder model. etc.

Weaknesses

- Some technical details are not clear to me. For example, how are the depth and width defined in decoder models? Is a self-attention block a layer? Or does the authors consider FFN and MHSA as two individual layers? In terms of the width, it is clear for FFN but is that the number of heads in MHSA? - What metric is used as the importance in subnets sampling? I assume that that importance is not derived from gradients, since DP is done layer-by-layer in forward manner. If so, will that metric work better than gradient-based metrics? - The authors introduce a new LoRA in the paper. Does the authors freeze the backbone LLM and only tune the adapters? If not, how does the authors avoid the co-adaptation of backbone and adapters? Since the original LoRA is proposed to preserve the learned knowledge and fine-tuning only low-rank adaptation for a small amount of new data. - Also are the adapters only added in MHSA? If so, can the authors provide some theoretical or experimental results to show why FFN layers do not suffer from gradient conflicts?

Questions

See above

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes

Authorsrebuttal2024-08-11

Dear Reviewer, We sincerely appreciate the time you dedicated to providing valuable feedback on our paper. In this author response, we have addressed all of your initial concerns. If you have any further questions or concerns, we are happy to discuss them with you. Additionally, we welcome any new suggestions or comments you may have! Best, The Authors of Paper #13852

Reviewer wNCb7/10 · confidence 3/52024-07-15

Summary

This paper proposes AmoebaLLM: a one-for-all fine-tuning and compression framework for delivering pruned and accurate subnets from a pre-trained LLM at various pruning (both depth and width) ratios without the need to fine-tune individual subnets. AmoebaLLM consists of two components: 1) a dynamic programming-based approach to layer selection for selecting subnets of certain depth that takes into account various layer interactions, as opposed to greedy-layer based approaches; 2) after the subnets are defined, a one-for-all fine-tuning strategy is proposed based on gating multiple LoRA heads that are subnet-specific. The training strategy uses a hybrid loss where the largest subnet (always sampled) computes its loss with the ground truth, and all subsequently sampled subnets compute a distilled loss from the largest subnet. The final loss is computed by weighing each subnet's loss appropriately. To demonstrate the efficacy of their method, the authors perform experiments on both LLama 2 7B and Vicuna 7b v1.5 and compare against prior LLM pruning techniques.

Strengths

- The paper tackles a challenging and relevant problem: compressing pre-trained LLMs with various complexities while maintaining their performance - The paper is well written and easy to follow - The method overall is novel, and performs better than other pruning methods. - The paper provides interesting insights into how different LLM metrics are impacted by model compression, e.g., Sec 5.4 shows that MMLU is a better metric for calibration.

Weaknesses

- The claims that full model fine-tuning suffers from severe gradient conflicts appear rather weak. In line (233) the authors state "As analyzed in Sec. 2 and demonstrated in Sec. 5.3 ... ", whereas Section 2 simply conjectures that this is an issue as it uses phrases like" .. would likely fail .." and "as it may omit layers". I would suggest citing the results from Section 5.3 directly to motivate this issue. As for the results themselves in Table 3, the performance of the full model (second row) does not seem consistent. While it performs poorly on Wikitext2, the MMLU performance remains comparable with the SMoL baseline, and seems to improve when pruning more (24 to 20) as opposed to all other methods. It is not clear how these results lead to the conclusion that there is a "gradient conflict". How would the same experiments look when using the Vicuna model? and how would the other evaluation metrics look? - Using dynamic programming in the context of neural network pruning is not new (https://arxiv.org/pdf/2308.10438, https://ieeexplore.ieee.org/document/9116287 to name a few). The authors should cite and discuss these works, even if they do not tackle LLMs specifically. - In line 199, the authors claim that the DP complexity is O(MN). This is a bit misleading, as it ignores the complexity of computing P(n,m) which I argue is also dependent on M and N. Furthermore, the actual run-time of the DP step is not discussed, is it too negligible? - In Table 1, comparing with AmoebaLLM$^\dagger$ is not a fair as it is trained for longer (on top of the training performed in regular AmoebaLLM). How will the other methods compare with the same amount of training? - It would improve reading clarity if the authors explicitly mention how many layers do the models being studied contain, as simply mentioning how many layers are being pruned in depth shrinkage is hard to gauge without that context. This should also be mentioned in the Tables.

Questions

- AmoebaLLM produces different subnets, some may have the same parameter count due to different width/depth shrinkage ratios, but it is not clear how a model is selected for a certain experiment (e.g. Fig 3). Is it simply enumerating all the subnets and picking the fastest one? Are the profiling results obtained from Section 3 being used for this, or is that purely for motivating the problem? - What ranks are used for LoRA?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have addressed the limitations of their proposed work.

Authorsrebuttal2024-08-07

Author Response - Part 2

**5. Benchmark with baselines with the same amount of training time** We first clarify that in Table 1, we follow a standard criterion in the literature, i.e., ensuring the same number of fine-tuning tokens across both the baselines and our method. Following your suggestion, we further train the baselines with the same amount of training time as our method, which corresponds to 20k training iterations on Alpaca-gpt4. As shown in the table below, we report the MMLU accuracy as well as the average accuracy across 7 tasks, following the task list of Table 1. We observe that (1) the baselines fine-tuned with 20k iterations generally maintain comparable performance with those trained with 10k iterations, potentially because 10k iterations is sufficient for fine-tuning on Alpaca-gpt4; and (2) our method still outperforms all baselines across various tasks. | **Remaining Ratio** | **Method** | **Training Iterations** | **Training Time** | **MMLU (%)** | **Average Acc (%)** | |:---:|:---:|:---:|:---:|:---:|:---:| | 80% | FLAP | 10k | 15h | 40.21 | 60.98 | | | | 20k | 30h | 38.62 | 61.40 | | | Shortened LLaMA | 10k | 15h | 26.45 | 58.72 | | | | 20k | 30h | 26.55 | 59.63 | | | Ours | 10k | 30h | **40.70** | **62.29** | | 65% | FLAP | 10k | 15h | 33.28 | 56.12 | | | | 20k | 30h | 35.4 | 55.93 | | | Shortened LLaMA | 10k | 15h | 24.89 | 52.57 | | | | 20k | 30h | 24.70 | 53.95 | | | Ours | 10k | 30h | **36.00** | **56.96** | | 50% | FLAP | 10k | 15h | 27.67 | 51.12 | | | | 20k | 30h | 27.55 | 51.32 | | | Shortened LLaMA | 10k | 15h | 24.76 | 47.35 | | | | 20k | 30h | 25.10 | 49.22 | | | Ours | 10k | 30h | **30.60** | **52.19** | **6. The number of remained depth and width in Table 1** Thank you for the suggestion! The (depth, width scale) for AmoebaLLM with 80%/65%/50% remaining ratios in Table 1 are (30, 0.875)/(28, 0.75)/(22, 0.75), respectively. We will follow your suggestion to add this information to Table 1 to improve readability. --- **7. The detailed subnet selection strategy** The profiling results in Section 3 are purely for motivating the problem. Here is our current subnet selection strategy: we adopt a hierarchical search strategy to deliver subnets from our design space that satisfy the target efficiency constraint, e.g., 50% weight remaining ratios in Table 1 of our manuscript, while maximizing the achievable accuracy. Specifically, we first perform a coarse grid search across uniformly spaced depth and width settings based on a small calibration set, e.g., 20 samples from the MMLU dataset, to identify the subnets that satisfy the given efficiency constraint with maximized accuracy. Next, we perform a more fine-grained grid search within depth/width ranges surrounding the optimal subnet identified in the coarse grid search stage. This process typically evaluates 40 promising subnets and takes no more than 10 minutes on an NVIDIA A5000 GPU. We empirically find that the above strategy works well at the scale of our target problem and can already deliver high-quality subnets that outperform previous compression methods. We also note that more complex subnet selection strategies, such as the evolutionary search adopted by [11]-[15] cited in our manuscript, can also be employed, which will be our future work. We will clarify this in the final version. --- **8. The rank of LoRA** We follow QLoRA [1] and adopt a rank of 64. We will add this information to the final version. [1] “QLoRA: Efficient Finetuning of Quantized LLMs”, T. Dettmers et al., NeurIPS’23.

Reviewer wNCb2024-08-13

Thank you for the detailed response. The authors have addressed most of my concerns/questions. The additional data/clarifications provided in this rebuttal should be included in the final manuscript as the authors mentioned. I will increase my score (6->7).

Authorsrebuttal2024-08-13

Thank you for taking the time to review our rebuttal responses and for providing positive feedback! We are encouraged to hear that our rebuttal has addressed most of your concerns. Following your suggestion, we will include the additional data and clarifications from this rebuttal in our final manuscript.

Authorsrebuttal2024-08-11

Dear Reviewer, We sincerely appreciate the time you dedicated to providing valuable feedback on our paper. In this author response, we have addressed all of your initial concerns. If you have any further questions or concerns, we are happy to discuss them with you. Additionally, we welcome any new suggestions or comments you may have! Best, The Authors of Paper #13852

Reviewer LhJH5/10 · confidence 4/52024-07-21

Summary

The paper proposed a new framework, named AmoebaLLM, that adapts any LLMs to achieve optimal efficiency across different platforms and applications. In specific, the framework contains two stages. The first stage denoted as a knowledge-preserving stage, creates a subnet of the LLM by dynamic programming given an arbitrary depth or width. This second stage denoted as the one-for-all finetune stage is to fine-tune the obtained subnet to achieve the optimal performance on the given application (dataset). To finetune the subnet, the paper proposed shape-aware LoRAs to adapt the knowledge and a loss-magnitude balancing to ensure efficient distillation in the second stage. All combined, the method achieves SOTA performance among all the pruning methods and obtains the best accuracy in the same latency in different platforms.

Strengths

1. The target of the paper is practical and interesting. It is indeed a good open problem to implement different LLMs within different constraints of devices and platforms while achieving the optimal trade-off between efficiency and accuracy. 2. This method, the DP-based depth shrinking strategy, could construct an any-shape LLM with reasonable performance, which is promising in different applications. 3. The presentation and writing are good and easy to follow.

Weaknesses

1. About latency in section 5.2:The work claimed that it achieve the accuracy-efficiency frontier. In terms of accuracy, it is convincing as shown in Table 1. However, for efficiency, it seems only Fig.3 is related to this topic yet only shows a trade-off between latency and accuracy. It is not sufficient to evaluate the efficiency purely based on latency. 2. In my understanding, efficiency should also consider the resources consumed in the pruning method. First, the DP algorithm itself has a high computational complexity. Second, the proposed method is somewhat limited to its one-for-all finetune, which is a training process related to a specific dataset, and a manual selection of the subnet, the method is not perfectly efficient. In this case, the author may consider adjusting their claim in efficiency. 3. The selection of the subnet seems to be related to the profiling of the generation latency. It is purely empirical and may be error-prone when it is generalized to other platforms.

Questions

1. What is the definition of the latency? 2. Is it possible to provide the details of selecting the subnet shape that favors the hardware characteristics? 3. Which LLM is used across all the experiments?

Rating

5

Confidence

4

Soundness

2

Presentation

4

Contribution

3

Limitations

No negative societal impact is found.

Authorsrebuttal2024-08-07

Author Response - Part 2

**5. Detailed subnet selection strategy** Thank you for pointing this out! Currently, we adopt a hierarchical search strategy to deliver subnets from our design space that satisfy the target efficiency constraint, e.g., 50% weight remaining ratios in Table 1 of our manuscript, while maximizing the achievable accuracy. Specifically, we first perform a coarse grid search across uniformly spaced depth and width settings based on a small calibration set, e.g., 20 samples from the MMLU dataset, to identify the subnets that satisfy the given efficiency constraint with maximized accuracy. Next, we perform a more fine-grained grid search within depth/width ranges surrounding the optimal subnet identified in the coarse grid search stage. This process typically evaluates 40 promising subnets and takes no more than 10 minutes on an NVIDIA A5000 GPU. We empirically find that the above strategy works well at the scale of our target problem and can already deliver high-quality subnets that outperform previous compression methods. We also note that more complex subnet selection strategies, such as the evolutionary search adopted by [11]-[15] cited in our manuscript, can also be employed, which will be our future work. We will follow your suggestion and our promise in the abstract to open source all source code and the delivered subnets upon acceptance. --- **6. Which LLMs are used** As mentioned in Section 5.1 of our manuscript, we employ LLaMA 2 7B and Vicuna 7B v1.5.

Authorsrebuttal2024-08-11

Dear Reviewer, We sincerely appreciate the time you dedicated to providing valuable feedback on our paper. In this author response, we have addressed all of your initial concerns. If you have any further questions or concerns, we are happy to discuss them with you. Additionally, we welcome any new suggestions or comments you may have! Best, The Authors of Paper #13852

Reviewer LhJH2024-08-11

Thank you for your detailed response. My concerns have largely been addressed, and after considering the feedback from other reviewers, I am inclined to support the acceptance of this paper. I have a few additional suggestions: 1. The table provided in the first response should be included in the main paper. It would strengthen the argument regarding the “Frontier of accuracy and efficiency,”, as particularly the method demonstrates state-of-the-art performance across various efficiency metrics. 2. It would be beneficial to include a discussion of the complexity of the DP algorithm and its corresponding running time in the main paper. This addition could alleviate any concerns readers might have about the

Authorsrebuttal2024-08-11

Thank you for taking the time to review our rebuttal responses and for providing positive feedback! We are encouraged to hear that our rebuttal has addressed most of your concerns. Following your suggestion, we will include (1) the real-device efficiency table and (2) the complexity and runtime of our DP algorithm, both in our first response, in the final version of our manuscript to strengthen the coherence of our work.

Reviewer TKvU5/10 · confidence 5/52024-07-21

Summary

To address the problems of diverse resource constraints and deployment flows while using LLM for multiple real-world applications, this paper proposes an AmoebaLLM, featuring a knowledge-preserving subnet selection strategy, a shape-aware mixture of LoRAs and a distillation scheme with loss-magnitude balancing.

Strengths

1. The paper is well-organized. 2. The method achieves SOTA performance on most metrics. 3. The experimental setup is easy to follow.

Weaknesses

Overall, the motivation and ideas of this paper are commendable. However, I have a few concerns: 1. DP-based depth shrinking aims to select suitable layers that can achieve the optimal target metric. Does this mean that the proposed method needs to evaluate the performance for each selection strategy? If so, how long is required for each evaluation? What is the relationship between the number of layers and the overall search time? 2. Another concern is about the optimal subnet. I wonder if the optimal subnet identified in the first phase is truly “optimal.” I believe that the selected subnet is the optimal choice for the initial parameters. However, the proposed method consists of two stages. After the fine-tuning stage, some subnets that were not optimal in the first phase might achieve the best target metrics. I understand that searching subnets with fine-tuning is very expensive, but have the authors considered this issue? Is this correct? 3. I am particularly interested in the training efficiency comparisons, as the authors mentioned that existing methods are not efficient. However, the experimental section does not show such comparisons. Could the authors provide relevant comparisons in Table 1 and Table 2?

Questions

Please see the weaknesses part.

Rating

5

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes.

Authorsrebuttal2024-08-11

Dear Reviewer, We sincerely appreciate the time you dedicated to providing valuable feedback on our paper. In this author response, we have addressed all of your initial concerns. If you have any further questions or concerns, we are happy to discuss them with you. Additionally, we welcome any new suggestions or comments you may have! Best, The Authors of Paper #13852

Authorsrebuttal2024-08-13

Thank you for taking the time to review our rebuttal responses and provide feedback! Any further suggestions or comments you have are welcome!

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC