logits 대신 마지막 레이어의 hidden states를 캐싱하기: 전체 어휘(full-vocab) logits는...

요약

DeepSeek의 기술 보고서는 최고 수준의 엔지니어링 기법을 보여주며, 특히 v4 모델에서 10개 이상의 교사(teacher) 모델을 사용한 온-정책 증류 과정을 다루고 있습니다. 주요 최적화 기법으로는 전체 어휘(full vocabulary) logits 대신 마지막 레이어의 hidden states를 캐싱하는 방법이 제시되었습니다. 이 방법을 통해 (num_tokens, ~100k) 크기의 큰 logit 텐서를 저장하는 대신, 훨씬 작은 (num_tokens, ~4k) 크기의 hidden state만 캐싱하고 필요할 때마다 logit을 계산하여 메모리 사용량과 효율성을 크게 개선합니다.

핵심 포인트

• 여러 교사 모델(teacher models)을 사용하는 온-정책 증류(on-policy distillation) 관리 방법론 제시
• 교사 인덱스에 따라 미니 배치(mini-batch)를 정렬하여 여러 교사 모델의 헤드 로딩 횟수를 최소화함
• 전체 어휘 logits 대신 마지막 레이어 hidden states를 캐싱하는 최적화 기법을 적용함
• 큰 logit 텐서 (num_tokens, ~100k) 대신 작은 hidden state (num_tokens, ~4k)만 캐싱하여 메모리 효율성을 극대화함