메모리 아끼면서 Cross Entropy Loss 계산하기

요약

LLM 학습 시 긴 컨텍스트와 큰 어휘 사전으로 인해 발생하는 Cross Entropy Loss의 메모리 문제를 분석합니다. FLCE 방식을 통해 거대한 Logits 텐서가 메모리를 점유하지 않도록 효율적으로 계산하는 방법을 제안합니다.

핵심 포인트

• 긴 컨텍스트 학습 시 Logits 텐서가 모델 가중치보다 큰 메모리를 점유할 수 있음
• 단순한 Chunk 분할 방식은 Autograd의 그래프 유지 문제로 Peak Memory를 낮추지 못함
• FLCE는 Forward pass 내에서 Gradient를 즉시 계산하여 메모리 효율을 극대화함
• 메모리 사용량과 연산 지연 시간(Latency) 사이의 트레이드오프 분석 제공