[논문 세미나] LLM in a flash: Efficient Large Language Model Inference with Limited Memory

https://arxiv.org/abs/2312.11514

 

애플에서 발표한 논문

제한된 메모리에서 LLM 모델을 효과적으로 다루는 것에 대한 논문

플래시 메모리에 저장해서 효과적으로 늘린다.

 

서론

대규모 언어 모델(LLM)은 다양한 작업에서 탁월한 성능을 제공하며 현대 자연어 처리의 중심이 되고 있습니다.

그러나 이들의 상당한 계산 및 메모리 요구사항은 특히 DRAM 용량이 제한된 장치에서 과제가 되고 있습니다.

본 논문은 사용 가능한 DRAM을 초과하는 LLM을 실행하기 위해 모델 매개변수를 플래시 메모리에 저장한 다음 필요에 따라 DRAM으로 가져오는 과제를 다룹니다. 우리의 방법은 플래시 메모리의 특성을 고려한 추론 비용 모델을 구성하여 두 가지 중요 영역에서 최적화를 안내합니다:

 

플래시에서 데이터 전송 량 감소 및 더 큰 연속 덩어리로 데이터 읽기. 이 하드웨어 중심 프레임워크 내에서 우리는 두 가지 주요 기술을 소개합니다.

 

첫째, "윈도잉"은 전략적으로 이전에 활성화된 뉴런의 재사용을 통해 데이터 전송을 줄입니다.

 

둘째, 플래시 메모리의 순차적 데이터 접근 강점에 맞춘 "행-열 번들링"은 플래시 메모리에서 읽는 데이터 덩어리 크기를 늘립니다. 이 기법들은 통합적으로 사용 가능한 DRAM의 두 배 크기 모델을 실행할 수 있게 하며, 기존의 단순 로딩 방식에 비해 CPU에서 4-5배, GPU에서 20-25배 빠른 추론 속도를 제공합니다. 희소성 인식, 컨텍스트 적응형 로딩 및 하드웨어 중심 설계의 통합을 통해 메모리가 제한된 장치에서 LLM의 효과적인 추론을 위한 길을 열어줍니다.

 

온디바이스에 LLM을 적용하기

 

개인정보 보호, 오프라인 사용, 실시간 처리가 가능.

다만 LLM을 적용하기에는 하드웨어 한계와, 모델 최적화가 한계가 있다.

 

memory hierachy

 

chaches 메모리랑 registers는 용량이 적고, 

main Memory가 Dram, Flash disk가 ssd = dram과 ssd를 활용해서, 가격 방어를 위해 그런 방법을 제안한다.

 

거의 100ms 이하로 latency를 줄였다.

compute latency, load from flash latency,memory management를 전부 줄였다.

 

7B 모델을 BF16을 이용해 표현하면 무려 14GB나 필요하다. (7*10^9*2Byte = 14GB)

 

사용가능한 DRAM보다 모델의 가중치가 더 크기 때문에, 모델의 전체 가중치를 flash memory(SSD)에 저장해 로드해야 한다.

 

물론 SSD는 용량은 크지만, 속도가 느리다.

LLM을 돌리기 위해서는 충분한 용량과 속도가 필요하다.(충분한 Bandwidth 필요)

 

chunk size가 커져도 여전히 느린 처리속도가 나온다.

5000mb/s => 말도 안되는 사이즈

 

1. 필요한 부분을 작은 chunk로 읽기보다, 필요 이상의 데이터를 큰 chunk로 읽고 버리는 것이 더 나을 수 있다.

2. Flash memory의 병렬 읽기 특성을 활용하여 읽기 연산을 병렬화 한다.

 

Load From Flash

3가지 방법을 사용했다.

 

predictor

sliding window

 

predictor

 

1. FFN의 특성을 활용하여 데이터 전송 속도를 개선했다.

활성화 되는 것만 사용한다.

 

2. Low-Rank Predictor를 사용해 활성화 여부 예측 -> 일단 활성화 함수가 되기 전에 예측한다.

저랭크 예측기

3. ReLU 활성함수에서 활성화된 부분만을 Load

4. 이 방법은 크고 작은 모델 모두에서 효과적으로 사용할 수 있다.

 

물론 예측이기 때문에 FN의 범위가 있긴 하나, 크게 성능에 영향을 미치지는 않았다.

 

sliding window

기본 방식은 모든 뉴런을 전부 가져온다.

사용하지 않는 뉴런은 unload하고 사용하는 뉴련은 load하여 전송 데이터 양을 감소한다. 

 

슬라이딩 윈도우로 이동하면서 자주 사용되는 부분만 사용한다.

새로운 걸 불러올 때, Dram에 이미 있는 부분에서 중요하지 않은 부분에서 하나만 버리고, 다른 것을 하나 붙여서 추론한다.

 

Row-Column Bundling

up project : 여러 행을 묶어서 처리. 데이터를 전방향으로 전달한다.

Down Project : 여러 열을 묵어서 처리. 데이터를 역방향으로 전달

 

행렬을 함께 flash memory에 저장해서, 데이터를 더 큰 chunk로 읽을 수 있도록 한다.

 

동시에 여러 데이터를 로드함으로서, 여러 데이터를 함께 처리하여 데이터 처리량 향상

Flash memory가 큰 chunk로 읽을 때 더 빠르고, 병렬 처리가 가능한 특성을 살려 사용했다. 

 

공동으로 활성화된 뉴런들의 빈도 그래프

- 가까이 있을 수록 연관이 있다면, 특정 뉴런과 함께 활성화되는 뉴런 그룹이 있을 것이다.

활성화 빈도가 높을수록 숫자가 감소하는 걸 볼 수 있다.

가장 자주 활성화 하는 뉴런을 파악하는 게 중요하다.

 

closest friend라고 한다.

가장 가까운 녀석은 95% 정도 같이 활성화 된다.

 

가장 가까운 뉴런의 공동 활성화 빈도가 높기 때문에, Row-columns bunding의 근거가 되었다.

공동 활성화 기반 번들링을 했을 때 특정 뉴런이 자주 활성화되어 여러번 로드되는 문제 발생

Row-column Bunding은 공동 활성화 패턴을 고려하지 않고 사용할 수 있어 효율을 높일 수 있다.

 

 

결과

3가지 방법을 사용 했을 때, 데이터 전송량이 0.2GB로 감소, I/O latency가 87ms로 빨라졌다.

 

1. on-deivce에서 LLM을 사용하기 위해 Flash Memory와 Dram을 함께 사용하는 방식을 제안

2. predictor / sliding window / row-column bundling을 사용

3. I/O 요청수와 Latency를 줄였다.