AI 기초

KV 캐시에는 비트 문제가 아니라 기하학적 문제가 있습니다.

mm
Unite.AI를 Google의 선호 소스에 추가

동일한 2비트 정밀도에서 축을 Quantize하는 한 가지 결정이 벤치마크 점수를 2.88에서 63.53으로 변경합니다. 키와 값은 반대되는 처리가 필요합니다. 이유는 하드웨어가 아니라 주의 방정식에 있습니다.

Llama-2-13B를 고려해 보십시오. 키-값 캐시를 Quantization 그룹 크기 32로 그룹화하고, 나머지는 그대로 유지합니다. 동일한 모델, 동일한 비트 예산, 동일한 그룹 크기, 동일한 벤치마크입니다.

한 가지 구현 결정 요소에 따라 CoQA 정확도는 2.88 또는 63.53이 됩니다. 전체 정밀도에서 사용하는 점수는 66.37입니다.

결정은 사용하는 총 비트 수에 관한 것이 아닙니다. 질문은 간단합니다. 각 스케일 팩터를 계산할 때 어떤 축을 그룹화 축으로 사용할지 선택하는 것입니다. 채널을 키의 그룹화 차원으로 사용하고 토큰을 값의 그룹화 차원으로 사용하면 전체 정밀도 성능에서 4점 이내에서 끝나게 됩니다. 이러한 선택 중 하나를 뒤집으면 품질이 손상됩니다. 두 가지 선택을 모두 뒤집으면 모델이 더 이상 작동하지 않습니다.

동일한 2비트를 동일한 캐시에 사용하는 4가지 방법. Llama-2-13B에서 KIVI ablation의 결과입니다.

Quantization은 일반적으로 하나의 다이얼로 생각됩니다. 8비트, 4비트, 2비트와 함께 부드러운 정확도 비용이 있습니다. 그러나 KV 캐시 내에서는 그렇지 않습니다. 좌표계를 선택하는 것입니다. 키와 값에 대해 다른 좌표계를 사용합니다. 이 기사는 이유를 설명합니다. 간단히 말해서, Quantization 오차는 그룹 내의 값 범위에 따라 달라지며, 키와 값은 매우 다른 구조를 가지고 있습니다. 또한 값 분포에서 올바른 축을 유도할 수 없습니다. 주의가 소비하는 방식으로 오차가 변경되는 방식을 살펴보아야 합니다. 이는 중간 활성화를 압축하는 일반적인 원칙을 제공하며, 재구성 오차를 품질의 대리자로 의심하는 좋은 이유가 됩니다.

KV 캐시에서 문제가 발생하는 이유

생성 단계에서 트랜스포머는 이전에 처리한 토큰의 모든 키-값 투영(KV) 데이터를 캐시에 저장하여 데이터를 다시 계산할 필요가 없습니다. 캐시는 컨텍스트 길이와 배치 크기와 함께 선형적으로 증가합니다. 결국 캐시는 모델 자체보다 더 커질 것입니다.

이 증가를 모델의 다양한 부분의 메모리 사용량을 살펴보면 쉽게 식별할 수 있습니다. LLaMA-7B의 KVQuant 분석에서 가중치는 약 98%의 메모리를 차지하며, 활성화는 2%를 차지합니다. 컨텍스트 길이가 128K인 경우, 비율은 약 16%의 가중치와 84%의 KV 캐시로 반전됩니다. KIVI 저자들이 인용한 OPT-175B 분석에서도 유사한 결과가 나타났습니다. 구체적으로, 배치 크기가 512이고 512토큰 프롬프트인 경우, KV 캐시는 1.2TB에 도달하여 모델 가중치의 여러 배가 됩니다.

그러나 용량은 문제의 절반에 불과합니다. GPU는 각 토큰을 생성할 때마다 전체 KV 캐시를 디바이스 메모리에서 읽어야 합니다. 이는 GPU가 KV 캐시를 읽는 동안 컴퓨팅 코어가闲置한다는 것을 의미합니다. 따라서 캐시의 전체 크기를 줄임으로써 처리 가능한 처리량을 증가시키고 데이터 전송을 위해 대기하는 시간을 줄일 수 있습니다.

Quantization 오차는 무엇으로 구성되어 있는가

정수 Quantization은 수학적으로 간단합니다. 숫자 그룹에 대해 가장 작은 숫자를 기준점으로 기록하고, 표현할 수 있는 수준의 수로 그룹의 범위를 나누어 단계 크기를 얻습니다. 그런 다음 각 요소를 가장 가까운 단계로 반올림합니다. 두 가지 즉각적인 결과가 따릅니다. 첫째, 각 요소의 오차는 반 단계 크기 이내로 제한됩니다. 둘째, 단계 크기는 2^B – 1로 그룹의 범위를 나눈 값입니다. 2비트에서는 4개의 수준만으로 존재하는 범위를 커버해야 하므로, 이웃보다 100배 더 큰 요소는 단순히 나쁘게 수행되는 것이 아니라, 같은 그룹을 공유하는 모든 요소의 단계 크기를 늘립니다. 그룹은 손상 단위입니다. 축을 선택하는 것은 어느 요소가 함께 손상되는지를 결정하는 것입니다. 질문을 다르게 프레이밍하면, “몇 비트를 희생할 수 있나요?”가 아니라 “극단적인 값은 어디에 있으며, 그것들을 분리할 수 있나요?”가 됩니다.

키: 고정 채널에 있는 아웃라이어

대규모 언어 모델에는 대부분의 활성화와 비교하여 매우 큰 활성화가 포함되어 있습니다. Sun과 동료들은 다양한 모델 패밀리에서 이러한 매우 큰 활성화를 카탈로그화했습니다. Mixtral 8x7B의 경우 가장 큰 크기는 7000 근처이며, 특징 크기의 중간값은 약 0.3입니다. 즉, 약 4개의 크기 차이가 납니다. 이러한 활성화는 매우 희박하며, 입력과 거의 변경되지 않는 차원에 고정되어 있습니다. 또한 우연히 발생하는 것이 아닙니다. 그것들은 암시적인 편향으로 작용하며, 주의를 몇 개의 토큰에 집중시키는 것입니다. 키 캐시에서는 이러한 구조가 매우 명확합니다. 특정 채널은 일관되게 시퀀스 내의 모든 토큰에서 매우 큰 크기를 갖습니다. 토큰으로 그룹화하면 각 그룹에는 이러한 아웃라이어 채널이 포함되므로, 각 그룹의 단계 크기는 아웃라이어에 의해 설정되고, 모든 정상 채널이 그것을 지불합니다. 채널로 그룹화하면 아웃라이어 채널이 자신의 그룹을 형성합니다. 내부 범위는 크지만 자체적으로 포함되어 있으며, 정상 채널은 혼자 남겨집니다. 결과는 일치합니다. Llama-2-13B에서 평균화된 키 재구성 오차는 토큰당 그룹화에서 13.67, 채널당 그룹화에서 4.55입니다. 또한 주의 점수 오차는 47.00 대 9.60입니다. 토큰당 Quantization 키는 점수 오차가 약 5배 더 많습니다. 채널 Quantization은 의미 있는 키 메트릭에서 두 가지 측면에서 모두 우수합니다.

값: 직관이 깨지는 곳

값 캐시는 채널-아웃라이어 패턴을 표시하지 않습니다. 그것은相当 평평한 것으로 보입니다. 범위에 대한 논증으로, 우리는 두 축이 유사한 품질의 압축을 생성할 것으로 기대할 수 있습니다.

그러나 그렇지 않습니다. 키 관리가 구현되는 방식에 관계없이(2.80 및 2.88 결과), 채널당 값 압축은 모델을 붕괴시킵니다.

여기에는 함정이 있습니다. 원래 텐서에 대한 원시 재구성 오차를 사용하여 이 손실을 측정하면, 채널당 값 Quantization은 실제로 3.73 대 4.57로 약간 더 나은 것으로 보입니다. 압축을 가장 명백한 방법으로 검증한다면, 모델을 파괴하는 구성으로 선택할 것입니다.

Llama-2-13B에서 값 캐시 Quantization 오차, 두 가지 방법으로 측정됨. 저장된 텐서 메트릭과 소비된 출력 메트릭은 한 자릿수 이상으로 다릅니다.

해결책은 값 캐시가 직접 읽히지 않는다는 것입니다. 그것은 행렬곱으로 소비됩니다. 주의 출력은 토큰에 걸친 값 벡터의 가중 합이며, 소프트맥스 주의 점수가 가중치입니다. 이러한 이유로 관련 오차는 텐서 자체 내에서 발생하는 것이 아니라, 이 과정에서 발생하는 것입니다. 주의 출력을 기준으로 측정한 상대 오차는 KIVI에 의해 보고되었으며, 토큰당 값 벡터 Quantization의 경우 3.55, 채널당 Quantization의 경우 49.89로, 후者の 경우 약 14배 더 높습니다.

해설은 주의 희박성입니다. KIVI에서 측정한 희박성은 84.3%입니다. 출력에 포함된 대부분의 정보는 몇 개의 매우 중요한 토큰에 기인합니다. 토큰당 Quantization은 각 토큰의 오차를 해당 토큰으로 제한합니다. 중요하지 않은 토큰의 오차는 거의 0에 가까운 주의 가중치로 곱해져서 실제로 사라집니다. 채널당 Quantization은 모든 토큰의 오차를 공유 채널 스케일로 퍼뜨립니다. 잘 표현되지 않는 토큰은 중요한 토큰의 표현을 오염시킵니다. 주의가 효율적이게 하는 희박성은 토큰당 Quantization이 안전하다는 것을 의미합니다.

전이 가능한 교훈은 KV 캐시를 넘어서는 것입니다. 텐서가 소비되는 곳에서 압축 오차를 측정하십시오. 재구성 오차는 텐서의 모든 구성 요소가 최종 출력에 동등한 가중치를 갖는다는 것을 가정합니다. 주의는 명시적으로 그렇지 않습니다. 입력을 가중, 게이트 또는 희박하게 하는 모든 다운스트림 작업은 이러한 가정을 깨뜨립니다. 이전에 검색 시스템의 평가 지표에 대한 블라인드 스팟에 대한私の 이전 기사와 관련된 독자는 이러한 결과가 이전에 설명된 실패와 유사하다는 것을 알 것입니다. 쉽게 계산할 수 있는 메트릭이 의도한 것과 다른 것을 보고합니다.

회전 임베딩이 키를 복잡하게 만듦

회전 위치 임베딩(RoPE)을 사용하는 데 몇 가지 문제가 있습니다. RoPE는 각 토큰의 상대적 위치에 따라 채널 쌍을 회전시킵니다. 이러한 혼합은 키 Quantization이 작동하도록 만든 고정 채널 구조를 부분적으로 해체합니다. 아웃라이어 채널이 회전되어 이웃에게 범위가 전파됩니다. KIVI의 해결책은 순서입니다. 회전이 적용되기 전에 키를 Quantize하고, 역Quantization 후에 RoPE를 적용합니다. 채널당 키 Quantization, 비균일 데이터 유형 및 아웃라이어의 작은 부분을 분리하면, 3비트에서 0.1 퍼플렉시티 저하 아래로 들어가고, 단일 A100-80GB에서 최대 100만 토큰 컨텍스트에서 LLaMA-7B를 제공할 수 있습니다.

또한 RoPE의 영향 수준을 이해하는 것이 중요합니다. “RotateKV”라는 논문의 저자는 RoPE를 추가하면 Quantization 오차가 145% 증가한다고 보고했으며, 아웃라이어 채널이 주의 헤드 전체에 걸쳐 다르므로, 모든 곳에 공유 회전 행렬을 적용하는 것은 불충분하다고 지적했습니다.

시스템 세금, 그리고 그것이 세부 사항이 아닌 이유

토큰당 Quantization은 디코딩에 적합합니다. 각 토큰이 도착하면, Quantize하고, 시퀀스에 추가합니다. 토큰 차원으로는 아무 것도 이동하지 않습니다.

그러나 채널당 Quantization은 맞지 않습니다. 채널의 통계는 아직 생성되지 않은 토큰을 포함하므로, 토큰이 들어오면 스케일 팩터를 계산할 수 없습니다. KIVI의 해결책은 최근에 생성된 토큰을 최대 128개까지 전체 정밀도로 유지하는 잔여 버퍼를 유지하고, 충분히 모인 후에 Quantization 그룹을 생성하는 것입니다.

그런데 잔여 버퍼는 부차적인 것이 아니라 부하를 지는 것입니다. Llama-2-7B에서 전체 정밀도 점수는 13.50입니다. 올바른 축을 사용하여 2비트로 완전히 Quantization하면 점수는 5.76입니다. 동일한 축과 비트를 사용하고, 최근에 생성된 토큰의 잔여 버퍼를 전체 정밀도로 유지하면 점수는 12.74입니다. 최근에 생성된 토큰의 滑動 창을 전체 정밀도로 유지하면, 공격적인 Quantization으로 인해 어려운 다단계 문제에서 잃어버린 것을 대부분 회복할 수 있습니다.

올바른 모든 것을 수행하면 상당한 이점이 있습니다. KIVI에 따르면, Llama-2-7B에서 피크 메모리 사용량이 2.6배 줄어들고, 배치 크기가 최대 4배까지 증가하며, 실제 서비스 작업에서 처리량이 2.35에서 3.47배까지 개선됩니다.

이것을 어떻게 사용할 것인가

  1. 하나의 Quantizer를 사용하지 마십시오. 키(채널당)와 값(토큰당)에 대해 다른 Quantizer를 사용하십시오. 단일 Quantizer를 “KV 캐시”에 적용하는 파이프라인은 이미 작은 수의 비트를 사용하여 각 값을 나타낼 때 대부분의 가능한 품질을 희생했습니다.
  2. RoPE 전에 키를 Quantize하십시오. 이것은 선호도 문제가 아니라 올바름 문제입니다.
  3. 최근에 생성된 토큰의 전체 정밀도 창을 저장하십시오. 캐시보다 매우 적은 메모리를 사용하는 이 영역은 어려운 작업에 대한 정확도 대부분을 생성합니다.
  4. 재구성 오차로 검증하지 마십시오. 주요 출력 또는 최종 작업 성능을 기준으로 항상 검증하십시오. 저장 메트릭은 그냥 시끄럽지 않습니다. 값의 경우 완전히 반대 방향을 가리킵니다.
  5. 짧은 컨텍스트 다중 선택 벤치마크에서 검증하지 마십시오. KIVI 저자는 의도적으로 캐시를 시간이 지남에 따라 구축하고 생성에서 수행하는 단일 디코딩 단계만을 읽는 폐쇄형 작업을 피합니다. 캐시를 시간이 지남에 따라 구축하고 생성에서 수행하지 않는 모든 평가에서는 시스템 디자인의 내재된 실패를 관찰할 수 없습니다.

작업이 어디로 향하는가

기하학적 문제의 많은 부분이 여전히 남아 있지만, 많은 연구자들은 변압기 헤드 전체에 걸쳐 분산된 아웃라이어 채널과 하드웨어 제한이 어느 그룹화가 가장 저렴한지에 대한 방법을 계속 연구하고 있습니다. 예를 들어, InnerQ는 채널 방향 키 정규화를 키 및 쿼리 가중치에 대해 사전 채움 동안 접을 수 있습니다. 따라서 런타임에 추가 오버헤드가 발생하지 않습니다. 또한 InnerQ는 최근에 생성된 토큰과 주의 싱크 토큰 모두에 대해 높은 정밀도 창을 저장합니다.这样하면 싱크 채널의 아웃라이어가 중요한 채널의 표현을 오염시키는 기회가 제거됩니다.

다른 사람들은 캐시의 전체를 저장하는 대신, 키와/또는 값을 필요에 따라 다시 구축할 수 있는 정보만 캐시에 저장해야 한다고 제안합니다.

마지막으로, Quantization이 정확도만 영향을 미치는 것이 아니라는 것을 기억하는 것이 중요합니다. 최근에 발표된 연구에서는 KV 캐시를 Quantization하는 것이 정렬退化를 초래한다는 것을 보여주었습니다. 또한 FP8 캐시와 함께 훈련 없는 복구 프로토콜을 사용하여 최대 97%의 손실을 복구하는 생산 환경의 vLLM에서 이러한 정렬退化가 발생한다는 것을 문서화했습니다. 따라서 구성이 벤치마크 결과를 유지하더라도, 관심 있는 모든 다른 매개변수를 유지한다는 것을 의미하지는 않습니다.

일반 원칙

Quantization의 아이디어는 “정밀도 예산”으로 프레이밍되었습니다. 몇 비트를 희생할 수 있나요? 그러나 KV 캐시에서는 구조적인 질문입니다. 정밀도는 그룹화로 할당되며, 그룹은 손상 단위입니다. 축을 선택하는 것은 어느 요소가 함께 손상되는지를 결정합니다. 올바른 축은 텐서가 소비되는 방식입니다. 키는 닷 프로덕트 컴퓨테이션을 통해 사용됩니다. 하나의 손상된 채널이 모든 점수를 오염시킵니다. 값은 희박한 가중 평균 컴퓨테이션을 통해 토큰에 걸쳐 소비됩니다. 따라서 하나의 손상된 토큰은 단순히 가중되어 사라집니다.

두 개의 텐서가 동일한 차원을 가지며, 두 개의 연속적인 레이어에서 생성되더라도, 다르게 처리됩니다. 압축할 활성화가 계획된 경우, 어떤 작업이 이를 계약하는지, 그리고 그룹화가 이를尊重하는지 묻는 것이 가치 있습니다.

히만수 고엘은 생물의학, 금융, 및 규제 문서 워크플로우를 포함하는 고위험 도메인에서 검색 증강 생성을 전문으로 하는 AI/ML 연구자입니다.