torch.utils.checkpoint — 활성화 메모리를 줄이는 기법
torch.utils.checkpoint — 활성화 메모리를 줄이는 기법
PyTorch의 torch.utils.checkpoint 는 활성화 체크포인팅(activation checkpointing)을 위한 API예요. 기본적으로 순전파 중 계산된 텐서는 역전파에서 그래디언트를 계산할 때까지 살아있는데, 이는 큰 메모리를 요구해요. 체크포인팅은 이 텐서를 살려두지 않고, 필요할 때 함수를 다시 호출해 재계산하는 방식이에요.
핵심 아이디어
전달된 입력 텐서(args)는 유지하고, 저장하지 않은 텐서는 역전파에서 필요할 때 함수를 재호출해 다시 계산해요. 이 '계산 대 메모리' 트레이드오프는 모델의 어느 부분이라도 적용할 수 있어요.
두 가지 구현 (use_reentrant)
use_reentrant=True: 재진입 autograd 기반. 권장하지 않아요.use_reentrant=False: 비재진입 구현. 재계산을 필요한 만큼만 하고 일찍 멈춰 더 빠른 경우가 많아요. 권장이에요.
동작 원리
checkpoint(fn, *args, use_reentrant=False) 처럼 함수와 인자를 넘겨요. 순전파 중간 결과를 버리고, 역전파에서 '체크포인트된' 경계까지 다시 순전파를 돌려 필요한 중간 활성화를 복원하는 방식이에요.
주의점
재현성을 위해 기존 재진입 변형은 기본적으로 RNG 상태를 저장·복원해요. 순전파와 역전파 호출이 달라지면(전역 변수 등) 잘못된 그래디언트가 생길 수 있어 주의해야 해요.