활성값 재계산의 숨은 상태를 계약으로 고정하고 보존 비용까지 함께 계획한다

무엇인가

활성값 재계산(activation recomputation, 흔히 gradient checkpointing)은 역전파 중 순전파 구간을 다시 실행해 액티베이션 메모리를 아낀다. 문제는 재실행이 "원래 순전파가 본 상태"를 봐야만 한다는 것이다. 뮤터블 버퍼는 재실행 전에 바뀔 수 있고, 별칭(alias)은 서로 다른 인덱스로 같은 바이트를 노출하며, 난수 추출은 텐서 의존 그래프에 나타나지 않는다. 그래서 코드를 그대로 다시 실행해도 미분 대상이 다른 계산이 되어버릴 수 있다. 저자들은 정확성과 메모리 계획이 서로 맞물려 있다고 본다. 전체 상태 스냅샷은 재실행을 항상 유효하게 만들지만 비싸고, 그 비용을 계획 모델에서 빼면 플래너가 자기 예산을 넘겨버린다.

어떻게 동작하나

논문의 핵심 장치는 "버전이 붙은 리플레이 계약(replay contract)"이다. 논리적 값은 주소가 아니라 그 값을 만들어낸 쓰기(write)로 식별한다. 연산은 자신의 쓰기 이전에 유효했던 버전을 읽으므로, 소스와 목적지 집합이 겹칠 때 최종 버퍼를 보고 답을 낼 수 없다. 각 쓰기마다 새 노드를 만들고 피연산자의 버전 식별자·계수·샘플링된 노이즈·비선형성을 보존하며, 역방향 미분은 v̄_p(k) += v̄_k · f′_k(v_p(k)) 형태로 누적된다. 여기에 별도로 바이트 수준 요구가 붙는다. 재실행 전후의 저장 바이트와 난수 생성기 상태가 같아야 한다는 식 P(B0,R0) = (B_T,R_T)이다. 최종 바이트가 같아도 생성기가 다른 지점까지 진행돼 있으면 이후 실행이 달라지고, 타입이 붙은 별칭은 물리 저장소를 공유하므로 한 뷰의 수치적 일치만으로는 충분하지 않다.

무엇과 다른가

보존 방식은 여러 엔진으로 구현된다. 선언된 버퍼를 통째로 뜨는 전체 스냅샷, 덮어쓴 값을 기록해 두고 기록된 순서대로 undo/redo 하는 영역 저널(region journal), 과거 피연산자를 명시적으로 저장하는 read-set 보존, 가드된 영역 등이다. 영역 저널에서 주의할 점은 자기 세그먼트가 덮어쓴 값만 저장하면 부족하다는 것이다. 이후 세그먼트가 또 다른 필요한 영역을 파괴할 수 있어서, 자신과 이후의 파괴적 쓰기 전체의 합집합을 보수적으로 보존한다. 플래너는 세그먼트 i마다 옵션 o를 골라 유지 바이트 r_i,o, 임시 바이트 q_i,o, 시간 c_i,o 를 청구하고, P + Σr + max q ≤ M 제약 아래 Σc를 최소화한다. 동적 계획법은 (유지 메모리, 최대 임시 메모리, 비용)의 파레토 프런티어를 잘라내지 않고 유지하는데, 임시 메모리는 뒤 세그먼트가 최댓값을 바꿀 수 있어 합산 메모리만으로 지배 관계를 판단하면 위험하기 때문이다. 실험에는 두 개의 대조군이 붙는다. 보존 비용을 목적함수에서 빼는 상태 무관(state-agnostic) 플래너와, 같은 정확한 비용을 쓰되 선택지를 eager와 전체 스냅샷으로 제한한 "정확히 계산된 eager/whole" 플래너가 그것이다.

어떻게 쓰나

의미론 검증은 상당히 촘촘하다. 1,024개 프로그램 수치 스위트에서 스칼라 출력은 정확히 재현되고, 정방향 야코비안 기준 최대 gradient 불일치는 8.88×10⁻¹⁶이었다. 최종 저장소를 읽는 대조군은 모든 프로그램에서 검출됐고, 별도의 1,024개 바이트 스위트에서는 스냅샷 재실행과 영역 undo 두 경로 모두 모든 입력과 일치했으며 잘못된 RNG 상태에서의 재실행은 전 입력에서 검출됐다. Torch/JAX 코호트는 시드 50개(420000–420049)로 250개 프로그램, 5,900개 실행 기록(Torch 3,500, JAX 2,400)을 남겼다. Torch의 네이티브 체크포인트 두 변형은 로컬 뷰 조건 200개는 모두 맞혔지만 히든 버퍼 조건 200개에서는 전부 어긋났고, 선언 상태 트랜잭션은 두 계열 모두 일치했다. 로우 오버래핑 카피 100개 조건은 eager 실행조차 거부했으며, RNG 보존을 끄면 비카피 조건 400개 전부에서 발산했다. JAX는 함수형 조건 500개를 모두 맞혔고 인플레이스 실행과 impure-holder 체크포인팅은 거부했다. 별도로 1,200개 프로그램·5개 보존 엔진에서 최대 gradient 오차 1.56×10⁻¹⁵, 96개 프로그램 4,544개 좌표의 complex-step 검사에서 최대 불일치 1.00×10⁻¹⁵, 조인트 별칭 스트레스(시드 50, 워크로드 3종, 300개 입력, 4,800개 저장 연산)에서는 지원되는 1,500개 복구가 바이트·난수 상태·별칭 저장소·반복 역전파 결과에서 모두 일치했다. 정수 쓰기에 float gradient를 요구하는 300개 결과는 값을 지어내는 대신 "미지원"으로 분류해 남겼다.

전제와 한계

성능은 워크로드에 따라 갈린다. H100 캘리브레이션 24개 형상에서 1,024MiB 버퍼, 쓰기 비율 0.001, 라이브 버전 64일 때 영역 보존은 전체 스냅샷 피크 메모리의 5.05%, 15.04%, 6.48%만 썼고 스텝 시간도 각각 38.99%, 60.79%, 48.82%에 그쳤다(큐 워크로드 기준 피크 94.95% 감소, 스텝 시간 61.01% 감소). 그러나 쓰기 비율이 1이면 영역 보존과 스냅샷의 피크 메모리가 같아지고 부기(bookkeeping) 시간만 늘었으며, 캘리브레이션한 모든 형상에서 eager 실행이 영역 보존보다 빠르고 피크 메모리도 작았다. 20GiB 예산에서는 두 플래너 모두 전 구간 eager를 선택했다. 압력이 걸린 조건에서는 결과가 달라진다. 32GiB에서 조인트 플래닝은 29.936GiB로 고정 영역(26.200GiB)보다 메모리를 더 쓰면서 4.98% 빨랐고(시간 비율 0.950163, 구간 [0.949090, 0.951220]), eager/whole 플래너는 이 예산에서 실행 가능한 정적 정책 자체가 없었다. 44GiB에서는 정확히 비용을 계산한 eager/whole 대조군 대비 50개 신규 프로세스 실측으로 평균 스텝 시간 18.16% 감소(비율 0.818351, 구간 [0.817193, 0.819465]), 피크 할당 1.08% 감소(비율 0.989239)를 기록했다. 다만 1,024MiB·쓰기 비율 0.001·라이브 세그먼트 64의 3워크로드 150개 독립 실행 코호트에서는 eager 유지가 모든 균일 재실행 방식보다 빠르고 작았고, 두 플래너 모두 전부 eager를 골랐다.

응용 워크로드 실험이 이 경계를 더 분명히 보여준다. CIFAR-100에 맞춘 ResNet18 쿼리·모멘텀 인코더를 2뷰 InfoNCE와 64MiB FP32 네거티브 큐로 학습시키고, 7개 방법을 6개 시드(2701–2706)에서 100 에폭 인코더 + 25 에폭 프로브로 돌려 42개 조건을 만들었다. 7개 방법의 테스트 점수 파일은 시드 내에서 바이트 단위로 동일했고, 독립 CPU 재계산 정확도는 27.16%, 23.70%, 26.62%, 27.40%, 23.75%, 26.44%였다. 그런데 자원 쪽은 역전된다. 관측한 5개 시드에서 조인트 플래닝은 1,229.74MiB를 할당해 지연 eager(1,389.41MiB)보다 11.49% 적게 썼지만, 전체 스냅샷은 1,004.62MiB로 조인트 플래닝보다 22.41% 적었다. 전체 인코더 시간 비율은 기하평균 1.008024(시드 범위 [0.984483, 1.019058])로, 네 시드에서는 더 느리고 시드 2706에서만 1.55% 빨랐다. 즉 조밀한 큐 응용에서는 영역 보존이 이점을 주지 못했다.

개발자 입장에서 이 논문의 실용적 규칙은 명확하다. 액티베이션이 원하는 eager 예산을 넘고, 쓰기가 희소하며, 과거 버전이 여러 개라 전체 스냅샷이 비쌀 때 영역 보존이 값을 한다. eager 유지가 예산 안에 들어오면 대개 그것이 여전히 빠르고 작다. 그리고 재실행을 켤 때는 프레임워크가 무엇을 복원한다고 약속하는지 확인해야 한다. Torch 네이티브 체크포인트가 히든 버퍼 조건 200개에서 전부 어긋났다는 결과는, 프레임워크가 문서화된 순수 함수 가정을 지키는지의 문제이지 프레임워크 버그라는 주장이 아니라는 점도 함께 읽어야 한다. 뮤터블 상태나 인플레이스 연산이 섞인 모델이라면 재실행 경계를 선언적으로 관리하는 편이 안전하다.

한계도 저자들이 직접 밝힌다. 리플레이는 선언된 상태로 제한되며 임의의 별칭과 히든 버퍼를 자동으로 발견하지 못한다. Coop·Enzyme 어댑터는 구현되지 않았고, ZeRO-Offload 같은 오프로딩은 정적 정책 공간 밖이다. 예산은 논리적 제약일 뿐 실제 디바이스 상한이 아니어서 예산 초과를 OOM으로 서술하지 않는다고 명시한다. 정적 모델은 할당자 수준 최적성을 보장하지 않고, 예측 피크는 실측보다 약 1GiB 높게 나오는 근사다. 응용 평가의 독립 단위는 시드 6개이며 완전 프로세스 쌍별 관측은 5개뿐이라 교차 시드 신뢰구간이나 유의성 주장을 하지 않는다. 44GiB·32GiB 압력 코호트와 150 프로세스 상태 비용 코호트는 서로 다른 체제라 합산하지 않는다.