BitNet 1-bit LLM 추론 프레임워크: CPU에서 대규모 언어 모델 실행하기

import torch
import torch.nn as nn
import torch.nn.functional as F

class BitLinear(nn.Module):
    """BitNet b1.58의 핵심 레이어 - 삼진 가중치를 사용하는 선형 변환"""

    def __init__(self, in_features: int, out_features: int, bias: bool = False):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        # 전체 정밀도 가중치 (훈련 중 사용)
        self.weight = nn.Parameter(torch.randn(out_features, in_features))
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features))
        else:
            self.bias = None

    def ternary_quantize(self, weight: torch.Tensor) -> tuple:
        """가중치를 -1, 0, +1로 양자화"""
        # absmean 양자화: 평균 절대값을 스케일 팩터로 사용
        gamma = weight.abs().mean()
        # Round-to-Nearest (RtN) 방식으로 삼진값 생성
        weight_ternary = torch.clamp(
            torch.round(weight / (gamma + 1e-8)),
            min=-1,
            max=1
        )
        return weight_ternary, gamma

    def activation_quantize(self, x: torch.Tensor, bits: int = 8) -> tuple:
        """활성화를 INT8로 양자화"""
        Qb = 2 ** (bits - 1)
        gamma = x.abs().max()
        x_quantized = torch.clamp(
            x * Qb / (gamma + 1e-8),
            min=-Qb,
            max=Qb - 1
        ).round()
        return x_quantized, gamma

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 1. 활성화 양자화 (INT8)
        x_quant, x_scale = self.activation_quantize(x)
        # 2. 가중치 삼진 양자화
        w_quant, w_scale = self.ternary_quantize(self.weight)
        # 3. 정수 행렬 연산 (곱셈 없이 덧셈/뺄셈만 사용)
        output = F.linear(x_quant, w_quant, None)
        # 4. 역양자화 (스케일 복원)
        output = output * (w_scale * x_scale / (2 ** 7))
        if self.bias is not None:
            output = output + self.bias
        return output

class StraightThroughEstimator(torch.autograd.Function):
    """STE: 순전파에서는 양자화, 역전파에서는 기울기 직접 전달"""

    @staticmethod
    def forward(ctx, weight, gamma):
        # 순전파: 삼진 양자화 적용
        weight_ternary = torch.clamp(
            torch.round(weight / (gamma + 1e-8)),
            min=-1, max=1
        )
        return weight_ternary

    @staticmethod
    def backward(ctx, grad_output):
        # 역전파: 기울기를 그대로 전달 (양자화 무시)
        return grad_output, None

class BitNetTransformerBlock(nn.Module):
    """BitNet b1.58 Transformer 블록"""

    def __init__(self, d_model: int, n_heads: int, d_ff: int):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)

        # 어텐션: Q, K, V 프로젝션을 BitLinear로 교체
        self.q_proj = BitLinear(d_model, d_model)
        self.k_proj = BitLinear(d_model, d_model)
        self.v_proj = BitLinear(d_model, d_model)
        self.o_proj = BitLinear(d_model, d_model)
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads

        # FFN: 게이트 프로젝션도 BitLinear
        self.gate_proj = BitLinear(d_model, d_ff)
        self.up_proj = BitLinear(d_model, d_ff)
        self.down_proj = BitLinear(d_ff, d_model)

    def attention(self, x: torch.Tensor) -> torch.Tensor:
        B, T, C = x.shape
        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)

        # 스케일드 닷-프로덕트 어텐션
        scale = self.head_dim ** -0.5
        attn = (q @ k.transpose(-2, -1)) * scale
        attn = F.softmax(attn, dim=-1)
        out = (attn @ v).transpose(1, 2).contiguous().view(B, T, C)
        return self.o_proj(out)

    def ffn(self, x: torch.Tensor) -> torch.Tensor:
        # SwiGLU 활성화 함수
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attention(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x

# 기존 PTQ 양자화 추론 (GPTQ 방식 - 의사 코드)
def ptq_forward(x_fp16, weight_int4, scale, zero_point):
    # 1. 역양자화: INT4 -> FP16
    weight_fp16 = (weight_int4 - zero_point) * scale
    # 2. 부동소수점 행렬 곱셈 (비용이 큼)
    output = torch.matmul(x_fp16, weight_fp16.T)
    return output

# BitNet 1-bit 추론 (의사 코드)
def bitnet_forward(x_int8, weight_ternary):
    # weight가 -1, 0, +1이므로:
    #   weight == +1 -> 활성화를 더함
    #   weight ==  0 -> 아무것도 안 함
    #   weight == -1 -> 활성화를 뺌
    # 곱셈이 완전히 제거됨!
    output = torch.zeros(weight_ternary.shape[0])
    output += (weight_ternary == 1).float() @ x_int8.float()
    output -= (weight_ternary == -1).float() @ x_int8.float()
    return output

# 시스템 요구사항
# - Python >= 3.9
# - CMake >= 3.22
# - Clang >= 18 (권장) 또는 GCC >= 12
# - x86_64 CPU: AVX2 이상 지원 필수 (AVX-512 권장)
# - ARM CPU: NEON 지원 필수

# CPU 기능 확인 (Linux)
lscpu | grep -E "avx2|avx512"

# CPU 기능 확인 (macOS)
sysctl -a | grep machdep.cpu.features

# 1. 저장소 클론
git clone --recursive https://github.com/microsoft/BitNet.git
cd BitNet

# 2. Python 의존성 설치
pip install -r requirements.txt

# 3. 모델 다운로드 및 빌드 (자동화 스크립트)
# 공식 지원 모델 목록:
#   - BitNet-b1.58-2B-4T (2B 파라미터, 4T 토큰 훈련)
#   - BitNet-b1.58-4B-8T (4B 파라미터, 8T 토큰 훈련)

python setup_env.py \
    --hf-repo microsoft/BitNet-b1.58-2B-4T \
    -q i2_s

# 4. 빌드 확인
ls build/bin/
# 출력: llama-cli, llama-bench, ...

# 텍스트 생성 실행
python run_inference.py \
    -m models/BitNet-b1.58-2B-4T/ggml-model-i2_s.gguf \
    -p "The future of artificial intelligence is" \
    -n 128 \
    -t 4 \
    --temp 0.7

# 옵션 설명:
# -m: 모델 파일 경로
# -p: 프롬프트
# -n: 생성할 토큰 수
# -t: 사용할 CPU 스레드 수
# --temp: 생성 온도 (낮을수록 결정적)

# 성능 벤치마크
python run_inference.py \
    -m models/BitNet-b1.58-2B-4T/ggml-model-i2_s.gguf \
    -p "Benchmark prompt for measuring inference speed" \
    -n 512 \
    -t 8 \
    --benchmark

# 출력 예시:
# Model: BitNet-b1.58-2B-4T (i2_s)
# Threads: 8
# Prompt eval: 245.3 tokens/s
# Generation: 32.7 tokens/s
# Peak memory: 1.2 GB

import subprocess
import json
from pathlib import Path

class BitNetRunner:
    """BitNet.cpp 래퍼 클래스"""

    def __init__(
        self,
        model_path: str,
        binary_path: str = "build/bin/llama-cli",
        threads: int = 4
    ):
        self.model_path = Path(model_path)
        self.binary_path = Path(binary_path)
        self.threads = threads

        if not self.model_path.exists():
            raise FileNotFoundError(f"모델 파일 없음: {model_path}")
        if not self.binary_path.exists():
            raise FileNotFoundError(f"실행 바이너리 없음: {binary_path}")

    def generate(
        self,
        prompt: str,
        max_tokens: int = 256,
        temperature: float = 0.7,
        top_p: float = 0.9,
    ) -> str:
        """텍스트 생성"""
        cmd = [
            str(self.binary_path),
            "-m", str(self.model_path),
            "-p", prompt,
            "-n", str(max_tokens),
            "-t", str(self.threads),
            "--temp", str(temperature),
            "--top-p", str(top_p),
            "--no-display-prompt",
        ]

        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=120
        )

        if result.returncode != 0:
            raise RuntimeError(f"추론 실패: {result.stderr}")

        return result.stdout.strip()

    def benchmark(self, prompt: str, n_tokens: int = 512) -> dict:
        """성능 벤치마크 실행"""
        cmd = [
            str(self.binary_path),
            "-m", str(self.model_path),
            "-p", prompt,
            "-n", str(n_tokens),
            "-t", str(self.threads),
        ]

        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=300
        )

        # 출력에서 성능 메트릭 파싱
        lines = result.stderr.split("\n")
        metrics = {}
        for line in lines:
            if "prompt eval" in line.lower() and "token" in line.lower():
                metrics["prompt_tokens_per_sec"] = self._parse_tps(line)
            elif "eval" in line.lower() and "token" in line.lower():
                metrics["gen_tokens_per_sec"] = self._parse_tps(line)
        return metrics

    @staticmethod
    def _parse_tps(line: str) -> float:
        """토큰/초 파싱"""
        import re
        match = re.search(r"([\d.]+)\s*tokens?/s", line)
        return float(match.group(1)) if match else 0.0


# 사용 예시
runner = BitNetRunner(
    model_path="models/BitNet-b1.58-2B-4T/ggml-model-i2_s.gguf",
    threads=8
)

response = runner.generate(
    "Explain the concept of 1-bit LLMs in simple terms:",
    max_tokens=256,
    temperature=0.7
)
print(response)

# 스레드 확장성 벤치마크 스크립트
for threads in 1 2 4 8 16 32; do
    echo "=== Threads: $threads ==="
    python run_inference.py \
        -m models/BitNet-b1.58-2B-4T/ggml-model-i2_s.gguf \
        -p "Performance benchmark with varying thread count" \
        -n 256 \
        -t "$threads" \
        --benchmark 2>&1 | grep "eval"
done

// BitNet.cpp LUT 커널 핵심 로직 (간략화)
// 실제 구현은 SIMD 인스트럭션으로 최적화됨

void bitnet_lut_matmul(
    const int8_t* activation,   // INT8 활성화
    const int8_t* weight_ternary, // 삼진 가중치 (-1, 0, 1)
    float* output,
    int M, int N, int K
) {
    // 2비트 패킹: 4개의 삼진값을 1바이트에 패킹
    // 00 = 0, 01 = +1, 10 = -1

    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            int32_t acc = 0;
            for (int k = 0; k < K; k++) {
                int8_t w = weight_ternary[j * K + k];
                int8_t a = activation[i * K + k];
                // w가 삼진값이므로 곱셈 대신 조건부 덧셈
                if (w == 1) acc += a;
                else if (w == -1) acc -= a;
                // w == 0이면 아무것도 안 함
            }
            output[i * N + j] = (float)acc;
        }
    }
}

# Raspberry Pi 5 (8GB) 배포 예시

# 1. 크로스 컴파일 환경 설정
sudo apt-get install cmake clang-18 python3-pip
pip3 install -r requirements.txt

# 2. ARM 최적화 빌드
python setup_env.py \
    --hf-repo microsoft/BitNet-b1.58-2B-4T \
    -q i2_s \
    --cmake-args "-DCMAKE_C_FLAGS=-mcpu=cortex-a76"

# 3. 메모리 제약 환경에서의 실행
python run_inference.py \
    -m models/BitNet-b1.58-2B-4T/ggml-model-i2_s.gguf \
    -p "Summarize this sensor data:" \
    -n 64 \
    -t 4 \
    --ctx-size 512

# Raspberry Pi 5 예상 성능:
# - 메모리: ~0.9GB (8GB 중)
# - 생성 속도: ~5-8 tok/s
# - 첫 토큰 지연: ~200ms

# FastAPI 기반 로컬 추론 서버

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import asyncio
from concurrent.futures import ThreadPoolExecutor

app = FastAPI(title="BitNet Local LLM Server")
executor = ThreadPoolExecutor(max_workers=2)

# 모델 런너 초기화
runner = BitNetRunner(
    model_path="models/BitNet-b1.58-2B-4T/ggml-model-i2_s.gguf",
    threads=8
)

class GenerateRequest(BaseModel):
    prompt: str
    max_tokens: int = 256
    temperature: float = 0.7

class GenerateResponse(BaseModel):
    text: str
    tokens_generated: int
    generation_time_ms: float

@app.post("/generate", response_model=GenerateResponse)
async def generate(req: GenerateRequest):
    import time
    start = time.monotonic()

    try:
        loop = asyncio.get_event_loop()
        text = await loop.run_in_executor(
            executor,
            lambda: runner.generate(
                req.prompt,
                max_tokens=req.max_tokens,
                temperature=req.temperature
            )
        )
    except RuntimeError as e:
        raise HTTPException(status_code=500, detail=str(e))

    elapsed_ms = (time.monotonic() - start) * 1000
    token_count = len(text.split())  # 근사 토큰 수

    return GenerateResponse(
        text=text,
        tokens_generated=token_count,
        generation_time_ms=round(elapsed_ms, 1)
    )

@app.get("/health")
async def health():
    return {"status": "ok", "model": "BitNet-b1.58-2B-4T"}


# 실행: uvicorn server:app --host 0.0.0.0 --port 8080

# 대규모 배치 추론 파이프라인

import multiprocessing as mp
from dataclasses import dataclass
from typing import List
import time
import json

@dataclass
class InferenceJob:
    job_id: str
    prompt: str
    max_tokens: int = 256

@dataclass
class InferenceResult:
    job_id: str
    output: str
    tokens_per_sec: float
    latency_ms: float

def worker_process(
    model_path: str,
    job_queue: mp.Queue,
    result_queue: mp.Queue,
    threads_per_worker: int
):
    """워커 프로세스: 독립적인 BitNet 런너 인스턴스"""
    runner = BitNetRunner(
        model_path=model_path,
        threads=threads_per_worker
    )

    while True:
        job = job_queue.get()
        if job is None:  # 종료 신호
            break

        start = time.monotonic()
        try:
            output = runner.generate(
                job.prompt,
                max_tokens=job.max_tokens,
                temperature=0.1  # 배치 추론은 낮은 온도
            )
            elapsed = time.monotonic() - start
            approx_tokens = len(output.split())

            result_queue.put(InferenceResult(
                job_id=job.job_id,
                output=output,
                tokens_per_sec=approx_tokens / elapsed,
                latency_ms=elapsed * 1000
            ))
        except Exception as e:
            result_queue.put(InferenceResult(
                job_id=job.job_id,
                output=f"ERROR: {str(e)}",
                tokens_per_sec=0,
                latency_ms=0
            ))


def run_batch_inference(
    model_path: str,
    jobs: List[InferenceJob],
    num_workers: int = 4,
    threads_per_worker: int = 4
) -> List[InferenceResult]:
    """멀티프로세스 배치 추론"""
    job_queue = mp.Queue()
    result_queue = mp.Queue()

    # 워커 프로세스 시작
    workers = []
    for _ in range(num_workers):
        p = mp.Process(
            target=worker_process,
            args=(model_path, job_queue, result_queue, threads_per_worker)
        )
        p.start()
        workers.append(p)

    # 작업 큐에 삽입
    for job in jobs:
        job_queue.put(job)

    # 종료 신호
    for _ in range(num_workers):
        job_queue.put(None)

    # 결과 수집
    results = []
    for _ in range(len(jobs)):
        results.append(result_queue.get(timeout=600))

    # 워커 종료 대기
    for p in workers:
        p.join()

    return results


# 사용 예시
if __name__ == "__main__":
    jobs = [
        InferenceJob(f"job_{i}", f"Translate to Korean: {text}", 128)
        for i, text in enumerate(open("inputs.txt").readlines())
    ]

    results = run_batch_inference(
        model_path="models/BitNet-b1.58-2B-4T/ggml-model-i2_s.gguf",
        jobs=jobs,
        num_workers=4,
        threads_per_worker=4  # 총 16스레드 사용
    )

    for r in results:
        print(f"[{r.job_id}] {r.tokens_per_sec:.1f} tok/s, {r.latency_ms:.0f}ms")

# 태스크별 성능 비교 스크립트 (lm-evaluation-harness 활용)

# 설치
# pip install lm-eval

# BitNet 모델 평가 (GGUF 포맷 사용)
# lm_eval 명령어 예시:
# lm_eval --model gguf \
#     --model_args path=models/BitNet-b1.58-2B-4T/ggml-model-i2_s.gguf \
#     --tasks arc_easy,arc_challenge,hellaswag,winogrande,piqa,boolq \
#     --batch_size 1 \
#     --output_path results/bitnet_2b_eval.json

# 결과 비교 분석
import json

def compare_eval_results(
    bitnet_path: str,
    baseline_path: str
) -> None:
    """평가 결과 비교"""
    with open(bitnet_path) as f:
        bitnet = json.load(f)
    with open(baseline_path) as f:
        baseline = json.load(f)

    print(f"{'Task':<20} {'Baseline':>10} {'BitNet':>10} {'Delta':>10}")
    print("-" * 50)

    for task in bitnet["results"]:
        b_acc = baseline["results"].get(task, {}).get("acc", 0)
        n_acc = bitnet["results"][task].get("acc", 0)
        delta = n_acc - b_acc
        sign = "+" if delta >= 0 else ""
        print(f"{task:<20} {b_acc:>10.3f} {n_acc:>10.3f} {sign}{delta:>9.3f}")

현재 공식 지원 모델 (2026년 3월 기준):
- microsoft/BitNet-b1.58-2B-4T (2B 파라미터)
  - 4T 토큰으로 사전 훈련
  - 영어 중심, 다국어 제한적

- 커뮤니티 모델:
  - 1bitLLM/bitnet_b1_58-large (0.7B)
  - 1bitLLM/bitnet_b1_58-3B (3B)
  - HF1BitLLM/Llama3-8B-1.58-100B-tokens (8B, 실험적)

# BitNet 배포 체크리스트

사전_검증:
  - CPU 기능 확인 (AVX2/AVX-512/NEON)
  - 메모리 여유 확인 (모델 크기 x 1.5 이상)
  - OS 호환성 (Linux/macOS 권장, Windows 실험적)

성능_최적화:
  - 스레드 수를 물리 코어 수에 맞출 것
  - NUMA 환경에서는 numactl로 메모리 지역성 확보
  - Hyper-Threading은 오히려 성능 저하 가능

안정성:
  - OOM 방지를 위한 메모리 제한 설정
  - 장시간 운영 시 메모리 누수 모니터링
  - 타임아웃 설정 필수 (프롬프트 길이에 비례)

품질:
  - 출력 품질을 반드시 태스크별로 검증
  - 동급 FP16 모델과 A/B 테스트 수행
  - 한국어 등 비영어 태스크는 별도 검증 필수

# llama.cpp로 동일 태스크 실행 (비교용)
./llama-cli \
    -m models/llama-3b-q4_k_m.gguf \
    -p "The future of artificial intelligence is" \
    -n 128 \
    -t 8

# BitNet.cpp로 동일 태스크 실행
python run_inference.py \
    -m models/BitNet-b1.58-2B-4T/ggml-model-i2_s.gguf \
    -p "The future of artificial intelligence is" \
    -n 128 \
    -t 8

전용 하드웨어 전망:

1. FPGA 기반 가속기
   - Xilinx/Intel FPGA에 삼진 MAC 유닛 구현
   - 프로토타입 수준에서 10x 에너지 효율 개선 가능
   - 소규모 배치 생산에 적합

2. ASIC 설계
   - 삼진 연산 전용 프로세서
   - 대규모 생산 시 100x 에너지 효율 가능
   - 2027~2028년 양산 가능성

3. In-Memory Computing
   - RRAM/STT-MRAM에 삼진 가중치 직접 저장
   - 메모리-연산 간 데이터 이동 제거
   - 연구 단계 (2028년 이후)

BitNet 도입 적합성 판단:

Q1: GPU가 가용한가?
  - Yes -> 기존 양자화 (GPTQ/AWQ) + vLLM 권장
  - No -> Q2로

Q2: 메모리가 4GB 이상인가?
  - Yes -> Q3으로
  - No -> 모델 실행 불가

Q3: 영어 태스크인가?
  - Yes -> BitNet 검토 가능 -> Q4로
  - No -> llama.cpp + GGUF 권장 (다국어 모델 풍부)

Q4: 실시간 응답이 필요한가?
  - Yes -> CPU 코어 8개 이상 확인 후 BitNet 배포
  - No -> 배치 추론으로 BitNet 활용

Q5: 에너지 효율이 최우선인가?
  - Yes -> BitNet 강력 추천
  - No -> 총 소유 비용(TCO) 비교 후 결정

# TCO 계산기

def calculate_tco(
    model_type: str,
    requests_per_day: int,
    avg_tokens_per_request: int,
    months: int = 12
) -> dict:
    """연간 총 소유 비용 추정"""

    configs = {
        "gpu_fp16": {
            "hardware": "A100 80GB",
            "hardware_cost_monthly": 2400,  # 클라우드 기준
            "power_watts": 400,
            "tokens_per_sec": 42,
        },
        "gpu_gptq4": {
            "hardware": "RTX 4090",
            "hardware_cost_monthly": 800,
            "power_watts": 350,
            "tokens_per_sec": 68,
        },
        "cpu_gguf": {
            "hardware": "EPYC 9654 Server",
            "hardware_cost_monthly": 400,
            "power_watts": 280,
            "tokens_per_sec": 25,
        },
        "cpu_bitnet": {
            "hardware": "EPYC 9654 Server",
            "hardware_cost_monthly": 400,
            "power_watts": 180,  # 곱셈 연산 없어 전력 절감
            "tokens_per_sec": 35,
        },
    }

    cfg = configs[model_type]
    daily_tokens = requests_per_day * avg_tokens_per_request
    daily_seconds = daily_tokens / cfg["tokens_per_sec"]
    daily_kwh = (cfg["power_watts"] * daily_seconds / 3600) / 1000
    monthly_energy_cost = daily_kwh * 30 * 0.12  # kWh당 0.12 USD

    total_monthly = cfg["hardware_cost_monthly"] + monthly_energy_cost
    total = total_monthly * months

    return {
        "model_type": model_type,
        "hardware": cfg["hardware"],
        "monthly_cost": round(total_monthly, 2),
        "annual_cost": round(total, 2),
        "energy_monthly_kwh": round(daily_kwh * 30, 1),
    }


# 비교 실행
for model_type in ["gpu_fp16", "gpu_gptq4", "cpu_gguf", "cpu_bitnet"]:
    result = calculate_tco(
        model_type=model_type,
        requests_per_day=10000,
        avg_tokens_per_request=200,
        months=12
    )
    print(f"{result['model_type']:>15}: "
          f"월 ${result['monthly_cost']:>8,}, "
          f"연 ${result['annual_cost']:>10,}, "
          f"에너지 {result['energy_monthly_kwh']:>6} kWh/월")