품질 손실 없이 LLM 가중치를 22% 압축하는 방법

Unweight는 모델 weights만 손보는 lossless compression으로, 품질 저하 없이 LLM 가중치를 15~22% 줄인다. 초기 결과로 Llama-3.1-8B의 MLP weights에서만 약 30% 압축을 달성했고, 전체 모델 기준으로는 약 3 GB VRAM을 절약했다.

핵심 문제는 계산이 아니라 memory bandwidth다. NVIDIA H100 GPU에서는 tensor cores가 메모리보다 훨씬 빨리 연산하므로, 토큰을 생성할 때마다 HBM에서 읽어오는 바이트 수를 줄이는 것이 성능 개선의 관건이다. Unweight는 압축된 weights를 on-chip shared memory에서 풀어 tensor cores에 바로 공급해, 느린 main memory로의 왕복을 피한다.

압축은 BF16 weight의 구조를 이용한다. sign과 mantissa는 사실상 랜덤에 가깝지만, exponent는 매우 치우쳐 있어 상위 16개 exponent가 일반적인 layer의 99% 이상을 차지한다. Unweight는 이 exponent byte만 Huffman coding으로 압축하고, sign과 mantissa는 그대로 둔다.

적용 범위도 선택적이다. attention, embeddings, layer norms는 건드리지 않고, decode 시 메모리 트래픽의 큰 비중을 차지하는 MLP weights의 gate, up, down projection에만 적용한다. 또한 64개 weight 단위의 row 안에 top-16 palette 밖의 rare exponent가 하나라도 있으면 그 row는 통째로 verbatim으로 저장해 hot path의 per-element branching을 없앤다.

실행 경로는 하나가 아니다. Unweight는 워크로드에 따라 다음 4가지 pipeline을 고른다.

• Full Huffman decode: 원래 BF16을 완전히 복원한 뒤 cuBLAS로 matmul
• Exponent-only decode: exponent byte만 풀어 preprocess traffic을 절반으로 줄임
• Palette transcode: 런타임에 4-bit palette index로 바꿔 traffic을 더 줄임
• Direct palette: model load 시점에 4-bit 형식으로 미리 전환하고 matmul이 즉시 소비

작은 batch size에서는 custom kernel overhead가 커서 full decode + cuBLAS가 유리하다. 반대로 256+ tokens처럼 batch가 커지면 reconstructive matmul이 길어져서, preprocessing이 가벼운 palette/exponent 경로가 더 빨라진다. 같은 layer 안에서도 gate, up, down projection은 shape와 연산 순서가 달라 서로 다른 pipeline이 최적일 수 있다.

Three of the four pipelines는 reconstructive matmul을 사용한다. 이 커널은 HBM에서 압축 데이터를 읽는 producer와, shared memory에서 BF16을 재구성해 Hopper의 WGMMA tensor-core instruction에 넘기는 consumer로 나뉜다. 재구성된 weights는 main memory에 존재하지 않고 shared memory에서 바로 계산으로 이어진다.

여기서도 자원 경쟁이 생긴다. Hopper의 각 SM은 shared memory가 228 KB인데, reconstructive matmul이 약 227 KB를 쓰고 decode kernel도 약 16 KB의 Huffman lookup table이 필요하다. 둘을 같은 SM에 함께 둘 수 없어서, decode에 더 많은 SM을 주면 preprocessing은 빨라지지만 matmul은 느려진다. 따라서 SM 분할 비율도 autotuner가 조정한다.

압축 이득은 transformer의 레이어 구조를 이용해 더 커진다. Unweight는 레이어를 hard와 easy로 나누고, easy layer가 계산되는 동안 별도 CUDA stream에서 다음 hard layer의 preprocessing을 백그라운드로 돌린다. 특히 down projection은 MLP 순서상 마지막에 소비되므로 decode를 가장 오래 미리 끝낼 수 있어 overlap 효과가 크다.

결국 Unweight의 설계 철학은 단순하다. 고정된 한 가지 해법을 고르는 대신, pipeline 선택, matmul variant, SM allocation을 모두 실측 기반으로 최적화해 workload마다 가장 빠른 조합을 쓴다. 그 결과 H100 환경에서 bit-exact 품질을 유지하면서도 더 많은 모델을 더 적은 VRAM에 올릴 수 있게 된다.