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

LLM 학습 중 128K context와 같은 긴 시퀀스를 처리할 때, Cross Entropy Loss 계산 단계에서 갑작스러운 **OOM(Out of Memory)**이 발생하는 원인을 분석합니다.

문제의 핵심은 Logit Materialization에 있습니다. PyTorch의 기본 F.cross_entropy를 사용할 경우, 계산 과정에서 거대한 크기의 Logit 텐서가 메모리에 실체화(Materialization)되어 메모리 점유율이 수직 상승하게 됩니다.

이를 해결하기 위한 핵심 전략은 다음과 같습니다:

• Memory-aware 알고리즘 도입: Logit 전체를 메모리에 올리지 않고, 계산에 필요한 부분만 처리하는 방식입니다.
• Chunking 기법: 시퀀스를 작은 단위(Chunk)로 나누어 계산함으로써 메모리 피크(Peak)를 억제합니다.
• Fused Kernel 활용: 연산을 하나로 합쳐 중간 단계의 대규모 텐서 생성을 방지합니다.

이러한 최적화는 대규모 모델의 Long Context 학습 시 메모리 효율성을 극대화하는 데 필수적입니다.