Gradient Checkpointing

그래디언트 체크포인팅은 순전파의 중간 활성화 중 일부만 저장하고 나머지는 역전파 중에 다시 계산하는 메모리 절약형 학습 기법입니다. 활성화 체크포인팅이라고도 하며, 추가 연산을 수행하는 대신 최대 메모리 사용량을 줄입니다. 이름과 달리 이 기법은 파라미터 그래디언트나 모델 파일이 아니라 활성화를 체크포인트하므로, 활성화 텐서 때문에 신경망이 사용 가능한 GPU 메모리에 맞지 않는 경우 특히 유용합니다.

그래디언트 체크포인팅 작동 방식#

일반적인 순전파 중에 신경망은 활성화라고 하는 중간 텐서를 계산합니다. 자동 미분 시스템은 나중에 그래디언트를 계산하는 데 필요한 활성화를 보존하며, 이에 대한 설명은 PyTorch 자동 미분 메커니즘에서 확인할 수 있습니다. 심층 네트워크, 대규모 배치, 고해상도 이미지 및 긴 입력 시퀀스에서는 저장된 텐서가 상당한 메모리를 차지할 수 있습니다.

그래디언트 체크포인팅은 네트워크를 여러 세그먼트로 나눕니다.

1. 순전파에서는 선택한 세그먼트의 입력 또는 경계 활성화를 저장합니다.
2. 해당 세그먼트 내부의 다른 중간 활성화는 버립니다.
3. 역전파 중에는 각 체크포인트 세그먼트의 순전파를 다시 실행하여 누락된 값을 재구성합니다.
4. 재구성된 활성화는 즉시 그래디언트 계산에 사용됩니다.

PyTorch 활성화 체크포인팅 API는 torch.utils.checkpoint을 통해 이 동작을 제공합니다. 이에 대응하는 개념으로는 JAX 그래디언트 체크포인팅 및 재물질화와 TensorFlow 테이프 체크포인팅이 있습니다.

체크포인트 위치에 따라 절충점이 달라집니다. 일반적으로 더 많은 영역에 체크포인팅을 적용할수록 메모리는 더 많이 절약되지만 더 많은 연산을 반복합니다. 전체 네트워크를 다시 계산하는 것보다 활성화가 많이 발생하는 블록 주변에 선택적으로 적용하는 편이 더 나은 균형을 제공할 수 있습니다.

절충점 및 관련 기법#

그래디언트 체크포인팅은 일반적으로 모델 파라미터, 그래디언트 및 옵티마이저 상태를 변경하지 않습니다. 주요 대상은 활성화 메모리이며, 일부 순전파 연산을 두 번 실행하기 때문에 학습 속도가 느려집니다. 정확한 결과는 아키텍처, 체크포인트 경계, 배치 형태 및 하드웨어에 따라 달라집니다.

다음과 같은 여러 관련 기법과는 차이가 있습니다.

• 그래디언트 누적은 모델 가중치를 업데이트하기 전에 여러 마이크로배치를 처리하여 모든 샘플을 한 번에 메모리에 올리지 않고도 유효 배치 크기를 키웁니다. 반면 그래디언트 체크포인팅은 각 마이크로배치에서 보존되는 활성화를 줄입니다.
• 혼합 정밀도는 선택된 연산에 더 낮은 정밀도의 데이터 타입을 사용합니다. PyTorch 자동 혼합 정밀도 워크플로는 메모리를 줄이고 호환되는 연산을 가속할 수 있는 반면, 체크포인팅은 의도적으로 연산량을 늘립니다.
• 배치 크기 축소는 한 번에 처리하는 샘플 수를 줄여 메모리 사용량을 낮춥니다. 체크포인팅을 사용하면 더 큰 배치나 입력 해상도를 적용할 수 있는 경우가 있습니다.
• 학습 체크포인트는 복구 또는 후속 추론을 위해 가중치와 옵티마이저 상태를 저장합니다. PyTorch 모델 체크포인트 가이드는 활성화 재계산과는 관련이 없는 이 저장 메커니즘을 설명합니다.

체크포인트된 코드는 원래 순전파와 재계산 순전파에서 기능적으로 일관되어야 합니다. 체크포인트 영역 내부의 변경 가능한 상태, 기기 간 전송 또는 제어되지 않은 난수 사용은 오류나 잘못된 그래디언트를 초래할 수 있습니다. 디코더형 Transformer에서는 캐시된 추론 상태가 순전파 그래프 재구성과 충돌할 수 있으므로, 체크포인트를 적용한 학습 중에는 키-값 캐시를 비활성화해야 할 수도 있습니다.

PyTorch 예시#

다음 예시는 한 번의 학습 단계에서 메모리를 많이 사용하는 블록에 체크포인팅을 적용합니다.

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

torch.manual_seed(0)

block = nn.Sequential(
    nn.Linear(1024, 4096),
    nn.ReLU(),
    nn.Linear(4096, 1024),
)
optimizer = torch.optim.AdamW(block.parameters())
inputs = torch.randn(8, 1024, requires_grad=True)
targets = torch.zeros_like(inputs)

optimizer.zero_grad(set_to_none=True)
outputs = checkpoint(block, inputs, use_reentrant=False)
loss = nn.functional.mse_loss(outputs, targets)
loss.backward()
optimizer.step()

블록의 입력과 필요한 경계 정보만 유지하고, 블록 내부 활성화는 loss.backward() 중에 재구성합니다. 실제 프로젝트에서는 PyTorch 최대 GPU 메모리 할당량과 같은 측정값을 사용해 체크포인팅 적용 여부에 따른 실행을 비교해야 합니다.

실제 적용 사례#

• 고해상도 컴퓨터 비전: 의료 영상 분할, 항공 영상 탐지 및 산업 검사는 특징 맵이 모델 가중치보다 더 많은 메모리를 사용하는 대형 이미지로 학습할 수 있습니다. 일부 백본 단계에 체크포인팅을 적용하면 입력을 크게 축소하는 대신 이미지 해상도를 유지하거나 배치당 샘플을 늘릴 수 있습니다.

• 긴 시퀀스 Transformer 학습: 시퀀스 길이와 레이어 수가 증가하면 활성화 저장량이 빠르게 늘어납니다. Transformer 블록을 재계산하면 동일한 가속기에서 더 긴 컨텍스트나 더 큰 마이크로배치를 처리할 수 있습니다. NVIDIA 활성화 재계산 가이드는 Transformer 레이어에 전체 및 선택적 재계산을 적용하는 방법을 보여줍니다.

실용적인 지침#

프로파일링 결과 파라미터나 옵티마이저 상태가 아니라 활성화가 메모리 사용량의 대부분을 차지하는 경우 그래디언트 체크포인팅을 사용합니다. 큰 반복 블록부터 시작하고, 적용 범위를 넓히기 전에 최대 메모리 사용량, 반복 시간 및 검증 동작을 벤치마크합니다.

문서화된 Ultralytics YOLO 학습 구성의 경우 우선 적용할 메모리 제어 방법으로 자동 혼합 정밀도, 이미지 크기 및 물리적 배치 크기가 있습니다. Ultralytics AutoBatch 레퍼런스는 사용 가능한 GPU 메모리를 기준으로 배치 크기를 자동 선택하는 방법을 설명합니다. 로컬 하드웨어의 성능이 여전히 부족한 경우 Ultralytics Platform 클라우드 학습을 통해 관리형 학습 실행에 사용할 클라우드 GPU를 구성할 수 있습니다. 사용자 지정 PyTorch 아키텍처에서 활성화 메모리를 더 세밀하게 관리해야 하는 경우 그래디언트 체크포인팅은 이러한 제어 방법을 보완할 수 있습니다.