QORA-LLM-2B: 순수 Rust로 구현된 BitNet b1.58 기반 3진법 추론 엔진

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

핵심 요약

기존 LLM 추론의 높은 연산 비용 문제를 해결하기 위해 Microsoft의 BitNet b1.58-2B-4T 모델을 활용한 순수 Rust 기반 추론 엔진이 공개됐다. 이 엔진은 가중치를 {-1, 0, +1}의 3진법으로 처리하여 행렬 곱셈을 단순 가감산으로 대체함으로써 CPU 환경에서도 효율적인 성능을 발휘한다. Python, CUDA 또는 외부 ML 프레임워크에 의존하지 않는 단일 실행 파일 형태로 제공되어 이식성이 뛰어나며, 시스템 리소스를 자동으로 감지해 최적의 설정을 적용한다. 1.13GB의 작은 모델 크기로 20억 개의 파라미터를 처리하며, 저사양 하드웨어에서도 독립적인 AI 구동이 가능함을 입증했다.

배경

Rust 프로그래밍 언어 기본 지식, LLM 양자화 및 Transformer 아키텍처에 대한 이해, CLI 도구 사용 경험

대상 독자

저사양 하드웨어에서 LLM을 구동하려는 개발자 및 Rust 기반 AI 인프라 구축에 관심 있는 엔지니어

의미 / 영향

이 프로젝트는 LLM 추론이 반드시 고성능 GPU와 복잡한 소프트웨어 스택을 필요로 하지 않는다는 점을 시사한다. 3진법 연산의 효율성과 Rust의 성능이 결합되어 향후 온디바이스 AI 및 임베디드 시스템에서의 LLM 활용 가능성을 크게 확장할 것으로 기대된다.

섹션별 상세

순수 Rust(2024 에디션)로 작성되어 Python 런타임이나 CUDA 프레임워크 없이도 단일 실행 파일로 작동한다. Windows, Linux, macOS를 모두 지원하며 외부 라이브러리 의존성을 최소화하여 높은 이식성과 보안성을 확보했다.

bash

# Chat mode (default)
qor2b --prompt "Explain how ternary neural networks work"

# With token limit
qor2b --prompt "Write a haiku about Rust" --max-tokens 100

# Raw text completion (no chat template)
qor2b --prompt "Once upon a time" --raw

QORA-LLM-2B 실행을 위한 주요 CLI 명령어 예시

BitNet b1.58 아키텍처를 채택하여 가중치를 2비트에 패킹된 3진법 값으로 관리한다. 행렬-벡터 곱셈(GEMV) 시 가중치가 +1이면 입력값을 더하고, -1이면 빼며, 0이면 연산을 건너뛰는 방식을 사용하여 부동 소수점 곱셈 연산을 완전히 제거했다.

표준 LLaMA 구조를 변형한 SubLN(Sub-Layer Normalization) 패턴을 적용했다. 레이어당 4개의 RMSNorm을 배치하고, 20개의 Query 헤드와 5개의 KV 헤드를 사용하는 GQA(Grouped Query Attention) 구조를 통해 메모리 효율과 추론 속도를 동시에 개선했다.

rust

// residual = x
// x = input_layernorm(x) # RMSNorm [2560]
// q, k, v = q/k/v_proj(x) # Ternary linear (add/sub only)
// q, k = apply_rope(q, k)
// attn = attention(q, k, v) # GQA: 20Q/5KV
// attn = attn_sub_norm(attn) # SubLN RMSNorm [2560]
// attn = o_proj(attn) # Ternary linear
// x = residual + attn

BitNet b1.58의 SubLN 패턴이 적용된 Attention 레이어 추론 로직

실행 시 시스템의 가용 RAM과 CPU 스레드를 실시간으로 감지하는 지능형 시스템 인식 기능을 갖췄다. 가용 메모리 용량에 따라 최대 생성 토큰 수를 2048에서 8192까지 동적으로 조정하여 시스템 안정성을 극대화한다.

자체적인 .qor2b 이진 포맷을 사용하여 모델 로딩 속도를 높이고 용량을 최적화했다. HuggingFace의 4.8GB bf16 Safetensors 모델을 약 2분 만에 1.13GB 크기의 3진법 바이너리로 변환하는 기능을 내장하고 있다.

i5-11500 CPU 환경에서 초당 약 2.5 토큰의 디코딩 속도를 기록했다. 이는 기존 4비트 양자화 모델인 QORA-3B(0.86 tok/s)보다 약 3배 빠른 속도이며, RAM 사용량도 1.5GB 수준으로 매우 낮다.

실무 Takeaway

GPU 자원이 부족한 엣지 컴퓨팅이나 저사양 CPU 환경에서 LLM을 배포해야 할 경우, 곱셈 연산을 제거한 BitNet 기반 3진법 엔진이 최적의 대안이다.
Python 의존성 없이 Rust로만 구현된 추론 엔진을 사용하면 배포 패키지 크기를 획기적으로 줄이고 런타임 오버헤드를 최소화할 수 있다.
BitNet b1.58 모델은 1.58비트 수준의 극단적인 양자화에도 불구하고 2B 파라미터 규모에서 실용적인 텍스트 생성 성능을 유지한다.

언급된 리소스

문서microsoft/bitnet-b1.58-2B-4T-bf16