Adam 옵티마이저 상태를 업데이트 오차 보증으로 8비트 압축한다

무엇인가

이 논문은 Adam으로 학습한 모델의 체크포인트에서 옵티마이저 상태, 즉 1차·2차 모멘트를 어떻게 안전하게 압축할지를 다룬다. Adam은 파라미터마다 FP32 모멘트 두 개를 저장해 가중치의 두 배 바이트를 쓰는데, 2차 모멘트가 업데이트의 제곱근 분모에 들어가므로 작은 복원 오차도 다음 업데이트를 크게 바꿀 수 있다. 저자들은 모멘트 자체의 오차가 아니라 복원된 모멘트가 만들어 내는 업데이트를 제약해야 한다고 주장한다. 구체적으로 |g|≤C인 모든 그래디언트에 대해 원래 FP32 Adam 업데이트와 복원된 업데이트의 차이가 학습률 단위로 τ 이하가 될 때만 해당 좌표의 압축 코드를 수용하고, 실패한 좌표는 FP32 원본 그대로 escape한다. 계약은 C=0.2, Torch 2.8 스칼라 FP32 Adam, β1=0.9, β2=0.999, ε=1e-8, weight decay 없음이라는 고정된 연산 모델을 전제로 한다.

어떻게 동작하나

인코더는 256개 좌표 블록마다 두 모멘트에 8비트 코드를 제안하고, 독자가 복원할 방식 그대로 디코드한 뒤 두 단계 guard를 적용한다. 먼저 이상적 실수 경계를 유향 반올림으로 계산하고, 다음으로 실제 FP32 연산 경로를 반영한 native 경계를 계산한다. 핵심은 원본과 복원 업데이트가 같은 그래디언트와 파라미터를 공유한다는 점이다. 두 Adam 몫의 차이를 직접 묶으면 분자 차이에서 αg가, 업데이트 차이에서 w가 상쇄되어 오차가 모멘트 복원 오차 |m−m~|, |v−v~|에 비례하게 된다. 여기에 각 FP32 연산의 반올림 반경 R(M)=8u32M+4nmin을 전파하고, 선형 경계와 2^-40 C까지의 로그 경계를 섞은 고정 파티션에서 최대 오차를 평가한다. 통과하지 못한 좌표는 64, 256, 1024개 구간으로 재확인한 뒤 exact FP32로 escape한다. 1차 모멘트는 대칭 선형 8비트, 2차 모멘트는 블록별 양수 범위의 로그 8비트 코드를 쓰며, 직렬화 페이로드는 헤더, 블록 테이블, 두 코드 배열, escape 비트마스크, escape된 FP32 쌍, SHA-256을 포함하고 독립 디코더가 전체 해시를 재현해야 한다.

무엇과 다른가

논문은 먼저 압축이 어디서 오는지를 100만 개 OPT-125M 좌표에서 분해한다. 두 업데이트의 범위를 따로 계산하는 전통적 구간 산술은 988,820개(98.9%)를 escape시켜 페이로드가 원본 FP32의 126.63%까지 커졌다. 공유 그래디언트를 보존한 차이 경계를 쓰면 escape가 439,362개(71.68%)로 줄고, 여기에 혼합 선형·로그 파티션을 더하면 5,324개(28.28%)까지 떨어진다. 코드와 메타데이터는 모든 설정에서 2,219,592바이트로 같고 escape 하나가 8바이트를 더하므로, 경계의 tightness가 페이로드를 결정한다. 따뜻한 인코딩 중앙값도 공유 경계가 0.161초로 separate-range의 0.435초보다 빨랐다. H100 감사에서는 1,048,576개 스칼라 사례와 131,072개 파티션 쌍에서 enclosure 실패가 0이었고, 반경을 일부러 빼면 899,311건이 검출됐다.

어떻게 쓰나

동일한 guard를 적용한 베이스라인 비교에서 8비트 로그 코드는 원본 FP32의 28.28%를 차지해 guarded BF16보다 45.16%, guarded blockwise 8비트보다 28.82% 더 작았다. guard 없는 blockwise 8비트는 더 작지만 계약을 위반한다. 선언된 영역의 합성 그래디언트 혼합에서 100만 좌표 중 2,469개가 τ를 넘었고 최대 정규화 오차가 0.1294였다. 같은 guard를 적용하면 127,396개가 escape되어 최대 오차가 0.0098로 내려가고 고정밀 점 witness도 통과했다. BF16은 이 부분집합에서 escape가 거의 필요 없지만 16비트 값 자체가 FP32의 절반 바이트를 차지한다. 독립 코드 선택과 공동 선택의 차이는 3바이트에 불과했다.

전제와 한계

완전한 체크포인트 실험은 OPT-125M 네 시드(701, 707, 708, 709)를 학습 10%, 50%, 100%(t=3,125, 15,625, 31,250)에서, ResNet-18 네 시드(701–704)를 CIFAR-10 224 해상도 20에폭(t=1,250, 6,250, 12,500)에서 저장해 총 120개 조건을 복원했다. 로그 코드는 모멘트 페이로드를 FP32 대비 OPT에서 71.67%, 69.44%, 67.35%, ResNet에서 72.18%, 71.85%, 71.56% 줄였다. 같은 guard의 BF16 대비로는 OPT 45.08%, 40.76%, 36.69%, ResNet 46.05%, 45.41%, 44.85% 작았고, 같은 guard의 blockwise 8비트 대비로는 OPT 27.33%, 29.08%, 31.49%, ResNet 26.90%, 42.64%, 43.28% 작았다. Escape율은 guarded blockwise 8비트가 OPT에서 12%에서 21%, ResNet에서 11%에서 23%였지만 로그 코드는 각각 0.6~4.9%, 0.08~0.7%였다. 모멘트가 아닌 전체 직렬화 컨테이너 기준 감소는 OPT 48.77%, 48.46%, 47.09%, ResNet 52.39%, 52.22%, 51.83%로 더 작다.

120개 조건 모두 두 모멘트 벡터를 복원했고 첫 업데이트 감사를 통과했다. 실제 첫 업데이트 차이의 최댓값은 OPT 0.009992, ResNet 0.009991 학습률로 τ=0.01보다 작았다. 이후 각 분기가 자체 그래디언트로 100회 업데이트한 뒤 개발 손실은 FP32 분기와 가까웠다. OPT에서 공동 선택과 FP32의 평균 차이는 세 비율에서 -1.41e-4, -4.18e-5, -3.68e-5였고 50%, 100%에서는 시드별 부호가 갈렸다. ResNet은 BF16을 포함한 모든 압축 방식에서 차이가 더 컸고 공동 선택은 평균 +0.0087, +0.0012, +0.0030이었다. 저자들은 이를 서술적 관찰로 한정한다. 공동 코드 탐색은 OPT에서 모멘트 바이트를 0.000526~0.001365%만 더 줄였고(시드 평균 0.000971%), ResNet에서는 체크포인트당 11~195바이트(평균 0.000370%)만 줄였으며, 코덱 시간은 1.060배, 1.059배로 늘어 저자들은 독립 반올림을 권장한다. 코덱 시간은 로그 코드 인코딩이 FP32 직렬화의 7.2배(OPT), 7.6배(ResNet)였고 복원은 FP32 복원의 0.88배, 0.77배였다. 인코딩과 복원 합계는 guarded BF16의 2.31배, 2.58배지만 guarded blockwise 8비트의 0.72배, 0.83배로 바이트와 시간 모두 앞선다. OPT 인코딩 피크 할당은 24.9GiB로 guarded BF16 19.2GiB, guarded 8비트 17.4GiB보다 높다.

가장 중요한 한계는 보증의 범위다. 이 보증은 고정된 제한적 연산 모델, 즉 Torch 2.8 스칼라 FP32 Adam의 특정 설정 아래에서 첫 복원 업데이트에만 적용되는 조건부 enclosure이며, Torch나 임의 생성 커널을 기계적으로 검증한 증명은 아니다. 복원 직후에는 두 분기가 같은 그래디언트와 메타데이터를 공유해 상쇄가 성립하지만, 이후에는 각자 그래디언트를 만들므로 궤적 전체에 대한 경계가 되지 못한다. 100회 연속 학습은 그 결과를 경험적으로 관찰한 것이다. 실험도 OPT-125M과 ResNet-18, 아키텍처당 네 시드에 한정되고 신뢰구간 대신 기술적 평균과 시드 범위를 보고한다. 타이밍은 코덱 경계만 분리해 전체 컨테이너 직렬화, 파일 전송, 재개 학습을 제외했고 호스트 공유 가능성도 명시했다. 또한 이 방식은 체크포인트 표현을 압축하는 것이지 학습 내내 옵티마이저 재귀 자체를 압축 상태로 유지하는 것은 아니다.