PyTorch 2.13 LinearCrossEntropyLoss: 거대한 logits를 만들지 않는 학습 손실

마지막 projection이 학습 메모리를 삼키는 순간

대규모 언어 모델의 학습 메모리를 이야기할 때 보통 parameter, optimizer state, activation, attention을 먼저 본다. 하지만 vocabulary가 커지면 마지막 language-model head도 독립적인 병목이 된다. hidden state를 vocabulary 전체에 투영한 logits의 모양은 대략 다음과 같다.

hidden:  [token_rows, hidden_size]
weight:  [vocab_size, hidden_size]
logits:  [token_rows, vocab_size]

예를 들어 한 step에서 loss를 계산하는 token row가 8,192개이고 vocabulary가 128,000개라면, logits 하나만 FP16/BF16 기준 약 1.95 GiB다. 이는 단순 계산 예시이며 전체 peak memory가 아니다. backward를 위한 중간값과 allocator 동작까지 더해지면 실제 peak는 더 커질 수 있다.

2026년 7월 8일 공개된 PyTorch 2.13은 nn.LinearCrossEntropyLoss와 torch.nn.functional.linear_cross_entropy를 추가했다. 목적은 마지막 linear projection과 cross entropy의 경계를 합쳐 전체 [token_rows, vocab_size] logits를 한 번에 materialize하지 않는 것이다. 공식 release blog는 대규모 vocabulary workload에서 peak GPU memory를 최대 약 4배 줄였다고 설명한다.

그러나 이 수치만 보고 nn.Linear + nn.CrossEntropyLoss를 기계적으로 바꾸면 안 된다. 2.13.0 소스를 읽으면 세 가지 중요한 조건이 보인다.

1. memory-saving chunked path는 기본값이 아니라 options를 명시해야 활성화된다.
2. 실제 구현은 vocabulary를 나누는 것이 아니라 flatten된 batch/token row를 나눈다.
3. chunked path는 일반 autograd와 torch.compile의 모든 기능을 보존하지 않는다.

이 장의 초점은 새 API를 소개하는 데 있지 않다. 어떤 tensor가 사라지고 무엇이 새로 생기는지, 조용한 reference fallback을 어떻게 탐지할지, 정확성·처리량·메모리를 어떤 순서로 검증할지를 다룬다.

이 장은 PyTorch 2.13.0 release note, 공식 blog, 2.13 API 문서와 tag 고정 소스를 기준으로 한다. 공개 후 약 5일 지난 API이며 공식 blog에도 API Unstable로 표시되어 있으므로, 이후 minor release에서 option과 제한이 바뀔 수 있다.

먼저 수식에서 없앨 수 있는 tensor를 찾는다

일반적인 마지막 layer와 cross entropy는 다음 순서로 계산된다.

Z = H @ W^T
P = softmax(Z)
loss = -log(P[target])

여기서 H는 [N, F], W는 [V, F], Z는 [N, V]다. N은 loss에 참여하는 token row 수, F는 hidden size, V는 vocabulary size다.

cross entropy의 gradient는 개념적으로 다음처럼 쓸 수 있다.

dZ = softmax(Z) - one_hot(target)
dH = dZ @ W
dW = dZ^T @ H

핵심은 dH와 dW를 계산할 때 모든 row의 Z가 동시에 존재할 필요가 없다는 점이다. row 범위를 여러 chunk로 나누면 각 chunk에서 다음 작업을 끝낸 뒤 logits buffer를 재사용할 수 있다.

1. Z_chunk = H_chunk @ W^T를 계산한다.
2. 안정적인 log-sum-exp와 target logit으로 loss를 누적한다.
3. dZ_chunk를 만든다.
4. dH_chunk를 해당 위치에 기록한다.
5. dW += dZ_chunk^T @ H_chunk로 누적한다.
6. 다음 chunk가 같은 logits buffer를 덮어쓴다.

전체 logits의 주된 크기는 N × V지만, chunked logits buffer는 B_chunk × V가 된다. B_chunk를 줄일수록 peak memory는 줄고 loop와 kernel launch 수는 늘어난다. 즉, 이 기능은 계산량을 없애는 마법이 아니라 큰 생존 범위를 짧은 반복 buffer로 바꾸는 scheduling 선택이다.

"vocabulary chunking"이 아니라 batch row chunking이다

PyTorch 2.13 공식 blog는 이 기능이 vocabulary dimension을 chunk로 처리한다고 설명한다. 하지만 2.13.0의 LinearCrossEntropyOptions 문서와 구현은 더 구체적이다.

• batch_chunk_size는 한 번에 처리할 batch row 수다.
• 내부 logits buffer의 모양은 [batch_chunk_size, num_classes]다.
• loop는 range(0, num_batches, batch_chunk_size)로 row를 순회한다.
• 각 chunk에서도 vocabulary 전체에 대한 matrix multiplication과 softmax denominator를 계산한다.

LLM에서는 hidden state를 [batch, sequence, hidden]에서 [token_rows, hidden]으로 flatten하므로, 여기서 batch row는 사실상 loss에 참여하는 token row다. vocabulary 전체를 처리한다는 점은 바뀌지 않는다. 달라지는 것은 모든 token row의 vocabulary logits를 동시에 보관하지 않는다는 점이다.

이 구분은 운영상 중요하다. vocabulary가 매우 크더라도 N이 작으면 chunk를 나눌 여지가 적다. 공식 option 문서는 대략 다음 조건에서 chunking이 유리하다고 설명한다.

num_batches >= in_features
num_classes > in_features

반대로 짧은 sequence, 작은 microbatch, 작은 vocabulary에서는 reference path가 더 저렴할 수 있다. chunked라는 이름만 보고 모든 모델에 적용하지 말고, 실제 N, F, V와 kernel launch 수를 측정해야 한다.

기본 생성자만 쓰면 메모리 절감 경로가 켜지지 않는다

새 module의 options 기본값은 None이다. 이 값은 linear projection과 cross entropy를 같은 API에서 호출하지만, 내부적으로는 reference implementation을 사용한다. 즉, full logits가 만들어진다.

import torch
from torch import nn

# 새 API이지만 reference path다. full logits를 만들 수 있다.
loss_fn = nn.LinearCrossEntropyLoss(
    in_features=4096,
    num_classes=128_000,
    options=None,
)

memory-saving path를 사용하려면 LinearCrossEntropyOptions instance를 넘겨야 한다.

loss_fn = nn.LinearCrossEntropyLoss(
    in_features=4096,
    num_classes=128_000,
    options=nn.LinearCrossEntropyOptions(
        chunking_method="auto",
        acc_policy="auto",
    ),
)

새 module은 내부에 self.linear를 생성한다. 기존 모델에 이미 lm_head가 있고 embedding과 weight를 묶었거나 checkpoint key를 유지해야 한다면 functional API가 더 명시적이다.

import torch
import torch.nn.functional as F

lce_options = torch.nn.LinearCrossEntropyOptions(
    chunking_method="auto",
    acc_policy="auto",
)

def language_model_loss(hidden, labels, lm_head):
    # hidden: [batch, sequence, hidden_size]
    # labels: [batch, sequence]
    hidden_2d = hidden.reshape(-1, hidden.shape[-1])
    target_1d = labels.reshape(-1)

    return F.linear_cross_entropy(
        hidden_2d,
        lm_head.weight,
        target_1d,
        linear_bias=lm_head.bias,
        ignore_index=-100,
        reduction="mean",
        options=lce_options,
    )

이 방식은 기존 lm_head.weight object를 그대로 전달하므로 tied weight와 checkpoint naming을 보존하기 쉽다. module API로 교체할 때는 기존 lm_head.weight가 새 loss_fn.linear.weight로 이동한다는 사실을 migration code와 optimizer parameter group에 반영해야 한다.

auto heuristic은 무엇을 맞추려 하는가

LinearCrossEntropyOptions()의 기본 chunking_method="auto"는 device와 dtype별 기본값으로 해석된다. 알려지지 않은 조합의 fallback은 aspect_ratio:2다.

aspect_ratio의 목표는 chunked logits buffer가 원래 input tensor와 비슷한 memory 크기가 되도록 batch_chunk_size를 고르는 것이다. 문서의 개념식은 다음과 같다.

rows_per_chunk ≈ next_power_of_two(
    ceil(num_batches / ceil(num_classes / in_features))
)

aspect_ratio:2는 이 크기를 다시 약 2분의 1로 줄인다. 따라서 대략 두 배 적은 logits scratch를 쓰지만 chunk 수는 약 두 배 늘 수 있다.

직접 batch_chunk_size를 정할 수도 있다.

options = torch.nn.LinearCrossEntropyOptions(
    chunking_method=None,
    batch_chunk_size=256,
    acc_policy="compact",
)

여기서 작은 값이 항상 좋은 것은 아니다.

프로덕션에서는 auto를 기준선으로 두고 OOM에 가까운 shape와 대표 shape를 함께 측정하는 편이 낫다. 가장 긴 sequence만 최적화하면 짧은 batch가 많은 실제 traffic에서 throughput을 잃을 수 있다.

precision policy는 속도 옵션이 아니라 gradient 계약이다

chunk별 weight gradient를 누적하면 dtype과 accumulator 선택이 결과에 영향을 준다. PyTorch 2.13은 acc_policy를 accurate, balanced, compact, auto로 나눈다.

compact가 줄이는 scratch는 num_classes × in_features × sizeof(acc_dtype) 크기다. vocabulary와 hidden size가 모두 크면 수백 MiB 이상이 될 수 있다. 대신 "loss 값이 비슷하다"만으로 채택하면 안 된다. 최소한 다음을 별도로 비교해야 한다.

• scalar loss의 절대·상대 오차
• hidden gradient의 norm, max error, non-finite 여부
• LM head weight gradient의 norm과 cosine similarity
• 동일 seed에서 여러 step 뒤 optimizer state와 parameter drift
• BF16, FP16, FP32별 결과
• gradient clipping 직전의 global norm

bitwise equality는 backend와 reduction order가 달라지면 현실적인 기준이 아닐 수 있다. 그렇다고 tolerance를 임의로 넓히지 말고, 기존 reference path의 반복 변동과 학습 안정성 지표를 근거로 정한다.

scalar reduction에서는 forward가 gradient까지 계산한다

chunked path가 일반 autograd와 다른 가장 큰 이유는 scalar mean·sum 경로가 forward에서 loss와 gradient를 함께 계산해 보관하기 때문이다. 공식 source comment는 backward가 upstream scalar gradient를 미리 계산된 gradient에 곱하는 역할을 한다고 설명한다.

이 선택은 chunk마다 만든 logits를 즉시 소비하고 버리기 위해 필요하다. 일반적인 backward 시점에 logits를 다시 요구하지 않도록, forward loop에서 dH와 dW를 함께 만든다.

대가는 autograd contract의 축소다.

• 기본 allow_retain_graph=False에서는 precomputed gradient buffer를 소비하므로 두 번째 backward가 실패한다.
• allow_retain_graph=True는 반복 backward를 허용하지만 gradient 크기의 allocation 하나를 추가한다.
• higher-order autograd와 forward-mode AD는 지원하지 않는다.
• torch.func.grad와 vmap(grad(...))를 사용할 수 없다.
• chunked op에 대한 Inductor lowering이 없어 torch.compile은 이 op에서 eager로 fallback한다.
• torch.compile 안에서는 scalar 경로의 정확한 second backward를 위해 allow_retain_graph=True가 강제되고 warning이 난다.

reduction="none"은 다르다. forward에서는 per-row loss만 계산하고, backward에서 chunked gradient를 다시 계산한다. 그래서 retain_graph=True 동작은 scalar 경로와 다르지만, higher-order AD 같은 전체 제한이 사라지는 것은 아니다.

이 차이는 activation checkpointing, meta-learning, per-sample gradient, gradient penalty, Jacobian·Hessian 계산을 쓰는 학습 코드에서 결정적이다. 일반적인 한 번의 loss.backward()만 사용하는 LM pretraining과 달리, 이들 workload는 memory-saving path를 채택할 수 없거나 reference path를 분리해야 할 수 있다.

지원한다고 쓰여 있어도 chunked path가 아닐 수 있다

LinearCrossEntropyLoss module의 표면 API는 class probability target, label smoothing, K-dimensional output을 받을 수 있다. 하지만 2.13.0의 chunked path는 더 좁다. 다음 조건을 벗어나면 warning을 내고 reference implementation으로 돌아간다.

따라서 "label smoothing을 지원한다"는 말은 module 전체 API에는 맞지만, memory-saving path가 label smoothing을 지원한다는 뜻은 아니다. 자동 fallback은 결과를 계속 계산해 가용성을 높이지만, canary가 실제 chunked path를 검증했는지를 가릴 수 있다.

CI와 benchmark에서는 해당 warning을 오류로 승격하는 편이 안전하다.

import warnings

with warnings.catch_warnings():
    warnings.filterwarnings(
        "error",
        message=r"linear_cross_entropy: ``options`` ignored.*",
    )
    loss = language_model_loss(hidden, labels, lm_head)

production에서도 warning count를 수집하거나 startup contract test를 둔다. label smoothing 설정 하나가 바뀐 뒤 full logits path로 조용히 돌아가 OOM이 발생하는 상황을 막아야 한다.

"최대 4배"는 채택 기준이 아니라 재현할 가설이다

공식 blog의 최대 약 4배 peak-memory 절감은 upstream이 제시한 결과다. 모든 모델에 적용되는 보장이 아니다. 실제 절감 폭은 다음 항목에 따라 달라진다.

• N × V × dtype_size가 전체 peak에서 차지하는 비율
• hidden size와 vocabulary의 비율
• batch_chunk_size와 chunk 수
• accurate·balanced·compact policy
• LM head weight와 gradient, optimizer state의 크기
• activation checkpointing과 FSDP·tensor parallel 적용 범위
• allocator fragmentation과 다른 stream의 동시 allocation
• reference path가 이미 compiler fusion이나 별도 fused loss를 사용하는지

예를 들어 full logits가 2 GiB여도 모델 parameter와 optimizer state가 수십 GiB라면 전체 peak가 4분의 1로 줄지 않는다. 반대로 logits와 그 backward intermediate가 peak의 대부분인 large-vocabulary fine-tuning에서는 절감 폭이 클 수 있다.

성능도 같은 방식으로 검증한다. chunking은 full logits allocation을 줄이지만 matrix multiplication을 여러 번 호출하고 weight gradient를 누적한다. 처리량이 느려져 gradient accumulation step을 더 오래 돌려야 한다면, OOM 회피의 가치는 있어도 비용 절감이라고 바로 말할 수 없다.

안전한 도입 순서

1단계: 현재 loss 경로를 inventory로 만든다

먼저 다음을 기록한다.

• flatten 뒤 N, hidden size F, vocabulary V
• target dtype과 ignore_index
• label smoothing, class weighting, z-loss 같은 추가 loss
• tied embedding 여부와 checkpoint key
• torch.compile, activation checkpoint, higher-order AD 사용 여부
• FSDP·tensor parallel에서 LM head weight의 shard 방식
• 현재 peak allocated·reserved memory와 step time 분포

z-loss나 custom masking을 logits tensor에 직접 적용한다면 단순 치환이 아니다. functional API가 동일한 objective를 표현하는지 먼저 확인해야 한다.

2단계: reference와 chunked 결과를 같은 process에서 비교한다

동일한 hidden, weight, target으로 두 경로를 실행한다.

reference = F.cross_entropy(
    F.linear(hidden_2d, lm_head.weight, lm_head.bias),
    target_1d,
    ignore_index=-100,
)

chunked = F.linear_cross_entropy(
    hidden_2d,
    lm_head.weight,
    target_1d,
    linear_bias=lm_head.bias,
    ignore_index=-100,
    options=torch.nn.LinearCrossEntropyOptions(),
)

loss뿐 아니라 hidden과 weight gradient를 각각 비교한다. 같은 tensor의 .grad를 두 경로가 공유하지 않도록 clone한 parameter와 별도 graph를 사용한다.

3단계: memory와 throughput을 shape matrix로 측정한다

CUDA 측정은 asynchronous execution을 고려해야 한다.

torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()
torch.cuda.synchronize()

loss = run_one_step()
loss.backward()

torch.cuda.synchronize()
peak_allocated = torch.cuda.max_memory_allocated()
peak_reserved = torch.cuda.max_memory_reserved()

한 shape만 보지 말고 다음을 나눈다.

• 짧은 sequence × 큰 microbatch
• 긴 sequence × 작은 microbatch
• gradient accumulation을 포함한 실제 effective batch
• 정상 vocabulary와 확장 vocabulary
• auto, accurate, balanced, compact
• eager와 torch.compile 전체 step

warm-up, compile, allocator cache의 영향을 분리하고 median뿐 아니라 step-time p95도 본다. chunk size가 작을 때 kernel launch jitter가 tail에 나타날 수 있다.

4단계: fallback과 autograd 경계를 고의로 깨 본다

• label_smoothing=0.1을 넣어 warning이 CI 실패로 전환되는지 확인한다.
• target을 probability tensor로 바꿔 fallback 탐지가 작동하는지 확인한다.
• second backward, torch.func.grad, torch.compile을 실제 training wrapper에서 실행한다.
• activation checkpointing이 loss module을 다시 호출하는 구조인지 확인한다.
• checkpoint save/load 뒤 tied weight identity가 유지되는지 확인한다.

5단계: canary의 성공 기준을 memory 하나로 두지 않는다

채택 gate는 최소한 다음을 함께 포함한다.

• reference 대비 loss·gradient tolerance 통과
• non-finite loss·gradient 증가 없음
• peak allocated와 reserved memory 개선
• tokens/s와 step-time p95 허용 범위
• recompilation·graph break·warning 증가 없음
• checkpoint resume와 optimizer state load 성공
• unsupported option이 reference fallback으로 숨지 않음

도입 체크리스트

정리

PyTorch 2.13의 LinearCrossEntropyLoss는 새로운 loss 수식을 만든 것이 아니다. 마지막 projection과 cross entropy 사이에서 커다란 logits tensor가 오래 살아야 했던 실행 경계를 바꾼다. flatten된 token row를 chunk로 나누고, [B_chunk, V] buffer에서 loss와 gradient를 계산·누적해 전체 [N, V] logits를 피한다.

이 구조는 large-vocabulary 학습에서 실질적인 OOM 여유를 만들 수 있다. 동시에 chunk 크기에 따른 반복 비용, precision policy, forward에서 gradient를 미리 계산하는 비표준 autograd 경로라는 대가가 있다.

운영자가 기억할 핵심은 네 가지다.

1. 새 API를 호출하는 것과 memory-saving path를 쓰는 것은 다르다. options가 필요하다.
2. 공식 blog의 표현과 달리 2.13.0 구현은 vocabulary가 아니라 batch/token row를 chunking한다.
3. label smoothing이나 probability target이 동작해도 chunked path가 아니라 reference fallback일 수 있다.
4. 최대 4배는 upstream의 최대치다. 우리 workload에서는 loss·gradient 동등성, peak memory, tokens/s, compile·autograd 경계를 함께 측정해야 한다.

좋은 채택 기준은 "OOM이 사라졌다"에서 끝나지 않는다. OOM을 피한 경로가 실제로 선택되었고, 학습 objective와 gradient 계약을 지키며, rollback 가능한 비용 구조를 만들었는가까지 확인해야 한다.

References

• PyTorch 2.13.0 Release Notes — released 2026-07-08
• PyTorch 2.13 Release Blog — 2026-07-08
• PyTorch 2.13 nn.LinearCrossEntropyLoss API
• PyTorch 2.13 LinearCrossEntropyOptions API
• PyTorch v2.13.0 linear_cross_entropy functional source and fallback conditions
• PyTorch v2.13.0 chunked loss implementation
• PyTorch v2.13.0 chunking and accumulation option source