★★★ ZeroQuant: Efficient and Affordable Post-Training Quantization for Large-Scale Transformers

Basic Information

• Zhewei Yao, Reza Yazdani Aminabadi, Minjia Zhang, Xiaoxia Wu, Conglong Li, Yuxiong He 저
• Microsoft Paper
• 논문 아카이브: Link
• 코드 아카이브: Link (Microsoft Official)

서론

• 본 논문은 큰 트랜스포머 기반 모델을 압축하기 위한 효율적인 후훈련 양자화 접근법인 ZeroQuant를 소개합니다. ZeroQuant는 세 가지 주요 구성 요소로 이루어져 있습니다:
  • 가중치와 활성화 모두에 대한 세밀하고 하드웨어 친화적인 양자화 체계
  • 원래의 훈련 데이터에 접근하지 않아도 작동하는 새로운 계층별 지식 증류 알고리즘 (LKD)
  • 양자화와 역양자화의 오버헤드를 제거하기 위한 최적화된 양자화 시스템 백엔드 지원
• 본 논문은 ZeroQuant가 다음과 같은 성과를 낼 수 있다고 주장합니다:
  • BERT와 GPT-3 스타일 모델 모두에 대해 가중치와 활성화의 정밀도를 INT8로 줄이면서 최소한의 정확도 영향을 미치며, FP16 추론에 비해 상당한 속도 향상을 이룹니다.
  • LKD와 함께, 완전 연결 모듈의 가중치를 INT4로 양자화하고, 주의 모듈에서는 INT8 가중치와 INT8 활성화를 사용하여 FP16 모델에 비해 메모리 사용량을 3배 줄일 수 있습니다.
  • GPT-J6B와 GPT-NeoX20B와 같은 가장 큰 오픈 소스 언어 모델에 직접 적용될 수 있으며, 여기서 INT8 모델은 FP16 모델과 유사한 정확도를 달성하지만 최대 5.2배 더 효율적입니다.

배경

• Post-Training Quantization(PTQ)은 보통 모델을 재훈련 없이 양자화하는 PTQ를 적용하기 때문에 Quantization-Aware-Training(QAT)에 비해 큰 압축 효율을 보입니다.
• PTQ의 일반적인 전략은 훈련 데이터를 네트워크에 공급하고 실행 중인 평균을 사용하여 스케일링 요소 S를 보정하는 것입니다.
• BERTbase 모델에 대해 INT8 가중치와 혼합된 INT8/FP16 활성화 양자화와 함께 일부 작업이 수행되었습니다.

• Dynamic Activation Range
  • INT8 활성화가 BERT와 GPT-3 스타일 모델 모두에 대해 중요한 정확도 감소를 초래하는 이유를 조사하기 위해, 우리는 GPT-3350M의 다른 트랜스포머 계층에 대한 각 활성화의 토큰별(즉, 각 토큰의 숨겨진 상태) 범위를 Figure 1(왼쪽)에 표시합니다.
  • 보시다시피, 다른 토큰들은 활성화 범위가 크게 다릅니다.
  • 예를 들어, 마지막 계층의 최대 범위는 약 35 근처이지만 최소 범위는 8에 가깝습니다.
  • 이 큰 활성화 범위의 변동성 때문에 모든 토큰에 대해 예측 정확도를 유지하기 위해 고정된 양자화 범위(보통 최대 값)를 사용하는 것이 어렵습니다.
  • 왜냐하면 작은 범위의 토큰에 대한 제한된 표현력이 정확도 성능을 저하시킬 것이기 때문입니다.
• Different Ranges of Neurons in Weight Matrices
  • 마찬가지로, 우리는 GPT-3350M의 주의 출력 행렬(Wo)의 행별(즉, 출력 차원) 가중치 범위를 Figure 1(오른쪽)에 표시합니다.
  • 다른 행들의 가장 큰 크기 사이에는 10배의 차이가 있으며, 이것은 INT8 가중치 PTQ의 성능 저하를 초래합니다.
  • 이것은 또한 INT4 양자화가 적용될 때 매우 도전적으로 만듭니다.
  • 왜냐하면 INT4는 오직 16개의 숫자만 가지고 있고, 10배 더 작은 범위는 그러한 작은 범위 행의 표현을 위해 2(또는 3)의 숫자로 이어집니다.
  • 이 분석 결과는 더 비싼 hidden-state Knowledge Distillation이 극도로 낮은 정밀도 양자화를 위해 정확도 격차를 줄이기 위해 사용되는 이유를 나타냅니다.
  • 그러나 대규모 모델에 대한 KD의 훈련 비용이 너무 높기 때문에, PTQ에 대한 경량이면서 효율적인 방법이 바람직합니다.

연구 방법

세밀한 하드웨어 친화적 양자화 체계

• BERT/GPT-3 스타일 모델에 INT8 PTQ를 적용하더라도 중요한 정확도 저하가 발생합니다.
• 주요 도전 과제는 INT8의 표현이 가중치 행렬의 다른 행들과 다른 활성화 토큰의 다양한 숫자 범위를 완전히 포착하지 못한다는 것입니다.
• 이 문제를 해결하는 한 가지 방법은 가중치 행렬(활성화)에 대해 그룹별(토큰별) 양자화를 사용하는 것입니다.

가중치에 대한 그룹별 양자화

• 해당 개념은 "Q-BERT: Hessian based ultra low precision quantization of bert"에서 처음으로 제안됐습니다. 가중치 행렬 W2 Rnm은 g 그룹으로 분할되며, 각 그룹은 별도로 양자화됩니다.
  • 그러나 해당 논문에서 저자들은 이를 양자화 인식 훈련에만 적용합니다. 더 중요한 것은, 하드웨어 효율성 제약을 고려하지 않고 시스템 백엔드 지원이 없습니다.
  • 따라서 실제 지연 시간 감소 혜택이 부족합니다.
• 우리의 설계에서는 GPU의 암페어 아키텍처에서의 하드웨어 제약을 고려합니다(예: A100).
• 여기서 계산 단위는 Warp Matrix Multiply와 Accumulate (WMMA) 타일링 크기를 기반으로 하여 최상의 속도 향상을 달성합니다.
• 나중에 우리의 그룹별 양자화가 더 세밀한 양자화로 인해 단일 행렬 양자화에 비해 훨씬 더 나은 정확도를 달성하면서도 큰 지연 시간 감소를 달성한다는 것을 보여줄 것입니다.

활성화에 대한 토큰별 양자화

• 기존 PTQ 작업에 대한 일반적인 방법은 활성화에 대해 정적 양자화를 사용하는 것입니다.
• 여기서 최소/최대 범위는 오프라인 보정 단계에서 계산됩니다.
  • 이러한 방법은 활성화 범위의 변동이 작은 소규모 모델에는 충분할 수 있습니다. 그러나 GPT-3350M 및 BERTbase와 같은 대규모 트랜스포머 모델에는 활성화 범위에서 큰 변동이 있습니다.
  • 따라서 정적 양자화 체계(대개 모든 토큰/샘플에 적용됨)는 중요한 정확도 감소를 초래할 것입니다.
• 이 문제를 극복하기 위한 자연스러운 아이디어는 더 세밀한 토큰별 양자화를 채택하고 각 토큰에 대한 최소/최대 범위를 동적으로 계산하여 활성화로부터의 양자화 오류를 줄이는 것입니다.
• 결과적으로 활성화에 대한 토큰별 양자화가 GPT-3 스타일 및 BERT 모델의 정확도를 크게 향상시킨다는 것을 보여줍니다.
• 그러나 기존 DL 프레임워크, 예를 들면 PyTorch 양자화 스위트를 사용하여 토큰별 양자화를 직접 적용하면, 토큰별 양자화가 GPU 계산 단위와 주 메모리 사이에서 비싼 데이터 이동 오버헤드를 초래하는 추가 작업을 도입하기 때문에 상당한 양자화 및 역양자화 비용이 발생합니다.
• 이 문제를 해결하기 위해, 우리는 트랜스포머 모델의 토큰별 양자화를 위한 매우 최적화된 추론 백엔드를 구축합니다.
• 예를 들어, ZeroQuant의 추론 백엔드는 커널 융합 기술을 사용하여 양자화 연산자를 이전 연산자, 예를 들면 계층 정규화와 같은 것,와 융합하여 토큰별 양자화로부터의 데이터 이동 비용을 완화합니다.
• 마찬가지로, 다른 GeMMs의 출력의 역양자화 비용은 가중치 및 활성화 양자화 스케일을 사용하여 INT32 누적을 스케일링함으로써 완화되며, 다음 FP16 연산자(예: GeLU)에 대한 주 메모리로 최종 FP16 결과를 다시 작성하기 전에 완화됩니다.
• 토큰별 양자화는 양자화된 활성화에 대한 표현 오류를 크게 줄일 수 있습니다.
• 또한 활성화 범위를 보정할 필요가 없기 때문에, 추후 ZeroQuant에 대해서 활성화 범위 보정과 같은 양자화 관련 비용이 없다는 것을 보여줄 것입니다(예: INT8 가중치와 INT8 활성화를 사용하는 중간 양자화 체계).

합리적인 비용의 Layer-by-Layer Knowledge Distillation

• 지식 증류(KD)는 모델 압축 후 정확도 저하를 완화하는 가장 강력한 방법 중 하나입니다.
• 그러나 KD에는 여러 가지 제한 사항이 있으며, 특히 대규모 언어 모델에서의 숨겨진 상태 KD에 대한 제한 사항이 있습니다:
  • (1) KD는 교육 중에 교사와 학생 모델을 함께 유지해야 하므로 메모리와 계산 비용이 크게 증가합니다.
  • (2) KD는 일반적으로 학생 모델의 전체 훈련을 필요로 합니다. 따라서 가중치 매개변수의 여러 복사본(기울기, 첫 번째/두 번째 순서 모멘텀)이 메모리에 저장되어 모델을 업데이트해야 합니다.
  • (3) KD는 일반적으로 원래의 훈련 데이터를 필요로 하며, 때로는 개인 정보/기밀 문제로 인해 접근할 수 없습니다.
• 이러한 제한 사항을 해결하기 위해, 우리는 계층별 증류(LKD) 알고리즘을 제시합니다.
• 양자화 대상 모델이 N개의 트랜스포머 블록, L1, ..., LN을 가지고 있으며, 접근 가능한 데이터셋이 입력 (X & Y)을 가지고 있다고 가정합니다. 이는 원래의 훈련 데이터 또는 다른 리소스에서의 데이터셋일 수 있습니다.
• 우리의 LKD는 네트워크를 계층별로 양자화하고 원래의(즉, 양자화되지 않은) 버전을 교사 모델로 사용합니다.
• 구체적으로, Lk 계층이 양자화될 것이라고 가정하고, 그 양자화된 버전은 bLk입니다. 그런 다음 Lk-1의 출력(즉, 처음 k-1 계층에서 X에 대한 추론을 실행함)을 Lk와 bLk의 입력으로 사용하고 차이를 측정하고 Lk에 모델 업데이트를 수행합니다.
• 즉,

• 여기서 MSE는 평균 제곱 오차이며, 다른 손실(예: KL 발산)로도 대체될 수 있습니다.
• 보시다시피,
  • (1) 우리의 LKD는 교사/학생 모델 모두에 대해 동일한 L1에서 Lk-1을 사용하기 때문에 별도의 교사를 유지할 필요가 없습니다. 따라서 우리가 가진 유일한 추가 모델 비용은 Lk입니다.
  • (2) 최적화 계층이 Lk뿐이므로 최적화 상태의 메모리 오버헤드가 크게 줄어듭니다.
  • (3) 우리는 끝에서 끝까지의 모델을 최적화하지 않기 때문에, 훈련은 더 이상 레이블에 의존하지 않습니다.
    

Quantization-Optimized 트랜스포머 커널

• 실제로 대규모 트랜스포머 모델을 서비스하는 데 있어 추론 지연 시간과 모델 크기를 최적화하는 것은 중요합니다.
• 추론 중에는 배치 크기가 상대적으로 작은 경우가 많으므로 모델의 추론 지연 시간은 주로 주 메모리에서 필요한 추론 데이터를 로드하는 시간에 의존합니다.
• 가중치와 활성화를 더 낮은 정밀도로 양자화함으로써 이러한 데이터를 로드하는 데 필요한 데이터 양을 줄이게 되어 메모리 대역폭을 보다 효과적으로 사용하고 더 높은 로딩 처리량을 달성할 수 있습니다.
• 그러나 단순히 가중치/활성화를 INT8로 변환하는 것만으로는 개선된 지연 시간을 보장하지 않습니다.
• 왜냐하면 Figure 2(빨간색 상자)에 나타나 있는 양자화/역양자화 작업과 관련된 추가 데이터 이동 오버헤드가 있기 때문입니다.
• 이러한 오버헤드는 비용이 많이 들며, 경우에 따라서는 저 정밀도 사용의 성능 이점을 초과하기도 합니다.
• 토큰별 양자화로부터의 정확도 향상을 얻으면서 개선된 지연 시간을 얻기 위해, 우리는 ZeroQuant에 대한 추론 지연 시간을 가속화하기 위해 메모리 대역폭 활용을 극대화하는 최적화를 소개합니다.

CUTLASS INT8 GeMM

• INT8 계산을 지원하기 위해, 우리는 다양한 배치 크기에 튜닝된 CUTLASS INT8 GeMM 구현을 사용합니다.
• cuDNN과 같은 표준 GPU 백엔드 라이브러리와 달리 CUTLASS를 사용하면 GeMM 전후의 양자화 작업을 더 유연하게 융합하여 커널 시작과 데이터 이동 오버헤드를 줄일 수 있습니다.

Fusing Token-wise Activation Quantization

• 토큰별 양자화/역양자화는 추가적인 데이터 이동 비용을 초래하는 많은 추가 작업을 도입합니다.
• 이러한 비용을 제거하기 위해, 우리는 커널 융합을 사용하여 Figure 2의 녹색 상자에 나타나 있는 것처럼 bias-add, GeLU, LayerNorm과 같은 이전의 요소별 및/또는 축소 작업과 함께 활성화를 위한 양자화 작업을 단일 연산자로 융합합니다.
• GeMM 연산자에서의 정수 출력을 역양자화하는 작업(예: dequantizing the integer output from the GeMM operator)에 대해, 우리는 마찬가지로 Figure 2의 파란색 상자에 나타나는 것처럼 주 메모리로의 추가 읽기/쓰기 액세스를 피하기 위해 사용자 정의 GeMM 일정과 함께 그것을 융합합니다.
• 위의 최적화를 수행함으로써, 우리는 BERT와 GPT-3 스타일 모델에 대한 중요한 지연 시간 감소를 보여줄 수 있습니다.
  

실험 결과

결론

• 대규모 모델 크기의 빠른 성장으로 인해, 우리는 실제로 이러한 모델들을 어떻게 서비스할 것인지 고려해야 하는 시점에 도달했습니다.
• 여러 연구에서는 BERT 모델에 후훈련 양자화를 적용할 수 있다는 것을 보여주고 있지만, 우리의 지식으로는
  • (1) 10억 규모의 GPT-3 스타일 모델
  • (2) 초저 정밀도의 후훈련 양자화
  • (3) 양자화된 모델을 온라인으로 효율적으로 제공
• 하는 End-to-End 솔루션에 대한 기존의 연구가 없습니다.
• 이 연구에서는 최대 20B 규모의 모델(GPT-NeoX20B)에 대한 INT8 양자화를 가능하게 하는 가중치와 활성화를 위한 세밀한 압축 체계를 제공합니다.
• 또한 우리는 초저 정밀도 양자화를 위한 새로운 합리적인 계층별 지식 증류를 제공하며, 이는 FP16 모델과 비교하여 3배의 모델 크기 감소를 가져오면서 최소한의 정확도 저하를 달성합니다.
• 더욱이, 우리는 시스템 백엔드 지원을 제공하며 BERT 모델에서 최대 5.19배의 속도 향상과 GPT-NeoX20B에서 5.2배 더 나은 효율성을 보여줍니다.