FlashOptim: 메모리 효율적인 학습을 위한 최적화 기법

FlashOptim: Optimizers for Memory Efficient Training

신경망의 표준 혼합 정밀도 학습은 각 모델 파라미터마다 상당한 양의 가속기 메모리를 필요로 합니다. 이러한 메모리는 파라미터 자체뿐만 아니라, 그라디언트와 하나 이상의 최적화 상태 변수를 포함합니다. 일반적으로 각 값이 4바이트를 필요로 하기 때문에, 70억 개의 파라미터를 가진 모델을 학습하는 것조차 가속기 메모리가 100GB 미만인 연구자에게는 비현실적일 수 있습니다. 저희는 모델 품질과 API 호환성을 유지하면서 파라미터당 메모리 사용량을 50% 이상 줄이는 최적화 모음인 FlashOptim을 소개합니다. 저희의 접근 방식은 두 가지 핵심 기술을 도입합니다. 첫째, 양자화 오차에 대한 엄격한 상한을 찾아 분할된 마스터 가중치를 개선합니다. 둘째, 8비트 최적화 상태 양자화에서 발생하는 오류를 크게 줄이는 컴판딩 함수를 설계합니다. 16비트 그라디언트와 함께 이러한 기술은 AdamW의 파라미터당 메모리 사용량을 16바이트에서 7바이트로 줄이거나, 그라디언트 해제 시에는 5바이트로 줄입니다. 또한 모델 체크포인트 크기를 절반 이상 줄입니다. SGD, AdamW, Lion에 FlashOptim을 적용한 실험 결과, Llama-3.1-8B 미세 조정 작업을 포함한 다양한 표준 시각 및 언어 벤치마크에서 측정 가능한 품질 저하가 발생하지 않았습니다.

Standard mixed-precision training of neural networks requires many bytes of accelerator memory for each model parameter. These bytes reflect not just the parameter itself, but also its gradient and one or more optimizer state variables. With each of these values typically requiring 4 bytes, training even a 7 billion parameter model can be impractical for researchers with less than 100GB of accelerator memory. We introduce FlashOptim, a suite of optimizations that reduces per-parameter memory by over 50% while preserving model quality and API compatibility. Our approach introduces two key techniques. First, we improve master weight splitting by finding and exploiting a tight bound on its quantization error. Second, we design companding functions that greatly reduce the error in 8-bit optimizer state quantization. Together with 16-bit gradients, these techniques reduce AdamW memory from 16 bytes to 7 bytes per parameter, or 5 bytes with gradient release. They also cut model checkpoint sizes by more than half. Experiments with FlashOptim applied to SGD, AdamW, and Lion show no measurable quality degradation on any task from a collection of standard vision and language benchmarks, including Llama-3.1-8B finetuning.