Gefen: 최적화된 확률적 옵티마이저 (Optimized Stochastic Optimizer)

AdamW는 현대 딥러닝의 기본 옵티마이저 (optimizer)이지만, 1차 및 2차 모멘트 상태 (first and second moment states)로 인해 훈련 메모리에 파라미터 크기의 약 2배에 달하는 버퍼를 추가합니다. 우리는 파라미터 블록 간에 2차 모멘트 추정치를 자동으로 공유하고 학습된 코드북 (codebook)을 사용하여 1차 모멘트를 양자화 (quantize)함으로써, 동일한 성능을 유지하면서 AdamW의 메모리 사용량을 약 8배 줄이는 메모리 효율적인 옵티마이저인 Gefen을 제안합니다. 이는 10억 개의 파라미터당 6.5 GiB의 감소에 해당합니다. 이 방법은 큰 혼합 헤시안 (mixed Hessian) 항목이 제곱 기울기 (squared gradients)의 비율을 1에 가깝게 제한한다는 이론적 결과에서 영감을 얻었으며, 이는 헤시안 정렬 파라미터 (Hessian-aligned parameters)가 2차 모멘트 통계량을 공유하기 위한 자연스러운 후보임을 시사합니다. 헤시안을 계산하는 것은 대규모 환경에서 비실용적이므로, Gefen은 초기 제곱 기울기로부터 블록 구조를 추론하며, AdamW 기본값 외에 아키텍처별 메타데이터나 하이퍼파라미터 (hyperparameters)를 요구하지 않습니다. Gefen은 정확한 히스토그램 기반 동적 계획법 (dynamic-programming) 양자화 코드북을 학습하며, 1차 모멘트 스케일링 (scaling)을 위해 동일한 블록을 재사용합니다. 다양한 실험을 통해 Gefen은 비교된 AdamW 유사 방법들 중 가장 낮은 피크 옵티마이저 메모리 (peak optimizer memory)를 달성하면서도 AdamW 수준의 성능을 유지합니다. FSDP 및 DDP 훈련에서 감소된 메모리 사용량은 더 큰 마이크로배치 (microbatches)를 가능하게 하고 AdamW 대비 처리량 (throughput)을 크게 향상시켜, 메모리 사용량을 낮추면서 처리량을 높이고 더 큰 모델을 훈련하거나 더 큰 배치 크기를 사용할 수 있게 하는 실용적인 드롭인 교체 (drop-in replacement) 수단을 제공합니다. 우리는 https://github.com/ndvbd/Gefen 에서 제공되는 융합된 CUDA 커널 (fused CUDA kernels)을 포함한 완전한 Python 구현을 제공합니다.