생성 프로그램 순열로 추론 피크 메모리를 바이트 단위로 인증한다

무엇인가

이 논문은 작은 기기에서 실시간 추론을 할 때 활성화 메모리를 늘리지 않고 줄이는 문제, 특히 연산자 실행 순서를 다룬다. 순서 변경은 시간 비용이 없어 가장 값싼 메모리 레버지만, 기존에는 컴파일러 그래프에서 더 좋은 순서를 찾아 코드를 다시 생성하면 예측한 만큼 좋아질 것이라고 가정했다. 논문은 이 가정이 PyTorch Inductor에서 깨진다고 측정으로 보인다.

어떻게 동작하나

제안 방법은 컴파일러가 만든 프로그램을 그대로 두고, 한 번 프로파일링해 각 문장이 어떤 저장소를 읽고 쓰는지 기록한다. 커널 호출 문장을 단위로 묶고, 할당 문장은 그 버퍼를 처음 만지는 단위에 붙이며, 뷰·튜플 투영·del은 방출 시 의존성에 맞게 배치한다. 저장소별 read-after-write, write-after-read, write-after-write 간선으로 의존성을 만들고, 같은 버퍼 세대의 읽기끼리는 재정렬을 허용한다. 이후 단위 순서를 바꿔 원문 문장을 각각 한 번씩 방출하고, 할당은 단위 직전에, 별칭과 del은 가능한 이른 곳에 둔다. 커널·인자·버퍼 재사용은 건드리지 않는다. 논문은 여섯 조건(효과 완전성, 소유권과 완료 기록, 결정적 단위, 할당 순서 기록, 프레임 조건, 방출 시 추가 이벤트 없음)이 성립하면 모델링한 피크가 바이트 단위로 실현되고 결과가 비트 단위로 같다는 정리 1을 제시한다.

무엇과 다른가

멈출 시점을 알기 위해 포워드와 리버스 트랩 하한을 준다. 각 하한은 O(n^2)번의 최소 컷으로 계산되고, 고전적 이벤트 컷 하한을 지배하며 서로 비교 불가능하다. 학습·추론 그리드의 122개 래퍼 그래프 중 115개(94%)에서 찾은 순서가 최적임을 1초 이내에 인증했고, 포워드만 99개, 리버스만 112개를 인증했다. 800개 랜덤 DAG에서 max(F,R)은 773개에서 정확했고 이벤트 컷은 671개였다. 최악 비율은 랜덤 1.23, 13정점 적대적 탐색 1.50이다. 검색은 최대 259개 단위에 대해 이상적 순서를 깊이 우선 탐색하고 최소 컷 패킷 폐포를 쓰는 정확 용량 탐색이다.

어떻게 쓰나

실험은 PyTorch 2.14.1 Inductor, CPU FP32 단일 스레드, GPU RTX 2080 Ti CUDA 12.6에서 했다. 더 좋은 순서를 컴파일러에 돌려 코드를 재생성하면 CPU 훈련 54개 구성 중 22개가 예측 피크를 놓쳤고 9개는 기존 프로그램보다 나빴다. GPU에서는 50개 중 19개가 기존보다 나빴고 수치 결과도 바뀌었다. 반면 생성 프로그램을 직접 순열하는 고정 래퍼 실현은 CPU 훈련 78회와 추론 22회 모두 바이트 단위로 정확했고 출력과 그래디언트가 비트 단위로 같았다. GPU에서는 CUDA 캐싱 할당자가 계약 조건을 깨므로 측정 가드를 두고, 측정 피크가 개선될 때만 배포해 회귀 없이 평균 0.92%를 절약했다.

전제와 한계

서버급 추론 22개 구성에서는 DeepLabV3-ResNet50 512×512만 10.0% 줄었고, 인증서는 20개 그래프에서 기존 순서가 최적임을 증명했다. 마이크로컨트롤러급 16개 모델은 모두 모델 순서가 최적이었고, 피크는 단일 연산자의 입력과 출력이 만드는 하한과 같아서 전체 텐서 재계산으로는 줄일 수 없었다. 컴파일 훈련에서는 최대 7.6%를 절약했지만, 모든 접근을 쓰기로 보수적으로 처리하면 4.6%, 4.7%, 2.3% 절약이 사라지고 인코더-디코더는 7.6%에서 6.2%로 줄었다. 절약은 임계 역방향 호출의 수명에서 나오며 할당량 자체는 같다.

인증된 최적보다 더 낮춰야 할 때는 연산자별 재계산 상한을 둔 선택적 재계산을 다룬다. 논문은 균일 상한 k에 대해 트리폭 2에서도 강한 NP-완전임을 증명하고, 시리즈-병렬 그래프에서는 배낭 문제로 정확히 인코딩되며, 유계 트리폭에서는 f(t,K,B) poly(n,L) 시간의 정확 알고리즘이 있지만 B에 의존해 실용적이지 않다고 밝힌다. 프로토타입은 폭 2, 6정점, K=2, B=8에서 3×10^6 상태를 넘는다. ReDimNet 7개 크기에서는 순서가 이미 최적이었지만, 작은 입력 특징을 남기고 스템 체인을 한 번 재계산하면 추가 MAC 0.02~0.34%로 피크를 5.2~6.1% 낮추고 출력은 비트 단위로 같았다. B3와 B6에서는 탐욕적 아레나 배치가 단편화로 이득을 잃었고, 내부 배포 모델에서는 최대 15% 절감을 문맥으로만 보고한다.

개발자에게 이 논문은 컴파일러의 피크 메모리 재정렬 결과를 그대로 믿지 말고, 실제 생성된 프로그램에서 측정하고 필요하면 문장 수준으로 순서를 바꿔 검증하라는 지침을 준다. 특히 분기 네트워크는 재정렬이 무료로 큰 이득을 주고, 대부분의 체인형 추론 그래프는 이미 최적이므로 인증서로 탐색을 멈출 수 있다. 작은 모델은 단일 연산자 피크가 지배해 전체 텐서 재계산이 무용하고, 타일 단위로 쪼개는 접근이 필요하다. ReDimNet처럼 늦은 재사용을 위해 비싼 값을 오래 들고 있는 구조는 선택적 재계산이 효과적이다.

한계로 정리 1의 정확성은 요청한 크기 그대로 할당하는 CPU와 정적 아레나에 성립하고, CUDA 캐싱 할당자는 측정 가드로 대체한다. 라이브러리 커널의 효과와 할당 순서는 증명이 아니라 프로파일링과 재측정으로 확인하며, 모든 모델은 하나의 입력 형상에 특화된다. 마이크로컨트롤러 연구는 펌웨어가 아니라 연산자 그래프와 탐욕적 아레나 플래너를 평가했다. 재정렬 이득은 분기 네트워크를 빼면 훈련에서 1~3%, 대부분 추론 그래프에서 0에 가깝다. OPT/max(F,R)의 유계성은 열린 문제이고, 유계 트리폭 알고리즘은 상태 폭발로 실용적이지 않다.