장문맥 대규모 언어 모델 학습에서 병목은 대부분 가중치가 아니라 LM head와 Cross Entropy Loss 계산에서 발생한다. 특히 context 길이가 128K 토큰을 넘어가면 logits 텐서 하나만으로 수십 GB에 육박해, 16B급 모델 학습에서 실제로 OOM이 보고되고 있다. 본문에서는 이 문제의 구조적 원리와 이를 풀기 위한 분리형 Cross Entropy 기법을 짚어본다.
3줄 핵심 요약
- LM head와 Cross Entropy Loss는 128K 토큰 같은 장문맥에서 모델 가중치만큼 혹은 그 이상으로 메모리 병목이 될 수 있다.
- 16B 모델과 128K context 조합에서는 logits 텐서 및 관련 버퍼가 약 40GB에 근접해 OOM(out-of-memory) 사례가 보고된다.
- forward와 backward에서 Cross Entropy를 분리해 연산하면 logits를 메모리에 상주시키지 않고 학습할 수 있다.
결론적으로 logits 비저장 전략과 backward 재계산 흐름을 결합하는 방식이 장문맥 LLM 학습 메모리 최적화의 현실적인 해법 중 하나로 제시되고 있다.
왜 Cross Entropy가 LLM 학습의 메모리 킬러인가
LLM 학습에서 메모리를 가장 크게 잡아먹는 항목이 무엇이냐는 질문에, 많은 실무자는 가중치와 옵티마이저 상태를 떠올린다. 하지만 context 길이가 길어질수록 또 다른 변수가 무게중심을 옮기는데, 그것이 바로 LM head를 통과한 뒤 생성되는 logits 텐서와 그 위에서 계산되는 Cross Entropy Loss다. 이 영역은 그동안 상대적으로 덜 주목받았지만, 장문맥 학습이 보편화되면서 무시할 수 없는 비용 영역으로 부상하고 있다.
logits 텐서가 가중치보다 커지는 구조적 이유
LM head는 마지막 hidden state를 vocab 크기만큼의 차원으로 투사하는 선형층이다. 따라서 한 시점에서 생성되는 logits은 형태가 (batch, seq, vocab)가 된다. 일반적인 LLM은 vocab 크기가 수만에서 수십만에 이르고, 여기에 장문맥의 seq 길이까지 곱해지면 텐서 크기가 폭발적으로 커진다.
예를 들어 batch가 1이고 seq가 128K, vocab이 128K라면, fp32 기준 한 시점의 logits 텐서 크기는 1 × 131072 × 131072 × 4byte로 계산되며 약 64GB에 근접한다. fp16 또는 bf16을 써도 약 32GB 수준으로, 가중치 전체보다 큰 단일 활성화가 등장한다. 이 logits는 보통 softmax, log, Cross Entropy 계산을 위해 메모리에 상주해야 하는데, 이 지점이 학습 시 메모리 피크의 주범이 된다.
| 구성 요소 | 형태 | 메모리 영향 |
|---|---|---|
| 모델 가중치 (16B) | 파라미터 | 옵티마이저까지 합쳐 100GB 이상 |
| 활성화 (은닉 상태) | (batch, seq, hidden) | 체크포인팅으로 통제 가능 |
| logits 텐서 | (batch, seq, vocab) | 128K 기준 약 40GB까지 치솟음 |
| Cross Entropy 보조 버퍼 | log-softmax, mask, label | logits와 동등한 크기로 함께 점유 |
16B 모델, 128K context에서 실제 OOM 사례
GeekNews에 정리된 바에 따르면, 16B 파라미터 규모 모델을 128K 토큰 context에서 학습하는 환경에서 logits 텐서 단독 크기가 약 40GB에 근접하는 사례가 보고됐다. 동일한 환경에서 가중치와 옵티마이저 상태 외에 활성화 체크포인팅을 적용해도, Cross Entropy 단계에서 OOM이 발생하는 것으로 알려져 있다. 이는 vocab 크기와 seq 길이의 곱이 모델 자체의 파라미터 수를 능가하는 사례가 실제로 존재함을 보여준다.
메모리를 줄이는 핵심 기법: 분리형 Cross Entropy
이 문제를 푸는 자연스러운 접근은 Loss 계산 자체를 두 단계로 쪼개는 것이다. 핵심 아이디어는 logits 전체를 메모리에 남겨두지 않고, 가능한 한 빨리 필요한 스칼라 또는 부분 텐서로 환원한 뒤 버리는 것이다. 대표적으로 forward에서는 대상 토큰 위치의 logit만 잘라내고, backward에서는 그라디언트를 다시 계산해 LM head 파라미터와 hidden state로 흘려보내는 방식이 사용된다.
forward 단계에서 logits 비저장 전략
가장 큰 메모리 점유를 차지하는 것은 전체 vocab 차원의 logit 행렬이다. 하지만 Cross Entropy를 계산하는 데 실제로 필요한 정보는 정답 토큰 위치의 logit과, 분모에 해당하는 log-sum-exp 정도다. 따라서 다음과 같은 흐름으로 단순화할 수 있다.
- LM head 출력을 dense한
(B, T, V)텐서로 펼쳐 저장하지 않고, 필요한 부분만 잘라낸다. - 온라인 softmax 또는 chunk 기반 log-sum-exp를 적용해 log-prob를 누적한다.
- 음의 로그 가능도를 더해 스칼라 Loss로 환원한 뒤, logits 자체는 해제한다.
이 과정에서 vocab 차원의 큰 텐서를 autograd 그래프에 통째로 남기지 않게 되므로, 메모리 피크가 크게 낮아지는 것으로 보고된다.
backward 단계에서 그라디언트 재계산 흐름
앞 단계에서 logits를 버렸다면, backward 시점에는 그라디언트가 hidden state와 LM head 가중치까지 도달하지 못한다. 이를 해결하기 위해 다음과 같은 절충이 자주 쓰인다.
- Loss에 대한 LM head 출력의 그라디언트를 다시 만든다.
- 이 그라디언트는 보통
(B, T, V)형태지만, 대상 토큰과 log-sum-exp에 쓰인 위치에 대해서만 비영(非零)이다. - 이를 다시 matmul 역연산으로 분해해 hidden state에 대한 그라디언트와 LM head 가중치에 대한 그라디언트로 환원한다.
핵심은 메모리에 떠 있는 logits 없이도 결과적으로 동일한 Loss와 그라디언트를 재현할 수 있다는 점이다. 이 패턴은 PyTorch의 일부 메모리 친화적 Cross Entropy 구현체와 DeepSpeed, Megatron 계열 학습 프레임워크에서 유사한 형태로 채택되고 있는 것으로 분석된다.
장문맥 LLM 학습으로의 함의와 실무 체크리스트
장문맥 학습은 점차 일반화되는 추세이며, 64K, 128K를 넘어 1M 토큰을 다루는 시도도 공개되고 있어, logits 메모리 이슈는 사전에 점검해야 할 항목으로 부상하고 있다.
vocab 크기와 context 길이의 메모리 영향 추정법
실무적으로는 다음 식으로 한 시점 logits 메모리를 빠르게 추정할 수 있다. logits_bytes ≒ batch × seq × vocab × dtype_size이다. 예를 들어 batch 1, seq 131072, vocab 131072, bf16이라면 한 시점에 약 32GB의 logits 텐서 메모리가 필요하다. 여기에 softmax, log, 보조 마스크 버퍼가 더해지면 실제 피크는 더 커질 수 있다. 따라서 vocab 크기를 조정하거나, 시퀀스를 chunk로 잘게 쪼개 gradient accumulation으로 처리하는 방식이 실질적 대안으로 자주 거론된다.
앞으로의 학습 인프라 최적화 방향
현재 LLM 학습 인프라 최적화는 크게 세 갈래로 움직이는 것으로 보인다.
- logits 비저장형 Cross Entropy와 같은 loss-level 최적화
- 활성화 및 layer 단위 체크포인팅 확대
- ZeRO, FSDP 등 옵티마이저 및 가중치 분산 전략
이 중 loss-level 접근은 모델 구조를 바꾸지 않고 적용할 수 있어 호환성이 높다는 점이 강점으로 평가된다. 다만 구현 난이도와 디버깅 비용이 있기 때문에, 공식 지원 여부와 프로파일링 도구 호환성을 함께 확인해야 안정적인 도입이 가능하다는 의견이 제기된다.
정리하면
- 장문맥 LLM 학습의 메모리 병목은 가중치보다 LM head와 Cross Entropy 단계에 있다.
- 16B와 128K 조합에서 logits 한 텐서가 약 40GB로 치솟아 OOM이 실제로 발생한다.
- forward에서 logits를 가능한 빨리 환원하고, backward에서 그라디언트를 재계산하는 분리형 Cross Entropy가 현실적 해법이다.
- vocab과 seq의 곱, batch 크기, dtype에 따른 logits 메모리 추정 공식을 표준 점검 항목으로 두는 것이 권장된다.
참고 자료: GeekNews 메모리 아끼면서 Cross Entropy Loss 계산하기, GeekNews 관련 토픽