[논문 리뷰] 8-bit Optimizers via Block-wise Quantization

Tim Dettmers, Mike Lewis외 2명|arXiv (Cornell University)|2021. 10. 06.

Advanced Neural Network Applications참고 문헌 58인용 수 4

한 줄 요약

이 논문은 블록 단위 동적 양자화를 사용하여 32비트 최적화기 성능을 유지하면서 메모리 사용을 크게 줄인 8비트 최적화기를 소개한다. 최적화기 상태(예: Adam, AdamW)를 독립적인 블록으로 나누어 동적 양자화와 안정적인 임bedding 레이어를 적용함으로써, 초과 15억 파라미터 모델에서 하이퍼파라미터 조정 없이도 전체 훈련 안정성과 속도를 확보한다.

ABSTRACT

Stateful optimizers maintain gradient statistics over time, e.g., the exponentially smoothed sum (SGD with momentum) or squared sum (Adam) of past gradient values. This state can be used to accelerate optimization compared to plain stochastic gradient descent but uses memory that might otherwise be allocated to model parameters, thereby limiting the maximum size of models trained in practice. In this paper, we develop the first optimizers that use 8-bit statistics while maintaining the performance levels of using 32-bit optimizer states. To overcome the resulting computational, quantization, and stability challenges, we develop block-wise dynamic quantization. Block-wise quantization divides input tensors into smaller blocks that are independently quantized. Each block is processed in parallel across cores, yielding faster optimization and high precision quantization. To maintain stability and performance, we combine block-wise quantization with two additional changes: (1) dynamic quantization, a form of non-linear optimization that is precise for both large and small magnitude values, and (2) a stable embedding layer to reduce gradient variance that comes from the highly non-uniform distribution of input tokens in language models. As a result, our 8-bit optimizers maintain 32-bit performance with a small fraction of the memory footprint on a range of tasks, including 1.5B parameter language modeling, GLUE finetuning, ImageNet classification, WMT'14 machine translation, MoCo v2 contrastive ImageNet pretraining+finetuning, and RoBERTa pretraining, without changes to the original optimizer hyperparameters. We open-source our 8-bit optimizers as a drop-in replacement that only requires a two-line code change.

연구 동기 및 목표

• 대규모 모델 훈련에서 32비트 최적화기 상태의 높은 메모리 사용량(총 메모리의 33–75%)을 해결한다.
• 10억 파라미터 초과 모델에서 8비트 최적화의 과제인 양자화 정확도, 계산 효율성, 훈련 안정성 문제를 해결한다.
• 32비트 대체품으로 쉽게 사용할 수 있는 실용적이고 고성능 8비트 최적화기를 구현한다.
• 모델과 기울기에는 16비트 혼합 정밀도를 사용함으로써, 양자화가 최적화기 상태에 미치는 영향을 격리한다.
• 기존 하이퍼파라미터를 수정하지 않고도 다양한 NLP 및 비전 작업에서 안정성과 성능을 확보한다.

제안 방법

• 블록 단위 양자화 적용: 최적화기 상태 텐서를 독립적인 블록으로 나누어 병렬 처리와 국소 정규화를 가능하게 한다.
• 동적 양자화 사용: 블록 별로 적응적인 지수를 가지는 비선형 양자화로, 작은 값과 큰 값 모두 정밀도를 유지한다.
• 안정적인 임bedding 레이어 도입: 극도로 비균일한 입력 토큰 분포를 정규화하여 기울기 분산을 감소시킨다.
• 효율적인 블록 단위 정규화 및 양자화/해양자화 연산을 위한 커스텀 CUDA 커널 구현.
• 32비트 성능 유지: 최적화기 업데이트를 32비트로 수행한 후 상태를 8비트로 양자화하여 저장한다.
• 역전파 단계에서 해양자화를 위한 룩업 테이블 사용으로, 빠르고 정확한 상태 복구를 가능하게 한다.

실험 결과

연구 질문

• RQ1 10억 파라미터를 초과하는 대규모 모델에서 하이퍼파라미터 조정 없이 8비트 최적화기 상태가 32비트 훈련 성능과 안정성을 유지할 수 있는가?
• RQ2 텐서 전체에 대한 양자화와 비교해 블록 단위 양자화가 안정성과 효율성에 어떻게 기여하는가?
• RQ3 최적화기 상태의 극단적인 값 범위에서 양자화 오차를 줄이는 데 동적 양자화가 어떤 역할을 하는가?
• RQ4 비균일한 토큰 분포로 인한 기울기 분산을 안정적인 임bedding 레이어가 어느 정도 감소시키는가?
• RQ5 8비트 최적화기는 다양한 NLP 및 비전 작업에서 32비트 최적화기의 즉각적인 대체로 구현될 수 있는가?

주요 결과

• 제안된 8비트 최적화기는 15억 및 3억 5,500만 파라미터 언어 모델링, GLUE 피니팅, ImageNet 분류, WMT’14/16 번역, MoCo v2 프리트레인 + 피니팅, RoBERTa 프리트레인 등 다양한 작업에서 32비트 성능을 유지한다.
• 기존 32비트 대비 최적화기 상태의 메모리 사용량이 최대 80% 감소했으며, 하이퍼파라미터는 그대로 유지되었다.
• 블록 단위 양자화는 블록 내에서 이질적인 값(이상치)을 격리함으로써 전반적인 양자화에 미치는 영향을 줄여 훈련 안정성을 향상시킨다.
• 동적 양자화는 최적화기 상태 값의 5자리 수준의 범위에서 고정밀도 표현을 가능하게 하여 대규모 훈련에 필수적이다.
• NLP 작업에서는 안정적인 임bedding 레이어가 비균일한 토큰 입력으로 인한 기울기 분산을 감소시켜 8비트 훈련의 안정성을 확보하는 데 핵심적이다.
• 제거 실험 결과, 블록 단위 양자화, 동적 양자화, 안정적인 임bedding 레이어 모두가 전체 성능과 안정성 확보에 필수적임을 확인했다.