파이토치는 모델을 실행하기 전에 순전파와 역전파를 합친 계산 그래프를 만들고, 그 위에서 속도와 메모리 효율을 높이기 위한 여러 최적화를 시도한다. 그중 하나가 활성화 체크포인팅과 맞닿아 있는 결정이다. 각 연산의 결과를 메모리에 저장해두고 나중에 재사용할지, 아니면 저장하지 않고 필요할 때 다시 계산할지를 고르는 작업이다. 저장하면 재계산 시간을 아끼지만 메모리를 차지하고, 쓸 수 있는 메모리는 한정돼 있다. 결국 '정해진 메모리 예산 안에서 어떤 연산을 저장해 런타임 시간을 최대로 아낄 것인가'라는 최적화 문제가 된다.
왜 냅색 문제인가
이 구조는 고전적인 0/1 냅색 문제와 정확히 같다. 용량이 정해진 배낭에 가치와 무게가 각기 다른 물건들을 넣어 총 가치를 최대화하되 용량은 넘기지 않아야 하는 문제 말이다. 파이토치에서는 메모리 예산이 배낭 용량, 각 연산이 물건, 저장했을 때 아끼는 시간이 가치, 저장에 드는 메모리가 무게에 대응한다. 무게가 작은 것만 고르는 식의 휴리스틱으로도 일부 이득은 얻지만 최적해를 보장하지는 못한다. 그래서 정확한 최적 조합을 구하는 냅색 솔버가 쓰인다.
파이토치의 기본 솔버는 동적 계획법 기반의 dp_knapsack이다. 품목 수를 행, 최대 무게를 열로 하는 2차원 DP 테이블을 채우며 각 용량에서 얻을 수 있는 최대 가치를 계산한 뒤, 테이블을 거꾸로 거슬러 올라가며(백트래킹) 실제로 어떤 품목이 선택됐는지를 복원한다. 파이토치 입장에서는 이 복원 과정이 '어떤 연산을 메모리에 저장하고 어떤 연산을 재계산할 것인가'라는 최종 답에 해당한다. 이 솔버는 정확한 해를 내놓고 대부분의 경우 충분하다.
문제는 테이블의 크기다
한계는 메모리다. dp_knapsack은 (품목 수 × 최대 무게) 크기의 완전한 2차원 테이블을 통째로 할당한다. 애초에 메모리를 아끼려는 계획 단계가 거꾸로 막대한 메모리를 먹는 셈이다. 파이토치 메모리 플래닝의 상당 부분을 설계한 Horace He(Chillee)는 솔버 코드에 개선 아이디어를 메모로 남겨두었는데, 핵심은 두 가지 기법을 결합하는 것이다. 하나는 '윈도 트릭'으로, 전체 테이블 대신 직전 행과 현재 행만 유지하며 미끄러지듯 계산을 진행한다. 다음 행으로 넘어가면 현재 행이 직전 행이 된다. 최대 가치만 구할 거라면 이것만으로 충분하지만, 우리는 선택된 품목까지 알아야 하므로 한 가지가 더 필요하다.
그 보완책이 Hirschberg 트릭이다. 퀵소트처럼 분할 정복으로 문제를 반으로 쪼개 나가는 방식인데, 이 문제에서 가장 큰 장점은 백트래킹 자체를 없앤다는 점이다. 글쓴이는 재귀 대신 파이썬 리스트를 명시적 스택으로 써서 품목 인덱스 범위와 가용 용량을 저장한다. 스택에서 하나를 꺼내 품목이 하나뿐이면 가용 용량에 들어가는지 보고 저장 목록 또는 재계산 목록에 바로 넣는다. 범위가 둘 이상이면 품목을 둘로 나눠 각각 윈도 트릭으로 DP 행을 구하고, 왼쪽과 오른쪽 행을 더해 가장 값이 높은 위치를 찾는다. 그 위치가 메모리를 가르는 새 분할점이 되고, 메모리와 품목을 같이 쪼개 알고리즘을 재귀적으로 반복한다. 분할 결과 한쪽이 비는 경우도 정상이며 그냥 건너뛴다. 저장·재계산 목록이 알고리즘 진행 중에 완성되므로 사후 복원 단계가 사라진다.
실측된 효과와 선택지
두 기법을 합치면 DP 테이블 크기가 (품목 수 × 최대 무게)에서 (2 × 최대 무게)로 줄어든다. 글쓴이의 64GB RAM 환경에서 dp_knapsack은 품목 100개짜리 문제에서 이미 메모리 부족으로 죽지만, dp_knapsack_sliding_hirschberg는 같은 예산에서 품목 2000개까지 처리한다. 피크 메모리를 약 20배 낮춘 결과다. 런타임도 대략 37% 빨라졌다고 하는데, 이 수치는 엄밀하게 측정한 값이 아니라는 단서가 붙어 있다. 적용은 간단해서 torch._functorch.config의 activation_memory_budget_solver 값을 바꿔주면 된다.
선택의 기준은 명확하게 정리된다. 정확성보다 속도가 우선이라면 greedy_knapsack을, SciPy 의존성을 추가해도 괜찮다면 ilp_knapsack을 쓰는 편이 낫다. ilp_knapsack은 기존 어떤 DP 구현보다 훨씬 빠르면서도 그리디와 달리 정확한 해를 보장하고, 그리디보다 약간 느릴 뿐이다. 패키지 매니저로 설치한 파이토치를 그냥 쓴다면 기본값인 dp_knapsack이 적용되므로 아무것도 할 필요가 없다. 그 외의 경우, 즉 외부 의존성은 피하면서 메모리 한계를 늘리고 싶을 때 dp_knapsack_sliding_hirschberg가 제자리를 찾는다.
실무적으로 보면 이 변화는 거대 모델을 다루며 컴파일·메모리 플래닝 단계에서 OOM을 겪어본 사람에게 특히 와닿는다. 모델이 커질수록 저장·재계산을 가를 연산의 수가 늘고 냅색 품목 수도 함께 불어나는데, 기본 솔버는 바로 그 지점에서 무너진다. 다만 ilp_knapsack이라는 더 빠른 대안이 존재하는 만큼, SciPy 도입 가능 여부가 실제 선택을 가르는 기준이 된다는 점은 염두에 둘 만하다. 이 변경은 Edward Yang의 코드 리뷰를 거쳐 병합됐고, 작성자는 사용 중 버그를 발견하면 이슈로 등록해 자신(@jmaczan)을 태그해달라고 덧붙였다.