대형 모델 학습에서 GPU 메모리를 압박하는 대표 원인 중 하나는 활성화다. 활성화는 순전파 중 만들어지고, 역전파 계산을 위해 저장되는 중간 텐서다.1 활성화 체크포인팅, 또는 그래디언트 체크포인팅은 이 활성화의 일부를 저장하지 않고 역전파 때 다시 계산해 메모리를 줄이는 기법이다.1
문제는 체크포인팅이 공짜가 아니라는 점이다. 더 많은 활성화를 버릴수록 메모리는 더 아낄 수 있지만, 그만큼 역전파 중 다시 수행해야 하는 순전파 계산도 늘어난다.1 따라서 실전에서 중요한 질문은 “체크포인팅을 쓸까?”가 아니라 “무엇을 저장하고, 무엇을 다시 계산할까?”에 가깝다.
블록 단위에서 op 호출 단위로
기존 PyTorch의 일반적인 활성화 체크포인팅은 보통 함수나 블록 단위로 적용한다. 특정 구간의 활성화를 저장하지 않고, 필요할 때 그 구간을 다시 실행하는 방식이다. PyTorch의 Selective Activation Checkpointing(SAC)은 여기서 한 단계 더 나아가 op 종류별 정책 함수로 저장 또는 재계산 여부를 정할 수 있다.2
torch_remat은 이 선택을 더 세밀하고 명시적으로 다루려는 새로운 활성화 체크포인팅 API다. 공개 저장소는 meta-pytorch/remat이며, 기본 설계 방향은 “일단 모두 재계산하고, 꼭 저장할 것만 지정한다”에 가깝다.3
핵심 차이는 선택의 단위다. SAC가 op 종류 중심의 정책을 사용한다면, torch_remat은 개별 op 호출 단위로 저장할지 재계산할지 지정하는 데 초점을 둔다.3 즉 같은 matmul이라도 어떤 호출은 저장하고, 다른 호출은 다시 계산하도록 다르게 다룰 수 있다.
왜 이런 제어가 필요한가
체크포인팅의 효과는 버리는 활성화의 양에 비례하고, 비용은 추가 순전파 재계산에 비례한다.1 결국 좋은 정책은 “많이 버리는 것”이 아니라 “저렴하게 다시 계산할 수 있는 것은 버리고, 다시 계산하기 비싼 것은 저장하는 것”에 가깝다.
torch_remat의 관점은 이 판단을 코드에 더 직접적으로 반영하게 해준다. 특히 custom autograd Function을 많이 사용하는 PyTorch 코드베이스에 맞춰 설계됐다는 점이 중요하다.3 이런 코드에서는 기본 autograd 흐름에 맡기기보다, 역전파를 위해 어떤 값을 남길지 더 세밀하게 제어해야 하는 경우가 많다.
적용 전에 확인할 세 가지
torch_remat을 검토할 때는 먼저 병목을 분명히 해야 한다. 현재 문제는 파라미터 크기 자체가 아니라, 순전파 중 쌓이는 활성화 메모리인가? 활성화가 peak memory의 주요 원인일 때 체크포인팅의 가치가 커진다.
다음으로 op별 판단이 가능해야 한다. 다시 계산해도 부담이 작은 연산과, 저장하는 편이 나은 비싼 연산을 구분할 수 있어야 한다. torch_remat의 장점은 바로 이 구분을 op 호출 단위로 표현할 수 있다는 데 있다.
마지막으로 반드시 측정해야 한다. 체크포인팅은 메모리와 시간의 교환이므로, 변경 전후의 peak memory와 step time을 함께 봐야 한다. 메모리만 줄고 학습 시간이 지나치게 늘어난다면 좋은 선택이 아닐 수 있다.
결론
torch_remat은 체크포인팅을 더 똑똑하게 자동화하는 마법이라기보다, 메모리 예산을 코드 수준에서 더 명확히 표현하게 해주는 도구에 가깝다. 기본적으로 모두 재계산하고, 필요한 op 호출만 저장하도록 지정하는 방식은 활성화 메모리와 재계산 비용 사이의 균형을 더 세밀하게 조정할 수 있게 한다.
따라서 torch_remat의 가치는 단순히 “메모리를 줄인다”가 아니라, “무엇을 저장할지에 대한 판단을 더 정확한 단위로 옮긴다”에 있다. 대형 모델 학습에서 활성화 메모리가 병목이고, custom autograd Function을 포함한 복잡한 PyTorch 코드베이스를 다룬다면 특히 눈여겨볼 만하다.
1Gradient Checkpointing for Memory-Efficient PyTorch Training
2Current and New Activation Checkpointing Techniques in PyTorch
3Meta PyTorch Unveils Torch_Remat API For Flexible Activation Checkpointing
이 노트는 보도·비평·교육·연구 목적으로 작성되었습니다. 저작권 문의가 있으시면 에서 알려주세요.