The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits

"The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits" 논문에서 소개된 방법론이 소개되었습니다.

1비트 LLM(1-bit LLM)은 자연어 처리(NLP) 모델의 성능을 개선하려는 시도 중 하나로, 기존의 32비트나 16비트 LLM과는 다른 접근 방식을 취합니다. 이 모델은 모든 파라미터나 가중치를 -1, 0, 1 중 하나의 값으로 양자화하여 구성됩니다. 이렇게 함으로써, 각 파라미터가 더 적은 비트로 표현되어 전체적인 모델 크기를 줄이고 메모리 사용량을 최적화할 수 있습니다.

1비트 LLM은 절대평균 양자화(absolute mean quantization)라는 특정한 양자화 기법을 사용하여 수치를 양자화합니다. 이 양자화 기법은 평균 값과의 차이를 계산하여 각 수치를 양자화된 범주(-1, 0, 1) 중 하나로 변환합니다. 이 과정은 모델의 가중치를 간단한 형태로 만들어주며, 모델의 학습 및 추론 과정에서 계산 비용을 줄여줍니다.

BitNet b1.58은 absmean Quantization을 사용합니다. 이는 기존의 양자화보다 오차를 줄이면서도 하드웨어 성능을 요구하는 정확한 결과를 얻기 위해 고안된 방법입니다.

 

여기서 W는 양자화되는 값이고, gamma는 가중치의 절대값 평균에 비례한 상수입니다. RoundClip 함수는 값 x를 범위 [a,b] 내의 가장 가까운 정수로 반올림하고, 벗어나는 경우에는 각각 a 또는 b 클리핑합니다.

또한, r+e는 값이 튀는 것을 방지하기 위한 상수입니다. RMSNorm 과정에서는 분모가 0이 되어 계산이 불가능한 경우를 방지하기 위해 매우 작은 값을 추가하여 정규화에 영향을 미치지 않으면서도 공식의 불완전성을 보완합니다.

이러한 변화는 추론 과정에서 특히 유용합니다. 기존의 32비트나 16비트 LLM 모델에서는 많은 계산이 필요했지만, 1비트 LLM에서는 각 가중치가 -1, 0, 1 중 하나로 제한되어 있기 때문에 단순한 연산만으로도 원하는 결과를 얻을 수 있습니다. 따라서 추론 과정에서의 계산 비용이 크게 줄어들어 더 효율적인 추론이 가능해집니다.

또한, 1비트 LLM은 특징 필터링(feature filtering)을 명시적으로 지원하여 모델의 용량을 강화합니다. 각 가중치에 0을 포함시킴으로써, 모델이 불필요한 특징을 무시하고 중요한 특징에 집중할 수 있게 됩니다. 이는 모델의 성능을 향상시키는 데 도움이 됩니다.

1비트 LLM은 적은 비트로 모델을 효율적으로 표현하여 메모리 사용량을 줄이고, 추론 과정에서의 계산 비용을 최적화합니다. 이러한 기술적 혁신은 자연어 처리 분야에서 모델의 성능과 효율성을 향상시키는 데 중요한 역할을 할 것으로 기대됩니다.

NousResearch의 구현 사례

OLMo-Bitnet-1B은 1B 개의 파라미터를 가진 모델로 훈련되었습니다.

이 모델은 Dolma 데이터셋의 처음 60B 토큰을 사용하여 훈련되었으며, 이는 단순히 방법론을 실험하기 위한 연구 개념 증명입니다.

동일한 하이퍼파라미터를 사용하여 표준 fp16 가중치를 사용하여 별도의 훈련 실행이 있었고, 이에 대한 비교는 해당 wandb 보고서에서 확인할 수 있습니다.

아래는 샘플 추론 코드입니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, TextStreamer

tokenizer = AutoTokenizer.from_pretrained("NousResearch/OLMo-Bitnet-1B")
model = AutoModelForCausalLM.from_pretrained("NousResearch/OLMo-Bitnet-1B",
    torch_dtype=torch.bfloat16, trust_remote_code=True, device_map="auto")

streamer = TextStreamer(tokenizer)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, pad_token_id=tokenizer.eos_token_id,
    temperature=0.8, repetition_penalty=1.1, do_sample=True,streamer=streamer)
pipe("The capitol of Paris is",  max_new_tokens=256)

위 코드를 사용하여 OLMo-Bitnet-1B 모델을 활용할 수 있습니다.