BitNet

BitNet은 Multi-Head Attention과 feed-forward network의 기존 linear layer를 전용 BitLinear layer로 대체합니다. BitLinear layer는 가중치를 3원 정밀도(값이 -1, 0, 1)로 양자화하고, 활성화 값을 8비트 정밀도로 양자화합니다.

BitNet 모델은 추론 시점에 양자화(on-the-fly quantization)할 수 없습니다. 양자화 인식 학습(Quantization-Aware Training, QAT) 기법이기 때문에 사전 학습 또는 파인튜닝 단계에서 양자화해야 합니다. 학습 중에는 대칭 텐서별(per-tensor) 양자화를 통해 가중치를 3원(ternary) 값으로 양자화합니다.

1. 가중치 행렬의 절댓값 평균을 계산하여 스케일로 사용합니다.
2. 가중치를 scale로 나누고, 값을 반올림한 뒤, -1과 1 사이로 제한하고, 다시 원래의 정밀도로 되돌려 계속 진행합니다.
3. Activations는 absmax 양자화(채널별 대칭 양자화)를 사용하여 지정된 비트 폭(8비트)으로 양자화됩니다. 이 과정은 activations를 [−128,127] 범위로 스케일링하는 것을 포함합니다.

1.58-bit 모델을 Nanotron으로 사전 학습하거나 파인튜닝하려면 이 PR을 참조하세요. 파인튜닝을 위해 Hugging Face에서 Nanotron 형식으로 모델을 변환해야 합니다. 변환 단계는 이 PR에서 확인할 수 있습니다.

from_pretrained()로 BitNet 양자화 모델을 로드합니다.

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

커널

@torch.compile은 가중치를 언팩하고 forward pass를 수행하는 데 사용됩니다. 구현이 매우 간단하며 상당한 속도 향상을 제공합니다. 향후 버전에서는 추가적인 최적화 커널이 통합될 예정입니다.

리소스

BitNet 모델이 어떻게 학습되고 파인튜닝되는지 자세히 알아보려면 Fine-tuning LLMs to 1.58bit: extreme quantization made easy를 읽어보세요.