BitNet

BitNet

BitNet은 Multi-Head Attention과 피드포워드 네트워크의 일반적인 선형 레이어를 전용 BitLinear 레이어로 대체합니다. BitLinear 레이어는 가중치를 삼진 정밀도(ternary precision, 값이 -1, 0, 1)로 양자화하고 활성화는 8비트 정밀도로 양자화해요.

출처: 문서

본문

BitLinear 레이어를 갖춘 BitNet의 아키텍처.

BitNet 모델은 그 자리에서(on the fly) 양자화할 수 없어요. 양자화 인지 훈련(Quantization-Aware Training, QAT) 기법이기 때문에 사전 훈련(pre-training)이나 미세 조정하는 동안 양자화해야 합니다. 훈련 중에는 가중치가 대칭 텐서별(symmetric per tensor) 양자화로 삼진 값에 양자화돼요.

  1. 가중치 행렬 절댓값의 평균을 계산해 스케일(scale)로 사용합니다.
  2. 가중치를 스케일로 나누고 값을 반올림한 뒤 -1과 1 사이로 제약하고, 다시 스케일링해 전체 정밀도(full precision)로 계속 진행해요.
  3. 활성화는 absmax 양자화(대칭 채널별 양자화)로 지정된 비트 폭(8비트)에 양자화합니다. 여기서 활성화를 [−128,127] 범위로 스케일링하는 작업이 포함돼요.

Nanotron으로 1.58비트 모델을 사전 훈련하거나 미세 조정하려면 이 PR을 참고하세요. 미세 조정의 경우 Hugging Face 모델을 Nanotron 형식으로 변환해야 하는데, 변환 단계는 이 PR에서 확인할 수 있습니다.

BitNet으로 양자화된 모델은 from_pretrained()로 불러옵니다.

from transformers import AutoModelForCausalLM
path = "/path/to/model"
model = AutoModelForCausalLM.from_pretrained(path, device_map="auto")

커널 (Kernels)

가중치를 풀어내고(unpack) 포워드 패스를 수행하는 데 @torch.compile이 사용됩니다. 구현이 아주 간단하면서도 상당한 속도 향상을 가져와요. 앞으로의 버전에는 추가로 최적화된 커널이 통합될 예정입니다.

더 알아보기 (Learn more)

Fine-tuning LLMs to 1.58bit: extreme quantization made easy를 읽으면 BitNet 모델이 어떻게 훈련되고 미세 조정되는지 더 자세히 알 수 있어요.