지식 증류를 대규모로 실행할 수 있을 만큼 저렴하게 만들기

작은 학생(student) 모델이 더 큰 교사(teacher) 모델의 출력을 모방하도록 학습시키는 지식 증류(knowledge distillation)는 대형 언어모델(large language model, LLM)을 경량화하는 표준 기법이다. Kimi-K3처럼 2.8조 개의 파라미터를 가진 최근 오픈소스 LLM은 약 3TB의 VRAM을 필요로 하기 때문에, 이러한 압축 기법이 필수적이다. 논문 'Efficient Knowledge Distillation for LLMs: Offline Top-K Logits and a Fused Chunked KL Loss'는 교사 모델의 상위 K개 로짓(logits)을 오프라인으로 캐싱함으로써 증류 과정에서 발생하는 높은 비용 문제를 해결한다. 이 방식은 교사 모델을 메모리에 계속 유지할 필요를 없애준다. 또한 이 논문은 손실을 청크 단위로 계산하여 전체 어휘×시퀀스 행렬을 만들지 않아도 되는 융합 청크 KL 손실(fused chunked KL loss, KL은 쿨백-라이블러(Kullback-Leibler)의 약자)을 제안한다. 이를 통해 단일 H200 GPU에서 최대 VRAM 사용량을 250GB에서 128GB로 줄였으며, 4개의 노드가 아닌 1개의 GPU만으로도 32K 컨텍스트에서 증류를 수행할 수 있게 되었고, 스텝 처리 속도도 5배 빨라졌다. Llama 3.1 8B Instruct 모델에서 증류되어 만들어진 32억(3.2B) 파라미터 규모의 학생 모델은 BoolQ와 HellaSwag 같은 벤치마크에서 대부분의 정확도를 유지했으며, MMLU에서도 9점 이내의 격차를 보였다.