메모리 문제를 대충 보면 전략 선택도 틀린다

LLM 학습에서 메모리를 이야기할 때 흔히 "모델이 몇 GB냐"부터 묻는다. 하지만 실제 학습 중 메모리 사용량은 파라미터만으로 결정되지 않는다.

주요 항목은 다음 네 가지다.

  • parameters
  • gradients
  • optimizer state
  • activations

추가로 CUDA allocator fragment, temporary buffer, communication buffer도 무시할 수 없다.

어떤 항목이 가장 큰가?

상황에 따라 다르지만, 학습에서는 optimizer state와 activation이 생각보다 훨씬 크다.

parameters

모델 자체의 가중치다. inference에서는 이 항목이 핵심이지만, training에서는 이것만 보고 판단하면 거의 항상 낙관적으로 계산하게 된다.

gradients

backward 중 생성되며 optimizer step 전까지 유지된다. data parallel에서는 rank마다 전체 gradient 복사본을 가진다.

optimizer state

Adam류 optimizer는 모멘트 통계 때문에 파라미터보다 더 큰 메모리를 쓴다. 이 항목 때문에 ZeRO나 FSDP 같은 전략의 필요성이 커진다.

activations

forward 결과를 backward를 위해 저장해야 하므로 sequence length, hidden size, batch size에 따라 빠르게 커진다. 특히 긴 context에서는 activation이 메모리 병목의 중심이 되기 쉽다.

왜 이 구분이 중요할까?

어떤 기법이 무엇을 줄이는지 이해하려면, 메모리 항목을 분해해서 봐야 한다.

  • activation checkpointing: activation 저장량을 줄인다
  • ZeRO stage 1: optimizer state를 나눈다
  • ZeRO stage 2: gradient도 나눈다
  • ZeRO stage 3 / FSDP: parameter까지 나눈다
  • tensor parallel: 모델 계산과 parameter 배치를 분산한다

즉 같은 "메모리 최적화"라도 줄이는 대상이 다르다.

실무에서 자주 보는 오해

많은 경우 OOM이 발생하면 바로 model parallel을 떠올리지만, 실제로는 activation checkpointing이나 micro-batch 조정만으로 충분한 경우도 많다. 반대로 parameter 규모가 너무 크면 checkpointing만으로는 답이 안 나온다.

그래서 먼저 확인해야 한다.

  • 현재 peak memory의 주범이 activation인가
  • optimizer state가 메모리를 먹는가
  • parameter replication이 이미 너무 큰가

이 질문에 따라 다음 선택지가 달라진다.

다음 단계

이제부터는 "여러 GPU를 어떻게 연결해 통신하는가"를 더 구체적으로 봐야 한다. 같은 분산 전략도 NVLink, PCIe, InfiniBand 환경에 따라 체감 성능이 크게 다르기 때문이다.

다음 글에서는 NCCL과 하드웨어 토폴로지가 왜 분산 학습 성능의 핵심인지 정리한다.