KVSharer Efficient Inference via Layer-Wise Dissimilar KV Cache Sharing

논문 : https://arxiv.org/abs/2410.18517

아래글은 Scholar GPT와 대화를 나눈 내용입니다.

Q : 논문의 핵심 내용과 강점, 알고리즘 설명, 그리고 한계점

논문의 강점과 독창성

1. 독창적인 KVSharer 방식: 이 논문은 LLM(대형 언어 모델)에서 메모리 사용량을 줄이기 위해 KVSharer라는 새로운 계층 단위의 KV(키-값) 캐시 공유 기법을 제안합니다. 기존 방식은 유사한 KV 캐시를 공유하여 메모리를 절약하려고 했지만, KVSharer는 비유사한 계층 간 KV 캐시를 공유하는 방법이 성능을 더 잘 유지한다는 역설적이고 직관에 반하는 현상을 발견했습니다.

2. 효율적인 메모리 사용과 속도 개선: KVSharer는 GPU 메모리 사용량을 30% 줄이고 생성 속도를 1.3배 이상 향상시키면서도 모델 성능을 95% 이상 유지합니다. 이는 기존 모델 학습 없이도 기존 LLM에 적용할 수 있는 플러그 앤 플레이 방식입니다.

3. 다양한 모델과 호환 가능: KVSharer는 기존의 계층 내 KV 캐시 압축 방식과 호환되어 추가 메모리 절감 및 속도 개선을 가능하게 합니다. 또한 일반적인 교정 데이터셋을 사용하여 학습하기 때문에 특정 작업에 구애받지 않고 다양한 작업에 적용할 수 있습니다.

핵심 알고리즘 및 과정 설명

KVSharer의 핵심 알고리즘은 유클리드 거리 기반의 KV 캐시 공유 전략 검색으로, 주요 단계는 다음과 같습니다:

  1. 교정 데이터셋으로 KV 캐시 유사도 계산: 모델이 교정 데이터셋에서 예측을 수행하며 각 계층의 KV 캐시 간의 유클리드 거리를 계산합니다. 유사도가 낮은 계층 쌍을 우선하여 높은 비유사성을 가진 계층 간 KV 캐시를 공유 대상으로 선정합니다.

  2. KV 캐시 공유 전략 검색:
    • 계층 쌍을 순서대로 공유 목록에 추가하고, 특정 계층의 KV 캐시를 다른 계층에 복사하여 공유합니다.
    • 모델의 최종 출력이 기준 유사도 임계값(T)을 초과하는 경우 공유 쌍을 유지하고, 그렇지 않으면 제거합니다.
    • 이 과정을 통해 최적의 공유 전략을 도출합니다.
  3. 실제 추론 단계에서의 KV 캐시 복사:
    • 추론 시 각 계층이 미리 계산된 KV 캐시를 필요할 때마다 복사하여 사용합니다. 이렇게 하면 계층별로 중복된 계산을 줄여 메모리를 절약하고 속도를 높입니다.

한계점

  1. 특정 데이터셋에 의존: 교정 데이터셋을 사용하여 최적의 KV 캐시 공유 전략을 찾는 방식이므로, 교정 데이터셋의 특성에 따라 성능이 다를 수 있습니다. 다만, 연구에서는 데이터셋 크기에 따른 큰 영향을 발견하지는 않았습니다.

  2. 성능 저하 가능성: 비유사한 KV 캐시를 공유할 경우 특정 작업에서 성능이 약간 저하될 수 있으며, 압축률이 높을수록 성능에 더 큰 영향을 미칠 수 있습니다.

  3. KV 캐시 공유에 대한 추가 연구 필요: 비유사성을 이용한 KV 캐시 공유 방식은 효율적이지만, 다른 설정에서의 성능 검증과 추가 최적화가 필요합니다.

KVSharer는 모델 성능을 크게 저하시키지 않으면서 메모리 사용을 줄이고 속도를 높이는 혁신적인 접근으로, 대규모 언어 모델의 효율적 추론에 기여할 수 있는 잠재력을 보여줍니다【17†source】.

Q: 유클리드 거리 기반의 KV 캐시 공유 전략 검색을 통한 알고리즘을 llama 3 8b 모델을 32K까지 decoding하는 상황에서 어떻게 사용되는지 구체적인 예시를 들어서 설명해주겠니?

A :

Llama 3 8B 모델이 32K의 디코딩 길이로 추론할 때, KVSharer 알고리즘의 유클리드 거리 기반 KV 캐시 공유 전략을 어떻게 적용하는지 예시를 통해 설명하겠습니다.

예시 상황 설명

Llama 3 8B 모델을 32K의 디코딩 길이로 사용할 때, 메모리 사용량이 매우 높아 효율적인 KV 캐시 관리가 필수적입니다. KV 캐시는 모델의 각 계층에서 생성되는 키-값 쌍을 저장하여, 긴 시퀀스를 처리할 때 속도와 메모리 효율성을 높입니다. 이때, KV 캐시는 디코딩 길이가 길어질수록 많은 GPU 메모리를 차지하게 됩니다.

KVSharer 알고리즘 적용 과정

1. 교정 데이터셋을 통한 KV 캐시 유사도 계산

먼저, KVSharer는 교정 데이터셋을 사용하여 각 계층의 KV 캐시 간 유사도를 계산합니다. 구체적인 과정은 다음과 같습니다.

  1. Llama 3 8B 모델을 교정 데이터셋 (예: 30개의 문장으로 구성된 64 토큰의 문장들)에서 추론 실행합니다.
  2. 각 계층에서 생성된 KV 캐시를 저장하고, 이를 평균 KV 캐시 벡터로 표현하여 계층 간의 유사도를 계산합니다.
  3. 각 계층의 KV 캐시를 1차원 벡터로 평탄화한 후, 유클리드 거리를 계산하여 계층 간 유사도를 파악합니다.
  4. 계산된 유클리드 거리 값에 따라, 유사도가 낮은(즉, 서로 다른 특성을 가진) 계층 쌍을 우선 순위로 선정합니다.

2. KV 캐시 공유 전략 검색

이제 유클리드 거리가 높은 계층 쌍을 기반으로, 가장 적합한 KV 캐시 공유 전략을 설정합니다.

  1. 유클리드 거리 순으로 각 계층 쌍을 나열하고, 거리가 가장 큰 쌍부터 순차적으로 후속 계층의 KV 캐시를 이전 계층의 KV 캐시로 대체합니다.
  2. 매 단계에서, 해당 공유 전략으로 모델의 마지막 계층의 출력 유사도를 원래 모델의 출력과 비교합니다. 이때, 출력 유사도가 설정된 임계값(T)을 넘는지 확인합니다.
    • 예를 들어, Llama 3 8B 모델이 10번째 계층과 18번째 계층 간에 유사도가 낮아 이 계층 쌍이 선택되었다면, 18번째 계층의 KV 캐시를 10번째 계층의 것으로 대체합니다.
  3. 만약 유사도가 임계값을 넘지 못하면, 해당 계층 쌍을 공유 전략에서 제외하고, 다음으로 유사도가 낮은 쌍을 평가하여 전략을 수정합니다.
  4. 이러한 과정을 반복하여 사전에 설정된 압축률에 맞는 최적의 계층 쌍 공유 전략을 도출합니다.

3. 디코딩 시에 최적화된 KV 캐시 공유 전략 적용

32K 길이의 디코딩 중, 이전 단계에서 도출된 최적의 KV 캐시 공유 전략을 적용하여 메모리를 절약합니다.

  • KV 캐시 복사: 각 계층이 공유 전략에 따라 미리 계산된 KV 캐시를 복사해 사용합니다. 예를 들어, 18번째 계층의 KV 캐시가 10번째 계층의 캐시로 대체되었다면, 18번째 계층에서는 별도의 KV 캐시 계산을 하지 않고 10번째 계층의 캐시를 그대로 사용합니다.
  • 메모리 절감 및 속도 향상: 이를 통해 Llama 3 8B 모델은 더 적은 메모리로 32K 길이의 시퀀스를 디코딩할 수 있으며, 계산 자원을 절약하여 디코딩 속도도 향상됩니다.

요약

KVSharer 알고리즘을 적용함으로써, Llama 3 8B 모델은 32K 디코딩 과정에서 GPU 메모리 사용량을 줄이고 디코딩 속도를 향상시킵니다. 교정 데이터셋을 통해 계산된 계층 간 유사도에 따라, 유사도가 낮은 계층 간의 KV 캐시를 공유하여 모델 성능을 유지하면서 메모리 효율성을 극대화할 수 있습니다.

Q : 위 사례에서 얼마만큼의 memory 사용량 이득과 디코딩 속도 이득이 있을지 A100 gpu를 기준으로 설명해주겠니?

A :

Llama 3 8B 모델을 A100 GPU에서 32K 디코딩 길이로 실행할 때, KVSharer 알고리즘을 적용했을 때의 메모리 사용량 절감과 디코딩 속도 이득을 구체적으로 설명하겠습니다. 이 설명은 Llama2-13B와 KVSharer 실험 결과를 바탕으로 추정됩니다.

1. 메모리 사용량 이득

Llama2-13B 모델에서 KVSharer를 적용한 결과, 25%의 압축률을 사용했을 때 메모리 사용량이 최대 30% 절감되는 것을 확인할 수 있습니다. 이 결과를 Llama 3 8B 모델에 적용하면 다음과 같은 메모리 절감이 기대됩니다:

  • 32K 토큰을 디코딩할 때의 메모리 사용량:
    • Llama 3 8B는 대략 Llama2-13B보다 작은 모델이므로 A100 GPU에서 기본적으로 약 40~50GB의 메모리를 사용합니다.
    • KVSharer 적용 시, 최대 30%까지 메모리 절감이 가능하므로, 약 12~15GB의 메모리 절감이 예상됩니다.
    • 따라서, 32K 디코딩을 위한 A100 메모리 사용량은 약 28~35GB로 줄어들게 되어, GPU 메모리 사용에 여유가 생기고 더 긴 입력 시퀀스 또는 추가적인 모델을 병렬로 실행할 수 있는 가능성이 높아집니다.

2. 디코딩 속도 이득

KVSharer 알고리즘은 디코딩 속도에서 최대 1.6배의 속도 향상을 보여주었습니다. 이 속도 향상은 특히 디코딩 길이가 길고, KV 캐시로 인해 메모리 대역폭과 계산이 크게 증가하는 경우에 효과적입니다.

  • 32K 디코딩 시의 예상 속도 이득:
    • Llama2-13B 모델의 실험 결과에 기반하여, Llama 3 8B에서도 비슷한 디코딩 속도 향상이 기대됩니다.
    • KVSharer 적용 시 디코딩 속도는 최대 1.6배까지 빨라질 수 있으며, 예를 들어 기본 디코딩 속도가 초당 15 토큰이었다면, 적용 후에는 초당 약 24 토큰으로 증가할 수 있습니다.
    • 이로 인해 전체 디코딩 시간이 30~40% 정도 단축되며, 긴 문장 또는 복잡한 작업을 수행할 때 효율성이 크게 향상됩니다.

요약

Llama 3 8B 모델을 A100 GPU에서 32K 디코딩 길이로 실행할 때, KVSharer 알고리즘을 적용하면 다음과 같은 이점을 얻을 수 있습니다:

  • 메모리 절감: 최대 30% 절감되어, 약 12~15GB의 메모리 이득.
  • 디코딩 속도 향상: 최대 1.6배 향상되어, 디코딩 속도가 약 24 토큰/초로 증가 가능.

이러한 이점은 A100과 같은 GPU의 메모리와 처리 성능을 더 효율적으로 활용할 수 있게 하며, 특히 메모리 집약적인 작업이나 긴 시퀀스 처리에서 큰 효율성을 제공합니다.