그래디언트 체크포인팅

Gradient checkpointing, activation checkpointing, activation recomputation; 그래디언트 체크포인팅
순전파의 중간 활성값을 일부만 저장하고 나머지는 역전파 때 다시 계산해, 연산을 조금 더 쓰는 대신 학습 메모리를 크게 줄이는 기법

그래디언트 체크포인팅(gradient checkpointing)은 신경망 학습에서 순전파 중간 결과(활성값)를 모두 메모리에 들고 있지 않고, 정해 둔 지점(체크포인트)의 값만 저장했다가 역전파 때 필요한 구간을 다시 순전파해 복원하는 기법이다. 활성값 체크포인팅(activation checkpointing), 활성값 재계산(activation recomputation), 재물질화(rematerialization)라고도 한다. 이름에 "그래디언트"가 붙지만 저장하는 것은 기울기가 아니라 활성값이다.

깊은 신경망을 학습할 때 메모리의 큰 몫은 가중치보다 활성값이 차지한다. 역전파는 각 층의 입력이나 중간 결과를 다시 써야 하므로, 보통은 순전파가 끝날 때까지 모든 층의 활성값이 메모리에 남는다. 그래디언트 체크포인팅은 이 메모리를 층 수에 비례하는 양에서 그 제곱근 수준까지 줄이고, 대가로 순전파를 대략 한 번 더 한다. 대형 언어 모델 학습과 미세 조정에서 긴 시퀀스나 큰 배치를 GPU 메모리에 넣기 위한 기본 도구이며, 혼합 정밀도 학습, ZeRO, 텐서 병렬화와 함께 쓴다.

원리

$n$ 개 층이 사슬처럼 이어진 신경망을 생각하자. 층 $i$ 의 출력을 $a_{i} = f_{i} (a_{i - 1})$ 라 하면, 역전파에서 $\partial L / \partial a_{i - 1}$ 과 가중치 기울기를 구하려면 $a_{i - 1}$ (과 $f_{i}$ 내부 중간값)이 필요하다.

전략	저장하는 활성값	추가 순전파 연산	활성값 메모리
모두 저장(기본)	$a_{0}, \dots, a_{n}$ 전부	없음	$O (n)$
아무것도 저장 안 함	입력만	층마다 처음부터 재계산	$O (1)$ , 연산 $O (n^{2})$
$k$ 개 구간으로 나눠 경계만 저장	구간 경계 $k$ 개 + 재계산 중인 구간 하나	약 순전파 1회	$O (k + n / k)$
$k = \sqrt{n}$	경계 $\sqrt{n}$ 개	약 순전파 1회	$O (\sqrt{n})$
재귀적 분할	로그 개수	$O (n \log n)$	$O (\log n)$

구간 분할 방식은 다음과 같다.

순전파: 신경망을 $k$ 개 구간으로 나누고, 각 구간의 입력만 저장한다. 구간 안의 중간 결과는 계산 뒤 버린다.
역전파: 마지막 구간부터 거꾸로, 저장된 구간 입력에서 그 구간만 다시 순전파해 중간 결과를 만들고, 곧바로 그 구간의 역전파를 한다. 끝나면 그 구간의 중간 결과를 버린다.

메모리는 경계 저장분 $O (k)$ 와 한 구간을 펼칠 때 드는 $O (n / k)$ 의 합이므로 $k = \sqrt{n}$ 에서 $O (\sqrt{n})$ 이 최소다. 각 구간을 정확히 한 번씩 다시 순전파하므로 추가 비용은 순전파 한 번이다. 역전파 연산이 순전파의 약 두 배이므로(순전파 1 + 역전파 2) 순전파 1이 더해지면 전체 학습 연산은 약 3분의 1 늘어난다.

이 아이디어는 자동 미분 분야에서 오래전부터 쓰인 체크포인팅(그리반크와 발터의 Revolve 알고리즘 등)에서 왔고, 2016년 천톈치(Tianqi Chen) 등이 이를 딥 러닝 계산 그래프에 맞게 정리했다. 이 논문은 $n$ 층 신경망을 $O (\sqrt{n})$ 메모리와 미니배치당 순전파 한 번의 추가 비용으로 학습하는 알고리즘을 내놓았고, 극단적으로는 $O (n \log n)$ 추가 연산으로 메모리를 $O (\log n)$ 까지 줄일 수 있음을 보였다. 실험에서는 1000층 잔차 신경망의 ImageNet 학습 메모리를 48GB에서 7GB로 줄였다.^[1]

트랜스포머에서의 활성값 메모리

엔비디아 연구진은 트랜스포머 층 하나가 16비트 활성값으로 저장하는 메모리를 바이트 단위로 다음과 같이 계산했다. $s$ 는 시퀀스 길이, $b$ 는 마이크로배치 크기, $h$ 는 은닉 차원, $a$ 는 어텐션 헤드 수, $t$ 는 텐서 병렬 크기다.^[2]

설정	층당 활성값 메모리(바이트)
병렬화 없음	$s b h (34 + \frac{5 a s}{h})$
텐서 병렬	$s b h (10 + \frac{24}{t} + \frac{5 a s}{h t})$
텐서 병렬 + 시퀀스 병렬	$s b h (\frac{34}{t} + \frac{5 a s}{h t})$
텐서·시퀀스 병렬 + 선택적 재계산	$s b h \cdot \frac{34}{t}$
층 전체 재계산(층 입력만 저장)	$2 s b h$

마지막 항 $5 a s / h$ 는 어텐션 점수 행렬( $Q K^{⊤}$ ), 소프트맥스 출력, 드롭아웃 마스크처럼 시퀀스 길이의 제곱에 비례하는 부분이다. GPT-3(175B, $a = 96, s = 2048, h = 12288$ )에서는 이 값이 80으로 나머지 34보다 크다. 이 부분은 메모리를 많이 먹지만 다시 계산하는 데 드는 연산은 적다.

선택적 재계산

층 전체를 재계산하면 메모리는 가장 적지만 순전파 한 번만큼의 시간이 든다. 위 논문은 층 전체 재계산의 실행 시간 부담을 30~40%로 측정했다. 대신 메모리는 크고 재계산은 싼 어텐션 내부( $5 a s / h$ 항)만 다시 계산하는 선택적 활성값 재계산(selective activation recomputation)을 제안했다. GPT-3 규모에서는 활성값 메모리를 70% 줄이면서 추가 연산은 2.7%에 그쳤고, 시퀀스 병렬화와 함께 써서 활성값 메모리를 5배 줄이고 재계산으로 인한 실행 시간 부담을 90% 넘게 없앴다. 엔비디아 A100 2240장으로 5300억 매개변수 모델을 학습할 때 모델 연산 활용률(MFU)이 42.1%에서 54.2%로 올랐다.^[2]

플래시 어텐션은 어텐션 점수 행렬을 메모리에 만들지 않고 역전파 때 블록 단위로 다시 계산하므로, 그 자체가 어텐션 부분에 대한 선택적 재계산이라 볼 수 있다. 플래시 어텐션을 쓰면 $5 a s / h$ 항이 사실상 사라지고, 남는 선택지는 층 전체 재계산을 얼마나 할지다.

PyTorch 구현

파이토치는 torch.utils.checkpoint 모듈을 제공한다. 감싼 함수 안에서 만들어진 텐서는 역전파 때까지 보관하지 않고, 입력만 보관했다가 역전파 때 함수를 다시 실행한다.^[3]

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint

class Block(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.norm = nn.LayerNorm(d)
        self.ff = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))

    def forward(self, x):
        return x + self.ff(self.norm(x))

class Net(nn.Module):
    def __init__(self, d=1024, n_layers=24, use_ckpt=True):
        super().__init__()
        self.blocks = nn.ModuleList(Block(d) for _ in range(n_layers))
        self.use_ckpt = use_ckpt

    def forward(self, x):
        for blk in self.blocks:
            if self.use_ckpt and self.training:
                # 블록 입력만 저장, 블록 안의 중간값은 역전파 때 다시 계산
                x = checkpoint(blk, x, use_reentrant=False)
            else:
                x = blk(x)
        return x

위 코드를 $d = 1024$ , 24블록, 입력 (8, 512, 1024)로 RTX 3060에서 FP32로 돌려 보면, 순전파·역전파 중 최대 메모리 증가량이 체크포인팅 없이 약 3.9GiB, 블록 단위 체크포인팅으로 약 1.0GiB였고(이 중 약 0.8GiB는 기울기 자체) 기울기 값은 같았다. 수치는 환경에 따라 다르다.

use_reentrant 인자

checkpoint에는 두 구현이 있고 use_reentrant 인자로 고른다. PyTorch 문서는 use_reentrant=False를 권장하며, 이 인자를 명시적으로 넘기라고 경고한다.^[3]

항목	use_reentrant=True (예전 방식)	use_reentrant=False (권장)
재계산 범위	함수 전체를 항상 다시 실행	필요한 중간값이 다 나오면 멈춤(early stop)
`torch.autograd.grad`	지원 안 함	지원
입력 조건	입력·출력 중 하나 이상이 requires_grad여야 함	제약 없음
키워드 인자, 중첩 구조 텐서	지원 안 함	지원

드롭아웃처럼 난수를 쓰는 연산이 있으면, 기본 설정(preserve_rng_state=True)은 순전파 때 난수 상태를 저장했다가 재계산 때 되살려 같은 마스크가 나오게 한다. 재계산 함수의 동작이 순전파와 달라지면(전역 변수에 따라 분기하는 경우 등) 기울기가 조용히 틀릴 수 있다.

선택적 체크포인팅(SAC)

최신 PyTorch는 연산 단위로 저장할지 재계산할지 정하는 정책을 줄 수 있다. 아래는 행렬 곱 결과는 저장하고 나머지(활성화 함수, 정규화 등 싼 연산)는 다시 계산하는 예다.

import functools
from torch.utils.checkpoint import (
    checkpoint, create_selective_checkpoint_contexts, CheckpointPolicy)

ops_to_save = [torch.ops.aten.mm.default, torch.ops.aten.addmm.default]

def policy_fn(ctx, op, *args, **kwargs):
    if op in ops_to_save:
        return CheckpointPolicy.MUST_SAVE
    return CheckpointPolicy.PREFER_RECOMPUTE

context_fn = functools.partial(create_selective_checkpoint_contexts, policy_fn)
y = checkpoint(blk, x, use_reentrant=False, context_fn=context_fn)

라이브러리별 사용법

허깅 페이스 Transformers: model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) 또는 TrainingArguments(gradient_checkpointing=True). 체크포인팅 중에는 KV 캐시를 쓸 수 없어 use_cache=False가 된다.
순차 모델: torch.utils.checkpoint.checkpoint_sequential(model, segments, input, use_reentrant=False)는 nn.Sequential을 segments개 구간으로 나눠 위 원리대로 처리한다.
Megatron-LM: --recompute-granularity를 full 또는 selective로 고르고, full일 때 --recompute-method(uniform·block)와 --recompute-num-layers로 범위를 정한다.
JAX: jax.checkpoint(별칭 jax.remat)과 저장 정책(jax.checkpoint_policies)으로 같은 일을 한다.

언제 쓰나

메모리 부족(OOM)으로 배치나 시퀀스 길이를 줄여야 할 때. 배치를 키워 얻는 처리량 이득이 재계산 비용보다 크면 오히려 빨라진다.
메모리가 남으면 끈다. 남는 메모리를 버려 두고 재계산 비용만 내는 셈이 된다. 일부 층만 체크포인팅해 중간 지점을 찾을 수 있다.
파라미터 효율 미세 조정(LoRA, QLoRA)에서는 가중치·옵티마이저 메모리가 작아 활성값이 메모리의 대부분을 차지하므로 효과가 크다.
파이프라인 병렬 학습에서는 여러 마이크로배치의 활성값이 동시에 메모리에 있으므로 재계산과 자주 함께 쓴다. GPipe도 활성값 메모리를 줄이려고 재계산을 지원한다.^[4]
활성값을 버리지 않고 CPU 메모리로 내리는 방법(활성값 오프로딩, PyTorch torch.autograd.graph.save_on_cpu)도 있다. 재계산 대신 PCIe 전송 시간을 쓴다.

같이 보기

각주

↑ Tianqi Chen, Bing Xu, Chiyuan Zhang, Carlos Guestrin, 〈Training Deep Nets with Sublinear Memory Cost〉, 2016. arXiv:1604.06174
↑ ^2.0 ^2.1 Vijay Korthikanti 외, 〈Reducing Activation Recomputation in Large Transformer Models〉, 2022. arXiv:2205.05198
↑ ^3.0 ^3.1 PyTorch, 〈torch.utils.checkpoint〉, 2.14 문서. docs.pytorch.org
↑ Yanping Huang 외, 〈GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism〉, NeurIPS 2019. arXiv:1811.06965

참고 문헌

PyTorch 2.14 문서: torch.utils.checkpoint
Andreas Griewank, Andrea Walther, 〈Algorithm 799: Revolve: An Implementation of Checkpointing for the Reverse or Adjoint Mode of Computational Differentiation〉, ACM TOMS, 2000. doi:10.1145/347837.347846

[chen16-1] Tianqi Chen, Bing Xu, Chiyuan Zhang, Carlos Guestrin, 〈Training Deep Nets with Sublinear Memory Cost〉, 2016. arXiv:1604.06174

[kor22-2] 2.0 ^2.1 Vijay Korthikanti 외, 〈Reducing Activation Recomputation in Large Transformer Models〉, 2022. arXiv:2205.05198

[ptckpt-3] 3.0 ^3.1 PyTorch, 〈torch.utils.checkpoint〉, 2.14 문서. docs.pytorch.org

[gpipe-4] Yanping Huang 외, 〈GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism〉, NeurIPS 2019. arXiv:1811.06965

[1]

[2]

[3]

[4]