TL;DR
교사의 전체 확률 분포를 매 스텝 재생성하는 전통적 온라인 KL 증류는 어휘×시퀀스 그리드 때문에 VRAM 비용이 폭증한다. 교사 출력의 상위 100개 로짓을 한 번 캐시하고 손실 계산을 청크 단위로 퓨전하면 교사를 메모리에서 제거하면서도 학습 손실을 유지할 수 있다. 특히 Fused chunked KL은 시퀀스 길이에 대해 메모리 성장이 선형이 되도록 해 32K 문맥에서 메모리를 수십 기가바이트에서 단일 기가바이트 수준으로 줄여 장문 증류를 싱글 GPU에서 실현 가능하게 만든다.
섹션별 상세
이미지 분석
왼쪽 그래프에서 Dense KL은 시퀀스 길이가 커지며 H200 용량(≈141GB)을 넘는 OOM 구간을 만들지만 Fused chunked KL은 같은 길이에서 훨씬 낮은 피크 메모리를 유지한다. 오른쪽 반복률 그래프는 처리량이 길이 증가에 따라 떨어지지만 청크화된 방법들이 Dense 대비 더 나은 균형을 제공함을 시사한다. 그래프의 주요 수치로 32K에서 15.6× 메모리 절감이 보고되어 장문 문맥에서의 실효성을 정량적으로 뒷받침한다.
시퀀스 길이별 피크 VRAM과 반복률을 비교한 차트로, Dense KL이 시퀀스 길이에 따라 메모리가 급격히 증가하는 반면 Fused chunked KL은 완만하게 증가함을 보여준다.
왼쪽의 Dense KL 패널은 Loss/logits(파란 대각 텍스처)가 큰 피크를 형성하며 H200의 141GB 한계를 넘는 반면 오른쪽 Fused Chunked KL 패널은 스파이크 없이 약 128GB에 머문다. 이 시각화는 Dense 방식이 전체 어휘×시퀀스 그리드를 한꺼번에 보관하기 때문에 일시적이지만 큰 메모리 폭증을 유발함을 직관적으로 보여준다. H200 용량선과 피크 라벨이 장비 제약과 실용적 한계를 연결해 준다.
32K 컨텍스트 한 반복에서 Dense KL의 메모리 스파이크(약 250GB)와 Fused Chunked KL의 낮은 피크(약 128GB)를 누적 영역 그래프로 비교한 그림이다.
애니메이션은 Dense 접근이 왜 메모리 스파이크를 만드는지 구조적으로 보여주며, Fused chunked는 한 번에 한 조각만 투영·합산하고 폐기하는 흐름을 통해 메모리 상한을 낮춘다는 핵심 아이디어를 직관적으로 전달한다. 역전파에서의 재계산 비용이 존재하지만 메모리 절감이 더 큰 상황에서 이 트레이드오프가 실용적임을 시각적으로 보강한다.
밀집 Dense KL이 전체 어휘×시퀀스 격자를 한꺼번에 구성해 유지하는 반면 Fused chunked는 격자를 청크 단위로 만들고 버리는 과정을 애니메이션으로 보여준다.
손실 곡선은 네 방식이 거의 동일하게 겹치며 오프라인 top-100 캐시가 손실 관점에서 무손실 근사가 가능함을 보여준다. 메모리 분해와 막대그래프 라벨은 Online이 102.8 GB로 가장 큰 피크를 보였고 Fused Chunked KL이 58.3 GB로 가장 낮았음을 명확히 제시하며 반복 시간과 TFLOP/s의 trade-off도 함께 제시한다. 이 그림은 실무자가 메모리 절감 대비 속도 손실을 정량적으로 비교할 수 있게 해준다.
훈련 손실, 반복 시간·처리량 막대, 그리고 피크 GPU 메모리 분해를 한 페이지에 정리한 복합 차트로 네 방식(Online, Dense KL, Chunked KL, Fused Chunked KL)을 비교한다.
그래프는 학생이 BoolQ와 HellaSwag에서 교사 성능에 매우 근접하고 MMLU에서는 약 9포인트 차이를 보이며 전체적으로 교사보다 성능 저하가 크지 않음을 나타낸다. 이는 오프라인 Top‑100 캐시와 청크 손실 전략이 메모리를 크게 줄이면서도 핵심적인 짧은 문맥 성능을 대부분 보존했다는 실험적 증거를 제공한다.
Llama 3.1 8B 교사와 증류된 약 3.2B 학생의 단문(context) 성능을 BoolQ, WinoGrande, MMLU, HellaSwag, GSM8K에서 비교한 막대그래프다.
용어 해설
- 지식 증류(Knowledge Distillation)
- — 큰 교사 모델의 예측을 작은 학생 모델이 모방하도록 학습시키는 방법으로, 교사의 확률 분포를 학생이 재현하도록 KL 손실 등으로 맞추며 모델 압축과 배포 비용 절감에 사용된다.
- KL 손실(KL loss (Kullback‑Leibler))
- — 교사와 학생의 확률 분포 차이를 측정하는 손실로, 전체 어휘 대 시퀀스 그리드를 만들면 메모리 사용량이 어휘 크기×시퀀스 길이만큼 급증하므로 구현 방식이 학습 가능성에 큰 영향을 준다.
- Top‑K 로짓 캐시(Top‑K logits cache)
- — 교사 모델 출력에서 각 위치별로 확률이 높은 상위 K개 토큰의 로짓만 저장하는 방식으로, 교사를 매 스텝 실행하지 않고도 근사된 분포를 재사용해 VRAM과 계산을 줄인다.
- 청크 단위 KL 손실(Chunked KL loss)
- — 시퀀스를 작은 청크로 나눠 각 청크를 순차 처리하고 필요 시 역전파에서 재계산하여 전체 어휘×시퀀스 그리드를 메모리에 올리지 않는 손실 계산 기법으로, 메모리 성장률을 선형으로 제한한다.
근거 모음
- 32K 토큰에서 Dense KL의 피크 메모리 85.2 GiB가 Fully chunked로 5.45 GiB로 줄어들어 15.6× 절감이 관찰되었다. — 그림(시퀀스 길이별 메모리/반복률 스케일링)과 본문(32K 결과 요약).
- 단일 H200, 8K 컨텍스트 실험에서 Offline fused chunked KL은 피크 58.3 GB, 반복 시간 20.2 s, 처리량 304 TFLOP/s로 측정되었다. — 본문의 표(8K context, single H200)와 해당 막대그래프라벨.
- 오프라인 Top‑100 로그잇 캐시와 세 가지 오프라인 손실(출력 참조) 모두 온라인 증류와 거의 동일한 학습 손실 곡선을 보였다(무손실 근사). — 훈련 손실 그래프와 본문 'loss curves overlap — losses match across methods (lossless)' 문구.
- GPT‑OSS 20B를 32,768 토큰 문맥으로 증류할 때 fused chunked 손실은 전체 설정을 4 GPU 노드에서 1 GPU로 줄였고 스텝 시간을 57.0초에서 12.23초로 감소시켰다. — 본문의 확장 사례(32,768-token distillation) 숫자 및 처리량 변화 설명.
기술
- Llama 3.1 8B Instruct
- H200 GPU
- bfloat16
- TFLOP/s 측정 지표
- GitHub: CompactifAI/Full-Chunked-KL-Loss
활용 사례
- 긴 문맥을 필요로 하는 지식 증류 작업에서 단일 GPU로 대규모 증류를 수행하는 엔지니어링 파이프라인 구축에 유용하다.
- 대규모 교사 모델을 여러 학생 모델로 대규모로 실험하면서 교사 재실행 비용을 줄여 반복적인 ablation을 빠르게 수행할 때 비용을 절감한다.
- 어휘가 크고 시퀀스가 긴 생성 모델의 손실 계산을 메모리 제한 내에서 확장할 필요가 있는 연구·상용 환경에서 적용 가능하다.
언급된 리소스
AI 요약 · 북마크 · 개인 피드 설정 — 무료
출처 · 인용 안내
인용 시 "요약 출처: AI Trends (aitrends.kr)"를 표기하고, 사실 확인은 원문 보기 기준으로 진행해 주세요. 자세한 기준은 운영 정책을 참고해 주세요.
