LLM을 위한 효율적인 지식 증류: Offline Top-K Logits 및 Fused Chunked KL Loss

Multiverse Computing은 VRAM 사용량과 학습 비용을 크게 줄여 단일 GPU에서 긴 문맥(long-context) 증류를 가능하게 하는 새로운 지식 증류 방법을 도입했습니다. 연구진은 offline logit caching과 fused chunked KL-divergence loss를 결합함으로써, 전체 어휘 확률 분포와 관련된 막대한 메모리 오버헤드 없이 대규모 teacher 모델을 더 작은 student 모델로 증류할 수 있게 되었습니다.

표준 지식 증류의 높은 비용

Kullback-Leibler (KL) divergence loss를 사용하는 전통적인 "online" 지식 증류는 teacher 모델과 student 모델을 동시에 VRAM에 로드해야 하므로 계산 비용이 많이 듭니다. 모든 학습 단계에서 teacher는 모든 토큰 위치에 대해 전체 어휘에 대한 확률 분포를 생성하기 위해 전체 forward pass를 수행해야 합니다.

대규모 모델의 경우, 이는 엄청난 메모리 병목 현상을 일으킵니다. 예를 들어, 어휘 수가 201,088개인 gpt-oss-120b 모델을 사용하여 sequence length 32K, batch size 4로 학습할 경우, teacher-probability tensor 하나만으로도 bfloat16 기준 약 50GB의 VRAM이 필요합니다. gradients, activations, optimizer states를 포함하면 단일 iteration에서 약 250GB의 VRAM 피크가 발생할 수 있으며, 이는 H200이나 B200과 같은 하이엔드 GPU의 용량도 초과합니다.

확장 가능한 증류를 위한 두 가지 시스템 변경 사항

이러한 메모리 제약을 해결하기 위해 Multiverse Computing은 두 가지 주요 기술적 변경 사항을 구현했습니다:

1. Top-K Logit Caching을 통한 Offline Distillation

매 단계마다 teacher의 출력을 다시 계산하는 대신, 시스템은 teacher의 출력을 한 번 계산하고 위치당 가장 가능성이 높은 top-100 토큰을 캐싱합니다. 그 다음 student는 이 캐시를 대상으로 학습됩니다. 이를 통해 학습 과정 중에 teacher 모델을 메모리에 유지할 필요가 없으며, 동일한 캐시를 다양한 실험적 ablation에 재사용할 수 있습니다.

2. Fused Chunked KL Loss

표준 KL loss 구현은 student와 teacher 사이의 불일치를 계산하기 위해 전체 어휘-시퀀스 그리드를 구축합니다. Multiverse Computing의 "fused chunked KL loss"는 모델의 output projection을 loss 계산에 직접 융합함으로써 이를 최적화합니다.

전체 logits 그리드를 실체화하는 대신, 프로세스는 다음과 같이 작동합니다:

• 시퀀스의 한 청크(chunk)를 한 번에 처리합니다.
• 해당 특정 청크에 대해서만 hidden states를 logits로 투영합니다.
• 결과를 실행 중인 loss에 통합하고 즉시 청크를 폐기합니다.
• backward pass는 각 청크를 실시간으로 다시 계산합니다.

이 접근 방식은 피크 메모리가 전체 어휘 크기에 따라 급증하는 대신 시퀀스 길이에 따라 선형적으로 증가하도록 보장합니다.

성능 및 메모리 벤치마크

단일 H200 GPU에서 Llama 3.1 8B Instruct (teacher)와 3.2B Llama 모델 (student)을 사용하여 8K 토큰 문맥으로 다양한 loss 구현을 비교한 결과, 연구진은 모든 방법이 거의 동일한 training loss에 도달함을 발견했습니다. 이는 top-100 캐시된 logits만을 사용하는 offline distillation이 online distillation에 비해 실질적으로 손실이 없음을 확인시켜 줍니다.

방법 (8K context, 단일 H200)	피크 메모리	Iteration 시간	Throughput
Online distillation	102.8 GB	25.9 s	237 TFLOP/s
Offline, dense KL	78.3 GB	18.5 s	331 TFLOP/s
Offline, forward-chunked KL	61.8 GB	18.4 s	335 TFLOP/s
Offline, fused chunked KL	58.3 GB	20.2 s	304 TFLOP/s

긴 문맥으로의 확장

Fused chunked loss의 장점은 문맥 길이가 길어질수록 더욱 두드러집니다. toy output-projection network를 사용한 격리된 벤치마크에서:

• 32K 토큰: 피크 메모리가 85.2 GiB (dense loss)에서 5.45 GiB (fully chunked)로 감소하여 15.6배 절감되었습니다.
• 64K 토큰: dense loss 구현은 완전히 실패했습니다.
• 256K 토큰: fully chunked loss는 11.6 GiB를 사용한 반면, 차선책인 chunked variant는 134.2 GiB를 사용하여 iteration당 3.3배 더 빨랐습니다.

32,768-토큰 문맥에서 GPT-OSS 20B 모델을 증류하는 실제 시나리오에서, fused loss는 필요한 하드웨어를 4개의 GPU 노드에서 1개로 줄였습니다. 이는 step time에서 5배의 속도 향상(57.0초에서 12.23초로)과 GPU당 throughput의 증가(74.2에서 345.7 TFLOP/s로)를 가져왔습니다.

Student 모델 역량

이 파이프라인의 효율성 덕분에 Llama 3.1 8B Instruct로부터 증류된 3.2B 파라미터 student 모델을 생성하는 대규모 증류 캠페인이 가능했습니다. 이 컴팩트한 student는 파라미터 수가 절반 미만임에도 불구하고 BoolQ 및 HellaSwag에서 teacher의 정확도를 대부분 유지하며, MMLU 벤치마크에서 teacher와 약 9점 차이 이내를 유지합니다.