극한의 다이어트 - VRAM을 스크래치 메모리로 사용하기

LLM을 훈련할 때 batch size는 쉽게 욕심나는 값이다. batch가 너무 작으면 gradient가 흔들리고, sequence 길이를 늘리거나 모델을 키울 때 throughput을 맞추기도 어렵다. 하지만 실제로 batch를 키우려 하면 가장 먼저 부딪히는 것은 계산량보다 메모리였다.

weight와 optimizer state도 크지만, 학습 중에는 forward에서 생긴 activation을 backward까지 붙잡고 있어야 한다. 특히 Transformer block 안에는 RMSNorm, QKV projection, attention score, output projection, FFN activation이 차례로 쌓인다. batch와 sequence가 커질수록 이 saved activation checkpoint들이 그대로 메모리를 먹는다.

일반적인 해결책은 gradient accumulation으로 effective batch를 흉내 내거나, activation checkpointing으로 일부 forward를 다시 계산하는 것이다. 하지만 나는 좀 더 비현실적인 방법을 실험해보고 싶었다.

내가 확인해보고 싶은 가능성은 saved activation을 충분히 작게 만들어 CPU로 오프로딩(offloading)하는 것이었다. 여기서 오프로딩이란 backward까지 보관할 activation을 GPU에서 host memory로 옮겨 GPU 메모리를 확보하는 방식이다. 다음 레이어의 forward를 계산하는 동안 이전 레이어의 saved activation을 압축해 오프로딩하고, backward 때에는 계산 순서에 맞춰 GPU로 미리 불러오는 prefetch 구조를 생각했다.

forward layer N
  compute activation on GPU
  compress saved activation
  offload to host memory while layer N+1 runs

backward layer N
  prefetch compressed saved activation while backward layer N+1 runs
  restore and use it for gradient computation

이 구상이 성립하려면 오프로딩과 prefetch에 걸리는 각각의 전송 시간이 그 사이에 실행되는 계산 시간보다 짧아야 한다. PCIe bandwidth가 느려도, 압축된 activation이 충분히 작고 forward layer N+1이나 backward layer N+1이 충분한 계산량을 갖고 있다면 전송은 계산 뒤에 숨는다. 반대로 GPU가 saved activation을 기다리기 시작하면 메모리는 아껴도 throughput은 무너진다.

다르게 말하면, GPU memory를 모든 activation을 오래 보관하는 공간이 아니라, 곧 계산에 쓸 activation만 올라오는 빠른 scratch tier처럼 쓰고 싶었다. batch를 키우면 계산량도 함께 커지므로, 공격적으로 압축된 saved activation은 그 계산 시간 사이에 GPU 안팎을 오갈 수 있을지 모른다.

그래서 이 실험은 "모델을 줄이자"가 아니라 "훈련 중 저장해야 하는 activation을 줄이자"에서 시작했다. 같은 모델과 같은 sequence 길이를 유지하면서, 더 큰 batch를 감당할 수 있는 activation 압축·오프로딩 방식을 찾고 싶었다.

이후의 질문은 자연스럽게 압축률로 좁혀졌다. saved activation을 얼마나 줄여야 PCIe 전송이 병목이 되지 않을까? 그리고 그렇게 줄인 activation으로도 backward가 버틸 수 있을까?