Chunked KL loss, 32K를 5GB로

Chunked KL loss는 Flash Attention과 유사하게 계산을 청크 단위로 나누고 forward·backward를 융합해 KL loss의 메모리 사용량을 quadratic에서 linear로 줄인다.

• 32K 컨텍스트 VRAM: 약 85GB → 5GB
• 긴 컨텍스트에서 기존 구현 대비 약 3배 빠른 처리
• PyTorch의 일반적인 KL loss와 수학적으로 동등
• 대형 teacher 모델의 지식을 소형 student 모델에 증류하는 로컬 학습을 지원

teacher 모델의 전체 logits 대신 top-100 logits를 캐시해 사용하는 방식도 지원하며, 논문에서는 전체 분포와 거의 동일한 loss를 얻는다고 설명한다. 다만 모델의 lm-head forward를 청크 처리하도록 패치해야 하므로 PyTorch KL loss의 직접적인 drop-in replacement는 아니다.

코드는 GitHub에 공개됐고, 구현 세부사항과 top-k logits 실험 결과는 arXiv 논문에 정리됐다.