그래디언트 체크포인팅은 Deep learning에서 신경망 훈련 시 메모리 사용량을 줄이기 위해 사용되는 기법이다. 표준 역전파 과정에서 네트워크는 순전파 중 계산된 모든 중간 활성화 값을 저장해야 역전파 시 그래디언트를 계산할 수 있다. 대규모 언어 모델 및 트랜스포머와 같은 매우 깊은 모델의 경우, 이러한 저장 공간이 사용 가능한 하드웨어의 메모리 용량을 초과할 수 있다. 그래디언트 체크포인팅은 모든 활성화 값을 저장하는 대신 일부만 유지하고, 역전파 중 필요할 때 버려진 값들을 다시 계산하여 이 문제를 해결한다. 이는 계산 비용 증가와 메모리 사용량의 상당한 감소를 맞바꾸는 것으로, 동일한 하드웨어에서 더 큰 모델을 훈련하거나 더 큰 배치 크기를 사용할 수 있게 한다.
이 기법은 2016년 Carnegie Mellon University와 OpenAI의 연구자들이 "Training Deep Nets with Sublinear Memory Cost"라는 논문에서 처음 소개했다. Tianqi Chen, Bing Xu, Chiyuan Zhang, Carlos Guestrin을 포함한 저자들은 특정 체크포인트(예: 몇 레이어마다)에서만 활성화 값을 저장하고 나머지를 다시 계산함으로써, n개의 레이어를 가진 네트워크의 훈련 메모리 비용을 O(n)에서 O(sqrt(n))으로 줄일 수 있음을 입증했다. 이는 대략 한 번의 추가 순전파 비용을 수반한다. 이 기초 작업은 이후 머신 러닝 커뮤니티에서 표준 도구가 되었으며, 특히 모델 크기가 급격히 증가함에 따라 더욱 중요해졌다.
표준 역전파의 메모리 사용 방식
기존 훈련 루프에서 순전파는 네트워크의 모든 레이어에 대한 활성화 값을 계산한다. 이러한 활성화 값은 역전파가 연쇄 법칙을 통해 그래디언트를 계산하는 데 필요하므로 메모리에 저장된다. L개의 레이어를 가진 네트워크의 경우, 각각 클 수 있는 L개의 활성화 값 세트를 저장해야 한다. 예를 들어, 수백 개의 레이어를 가진 Residual Network (ResNet)나 수십 개의 어텐션 블록을 가진 트랜스포머는 단일 훈련 예제에 대해 수 기가바이트의 활성화 데이터를 축적할 수 있다. 큰 배치 크기로 훈련할 때 메모리 요구량은 배치 크기에 선형적으로 증가하므로, 종종 주요 병목 현상이 된다.
체크포인팅 전략
그래디언트 체크포인팅은 네트워크를 세그먼트로 나누고, 각 세그먼트의 경계에 체크포인트를 설정한다. 순전파 중에는 체크포인트 경계의 활성화 값만 메모리에 저장된다. 세그먼트 내의 다른 모든 중간 활성화 값은 폐기된다. 역전파가 세그먼트에 도달하면 저장된 체크포인트 활성화 값을 사용하여 해당 세그먼트의 순전파를 다시 계산하고, 그래디언트 계산에 필요한 중간 활성화 값을 재생성한다. 이 재계산은 일반적으로 훈련 단계당 한 번의 추가 순전파에 해당하는 계산 오버헤드를 추가하지만, 최대 메모리 사용량을 크게 줄인다.
체크포인트 배치 선택은 절충점이다. 체크포인트가 많을수록 재계산이 줄어들지만 메모리 사용량이 높아지고, 적을수록 메모리는 낮아지지만 계산이 늘어난다. n개의 레이어를 가진 네트워크의 최적 체크포인트 수는 대략 sqrt(n)으로, 메모리와 계산의 균형을 맞춘다. 실제로 PyTorch와 TensorFlow 같은 프레임워크는 사용자가 체크포인트 간격을 지정하거나 자동 휴리스틱을 사용할 수 있게 한다.
변형 및 개선
원래 기법에 대한 여러 개선이 개발되었다. 일반적인 변형 중 하나는 선택적 체크포인팅으로, 특정 레이어 유형(예: 어텐션 블록 또는 합성곱 레이어)만 체크포인트하고 나머지는 정상적으로 저장하는 방식이다. 또 다른 접근 방식인 메모리 효율적 그래디언트 체크포인팅은 여러 수준의 세분성으로 활성화 값을 저장하는 더 정교한 일정을 사용하여 추가 재계산 비용으로 메모리를 더 줄인다. 일부 프레임워크는 체크포인트를 CPU 메모리나 디스크로 이동하는 "오프로딩"도 구현하지만, 이는 데이터 전송 오버헤드를 도입한다.
Transformer (architecture) 모델의 맥락에서 그래디언트 체크포인팅은 Mixed Precision Training 및 Gradient Clipping과 같은 다른 메모리 절약 기법과 자주 결합된다. 예를 들어, 1,750억 개의 매개변수를 가진 GPT-3 같은 모델을 훈련하는 것은 이러한 최적화 없이는 불가능할 것이다. 이 기법은 대규모 모델의 미세 조정에도 사용되며, 메모리 절약 덕분에 클러스터 대신 단일 GPU에서 실행할 수 있다.
실제 구현
현대 딥 러닝 프레임워크에서 그래디언트 체크포인팅은 일반적으로 간단한 API로 제공된다. 예를 들어, PyTorch의 torch.utils.checkpoint 모듈은 모듈이나 일련의 연산을 감싸는 checkpoint 함수를 제공한다. 감싸진 모듈이 실행될 때 활성화 값이 저장되지 않고 역전파 중에 다시 계산된다. TensorFlow는 tf.recompute_grad를 통해 유사한 기능을 제공한다. 이러한 구현은 자동으로 기록을 처리하므로 연구자가 모델 아키텍처를 수정하지 않고도 기법을 쉽게 채택할 수 있다.
그래디언트 체크포인팅의 계산 오버헤드는 무시할 수 없다. sqrt(n)개의 체크포인트를 가진 네트워크의 경우, 훈련 중 총 순전파 계산은 표준 훈련에 비해 약 30-40% 증가한다. 그러나 이 비용은 종종 허용 가능한데, 대안인 배치 크기나 모델 크기 축소는 수렴이나 모델 품질에 해를 끼칠 수 있기 때문이다. 많은 경우 더 큰 배치 크기 사용으로 인한 속도 향상이 재계산 오버헤드를 능가한다.
대규모 모델 훈련에 미치는 영향
그래디언트 체크포인팅은 매우 큰 모델 훈련의 초석이 되었다. OpenAI, Anthropic, Google DeepMind 같은 기업들은 수천억 개의 매개변수를 가진 모델을 훈련하기 위해 이 기법에 의존한다. 예를 들어, 8개의 GPU가 있는 단일 노드에서 700억 개의 매개변수 모델을 훈련하려면 체크포인팅 없이는 해당 GPU의 결합 메모리를 초과하는 활성화 값을 저장해야 한다. 그래디언트 체크포인팅을 사용함으로써 이러한 조직은 훈련 시간이 길어지더라도 사용 가능한 하드웨어에 훈련 작업을 맞출 수 있다.
이 기법은 문서 요약이나 코드 생성과 같은 긴 시퀀스를 포함하는 Generative AI 애플리케이션에도 필수적이다. 이러한 경우 활성화 메모리는 시퀀스 길이에 따라 증가하며, 체크포인팅은 메모리 한도를 초과하지 않고 더 긴 컨텍스트를 허용한다. 이는 100,000개 이상의 토큰 컨텍스트 창을 가진 모델 개발을 직접적으로 가능하게 했다.
다른 메모리 최적화와의 관계
그래디언트 체크포인팅은 종종 다른 기법과 함께 사용된다. Batch Normalization과 Layer Normalization은 메모리를 직접 줄이지 않지만 훈련 안정성을 개선하여 체크포인팅을 보완한다. Model Pruning은 매개변수 수를 줄이지만 활성화 값은 여전히 병목 현상이므로 체크포인팅이 여전히 필요하다. Data Augmentation은 효과적인 데이터 세트 크기를 늘리지만 활성화 메모리에는 영향을 미치지 않는다. 분산 훈련에서 그래디언트 체크포인팅은 파이프라인 병렬 처리와 결합할 수 있으며, 여기서 다른 레이어가 다른 장치에 할당되어 장치별 메모리 압력을 더 줄인다.
주목할 만한 대안 중 하나는 Gradient Accumulation으로, 여러 개의 작은 배치에 걸쳐 그래디언트를 축적하여 더 큰 배치 크기를 시뮬레이션한다. 이는 옵티마이저 상태의 메모리를 줄이지만 활성화 메모리는 줄이지 않으므로 체크포인팅의 대체재가 아니다. 또 다른 관련 아이디어는 일부 Residual Network (ResNet) 변형에서 사용되는 가역 레이어로, 활성화 값을 출력에서 재구성할 수 있지만 아키텍처 변경이 필요하며 체크포인팅보다 덜 일반적이다.
한계 및 절충점
그래디언트 체크포인팅의 주요 한계는 훈련 단계당 벽시계 시간이 증가한다는 점이다. 이미 계산에 바인딩된 모델의 경우 추가 순전파로 훈련이 20-40% 느려질 수 있다. 또한 이 기법은 대규모 모델에서 상당할 수 있는 모델 매개변수나 옵티마이저 상태의 메모리를 줄이지 않는다. 매우 큰 모델의 경우 실무자는 체크포인팅을 매개변수 오프로딩과 결합하거나 AWS Trainium이나 Groq와 같이 더 큰 온칩 메모리를 가진 특수 하드웨어를 사용해야 할 수 있다.
또 다른 미묘한 문제는 재계산이 수치적 차이를 도입할 수 있다는 점이지만, 실제로는 일반적으로 무시할 수 있다. 이 기법은 또한 동일한 활성화 값을 여러 번 다시 계산하여 오버헤드를 증가시키지 않도록 신중한 구현이 필요하다. 이러한 어려움에도 불구하고 그래디언트 체크포인팅은 널리 사용되는 신뢰할 수 있는 방법으로 남아 있다.
향후 방향
모델이 계속 성장함에 따라 연구자들은 더 효율적인 체크포인팅 전략을 탐구하고 있다. 최근 일부 작업은 모델 구조와 하드웨어 메모리 프로필을 기반으로 어떤 활성화 값을 저장할지 결정하는 학습된 휴리스틱을 사용한다. 다른 연구는 레이어 수준이 아닌 개별 연산 수준에서 체크포인팅을 조사하여 더 세밀한 제어를 허용한다. 또한 체크포인팅을 저정밀 훈련과 결합하여 메모리를 더 줄이는 데 관심이 있지만, 이는 정확성 절충점을 도입한다.
이 기법은 메모리가 딥 러닝의 제약 조건으로 남아 있는 한 계속 관련성을 유지할 가능성이 높다. 대규모 언어 모델과 멀티 헤드 어텐션 아키텍처의 부상으로 메모리 효율적 훈련 방법에 대한 수요는 더욱 증가할 것이다. 그래디언트 체크포인팅은 다른 최적화와 함께 Artificial intelligence의 발전을 가능하게 하는 핵심 요소로 계속될 것이다.
결론
그래디언트 체크포인팅은 간단하면서도 강력한 아이디어다. 활성화 값을 선택적으로 저장하고 역전파 중에 다시 계산함으로써 깊은 네트워크를 훨씬 적은 메모리로 훈련할 수 있다. 2016년에 소개된 이후 이 기법은 현장의 표준 관행이 되었으며, 사용 가능한 하드웨어에서 불가능했을 모델 훈련을 가능하게 했다. 계산 오버헤드를 추가하지만, 특히 대규모 모델의 경우 그 절충점은 종종 가치가 있다. 분야가 더 큰 아키텍처로 나아감에 따라 그래디언트 체크포인팅은 머신 러닝 도구 상자에서 필수 도구로 남을 것이다.