[논문리뷰] The Era of 1-bit LLMs:All Large Language Models are in 1.58 Bits

 

MoRA논문을 읽어보며 양자화에 대해 궁금해졌으며 관련 리서치들을 읽어보는 중https://jihoonjung.tistory.com/93

[논문리뷰] LLM.int8() : 8-bit Matrix Multiplication for Transformers at Scale

 

 

LLM.int8에 의하면, 극단의 값을 갖는 텐서 및 웨이트가 매우 중요함

-> 허나 1bit LLM에서는 어떻게 -1,0,1로 LM이 작동하는지....

-> -1,0,1로 Scratch부터 학습을 하게 되면 가능한건가? 싶었음.

 

0 Abstract

• 최근 연구인 BitNet[WMD+23]은 1-비트 대규모 언어 모델(LLM) 시대를 열고 있다.
• 이 연구에서는 모든 매개변수(또는 가중치)가 3진법 {-1, 0, 1}인 1-비트 LLM 변형인 BitNet b1.58을 소개
• 동일한 모델 크기와 학습 토큰 수를 가진 전체 정밀도(FP16 또는 BF16) Transformer LLM와 지연 시간, 메모리, 처리량, 에너지 소비 측면에서 유사한 성능을 보이면서도 상당히 비용 효율적

여기서 1.58Bits는 -> [-1,0,1] 세가지의 숫자만 사용함 -> 로그 2의 (3)은 -> 1.58이다.

 

 

1 The Era of 1-bit LLMs

• 1-비트 LLM은 크기와 성능이 빠르게 증가하는 추세에서 비용 효율성을 높이기 위한 해결책으로 주목받음.
• 8비트(XLS+23, FAHA23, CCKS23, TCS+24)연구, 4비트 변형(FAHA23, LTT+23) 연구가 있음
• 이 연구에서는 1-비트 LLM의 변형인 BitNet b1.58을 소개하며, 이는 모든 매개변수가 3진수 {-1, 0, 1} 값을 갖는다. 
  
  • BitNet의 행렬 곱셈은 정수 덧셈만 포함하므로 LLM의 에너지 비용을 크게 절약
• BitNet b1.58은 기존 1-비트 BitNet의 장점을 유지하면서도 모델링 능력이 향상되었고, 3B 크기부터 전체 정밀도 모델과 동등한 성능을 보인다

2 BitNet b1.58

• BitNet b1.58은 BitNet 아키텍처를 기반, 이는 nn.Linear를 BitLinear로 대체한 Transformer이다

BitNet이라는 논문이 있음

• 처음부터 1.58-비트 가중치와 8-비트 activation를 사용하여 학습된다. 

BitNet은 추후 블로그에서 다룰예정이지만, 가볍게 한번 살펴보자

BitLinear Block이 있는데, 이는 [-1,0,1]로 맵핑할 수 있는 방법이라 생각하자.

아키는 생각보다 간단하다. 큰그림으로 보면 기존 Transformers구조와 동일하다.

 

Quantization Function(해당 논문이 주장하는 것)

• 양자화를  [-1, 0, +1]로 하기위해 absmean 양자화 함수를 사용.
• 먼저 가중치 행렬을 절대값 평균으로 스케일링한 다음 각 값을 {-1, 0, +1} 중 가장 가까운 정수로 반올림

• γ : 가중치 행렬 W의 절대값 평균, 행렬의 모든 원소의 절대값을 더한 후 행렬의 총 원소 수(nm)로 나눈 값
• ϵ : 매우 작은 양수 값
• a,b : -1, 1로 설정하여 [-1,1] 사이의 값이 나오도록 한다.

 

• activation을 위한 양자화 함수는 BitNet의 구현과 동일
• 비선형 함수 전에 활성화를 [0, Q_b] 범위로 스케일링하지 않는다
• 대신 토큰별로 activation을 모두 [-Q_b, Q_b] 범위로 스케일링함(제로-포인트 양자화 제거)
  • 제로-포인트 양자화 : 양자화된 값의 범위를 이동, 양의 정수만 사용할 수 있는 경우에 유용
• 구현과 시스템 수준 최적화 모두에 더 편리하고 단순하며, 성능에 미미한 영향 미침

BitNet b1.5B의 모델 설계

• LLaMA와 유사 : RMSNORM, SwiGLU, Rotarty Embedding, No Bias Term

 

3 Results

BitNet b1.58 vs LLaMA(FP16)

• 공정한 비교를 위해 RedPajama 데이터셋에서 1000억 토큰으로 모델을 사전 학습
• 양한 언어 작업에 대해 제로샷 성능을 평가
• GPU메모리 비교

표 1은 perplexity와 비용 비교

• BitNet b1.58이 3B 모델 2.71배 빠르고 GPU 메모리를 3.55배 적게 사용
• 특히 3.9B 모델 크기의 BitNet b1.58은 2.4배 빠르고 메모리를 3.32배 적게 소모하지만 LLaMA LLM 3B보다 성능이 훨씬 우수

 

표 2는 엔드 태스크에 대한 제로샷 정확도

 

• 모델 크기가 커질수록 BitNet b1.58과 LLaMA LLM 간의 성능 격차가 좁혀짐
• BitNet b1.58이 3B 크기부터 전체 정밀도 기준 성능과 일치
• BitNet b1.58 3.9B가 더 적은 메모리와 지연 시간 비용으로 LLaMA LLM 3B보다 우수한 성능을 보임

 

Memory and Latency

• 모델 크기를 7B, 13B, 70B로 더 확장&비용을 평가
• BitNet b1.58 70B는 LLaMA LLM 기준보다 4.1배 빠름
  • nn.Linear에 소요되는 시간이 모델 크기에 따라 증가하기 때문

Throughout

• 두 개의 80GB A100 카드에서 70B 파라미터를 가진 BitNet b1.58과 LLaMA LLM의 처리량을 비교
• 시퀀스 길이를 512로 하여 GPU 메모리 한계에 도달할 때까지 배치 크기를 늘림
• 표 3은 BitNet b1.58 70B가 LLaMA LLM의 최대 11배 배치 크기를 지원할 수 있어 처리량이 8.9배 더 높음

 

BitNet b1.58 is enabling a new scaling law with respect to model performance and inference

• 13B BitNet b1.58은 지연 시간, 메모리 사용량, 에너지 소비 측면에서 3B FP16 LLM보다 효율적
• 30B BitNet b1.58은 지연 시간, 메모리 사용량, 에너지 소비 측면에서 7B FP16 LLM보다 효율적
• 70B BitNet b1.58은 지연 시간, 메모리 사용량, 에너지 소비 측면에서 13B FP16 LLM보다 효율적

 

4 Discussion and Future Work

 

1-bit Mixture-of-Experts (MoE) LLMs

• 전문가 혼합(MoE)은 LLM에 비용 효율적인 접근 방식임이 입증
• MoE는 계산 FLOPs를 크게 줄이지만, 높은 메모리 소비와 칩 간 통신 오버헤드로 인해 배포 및 적용이 제한
  • 이러한 문제는 1.58-비트 LLM으로 해결할 수 있다.

MoE가 궁금하다면 여기로...

https://jihoonjung.tistory.com/92

[논문리뷰] Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

 

Native Support of Long Sequence in LLMs

• 긴 시퀀스를 처리하는 능력은 중요
• 긴 시퀀스 추론의 주요 과제 중 하나는 KV 캐시로 인한 메모리 소비
  • BitNet b1.58은 활성화를 16비트에서 8비트로 줄여 동일한 리소스로 컨텍스트 길이를 두 배로 늘릴 수 있음.
  • 1.58-비트 LLM에 대해 4비트 또는 더 낮은 비트로 무손실 압축가능

LLMs on Edge and Mobile

1.58-비트 LLM의 사용은 엣지 및 모바일 기기에서 언어 모델의 성능을 크게 향상시킬 수 있다.