Memory-Efficient Activation Checkpointing with Sliding Window and Hirschberg's Algorithm for 0/1 Knapsack Solving in PyTorch

당신이 세계에서 가장 맛있고 복잡한 케이크를 구우려고 노력 중이라고 상상해 보세요. 하지만 당신에게는 아주 작고 비좁은 주방뿐입니다. 당신은 나중에 이 과정을 완벽하게 되돌려 케이크가 어떻게 만들어졌는지 확인하기 위해, 혼합한 모든 재료와 모든 온도 변화, 그리고 모든 휘젓는 동작 하나하나를 기록해야 하는 레시피를 가지고 있습니다. 문제는 당신의 주방 조리대(당신의 컴퓨터 메모리)가 이 모든 노트를 담기에는 너무 작다는 것입니다. 만약 모든 것을 다 적으려고 한다면 조리대가 넘쳐버릴 것이고, 결국 당신은 베이킹을 중단해야 합니다. 이것이 거대한 인공지능 모델을 훈련시키는 과학자들이 매일 겪는 고충입니다. 그들은 AI를 가르치기 위해 많은 단계를 기억해야 하지만, 그들의 컴퓨터는 공간이 부족합니다. 이를 해결하기 위해 그들은 "액티베이션 체크포인팅(activation checkpointing)"이라는 영리한 기술을 사용합니다. 모든 단계를 다 적는 대신, 가장 중요한 단계들을 골라 저장하고, 덜 중요한 것들은 나중에 다시 수행하기로 약속하는 것입니다. 이는 마치 작은 사진 앨범에 남길 사진을 결정하고, 어떤 사진은 나중에 잊어버리더라도 다시 찍을 수 있는 것들로 정하는 것과 같습니다. 목표는 그 작은 주방 안에 케이크를 굽는 전체 과정을 담아내면서도 레시피의 마법을 잃지 않는 것입니다.

오랫동안 많은 AI 과학자들이 모델을 구축할 때 사용하는 특정 프로그램인 PyTorch는 어떤 단계를 저장할지 결정하는 특정한 방식을 가지고 있었습니다. 그것은 이 결정을 "0/1 배낭 문제(0/1 Knapsack Problem)"라는 고전적인 퍼즐처럼 취급했습니다. 당신이 일정한 무게만큼만 담을 수 있는 배낭을 멘 등산객이라고 상상해 보세요. 당신에게는 각각의 무게와 가치(그것이 당신에게 얼마나 도움이 되는지)를 가진 물건 목록이 있습니다. 당신은 배낭을 망가뜨리지 않으면서 가장 높은 가치를 주는 물건들을 골라야 합니다. PyTorch의 기본 방식은 이 문제를 해결하기 위해 모든 가능한 조합을 거대한 종이에 적으려는 것과 같았습니다. 이 방식은 완벽하며 최선의 답을 찾아내지만, 그 종이가 너무 거대해져서 컴퓨터의 메모리가 폭발하고 프로그램이 충돌하게 만들었습니다. 연구자들은 만약 선택해야 할 항목이 단 100개뿐이라 하더라도, 필요한 종이의 크기가 너무 커서 그들의 기계에 있는 64기가바이트보다 훨씬 많은 304기가바이트의 공간이 필요하다는 것을 발견했습니다. 그것은 방 안에 들어갈 수 없는 완벽한 해결책이었습니다.

이 논문에서 저자는 이 퍼즐을 해결하는 더 똑똑한 방법인 dp_knapsack_sliding_hirschberg를 소개합니다. 모든 거대한 종이를 한꺼번에 적으려고 하는 대신, 그들은 "슬라이딩 윈도우(sliding window)" 기술을 사용합니다. 당신이 긴 책을 읽고 있지만, 한 번에 두 페이지씩만 보여주는 작은 돋보기를 가지고 있다고 상상해 보세요. 당신은 돋보기를 아래로 밀어 내려가며 두 페이지를 보고, 그다음 두 페이지를 보고, 또 그다음 두 페이지를 보는 식입니다. 이렇게 하면 당신은 매 순간 오직 두 페이지만을 머릿속에 담고 있으면 되므로, 엄청난 정신적 공간을 절약할 수 있습니다. 하지만 단순히 두 페이지를 보는 것만으로는 전체 이야기를 기억하기에 부족합니다. 어떤 특정 항목들을 골라야 하는지 알아야 하기 때문입니다. 이를 해결하기 위해, 그들은 슬라이딩 윈도우를 "히르슈베르크 알고리즘(Hirschberg's algorithm)"이라는 오래되고 영리한 전략과 결합합니다. 이것은 "분할 정복(divide and conquer)" 게임이라고 생각하면 됩니다. 배낭 문제를 한꺼번에 해결하는 대신, 항목 목록을 절반으로 나눕니다. 왼쪽 절반을 해결하고, 오른쪽 절반을 해결한 다음, 두 개의 최선책을 어떻게 결합할지 알아냅니다. 이 과정을 재귀적으로 반복하여, 문제를 아주 작은 조각들로 쪼개어 아주 쉽게 해결할 수 있을 때까지 나눕니다. 이 모든 과정은 아주 적은 양의 메모리만을 사용하면서 이루어집니다.

이 새로운 방식의 결과는 인상적입니다. 저자는 64기가바이트의 RAM을 가진 컴퓨터에서 이를 테스트했습니다. 기존 방식은 단 100개의 항목을 해결하려 할 때 충돌이 발생했지만, 새로운 방식은 2,000개의 항목을 가진 문제를 최대 58.4기가바이트의 메모리를 사용하여 성공적으로 해결했습니다. 이는 컴퓨터가 이전보다 20배 더 큰 문제를 처리할 수 있게 되었음을 의미합니다. 게다가, 이 새로운 방식은 단순히 메모리를 아끼는 것뿐만 아니라 더 빠릅니다. 저자는 특정 기기에서 동일한 퍼즐을 1,000번 실행하여 테스트했고, 새로운 해결사가 기존의 것을 일관되게 앞지른다는 것을 발견했습니다. 저자는 이 새로운 방식이 답을 추측하여 약간의 오류를 범할 수 있는 다른 "임시방편" 방식들과 달리, 매번 정확하고 완벽한 해답을 찾아낸다는 점을 강조했습니다. 이 방식은 기존 방식만큼 정확하면서도 훨씬 더 효율적입니다.

이 논문은 이 새로운 접근 방식이 단순한 이론이 아님을 확인시켜 줍니다. 이는 PyTorch 소프트웨어에 성공적으로 병합되었으며 버전 2.10에서 사용할 수 있습니다. 저자는 슬라이딩 윈도우와 분할 정복을 결합함으로써, AI 모델의 성장을 가로막던 메모리 병목 현상을 해결할 수 있음을 보여줍니다. 저자는 이것이 유일한 방법이라고 주장하거나 모든 종류의 컴퓨터 퍼즐에 작동한다고 제안하는 것이 아닙니다. 다만, AI 단계를 저장하기로 결정하는 특정 작업에 있어서는, 기존 방식의 완벽한 정확성을 유지하면서도 메모리 충돌을 제거하는 검증된, 정확하고 매우 효율적인 업그레이드임을 보여줍니다. 이 논문은 기존 방식이 대규모 모델에 충분하지 않다는 점을 입증하며, 항목 수가 너무 많아지면 실패한다는 것을 명확히 보여줍니다. 대신, 그들은 완벽한 정확성을 유지하면서 메모리 충돌을 제거하여, 과학자들이 그 작은 주방에서도 더 크고 복잡한 AI 케이크를 구울 수 있게 해주는 해결책을 제시합니다.