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

지난 3월, 128K context로 Gravity 16B를 학습시키려던 참이었습니다. batch size를 1까지 낮추고 gradient checkpointing을 전부 켰는데도, 매번 정확히 같은 지점에서 OOM이 났습니다.

INFO  step: 18420  loss: 2.7136  ntp: 2.4819  mtp: 2.9452  mem: 92%
...
  File "cyan/models/gravity/model.py", line 284, in forward
    mtp_loss += cross_entropy(mtp_logit, mtp_target)
  File "cyan/models/gravity/model.py", line 27, in cross_entropy
    return F.cross_entropy(pred.flatten(0, 1).float(), target.flatten(0, 1), ...)
                                              ^^^^^^^
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 37.00 GiB.
  GPU 7 has 178.36 GiB total, 33.25 GiB free.

계속 마지막 loss 계산에서 34GB가 한 번에 잡히면서 터졌습니다. 어느 지점 때문에 이렇게 터지는지 궁금해 Memory Viz를 열어봤습니다.

프로파일은 대부분의 구간에서 평탄했습니다. 그러다 딱 한 곳에서 수직으로 솟았고, 그게 cross entropy였습니다.

이 글에서는 왜 이런 솟구치는 현상이 발생하는지를 살펴보려고 합니다. 나아가 이 문제를 해결하기 위한 cross entropy 계산 방법들을 소개해보려고 합니다. 차근차근 따라가다 보면 Logit materialization이 왜 문제이고, memory-aware한 알고리즘이라는 게 무엇인지 감을 잡으실 수 있을 것입니다.

LM Head and Cross Entropy

LLM 학습에서 마지막 hidden state는 두 단계를 거쳐서 loss가 됩니다. 먼저, vocab projection (LM head)를 통해 logit이 되고, 그 다음 cross-entropy를 통해 스칼라 loss가 됩니다. 이 두 연산에 대해서 Forward, Backward를 한 번 살펴보겠습니다.

Notations

• $B$: batch size
• $S$: sequence length
• $V$: vocabulary size
• $D$: hidden dimension
• $C$: token chunk size
• $M$: number of chunks

Forward

마지막 hidden state $h \in \mathbb{R}^{1 \times D}$와 LM head weight $W \in \mathbb{R}^{V \times D}$ 가 있을 때, LM head를 통과시키면 전체 Vocab에 대한 점수인 Logits가 나오게 됩니다. 여기서는 PyTorch의 linear layer처럼 hidden state를 row vector로 두겠습니다.

이 때, logits은 숫자 벡터입니다. 이를 다음 토큰이 무엇일까의 확률 분포로 만들기 위해, softmax로 정규화합니다.

이제 정답 토큰 인덱스 $y$에 대해, 모델이 그 토큰에 부여한 확률 $p^{(y)}$ 를 보겠습니다. 이 값이 1에 가까울수록 좋고, 0에 가까울수록 나쁜 것이죠. Loss로 쓰려면 “확률이 높을수록 작은 값”이 되어야 하니까, $-\log$를 씌워 Cross Entropy loss를 계산하게 됩니다.

위 식에 softmax 정의를 대입하고 $\log$를 분배하면 다음과 같이 쓸 수 있습니다.

Equation 2는 모든 logit에 대한 합(두 번째 항, 흔히 log-sum-exp라고 부르기도 합니다.)과 정답 토큰의 logit(첫 번째 항) 사이의 차이입니다. 정답 logit을 끌어올리거나, 나머지 전부를 끌어내리면 loss가 줄어듭니다.

Backward: dL/dz

이제 $\partial \ell / \partial z$를 계산하겠습니다. 위에서 풀어쓴 식에서 시작하겠습니다.

첫째 항: $-z^{(y)}$를 $z^{(j)}$로 미분하면, $j = y$일 때만 $-1$이고 그 외엔 $0$입니다. Indicator 함수로 쓰면 다음과 같습니다.

둘째 항: $A = \sum_k \exp(z^{(k)})$로 두면 chain rule로 다음과 같이 계산됩니다.

두 항을 합치면:

Equation 3은 예측 분포 p에서 정답 분포 (one-hot)를 뺀 것입니다. 정답 자리에선 음수, 나머지 자리에선 양수이면서, gradient가 곧 “정답에서 얼마나 빗나갔나”를 나타냅니다.

Backward: dL/dh and dL/dW

logit gradient를 기반으로 LM Head의 Wgrad, Hgrad도 계산해보겠습니다.

Hidden Gradient

$g = \partial \ell / \partial z$ 라고 두면, $z = h W^T$ 이므로 hidden state로 돌아가는 gradient는 다음과 같습니다.

Weight Gradient

한 토큰만 봤을 때, $\ell$ 이 $W$ 에 의존하는 방식은 $z = h W^T$ 한 줄뿐입니다.

시퀀스로 확장

이제 시퀀스 전체로 확장해보겠습니다. 길이 $S$, batch $B$ 이면 전체 토큰 수는 $B \cdot S$ 입니다. 각 토큰 $i$에 대해 위 forward/backward 계산 방식을 그대로 적용하면 됩니다.

Loss

Weight gradient

• dw

Hidden gradient

• dx

Logits at Scale

앞 섹션에서 LM head + CE의 동작을 처음부터 끝까지 따라가봤습니다. 그런데 막상 학습을 돌려보면, 이 두 연산이 메모리를 가장 많이 잡아먹는 부분 중 하나입니다. 특히 가장 문제가 되는 부분은 logits 텐서와 logit gradient의 크기 $(B,S,V)$에 있습니다.

특히 Llama 3 이후부터 vocabulary size가 100K는 기본으로 넘고, context length도 엄청나게 늘어나면서 logit 계산 과정의 메모리 소모량이 훨씬 더 커지게 되었는데, 실제 수치로 한 번 살펴보겠습니다.

늘어나는 Vocabulary Size

요즘 frontier 모델들의 vocab 크기를 살펴보면, 대부분이 130K~260K 범위에 자리잡았습니다. GPT-2의 50K, Llama 2의 32K 대비 거의 10배 가까이 늘어난 것을 알 수 있습니다.

$S = 4096$, $B = 4$ 기준으로 logits 텐서 크기를 계산해보면 다음과 같습니다.

배치 4, 시퀀스 4K 정도의 작은 셋업에서도 logits이 4~9 GB로, hidden state 대비 적게는 18배에서 많게는 70배까지 큽니다. 가장 극단적인 GPT-OSS 120B 기준으로 보면, hidden state는 0.09 GB인데 logits은 6.59 GB 이므로 약 70배입니다. Hidden dimension이 2,880으로 비교적 작은 반면 vocab은 200K가 넘기 때문에 LM head 한 번에 텐서가 그만큼 부풀어 오릅니다.

늘어나는 Context Length

Vocab만 문제면 그나마 다행인데, sequence length 트렌드가 더 무섭습니다. 최근 모델들은 32K context length를 기본으로 지원하고, DeepSeek V4와 Qwen3-Next는 1M context까지 갑니다. Llama 2의 4K에서 거의 25배 늘어난 수치입니다.

Figure 2의 배치 1짜리 계산만 봐도 128K context에서 logits이 40 GB에 육박합니다. H100 한 장의 메모리 (80 GB) 의 절반이고, 모델 weight (16B × 2 bytes ≈ 32 GB) 보다 큽니다. 256K로 가면 80 GB 한 장에 logits 하나도 못 담고, 1M로 가면 320 GB까지 부풀어서 텐서 하나가 HBM을 다 채워버립니다.

Backward 때도 Logits이 필요하다

지금까지는 forward에서 만들어지는 logits 텐서 자체의 크기만 봤습니다. 그런데 더 큰 문제는, 이 텐서가 forward 끝나고 사라지지 않는다는 점입니다.

이전 섹션에서 유도한 logit gradient를 다시 보면 Equation 3과 같습니다.

backward를 계산하려면 확률 $p$가 필요하고, $p$는 logits에서 나옵니다. PyTorch autograd는 이 사실을 모르지 않습니다. 그래서 backward 시점에 필요한 값을 만들 수 있도록 logits 또는 logits에서 이어지는 큰 intermediate를 computational graph에 남겨둡니다. 즉 forward에서 만든 큰 텐서가 backward가 끝날 때까지 메모리 pressure로 남습니다.

그리고 backward 자체가 또 한 덩이를 더합니다. Logit gradient $\partial \mathcal{L} / \partial z$ 도 $(B, S, V)$ 형태의 텐서라, logits과 같은 사이즈로 메모리에 한 번 더 만들어집니다. 40GB 짜리 logit, logit gradient 두 개로 H100 GPU HBM이 꽉 차게 되고, 학습에 OOM이 나게 됩니다.

Where Can We Cut?

결국 문제의 핵심은 $(B,S,V)$ 텐서가 너무 크다는 것입니다. 그러면, 해결책은 단순합니다.

어느 축이든 forward에서 chunk 단위로 만들었다 버리면, backward gradient도 같은 축을 따라 자연스럽게 분산할 수 있을 것입니다. 그럼 어느 축을 쪼갤 수 있을까요?

Batch Size B

$B$는 사실상 답이 아닙니다. 앞 섹션에서 본 숫자들은 이미 $B=1$ 기준이었습니다. 128K Context에서 logits이 40 GB, 1M context에선 320 GB. Batch를 더 쪼갤 자리가 없습니다.

Sequence Length S

$S$축은 첫 번째 쪼갤 수 있는 후보입니다. Sequence를 $M$개의 chunk $\mathcal{C}_1, \ldots, \mathcal{C}_M$ 으로 나누면, loss가 chunk별 부분합으로 분해됩니다:

Equation 4처럼 묶으면 기존 loss와 값은 동일하지만, 메모리 측면에선 큰 차이가 있습니다. chunk별로 logit $(B, |\mathcal{C}_m|, V)$ 만 만들고 처리한 뒤 버리면, $(B, S, V)$ 전체를 한 번에 안 만들어도 됩니다.

Vocabulary V

$V$ 축도 가능합니다. Vocab 인덱스 ${1, \ldots, V}$ 를 $M$개의 chunk $\mathcal{V}_1, \ldots, \mathcal{V}_M$ 으로 나누면, 각 chunk는 자기 영역의 partial logit만 만들면 됩니다. Logit 텐서가 $(B, S, V/M)$ 크기로 줄어듭니다.

문제는 합 분해 후의 처리가 $S$축만큼 단순하지 않다는 것입니다. CE 식을 다시 보면:

둘째 항의 합 자체는 chunk별로 분해됩니다.

다만 그 다음 $\log$ 가 걸려있어서, 모든 chunk의 부분합 $A_m$을 모은 뒤에야 최종 loss를 계산할 수 있습니다. Chunk별로 독립 처리한 뒤 그냥 누적하면 되는 $S$ 축과 달리, chunk들을 한 번 묶어서 같이 처리하는 단계가 필요한 거죠. 그래서 $V$축 chunking은 log-sum-exp reduction까지 포함한 별도 kernel 설계가 필요합니다. 이 글에서는 이 방향은 다루지 않고, token axis를 따라 쪼개는 방법에 집중하겠습니다.

Loss를 효율적으로 계산하는 방법은 결국 Logit의 $S$ 축을 잘 쪼개는 방법으로 생각해볼 수 있습니다.

1. Chunked Cross-Entropy (Chunked CE)

가장 직관적인 시도부터 시작하겠습니다. Sequence 차원 $S$ 을 따라 $M$개의 chunk로 나눠서, logit chunk를 만들고 처리한 뒤 버리는 식이죠.

수식으로 보면 단순합니다. Sequence를 $M$개의 chunk $\mathcal{C}_1, \ldots, \mathcal{C}_M$ 으로 나누면, loss는 합의 결합법칙으로 자연스럽게 분해됩니다.

각 chunk별 부분합 $A_m$ 을 계산할 때는 그 chunk에 속한 logits $(B, |\mathcal{C}_m|, V)$ 만 있으면 됩니다. 다음 chunk로 넘어가기 전에 이전 chunk logits을 버리면, $(B, S, V)$ 전체가 한 번에 메모리에 살아있을 필요가 없습니다.

코드로 쓰면 다음과 같습니다. for-loop을 돌면서 chunk별로 loss를 계산하게 됩니다.

def chunked_cross_entropy(x, weight, target, n_chunks=8, ignore_index=-100):
    x = x.reshape(-1, x.shape[-1])
    target = target.reshape(-1)
    loss = x.new_zeros(())
    count = target.new_zeros((), dtype=torch.long)
 
    for x_i, y_i in zip(x.chunk(n_chunks, dim=0), target.chunk(n_chunks, dim=0)):
        loss = loss + F.cross_entropy(
            x_i @ weight.T,
            y_i,
            ignore_index=ignore_index,
            reduction="sum",
        )
        count = count + (y_i != ignore_index).sum()
 
    return loss / count.clamp_min(1)

그럼 Chunked-CE를 실제 환경에서 측정해보겠습니다. 세팅은 다음과 같습니다.

• 모델: Gravity-16B-A3B (V = 151,552, D = 2,048, bf16)
• B=1, S=8192
• Forward + Backward (1 step training)

Figure 3의 결과를 보면, forward 시점의 chunk logit은 예상대로 1/M 만큼 줄지만, step peak memory는 어느 순간부터 ~4.7 GB 아래로 더 내려가지 않습니다.

텐서를 16개, 64개, 1024개로 쪼개도 peak은 ~4.7 GB 아래로 내려가지 않습니다. step time은 그에 반해 계속 증가합니다. Latency를 희생하고 memory를 잡으려 했지만, training step 관점에서는 원하는 만큼 memory가 내려가지 않은 겁니다.

왜 그럴까요? Backward 계산 때문입니다. Listing 2의 loop는 forward에서 한 번에 chunk 하나씩 만들지만, 각 chunk의 loss를 모두 더한 뒤 마지막에 한 번만 backward()를 호출합니다. 그동안 각 chunk의 CE graph는 최종 loss scalar에 연결된 채 남아 있어야 합니다. Forward-time logits buffer는 작아졌지만, training step peak에서는 chunk별 graph가 backward 전까지 누적되는 겁니다. M개가 전부 동시에 살아있으면 결국 우리가 피하려던 그 전체 Logit과 같습니다. 이건 어떻게 해결할 수 있을까요?

2. Fused Linear Cross-Entropy (FLCE)

Chunked CE는 forward 순간의 logits buffer 크기는 줄입니다. 하지만 training step의 peak memory는 forward만으로 결정되지 않습니다. backward가 실행되기 전까지 각 chunk의 CE graph가 살아 있어야 하므로, chunk를 잘게 쪼개도 전체 step 관점에서는 큰 중간 텐서들이 누적됩니다.

Chunked CE가 고려하지 못한 것은 loss와 gradient가 서로 다른 시점에 계산된다는 점입니다. Forward는 지금 logits를 만들고, backward는 나중에 그걸 소비하며, 그 사이의 모든 것은 보관되어야 합니다. 그렇다면 각 chunk의 logits가 손에 들어오는 바로 그 순간 forward pass 안에서 gradient를 계산해 버리면 어떨까요? 이것이 Fused Linear Cross Entropy (FLCE)의 핵심입니다.

FLCE는 토큰을 chunk로 쪼개서 처리하되, 각 chunk의 gradient를 그 자리에서 다 계산하고 chunk를 버립니다. 그래서 큰 텐서가 한순간도 graph에 남지 않게 합니다.

FLCE는 각 토큰 chunk마다 네 가지 일을 수행하고, 그 chunk를 완전히 버립니다.

1. Matmul: LM head projection
   • 먼저 현재 token chunk에 대해서만 LM head projection을 수행합니다.
   • logits_chunk = x_chunk @ weight.T
   • 이때 전체 $(B \cdot S, V)$ logits을 만들지 않고, 현재 chunk 크기인 $(C, V)$ logits buffer만 사용합니다.
2. Cross Entropy forward + dlogits 생성
   • cross_entropy_fwd_out이 logits_chunk를 입력으로 받아 각 token의 CE loss를 계산합니다.
   • 동시에 backward에 필요한 dlogits도 계산합니다. 수식으로는 dlogits = softmax(logits) - one_hot(target) 입니다.
   • 이때 log-sum-exp를 안정적으로 계산하는 streaming/online normalizer 아이디어는 online softmax 계산과 같은 계열입니다 [2].
3. dx 계산
   • 만들어진 dlogits_chunk를 이용해 hidden state에 대한 gradient를 바로 계산합니다.
   • dx_chunk = dlogits_chunk @ weight
   • 이렇게 하면 각 chunk의 input gradient는 forward 안에서 미리 계산되어 dx에 저장됩니다.
4. dw 누적
   • weight gradient는 chunk별 contribution인 dlogits_chunk.T @ x_chunk를 누적해서 계산합니다.
   • 첫 chunk에서는 dw를 만들고, 중간 chunk들은 여기에 계속 더합니다. 수식으로는 dw = dlogits_chunk.T @ x_chunk 입니다.

결국 FLCE의 핵심은 logits을 오래 들고 있지 않는 것입니다. 각 chunk에서 logits을 만든 직후 loss와 dlogits를 계산하고, 같은 buffer를 재사용해 dx와 dw에 필요한 값을 소비합니다. backward에서는 이미 계산된 dx, dw를 가져와 상위 graph에서 내려온 scalar gradient dloss로 scaling하고, forward에서 미뤄둔 마지막 chunk의 dw만 추가로 완성합니다.

gradient가 forward pass 바로 직후 만들어지기 때문에, FLCE가 반환하는 loss는 gradient가 미리 계산되어 따로 보관된 평범한 scalar입니다. Autograd의 backward()는 들어온 grad_output으로 그걸 scaling만 하면 됩니다. (C, V) 크기의 그 무엇도 다음 for-loop iteration에는 살아남지 않습니다.

이제 Naive CE와 FLCE를 memory-latency plane 위에 같이 올려보겠습니다. 실험 조건은 앞과 같은 Gravity-16B-A3B, $B=1$, $S=8192$, $V=151{,}552$, $D=2048$, bf16입니다. Naive CE는 큰 logits/dlogits buffer를 한 번에 들고 있어 peak memory가 높습니다. FLCE는 chunk가 작아질수록 peak memory가 내려가지만, 너무 작게 만들면 작은 GEMM과 kernel launch가 많아져 latency가 늘어납니다.

Figure 4를 보면 FLCE는 Naive CE보다 낮은 memory 영역으로 내려갑니다. 큰 chunk에서는 latency까지 유리하고, chunk를 더 작게 가져가면 memory를 더 줄이는 대신 latency를 지불합니다. 즉 FLCE의 핵심은 단순히 “항상 빠르다”가 아니라, full logits를 materialize하지 않아서 memory와 latency의 균형점을 선택할 수 있다는 데 있습니다.

다음으로 같은 chunk 수에서 Chunked CE와 FLCE를 비교해보겠습니다. Chunked CE는 forward logits buffer는 줄이지만, backward 전까지 graph가 남기 때문에 step peak가 잘 내려가지 않습니다. FLCE는 각 chunk의 gradient를 즉시 계산하고 buffer를 버리므로 같은 chunk 수에서도 더 낮은 peak memory를 유지합니다.

Figure 5의 결과를 보면, Chunked CE는 chunk 수가 많아져도 최소 4.7 GB 메모리를 잡아먹는 반면에, FLCE는 같은 chunk 조건에서 더 낮은 peak memory를 유지합니다. 다만 chunk를 지나치게 작게 만들면 FLCE도 kernel launch와 작은 GEMM 비용 때문에 latency가 커지므로, 실제로는 memory와 latency의 균형점을 잡아야 합니다.

Inside the Implementation

이제 코드를 살펴보겠습니다. QuACK의 linear_cross_entropy.py를 참고해서 설명드리겠습니다 [1]. 앞서 설명드린 FLCE의 네 가지 단계를 생각해보면서 따라오시면 더 쉽게 이해하실 수 있습니다.

ctx.weight_dtype = weight.dtype
x, weight = linear_fwd_convert_type(x, weight)
batch_shape = x.shape[:-1]
x = x.reshape(-1, x.shape[-1])
 
loss, dx, dw, last_dlogits_chunk, last_x_chunk = (
    chunked_linear_cross_entropy_fwd(
        x, weight, target, chunk_size, ignore_index, tuned=tuned
    )
)

for i, (x_chunk, target_chunk, loss_chunk, dx_chunk) in enumerate(
    zip(*(t.split(chunk_size) for t in (x, target, loss, dx)))
):
    chunk_len = x_chunk.shape[0]
    logits_chunk = logits_chunk_preallocated[:chunk_len]
 
    torch.mm(x_chunk, weight.mT, out=logits_chunk)

dlogits_chunk = logits_chunk
cross_entropy_fwd_out(
    logits_chunk,
    target_chunk,
    None,
    loss=loss_chunk,
    lse=None,
    dx=dlogits_chunk,
    weight=None,
    ignore_index=ignore_index,
)

torch.mm(dlogits_chunk, weight, out=dx_chunk)

if i == num_chunks - 1:
    last_dlogits_chunk = dlogits_chunk
    last_x_chunk = x_chunk
elif i == 0:
    gemm(dlogits_chunk.T, x_chunk, out=dw, tuned=tuned)
else:
    gemm_add_inplace(dlogits_chunk.T, x_chunk, dw, tuned=tuned)

loss_sum = loss.sum()
loss_scale = None
if reduction != "sum":
    loss_scale = 1.0 / (target != ignore_index).sum().float()
 
ctx.save_for_backward(dx, dw, last_dlogits_chunk, last_x_chunk, loss_scale)
return loss_sum if loss_scale is None else loss_sum * loss_scale

dx, dw, last_dlogits_chunk, last_x_chunk, loss_scale = ctx.saved_tensors
if loss_scale is not None:
    dloss = dloss * loss_scale
 
dx.mul_(dloss)
dx = dx.reshape(*ctx.batch_shape, dx.shape[-1])
 
gemm_add_inplace(
    last_dlogits_chunk.T,
    last_x_chunk,
    dw,
    alpha=dloss,
    beta=dloss,
    tuned=ctx.tuned,
)
return dx, dw, None, None, None, None, None

Conclusion

LM head와 Cross Entropy는 단순한 마지막 연산처럼 보이지만, 긴 context와 큰 vocabulary에서는 학습 메모리를 크게 좌우하는 병목이 됩니다. 문제의 핵심은 $(B, S, V)$ 크기의 logits와 logit gradient가 한 번에 materialize된다는 점입니다.

위에서 봤던 것처럼, 단순히 sequence를 chunk로 나누는 Chunked CE는 forward logits 크기는 줄이지만, backward를 위해 중간 graph가 남기 때문에 peak memory를 충분히 낮추지 못합니다. 반면 FLCE는 chunk별로 logits을 만든 직후 loss와 gradient를 계산하고 buffer를 버리기 때문에, 거대한 logits 텐서를 끝까지 들고 있을 필요가 없습니다.

결국 FLCE는 cross entropy의 수식을 바꾸는 방법이 아닙니다. 같은 loss를 계산하되, 어떤 중간 텐서를 실제 메모리에 만들고, 언제 gradient로 소비하고, 무엇을 graph에 남길지 다시 설계하는 방법입니다.

Cross entropy는 수식으로는 한 줄이지만, long-context training에서는 어떤 tensor를 materialize할지 결정하는 systems problem이 됩니다. Context와 vocabulary가 계속 커질수록, 이런 memory-aware loss 계산은 마지막 연산의 최적화가 아니라 학습 자체를 가능하게 만드는 조건에 가까워질 것입니다.

References

[1] QuACK. CUDA/CuTe kernel library used as the implementation reference for the FLCE backend discussed in this post.

[2] Maxim Milakov and Natalia Gimelshein. Online Normalizer Calculation for Softmax. 2018.