Knowledge Distillation 비용 줄이기: offline top-K logits와 chunked KL 적용 기준

요약: Hugging Face에 공개된 “Making Knowledge Distillation Cheap Enough to Run at Scale” 글은 LLM 지식 증류 비용을 줄이는 실용적인 시스템 패턴을 보여줍니다. 핵심은 teacher model을 매 step 같이 올리지 않고 top-K logits를 offline cache로 만들며, KL loss 계산에서 vocabulary × sequence 전체 행렬을 만들지 않는 것입니다. 긴 컨텍스트 distillation을 고민하는 팀이라면 바로 검토할 만한 구조입니다.

문제가 되는 지점: teacher보다 loss가 더 비쌀 수 있다

Knowledge distillation은 큰 teacher model의 출력을 작은 student model이 따라 하도록 학습시키는 방식입니다. 배포 비용을 줄이거나 특정 도메인에 맞춘 작은 모델을 만들 때 자주 씁니다. 그런데 실무에서 막히는 지점은 학습 데이터보다 메모리입니다. online distillation은 teacher와 student를 동시에 GPU에 올리고, 매 step마다 teacher forward를 다시 수행합니다.

Hugging Face 글은 gpt-oss-120b 같은 모델의 vocabulary가 201,088 tokens일 때, sequence length 32K, batch size 4에서 teacher probability tensor 하나만 bfloat16 기준 약 50GB VRAM을 차지한다고 설명합니다. 여기에 gradient, activation, model weight, optimizer state가 붙으면 한 iteration이 약 250GB VRAM까지 튈 수 있습니다. 단일 H200이나 B200에서도 버거운 수치입니다.

즉 병목은 단순히 “큰 모델을 못 올린다”가 아닙니다. teacher distribution과 student log probability를 full vocabulary × sequence 형태로 만드는 손실 계산 자체가 메모리를 먹습니다. 긴 컨텍스트를 다루는 distillation이라면 이 지점부터 줄여야 합니다.

offline top-K logits 방식

첫 번째 개선은 offline distillation입니다. teacher를 매 step 실행하지 않고, teacher output을 한 번 계산해 top-K logits만 저장합니다. 글에서는 position마다 top-100 most likely tokens를 cache한다고 설명합니다. 이후 student 학습은 이 cache를 기준으로 진행합니다.

이 방식의 장점은 세 가지입니다.

• teacher model을 student 학습 중 GPU에 계속 올릴 필요가 없다.
• 같은 teacher cache를 여러 ablation에 재사용할 수 있다.
• teacher forward 비용을 학습 루프 밖으로 빼서 반복 비용을 줄인다.

물론 단점도 있습니다. teacher의 full distribution을 완전히 저장하는 것이 아니라 top-K만 저장하므로, long tail token 정보는 줄어듭니다. 하지만 공개된 benchmark에서는 top-100 cached logits 기준으로 online distillation과 거의 동일한 training loss curve를 보였다고 합니다. 실무에서는 top-K 크기를 품질과 저장 비용의 knob으로 보면 됩니다.

추천 시작점은 다음과 같습니다.

• instruction tuning: top-50 또는 top-100
• 코드 모델 distillation: top-100 이상부터 검토
• 긴 문서 요약 특화: top-50으로 시작 후 validation loss 비교
• 안전성·정책 응답 distillation: top-K보다 eval set 품질 우선

chunked KL loss가 하는 일

두 번째 개선은 KL loss 계산 방식입니다. 기본 dense KL은 cached top-K logits를 다시 dense teacher probability grid로 펼치고, student의 dense log-probability grid와 비교합니다. 이 방식은 correctness baseline으로는 좋지만, vocabulary × sequence 전체 행렬을 잡아먹습니다.

forward-chunked KL은 teacher를 sparse하게 유지하고 sequence position을 조각으로 나눠 loss를 계산합니다. 이러면 dense teacher grid는 만들지 않아도 됩니다. 다만 student logits는 여전히 전체로 계산되어 backward pass를 위해 남습니다.

fused chunked KL은 한 단계 더 들어갑니다. output projection을 loss computation 안으로 합쳐서 student의 full logits grid도 만들지 않습니다. sequence chunk마다 hidden state를 logits로 projection하고 loss에 반영한 뒤 버립니다. backward에서는 chunk를 다시 계산합니다. 계산은 조금 늘어나지만 peak memory를 크게 줄입니다.

공개 글의 benchmark에서는 32K tokens의 toy output-projection 기준 dense loss peak memory가 85.2 GiB인 반면 fully chunked version은 5.45 GiB였고, 15.6× reduction이라고 설명합니다. 256K tokens에서는 fully chunked loss가 11.6 GiB를 사용했다고 합니다. 긴 컨텍스트에서 차이가 커지는 구조입니다.

언제 이 방식을 써야 하나

모든 팀이 바로 fused chunked KL을 구현할 필요는 없습니다. 짧은 context, 작은 vocabulary, 작은 student라면 일반 distillation도 충분할 수 있습니다. 이 패턴은 다음 조건에서 효과가 큽니다.

• 8K 이상 context에서 distillation한다.
• teacher와 student를 동시에 올리기 어렵다.
• 여러 dataset 또는 hyperparameter ablation을 반복한다.
• GPU 수보다 실험 반복 속도가 병목이다.
• distillation 목적이 “최고 성능”보다 “배포 가능한 작은 모델”이다.

반대로 teacher response를 실시간으로 바꿔야 하거나, RL-style online feedback이 핵심인 경우 offline cache가 맞지 않을 수 있습니다. teacher가 고정된 supervised distillation에 먼저 적용하는 편이 안전합니다.

실무 적용 순서

먼저 baseline을 작게 잡아야 합니다. 큰 모델과 긴 context로 바로 들어가면 메모리 절감 효과는 보이겠지만, 품질 회귀를 추적하기 어렵습니다. 작은 subset에서 online distillation, offline dense KL, offline chunked KL을 비교하고 validation metric이 얼마나 달라지는지 봐야 합니다.

권장 실험 순서는 다음과 같습니다.

1. 대표 dataset 1~5% 샘플링
2. teacher top-K logits cache 생성
3. online distillation baseline 1개 실행
4. offline dense KL로 loss curve 비교
5. forward-chunked KL로 memory와 throughput 측정
6. fused chunked KL로 long context에서 peak memory 측정
7. 최종 모델을 task eval set으로 비교

cache 저장 포맷도 중요합니다. position별 top-K token id와 logit, temperature, teacher snapshot, tokenizer version, dataset hash를 같이 저장해야 합니다. tokenizer가 바뀌면 cache는 재사용하면 안 됩니다.

운영에서 놓치기 쉬운 부분

offline cache는 비용을 줄이지만 storage와 lineage 문제가 생깁니다. teacher output은 모델의 파생 산출물이므로 접근 권한을 분리해야 할 수 있습니다. 특히 proprietary data로 만든 cache라면 원본 dataset만큼 민감하게 다뤄야 합니다.

또한 top-K cache를 여러 실험에서 재사용하면 편하지만, cache 생성 오류가 모든 실험에 전파됩니다. cache validation 단계가 필요합니다. 예를 들어 token id 범위, sequence length, missing position, top-K 정렬, logit dtype, dataset row mapping을 검사해야 합니다.

마지막으로 throughput만 보면 안 됩니다. fused chunked 방식은 backward에서 projection을 다시 계산하기 때문에 짧은 context에서는 forward-chunked보다 느릴 수 있습니다. 긴 context를 실제로 쓸 때만 memory 절감 이득이 커집니다.

실행 체크리스트

• teacher snapshot, tokenizer version, dataset hash를 cache metadata에 남긴다.
• top-K는 50/100/200 중 2개 이상으로 품질 차이를 비교한다.
• online distillation baseline을 작은 샘플에서 먼저 만든다.
• dense KL, forward-chunked KL, fused chunked KL을 같은 seed와 dataset으로 비교한다.
• peak memory, iteration time, validation loss, downstream eval을 함께 본다.
• cache 파일은 원본 학습 데이터와 같은 보안 등급으로 관리한다.
• 8K 이하에서는 단순 구현을 먼저 쓰고, 32K 이상부터 fused chunked KL을 우선 검토한다.
• 최종 선택 기준을 “가장 빠름”이 아니라 “목표 context에서 안정적으로 반복 실험 가능한가”로 둔다.

Knowledge distillation 비용 절감의 핵심은 모델 크기만 줄이는 것이 아닙니다. teacher 계산과 loss 메모리 구조를 분리해야 반복 실험이 가능해집니다.