EDGE-LLM: �Enabling Efficient Large Language Model Adaptation on Edge Devices via Layer-wise Unified Compression and Adaptive Layer Tuning and Voting

Georgia Institute of Technology

Zhongzhi Yu, Zheng Wang, Yuhan Li, Haoran You, Ruijie Gao, Xiaoya Zhou, Sreenidhi Reedy Bommu, Yang (Katie) Zhao, Yingyan (Celine) Lin

Motivation: Tuning LLMs on the Edge is Highly Demanding

• LLMs are revolutionizing our life

Motivation: Tuning LLMs on the Edge is Highly Demanding

• LLMs are revolutionizing our life
• Stimulate a growing demand for tuning LLMs on the edge

Healthcare

Personal Assistant

Security

Challenge: Costly LLMs Prohibits Edge Tuning

• Computation Challenge:

• Memory Challenge:

Challenge: Costly LLMs Prohibits Edge Tuning

• Computation Challenge:
• Tens of A100 hours w/ the state-of-the-art (SOTA) efficient tuning [T. Dettmers, NeurIPS 2023]
• Memory Challenge:

Challenge: Costly LLMs Prohibits Edge Tuning

• Computation Challenge:
• Tens of A100 hours w/ the state-of-the-art (SOTA) efficient tuning [T. Dettmers, NeurIPS 2023]
• Memory Challenge:
• A 1.48x ~ 3.78x memory gap

Challenge: Costly LLMs Prohibits Edge Tuning

• Computation Challenge:
• Tens of A100 hours w/ the state-of-the-art (SOTA) efficient tuning [T. Dettmers, NeurIPS 2023]
• Memory Challenge:
• A 1.48x ~ 3.78x memory gap
• Edge device mem.: 8~12GB
• Tuning mem.: 17.7GB~30.2GB �for tuning Llama-7B with �different techniques

Where Are We: Existing Solutions for Efficient Tuning

• Category 1: Parameter-efficient tuning
• Mostly commonly used
• Add lightweight learnable modules
• Reduces storage
• Minimal impact on memory and computation

LoRA [E. Hu, ICLR 2022], one of the most commonly �adopted parameter-efficient tuning techniques

Where Are We: Existing Solutions for Efficient Tuning

• Category 1: Parameter-efficient tuning
• Category 2: Compression-then-tuning
• Alleviate computation challenge
• Marginal impact to memory overhead
• Compress the LLM to a lower bit-width
• Same compression setting for all layers

​

Where Are We: Existing Solutions for Efficient Tuning

• Category 1: Parameter-efficient tuning
• Category 2: Compression-then-tuning
• Alleviate computation challenge
• Marginal impact to memory overhead
• Compress the LLM to a lower bit-width
• Same compression setting for all layers

​

Open Question: How to compress?

Further reduce computation with minimal impact to accuracy

Where Are We: Existing Solutions for Efficient Tuning

• Category 1: Parameter-efficient tuning
• Category 2: Compression-then-tuning
• Category 3: Partial tuning
• Alleviate memory challenge
• Cannot help much with computation overhead
• Only update last layers to reduce backprop. depth
• Limit the achieved accuracy

Where Are We: Existing Solutions for Efficient Tuning

• Category 1: Parameter-efficient tuning
• Category 2: Compression-then-tuning
• Category 3: Partial tuning
• Alleviate memory challenge
• Cannot help much with computation overhead
• Only update last layers to reduce backprop. depth
• Limit the achieved accuracy

Open Question: How to Update?

Memory efficient approach to update early layers

Where Are We: Summary of Existing Solutions

Where Are We: Summary of Existing Solutions

• Two directions to improve the efficiency

Dir. 1: Compress the model

Dir. 2: Reduce backprop. depth

Where Are We: Summary of Existing Solutions and Open Questions

• Two directions to improve the efficiency

​

​

Dir. 1: Compress the model

Question 1:

How to compress?

Dir. 2: Reduce backprop. depth

Question 2:

How to update?

Where Are We: Summary of Existing Solutions and Open Questions

• Two directions to improve the efficiency

​

​

Dir. 1: Compress the model

Question 1:

How to compress?

Dir. 2: Reduce backprop. depth

Question 2:

How to update?

Marry the merit

Question 3:

How to combine?

Where Are We: Summary of Existing Solutions and Open Questions

• Two directions to improve the efficiency

​

​

Dir. 1: Compress the model

Question 1:

How to compress?

Dir. 2: Reduce backprop. depth

Question 2:

How to update?

Marry the merit

Question 3:

How to combine?

Edge-LLM: Overview

• A comprehensive tuning framework for edge devices with improved computation and memory efficiency

Edge-LLM: Key Enablers

How to compress the model?

Enabler 1:

Layer-wise unified compression

​

Edge-LLM: Key Enablers

How to compress the model?

Enabler 1:

Layer-wise unified compression

​

How to update all layers?�Enabler 2:

Adaptive layer tuning and voting

​

Edge-LLM: Key Enablers

How to compress the model?

Enabler 1:

Layer-wise unified compression

​

How to update all layers?�Enabler 2:

Adaptive layer tuning and voting

​

How to combine above techniques? �Enabler 3:

A hardware scheduling module

Edge-LLM: Key Enablers

How to compress the model?

Enabler 1:

Layer-wise unified compression

​

How to update all layers?�Enabler 2:

Adaptive layer tuning and voting

​

How to combine above techniques? �Enabler 3:

A hardware scheduling module

Edge-LLM: Enabler 1 - Layer-Wise Unified Compression

• Motivating observation: Different layers have varying sensitivities to compression

Edge-LLM: Enabler 1 - Layer-Wise Unified Compression

• Motivating observation: Different layers have varying sensitivities to compression
• Setting: Visualize MSE before and after compressing each layer in Llama-7B

Quantization

Pruning

Edge-LLM: Enabler 1 - Layer-Wise Unified Compression

• Motivating observation: Different layers have varying sensitivities to compression
• Setting: MSE before and after compressing each layer in Llama-7B

Quantization

Pruning

A layer-wise compression policy is needed

Edge-LLM: Enabler 1 - Layer-Wise Unified Compression

• Goal: Generate layer-wise compression policy based on sensitivity

​

Edge-LLM: Enabler 1 - Layer-Wise Unified Compression

• Goal: Generate layer-wise compression policy based on sensitivity
• For pruning ratio:
• Sensitivity gradually increases with depth
• Fine-grained assignment to preserve differences
• Pruning ratio inversely proportional to sensitivity

​

Pruning

Edge-LLM: Enabler 1 - Layer-Wise Unified Compression

• Goal: Generate layer-wise compression policy based on sensitivity
• For pruning ratio:
• Sensitivity gradually increases with depth
• Fine-grained assignment to preserve differences
• Pruning ratio inversely proportional to sensitivity
• For quantization bit-width:
• Two types of sensitivity values
• Same bit-width for most layers
• An extra bit for highly sensitive layers

​

Quantization

Edge-LLM: Key Enablers

How to compress the model?

Enabler 1:

Layer-wise unified compression

​

How to update all layers?�Enabler 2:

Adaptive layer tuning and voting

​

How to combine above techniques? �Enabler 3:

A hardware scheduling module

Edge-LLM: Key Enablers

How to compress the model?

Enabler 1:

Layer-wise unified compression

​

How to update all layers?�Enabler 2:

Adaptive layer tuning and voting

​

How to combine above techniques? �Enabler 3:

A hardware scheduling module

Edge-LLM: Enabler 2 - Adaptive Layer Tuning

• Goal: Update early layer with reduced backprop. depth

Edge-LLM: Enabler 2 - Adaptive Layer Tuning

• Goal: Update early layer with reduced backprop. depth
• Motivation: LLM’s intermediate layers can generate tokens [T. Schuster, NeurIPS 2022]
• Key idea: Tune intermediate layers to generate tokens

Edge-LLM: Enabler 2 - Adaptive Layer Tuning

• Goal: Update early layer with reduced backprop. depth
• Motivation: LLM’s intermediate layers can generate tokens [T. Schuster, NeurIPS 2022]
• Key idea: Tune intermediate layers to generate tokens
• Proposed technique:
• Add skip connections
• Randomly activate one and update its preceding layers

Can the tuned LLM do even better?

Edge-LLM: Enabler 2 - Adaptive Layer Voting

• Design knob: Many intermediate layers can generate outputs
• Have intermediate layers vote on the final output token

Edge-LLM: Enabler 2 - Adaptive Layer Voting

• Design knob: Many intermediate layers can generate outputs
• Have intermediate layers vote on the final output token
• Post-softmax probability indicates confidence [T. Pearce, arXiv 2021]
• Vote for the token with maximum confidence

Edge-LLM: Enabler 2 - Adaptive Layer Voting

• Design knob: Many intermediate layers can generate outputs
• Have intermediate layers vote on the final output token
• Post-softmax probability indicates confidence [T. Pearce, arXiv 2021]
• Vote for the token with maximum confidence
• A max-in-max approach
• Select token w/ the highest probability among all layers’ outputs

Edge-LLM: Key Enablers

How to compress the model?

Enabler 1:

Layer-wise unified compression

​

How to update all layers?�Enabler 2:

Adaptive layer tuning and voting

​

How to combine above techniques? �Enabler 3:

A hardware scheduling module

Edge-LLM: Key Enablers

How to compress the model?

Enabler 1:

Layer-wise unified compression

​

How to update all layers?�Enabler 2:

Adaptive layer tuning and voting

​

How to combine above techniques? �Enabler 3:

A hardware scheduling module

Edge-LLM: Enabler 3 - Hardware Scheduling Module

• Challenge: Difficult to schedule the irregular computation patterns
• Require efficient memory scheduling and offloading strategies
• Impossible to load everything to SRAM (512KB ~ 1MB)
• Need offloading to DRAM (8GB ~ 16GB) and SSD (128GB ~ 256GB)

Edge-LLM: Enabler 3 - Hardware Scheduling Module

• Proposed technique:
• Conceptualize LLM tuning with offloading as a graph traversal problem

Edge-LLM: Enabler 3 - Hardware Scheduling Module

• Proposed technique:
• Conceptualize LLM tuning with offloading as a graph traversal problem
• Objective: Identify a valid path through all blocks with minimized execution time

Edge-LLM: Enabler 3 - Hardware Scheduling Module

• Proposed technique:
• Conceptualize LLM tuning with offloading as a graph traversal problem
• Objective: Identify a valid path through all blocks with minimized execution time

Search space

Edge-LLM: Enabler 3 - Hardware Scheduling Module

• Proposed technique:
• Conceptualize LLM tuning with offloading as a graph traversal problem
• Objective: Identify a valid path through all blocks with minimized execution time
• Execution time model: �Max latency among
• Read/write latency across each �mem. hierarchy
• Computation latency

Search space

Evaluation Setup

• Target model: Llama-7B
• Algorithm implementation:
• Quantize with LLM-QAT [Z. Liu, arXiv 2023]
• Prune with Sparse-GPT [E. Frantar, ICML 2023]

​

• Evaluation methodology:
• Apply on top of TRETA [H. Shao, VLSI 2023] to validate hardware scheduling speedup
• Simulate with Scale-Sim [A. Samajdar, Github, 2019]

• Dataset: MMLU
• Hardware config:
• SRAM: 1MB
• DRAM: 8GB LPDDR4
• SSD: 128GB

Evaluation Results

• Achieves a 0.70%~1.29% higher accuracy, a ~4x memory reduction with the comparable theoretical computation overhead
• With the hardware scheduling module, achieves a 2.92x~3.38x overall speedup over the vanilla implementation

Scan Me

Project Page

EDGE-LLM: �Enabling Efficient Large Language Model Adaptation on Edge Devices via Layer-wise Unified Compression and Adaptive Layer Tuning and Voting

¹Georgia Institute of Technology, ²University of Minnesota, Twin Cities, ³University of California, Santa Barbara

Zhongzhi Yu¹, Zheng Wang¹, Yuhan Li¹, Haoran You¹, Ruijie Gao¹, Xiaoya Zhou³, Sreenidhi Reedy Bommu¹, Yang (Katie) Zhao², Yingyan (Celine) Lin¹