그레이디언트 체크포인팅 (Gradient Checkpointing)
쉽게 풀면
신경망을 학습할 때는 역전파 계산을 위해 순전파 과정에서 나온 중간 결과들을 모두 메모리에 기억해 두어야 한다. 모델이 깊어질수록 이 중간 결과의 양도 늘어나 GPU 메모리가 금방 부족해진다. 그레이디언트 체크포인팅은 이 모든 중간 결과를 다 저장하는 대신 일부 지점만 저장해 두고, 역전파 시 필요한 나머지 값은 그때그때 다시 계산해서 얻는 방식이다. 메모리는 절약되지만 다시 계산하는 만큼 학습 시간은 다소 늘어나는 trade-off가 있다.
왜 중요한가
대규모 언어모델이나 매우 깊은 신경망은 파라미터 수뿐 아니라 학습 중 저장해야 하는 중간 활성값 때문에 GPU 메모리가 병목이 되는 경우가 많습니다. 그레이디언트 체크포인팅은 제한된 하드웨어에서도 더 크거나 더 긴 시퀀스를 다루는 모델을 학습할 수 있게 해주기 때문에, 대규모 모델 학습 인프라를 다루는 논문에서 혼합 정밀도 학습이나 모델 병렬화와 함께 메모리 효율화 기법으로 자주 언급됩니다.
논문에서는 이렇게 쓰입니다
모델이 커서 GPU 메모리가 부족한 문제를, 중간 결과를 다시 계산하는 방식으로 메모리 사용량을 줄여 해결했다는 뜻이다.
메모리 절약 덕분에 한 번에 더 많은 데이터를 학습할 수 있게 되었지만, 그 대가로 학습 속도가 느려졌다는 결과를 보고한 문장이다.
여러 메모리 절약 기법을 함께 사용해 제한된 하드웨어에서도 대규모 모델 학습을 가능하게 했다는 의미다.
조금 더 깊게 보면
그레이디언트 체크포인팅은 신경망을 몇 개의 구간(세그먼트)으로 나눈 뒤, 각 구간의 경계 지점에서 나온 활성값만 저장하고 구간 내부의 중간 결과는 저장하지 않는 방식으로 동작합니다. 역전파 시 필요한 값은 저장해둔 경계값에서부터 순전파를 다시 실행해 그 구간만 재계산하므로, 층 수가 n일 때 구간을 √n개 정도로 잡으면 메모리 사용량이 대략 n의 제곱근에 비례하는 수준으로 줄어드는 대신 순전파 연산량은 늘어납니다. 트랜스포머 계열 모델에서는 흔히 각 층(레이어) 단위로 체크포인트를 두는 방식이 널리 쓰이며, 최근에는 이러한 재계산 비용을 고려해 최적의 체크포인트 배치를 자동으로 찾는 기법도 연구되고 있습니다.
주의할 점
그레이디언트 체크포인팅은 메모리를 아끼는 대신 계산량과 학습 시간이 늘어나는 trade-off가 있으므로, 메모리가 충분하다면 굳이 쓸 필요가 없다.