BitNet

논문 링크

BitLinear = tensor -> layernorm -> Binarize -> abs max quantization -> dequant

"BitNet 아키텍처의 구현은 상당히 간단하여 트랜스포머 내의 선형 투영(즉, PyTorch에서 nn.Linear)만 교체하면 됩니다." -- BitNet을 정말 쉽게 구현할 수 있으며, 단순히 라인어를 BitLinear 모듈로 바꿔주기만 하면 됩니다!

뉴스

• 새로운 반복 버전 (New Iteration) 🔥 논문 "The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits"에서 나온 완전히 새로운 반복 버전이 있으며, 현재 이를 구현하고 있습니다. Agora discord에 참여하여 기여해 주세요! 여기서 참여하기
• 새로운 최적화 (New Optimizations) 첫 번째 BitLinear가 최적화되었으며, 이제 BitAttention인 BitMGQA를 갖게 되었습니다. 이는 BitLinear를 어텐션 메커니즘에 구현합니다. Multi Grouped Query Attention은 또한 빠른 디코딩과 긴 컨텍스트 처리가 가능하여 최고의 어텐션으로 널리 인정받고 있으며, 사용하기 쉬운 구현을 제공해 준 Frank에게 감사드립니다!
• BitLinear 1.5 출시 🔥: 새로운 BitLinear 1.5는 아직 진행 중입니다 🔥 [파일 보기] 데퀀타이징(dequantization) 알고리즘과 같은 버그가 여전히 존재하며, 곱셈을 요소별 덧셈으로 대체해야 하므로 이 부분에 도움을 주시면 정말 감사하겠습니다.
• 공지 (NOTICE): BitLinear를 사용하려면 모델을 처음부터 파인튜닝(finetune)해야 합니다. 이미 학습된 모델의 선형 메서드를 단순히 변경하는 것만으로는 작동하지 않습니다. 처음부터 파인튜닝하거나 학습시켜야 합니다.

감사 인사 (Appreciation)

• 분석 및 코드 검토, 수정에 기여한 Dimitry, Nullonix 님께
• 학습용 4080을 제공해 준 Vyom 님께!

설치 (Installation)

pip3 install bitnet

사용법

여기에는 방대한 양의 예제 스크립트가 있으며, examples 폴더에서도 확인할 수 있습니다. 디스코드에서 특정 사용 사례에 대한 도움이 필요하면 알려주세요!

BitLinear

• 논문의 주요 혁신인 BitLinear 레이어의 예시입니다!

import torch

from bitnet import BitLinear
...

BitLinearNew

import torch
from bitnet import BitLinearNew

...

BitNetTransformer

• MHA와 BitFeedforwards를 포함하여 다이어그램에 설명된 트랜스포머가 완벽하게 구현되었습니다.
• 텍스트뿐만 아니라 이미지, 심지어 비디오나 오디오 처리에도 활용될 수 있습니다.
• 기울기 흐름을 위해 잔차 연결(residuals)과 스킵 연결(skip connections)이 모두 갖춰져 있습니다.

# 필요한 라이브러리 가져오기
import torch
from bitnet import BitNetTransformer
...

BitAttention

이 어텐션은 기본 선형 투영(linear projection) 대신 BitLinear를 사용하도록 수정되었습니다. 또한, 빠른 디코딩과 긴 컨텍스트 처리를 위해 일반적인 멀티 헤드 어텐션(multi-head attention) 대신 다중 그룹화 쿼리 어텐션(Multi-Grouped Query Attention)을 사용합니다.

import torch
from bitnet import BitMGQA

...

BitFeedForward

• 다이어그램에 표시된 BitLinear와 GELU를 사용한 피드포워드입니다:
• Linear -> GELU -> Linear
• 더 나은 ffn을 위해 드롭아웃(dropouts), 레이어노름(layernorms) 또는 다른 레이어를 추가할 수 있습니다.

import torch
from bitnet import BitFeedForward

...

추론 (Inference)

from bitnet import BitNetInference

bitnet = BitNetInference()
...

Huggingface 사용법

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

...

PyTorch 모델 대체 (Drop in Replacement for Pytorch Models)

import torch
from torch import nn
from bitnet import replace_linears_in_pytorch_model
...

최적화된 Cuda 커널

python setup.py build_ext --inplace

import torch
import gemm_lowbit_ext  # 컴파일된 모듈을 가져옵니다

...

BitLora

BitLora 구현입니다!

import torch
from bitnet import BitLora

...

BitMamba

import torch
from bitnet import BitMamba

...

BitMoE

import torch
from bitnet.bit_moe import BitMoE

...

1 비트 Vision Transformers

이 아이디어는 갑자기 떠올랐지만, 초압축을 위해 vision 작업에 bitlinear를 활용할 수 있어 매우 재미있어 보입니다. 스크립트를 작성해 주신다면 imagenet으로 학습시키는 것이 좋을 것 같습니다. 컴퓨팅 자원은 저희가 제공하겠습니다. 다음 단계는 join vision language model인 gpt-4o를 학습시키는 것입니다.

import torch
from bitnet import OneBitViT

...

라이선스

MIT

인용

@misc{2310.11453,
Author = {Hongyu Wang and Shuming Ma and Li Dong and Shaohan Huang and Huaijie Wang and Lingxiao Ma and Fan Yang and Ruiping Wang and Yi Wu and Furu Wei},
Title = {BitNet: Scaling 1-bit Transformers for Large Language Models},
...

할 일(Todo)

• BitLinear 구현을 재확인하고 논문과 정확히 작동하는지 확인하기
• BitNetTransformer 학습 스크립트 구현
• Enwiki8에서 학습시키기, Lucidrains 리포지토리의 코드와 데이터 복사 및 붙여넣기
• 성능 벤치마킹
• 미분 불가능한 역전파(backprop)를 위한 Straight Through Estimator 검토
• BitFeedForward 구현
• 코드베이스 정리하기
• 각 모듈에 대한 단위 테스트 추가
• 논문의 새로운 BitNet1.5b 구현
• Cuda에서 BitNet15b 구현하기
• 저비트 gemm cuda 커널 구현