그래디언트 체크포인팅

수치해석 소프트웨어 최적설계 마지막 수정: 2026-08-18 04:44:37

1. 개요[편집]

그래디언트 체크포인팅 Gradient / Activation Checkpointing
다른 이름	활성값 재계산, rematerialization, recompute
거래 내용	메모리를 줄이고 계산을 더 낸다
균등 분할	메모리 $O(\sqrt{L})$, 계산 약 1.3~1.5배 (Chen 외 2016)
최적 스케줄	Griewank–Walther revolve — 이항계수 동적계획
기원	자동미분·비정상 수반해석 (1990년대)
최대 함정	재계산 시 난수·통계가 처음과 달라지는 것
대안·조합	오프로딩, 선택적 재계산, FlashAttention, 가역 구조

그래디언트 체크포인팅(gradient checkpointing)은 역전파에 필요한 중간 활성값을 전부 저장하는 대신 일부 지점만 저장해 두고, 역방향 패스에서 필요해지는 순간 그 지점부터 순전파를 다시 돌려 나머지를 복원하는 기법이다. 이름에 “그래디언트”가 붙어 있지만 기울기를 저장하는 게 아니라 활성값을 안 저장하는 기법이라, 요즘은 활성값 체크포인팅(activation checkpointing)이나 재계산(recompute)이라는 이름이 더 정확하게 쓰인다.

거래 조건은 단순하다. 메모리를 사고 계산을 판다. 순전파를 한 번 더 도는 값을 치르면 활성값 메모리가 층 수의 제곱근 수준으로 줄어든다. “배치 크기 1도 OOM” 과 “배치 8이 돈다” 사이를 가르는 스위치이므로, 대형 모형 학습에서는 취향이 아니라 필수 옵션이다.

이 기법은 딥러닝이 발명한 것이 아니다. 자동 미분 커뮤니티가 1990년대에 비정상 문제의 수반해석을 위해 정립해 둔 것이고, 최적 스케줄에 대한 정리까지 이미 나와 있었다. 역전파 문서가 메모리 병목의 큰 그림과 요약을 다루므로, 이 문서는 스케줄의 수학, 자동미분 테이프와의 관계, 재계산이 조용히 틀린 결과를 내는 경우, 그리고 대안들과의 비교를 다룬다.

2. 왜 활성값이 문제인가[편집]

역방향 패스는 각 층에서 순방향의 입력을 필요로 한다. $\mathbf y = W\mathbf x$ 의 가중치 기울기가 $\boldsymbol\delta\mathbf x^\top$ 이므로 $\mathbf x$ 가 없으면 계산이 안 된다. 그래서 표준 구현은 순전파에서 만든 모든 중간 텐서를 역방향이 끝날 때까지 붙들고 있는다.

트랜스포머 학습을 예로 들면 층 하나당 저장해야 하는 텐서가 어텐션 입력·QKV·소프트맥스 출력·MLP 중간층 등 십수 개이고, 총량은 대략

M_{\text{act}} \;\sim\; c\,\cdot\, B\,\cdot\, S\,\cdot\, d\,\cdot\, L

( $B$ 배치, $S$ 시퀀스 길이, $d$ 모형 차원, $L$ 층 수, $c$ 는 구현에 따라 십 단위의 상수)로 붙는다. 파라미터·옵티마이저 상태는 배치와 무관하게 고정인 반면 이쪽은 배치와 시퀀스 길이에 곱으로 커지므로, 어느 순간부터 GPU를 먼저 터뜨리는 것은 항상 활성값이다. 추론은 층 하나 지나면 앞의 것을 버려도 되니 이 문제가 없다 — “추론은 되는데 학습이 안 된다”의 정체.

3. 균등 분할과 제곱근 트레이드오프[편집]

첸 외(2016)의 처방은 이렇다. $L$ 개 층을 길이 $L/k$ 짜리 구간 $k$ 개로 나누고, 구간 경계의 활성값만 저장한 채 순전파를 끝낸다. 역방향에서 마지막 구간에 도달하면 그 구간의 시작 체크포인트에서 순전파를 다시 돌려 구간 내부 활성값을 복원하고, 그걸로 역전파한 뒤 버린다. 그리고 앞 구간으로 넘어간다.

메모리 회계는 두 항의 합이다.

M(k) \;\approx\; \underbrace{k}_{\text{체크포인트}} \;+\; \underbrace{L/k}_{\text{현재 구간 내부}}

$k=\sqrt{L}$ 에서 최소가 되고 $M = O(\sqrt{L})$ 이다. 이 분할을 재귀적으로 한 번 더 적용하면 $O(\log L)$ 까지 내려가는데, 그 대신 재계산 스윕이 늘어난다.

계산 비용은 각 층이 정확히 두 번 순전파된다 — 원래 한 번, 재계산 한 번. 순전파를 1, 역전파를 2로 놓으면 원래 총비용 3에 1이 더해져 $4/3 \approx 1.33$ 배다. 실측에서 1.3~1.5배로 보고되는 이유는 커널 실행 오버헤드, 층별 비용의 불균일, 그리고 재계산 구간의 캐시 지역성 차이 때문이다.

여기서 자주 오해되는 지점 하나. 체크포인팅을 켜면 학습이 느려진다는 것은 스텝당 이야기지 전체 학습 시간 이야기가 아니다. 메모리가 남으면 배치를 2~4배 키울 수 있고, 큰 배치는 GEMM 효율과 통신/계산 겹침을 개선한다. 실측하면 처리량이 오히려 올라가는 경우가 흔하다.¹

4. 최적 스케줄 — revolve[편집]

균등 분할은 최적이 아니다. “체크포인트를 $c$ 개 쓸 수 있고 각 층의 재계산을 $r$ 번까지 허용할 때, 되돌릴 수 있는 최대 스텝 수는 얼마인가”라는 문제를 정확히 풀 수 있다. 그리바인크(1992)와 그리바인크·발터(2000, ACM TOMS Algorithm 799 revolve)의 답은 놀랍도록 깔끔하다.

L_{\max}(c,r) \;=\; \binom{c+r}{c}

즉 메모리와 재계산 횟수를 각각 조금만 늘려도 처리 가능한 스텝 수가 이항계수로 폭증한다. 반대로 읽으면, $L$ 스텝을 되돌리는 데 필요한 체크포인트와 스윕이 둘 다 $O(\log L)$ 이면 충분하다는 뜻이다.

체크포인트 $c$	스윕 $r$	되돌릴 수 있는 스텝 수
2	2	6
3	3	20
5	5	252
10	5	3,003
10	10	184,756

체크포인트 10개와 재계산 10회로 18만 스텝짜리 시간 적분을 되돌릴 수 있다는 얘기다. 균등 분할이었다면 $\sqrt{184756}\approx 430$ 개의 체크포인트가 필요했을 것이다.

revolve 는 이 최적 배치를 동적계획으로 찾는다. 상태를 (남은 스텝 수, 남은 체크포인트 수)로 두고 “다음 체크포인트를 어디에 놓을 것인가”를 재귀적으로 결정하며, 비용 함수는 총 재계산 스텝 수다. 스텝 비용이 균일한 사슬(chain) 구조에서는 이 스케줄이 증명 가능하게 최적이다. 관련 문제를 동적 계획법의 전형적 사례로 소개하는 교재도 있다.

가정이 깨지면 얘기가 달라진다. 층마다 비용과 메모리가 제각각인 실제 신경망은 균일 사슬이 아니고, 분기·병합이 있는 계산 그래프는 사슬조차 아니다. 이 일반화된 재구체화 문제를 혼합정수계획으로 풀어 최적해를 얻는 접근이 Checkmate(제인 외 2020)이고, 실무 프레임워크는 대개 그 중간 — 트랜스포머 블록 단위로 균등하게 자르는 단순한 규칙 — 을 쓴다. 블록 구조가 균일하니 균등 분할의 가정이 실제로 잘 맞기 때문이다.

5. 자동미분 테이프와의 관계[편집]

역방향 모드 자동 미분의 구현은 테이프(tape, Wengert list)를 기록한다. 순전파를 돌면서 실행한 원시연산과 그 국소 편미분에 필요한 값을 순서대로 적어 두고, 역방향에서 테이프를 거꾸로 읽으며 벡터-야코비 곱을 조립한다. 활성값 메모리는 곧 테이프 길이다.

이 언어로 보면 선택지가 두 극단 사이의 스펙트럼이다.

전부 기록(store-all). 표준 역전파. 시간 최소, 메모리 최대.
전부 재계산(recompute-all). 각 층의 역전파 직전에 입력에서부터 다시 순전파. 메모리는 $O(1)$ 이지만 계산이 $O(L^2)$ 로 폭발한다.
체크포인팅. 그 사이. 프로그램 상태의 스냅숏만 몇 군데 남기고, 테이프는 필요한 구간만 그때그때 다시 만든다.

자동미분 문헌에서는 이 조작을 순차 반전(joint/split reversal)이라 부르고, 컴파일러 쪽에서는 재구체화(rematerialization)라는 이름을 쓴다. 계산해서 다시 만들 수 있는 값을 굳이 저장하지 않는다는 점에서, 레지스터 할당에서 스필(spill) 대신 재계산을 택하는 고전적 최적화와 정확히 같은 결정이다. JAX의 jax.checkpoint(별칭 remat)라는 이름이 이 계보를 정직하게 드러낸다.

기원도 여기다. 비정상 편미분방정식의 수반해석 — 예컨대 시간 적분되는 전산유체역학 해석의 형상 민감도 해석이나 최적 제어 문제 — 에서는 수천~수만 스텝의 유동장을 전부 디스크에 남길 수 없어 1990년대부터 이 스케줄링이 필수였다. 딥러닝은 20년 뒤 같은 벽에 부딪혔고, 같은 답을 다시 발견했다.

6. 재계산의 함정 — 두 번째가 첫 번째와 다를 때[편집]

체크포인팅의 정확성은 재계산한 순전파가 원래 순전파와 같은 값을 낸다는 가정 위에 서 있다. 이 가정이 깨지면 기울기가 조용히 틀린다. 손실도 잘 내려가고 에러도 안 나므로 발견이 늦다.

드롭아웃의 난수. 재계산 때 다른 마스크가 뽑히면, 순전파에서 쓴 마스크와 역전파가 가정하는 마스크가 달라져 기울기가 다른 함수의 것이 된다. 프레임워크는 체크포인트 진입 시점의 RNG 상태를 저장했다가 재계산 직전에 복원하는 방식으로 막는다(PyTorch torch.utils.checkpoint 의 preserve_rng_state=True 가 기본값). 그런데 이 보호는 현재 장치의 기본 생성기까지만이다 — 커스텀 Generator, 다른 장치의 RNG, 파이썬 random 을 쓰는 사용자 정의 연산은 보호 밖이다. 확률적 깊이(DropPath)나 모형 안에 박아 넣은 랜덤 증강도 같은 함정.
배치 정규화의 이동평균. 재계산은 학습 모드 순전파를 한 번 더 도는 것이므로 이동평균이 두 번 갱신된다. 출력값 자체는 같은 배치 통계를 쓰니 문제가 없지만, 모멘텀이 두 번 적용된 버퍼로 추론하면 통계가 미묘하게 어긋난다. 체크포인팅을 켠 뒤 학습 정확도는 그대로인데 평가 정확도만 떨어진다면 여기를 먼저 본다. 층 정규화·RMSNorm은 버퍼가 없어 이 문제가 원천적으로 없고, 트랜스포머에서 체크포인팅이 마음 편한 이유 중 하나다.
비결정 커널. 원자 덧셈 기반 리덕션, cuDNN 알고리즘 자동선택, TF32/bf16 누산 순서 차이 때문에 같은 입력에 대해 재계산 결과의 마지막 비트가 달라질 수 있다. 대개 무해하지만, 기울기 검증이 갑자기 실패하거나 비트 단위 재현성이 요구되는 상황에서는 문제가 된다.
기울기가 필요 없다고 판단되는 경우. 재진입(reentrant) 방식 구현에서는 체크포인트 구간의 입력 중 하나라도 requires_grad 가 켜져 있어야 구간 출력에 그래프가 붙는다. 임베딩 앞단처럼 입력이 정수 텐서인 첫 블록에 체크포인팅을 걸면 그 구간이 통째로 미분에서 빠지고, 손실은 줄지만 앞부분이 학습되지 않는다. 요즘은 use_reentrant=False 구현이 권장되며 이 함정을 피한다.
부작용이 있는 연산. 카운터 증가, 로깅, 캐시 갱신 등 순전파에 부작용이 있으면 두 번 실행된다. 재계산 구간 안에는 순수 함수만 두는 것이 원칙이다.

검증 방법은 단순하다. 작은 모형에서 체크포인팅 켠 기울기와 끈 기울기를 직접 비교한다. 배정밀도로 상대오차 $10^{-10}$ 수준이면 정상이고, $10^{-3}$ 이상 벌어지면 위 목록 중 하나에 걸린 것이다.²

7. 대안과 조합[편집]

체크포인팅은 “계산으로 메모리를 산다”는 하나의 축일 뿐이고, 다른 축들과 조합하는 것이 실제 레시피다.

선택적 재계산. 층 전체를 다시 도는 대신 메모리는 크고 계산은 싼 부분만 재계산한다. 코르티카티 외(2022)는 트랜스포머에서 어텐션 행렬·소프트맥스·드롭아웃 마스크처럼 $S^2$ 로 커지는 텐서만 버리고 GEMM 출력은 저장하면, 전체 재계산의 대부분의 메모리 이득을 얻으면서 계산 오버헤드를 수 % 수준으로 낮출 수 있음을 보였다. 균등 재계산이 30% 이상을 내는 것과 비교하면 훌륭한 거래다.
플래시 어텐션. 같은 아이디어를 커널 안으로 밀어 넣은 것. $S\times S$ 어텐션 행렬을 아예 물리적으로 만들지 않고 타일 단위로 처리하며, 역방향에서는 저장해 둔 소프트맥스 정규화 통계만으로 필요한 블록을 재계산한다. 메모리가 $O(S^2)$ 에서 $O(S)$ 로 줄면서 속도까지 빨라진다 — 어텐션이 연산량이 아니라 메모리 대역폭에 묶여 있었기 때문. 재계산이 항상 느려지는 거래는 아니라는 반례다.
활성값 오프로딩. 재계산 대신 CPU 메모리(또는 NVMe)로 밀어냈다가 역방향에서 되가져온다. 파는 자원이 FLOPs가 아니라 PCIe 대역폭이므로, 계산이 빡빡하고 전송을 계산과 겹칠 여지가 있으면 유리하다. 반대로 배치가 커서 전송량이 크면 병목이 그대로 드러난다. ZeRO 계열이 이 축을 파고들었다.
가역 구조. RevNet 계열은 블록을 수학적으로 역함수가 존재하도록 설계해, 출력에서 입력을 정확히 복원한다. 활성값 저장이 원리적으로 $O(1)$ 이고 체크포인트도 필요 없다. 대신 아키텍처 자유도를 포기해야 해서 널리 쓰이지는 않는다.
혼합 정밀도. 활성값을 bf16으로 들면 그냥 절반이다. 가장 싼 처방이라 언제나 먼저 켠다.
분산 축. FSDP·ZeRO의 파라미터·옵티마이저 상태 샤딩은 다른 축이다(활성값이 아니라 모형 상태를 나눈다). 시퀀스 병렬화는 활성값을 장치 간에 쪼개므로 체크포인팅과 직접 경쟁하며, 둘을 같이 쓰는 것이 최근 대형 학습의 표준이다. 병렬 컴퓨팅·GPU 컴퓨팅 참고.

우선순위를 실무 순서로 적으면 이렇다. 혼합 정밀도 → 선택적 재계산 / FlashAttention → 블록 단위 전체 재계산 → 오프로딩. 앞의 것이 뒤의 것보다 싸고, 앞의 것으로 안 되면 그때 뒤로 간다.

8. 실무 감각[편집]

입도(granularity)가 전부다. 트랜스포머에서는 블록 하나를 한 체크포인트 단위로 잡는 게 기본이고, 메모리가 조금만 부족하면 ” $n$ 블록마다 하나씩만” 거는 부분 적용이 낫다. 모든 블록에 거는 것과 안 거는 것 사이의 이분법이 아니다.
잔차 연결 구조가 자연스러운 경계를 준다. 잔차 블록의 입력은 그 블록을 재계산하기 위한 완전한 상태이므로, 블록 경계가 곧 체크포인트 후보다. 프레임워크가 블록 단위 API를 제공하는 이유.
재계산해도 값이 같은지가 곧 안전 조건이다. 정규화 계층에 버퍼가 있는지, 난수를 쓰는지, 부작용이 있는지 — 이 세 질문만 통과하면 체크포인팅은 정확하다.
켜 보기 전엔 모른다. 메모리 절감량과 처리량 변화가 모형·배치·하드웨어에 따라 크게 달라진다. 스텝당 시간과 최대 메모리를 같이 로깅해 두고 배치 크기를 함께 조정하는 것이 유일한 방법이다.³

9. 관련 문서[편집]

“메모리를 아끼려고 켰는데 학습이 빨라졌다”는 제보가 정기적으로 올라오는 이유. 스텝당 33% 느려져도 배치를 4배 키워 스텝 수가 4분의 1이 되면 총 시간은 줄어든다. 물론 배치를 키우면 학습률과 스케줄을 다시 잡아야 하니 공짜는 아니다. ↩
체크포인팅 버그의 악질적인 점은 학습이 그럭저럭 된다는 것이다. 드롭아웃 마스크가 어긋난 기울기는 완전히 틀린 방향이 아니라 잡음이 심한 방향이라, 손실은 내려가고 성능만 조금 나쁘다. 몇 주를 태우고 나서야 “체크포인팅 껐더니 잘 되네요”를 발견하는 전개. ↩
그리바인크의 revolve 논문에는 재밌는 관점이 하나 있다. 되돌리기(reversal)를 시간과 공간의 교환으로 보면, 이건 물리 시뮬레이션의 시간 역행 문제와 형식적으로 같은 구조다. 정보를 저장하지 않고 과거를 복원하려면 다시 계산하는 수밖에 없다는 것 — 컴퓨터에서든 우주에서든. ↩