[논문 리뷰] Training and inference of large language models using 8-bit floating point

Sergio P. Perez, Yan Zhang외 7명|arXiv (Cornell University)|2023. 09. 29.

Parallel Computing and Optimization Techniques인용 수 4

한 줄 요약

이 논문은 무게, 활성화 및 기울기의 텐서별 동적 스케일링 인자를 조정하여 최대 700억 파라미터의 대규모 언어 모델을 8비트 부동소수점(FP8) 형식으로 훈련 및 추론하는 확장 가능한 방법론을 제시한다. FP8-AMAX 방법은 GPT 및 라마2 모델 전반에서 정밀도가 높은 FP16 기준선과 거의 동일한 정확도를 달성하며, 대규모 환경에서도 강건함을 입증한다.

ABSTRACT

FP8 formats are gaining popularity to boost the computational efficiency for training and inference of large deep learning models. Their main challenge is that a careful choice of scaling is needed to prevent degradation due to the reduced dynamic range compared to higher-precision formats. Although there exists ample literature about selecting such scalings for INT formats, this critical aspect has yet to be addressed for FP8. This paper presents a methodology to select the scalings for FP8 linear layers, based on dynamically updating per-tensor scales for the weights, gradients and activations. We apply this methodology to train and validate large language models of the type of GPT and Llama 2 using FP8, for model sizes ranging from 111M to 70B. To facilitate the understanding of the FP8 dynamics, our results are accompanied by plots of the per-tensor scale distribution for weights, activations and gradients during both training and inference.

연구 동기 및 목표

• 대규모 언어 모델의 훈련 및 추론 과정에서 FP8 형식의 제한된 동적 범위 문제를 해결한다.
• FP8에서의 부족 및 오버플로우로 인한 수치적 불안정성을 방지하기 위해 스케일링 인자를 동적으로 조정한다.
• 정확도 저하를 최소화하면서 최대 700억 파라미터의 모델에 대해 엔드 투 엔드 FP8 훈련 및 추론을 가능하게 한다.
• FP16 및 INT8에 대한 실용적이고 하드웨어 효율적인 대안으로서 FP8의 계산적 이점을 활용한다.
• 대규모 모델에서 정적 스케일링 전략보다 뛰어난 성능을 보이는 신뢰할 수 있는 FP8 양자화 방법을 확립한다.

제안 방법

• 무게, 활성화 및 기울기의 텐서별 스케일링 바이어스를 각 텐서의 최대 절대값 기반으로 계산하는 동적 스케일링 전략인 FP8-AMAX를 도입한다.
• 계산_바이어스 함수를 통해 동적 지수 바이어스 조정을 수행하며, 이는 형식 최대값과 텐서의 진폭 비율의 로그2 기반 스케일링을 계산하고, 설정 가능한 여유를 포함한다.
• FP8 연산 이전 및 이후에 스케일링을 적용한다: FP8로 캐스팅하기 전에 입력을 스케일링하고, 행렬 곱셈 이후에 복원한다. 이를 통해 정밀도를 유지한다.
• 무게 및 활성화에는 FP8_E4, 기울기에는 FP8_E5를 별도로 사용하여 범위와 정밀도의 균형을 이룬다.
• 훈련(기울기 스케일링 및 가중치 갱신 포함) 및 추론(훈련 후 양자화) 파이프라인에 이 방법을 통합한다.
• 수렴 안정성을 검증하기 위해 손실 스케일링 스윕을 실시하며, 고정 스케일링(FP8-CSCALE)과 대비하여 동적 스케일링(FP8-AMAX)의 성능을 비교한다.

(a) Forward pass for FP16 inference.

실험 결과

연구 질문

• RQ1 대규모 언어 모델의 훈련 및 추론에 있어 FP8가 정확도 저하 없이 효과적으로 사용될 수 있는가?
• RQ2 모델 크기에 따라 동적 텐서별 스케일링(FP8-AMAX)과 고정 스케일링(FP8-CSCALE) 간의 수렴성 및 안정성은 어떻게 비교되는가?
• RQ3 무게, 활성화 및 기울기에서 부족 및 오버플로우를 방지하기 위해 어떤 스케일링 전략이 수치적 안정성을 확보하는가?
• RQ4 훈련 중 스케일링 바이어스는 어떻게 변화하는가? 이는 갱신 빈도와 계산 비용에 어떤 영향을 미치는가?
• RQ5 FP8 양자화가 다양한 LLM 아키텍처 및 작업에서 FP16 수준의 성능을 달성할 수 있는가?

주요 결과

• FP8-AMAX는 MNLI, QQP 및 SST-2 작업 전반에서 1억 1,100만에서 130억 파라미터의 GPT 및 라마2 모델을 성공적으로 피나이팅하며, 검증 정확도가 FP16 기준선의 99.5% 이내를 유지한다.
• 130억 파라미터 모델의 경우, FP8-CSCALE는 스케일링 범위 [-1, 0] 내에서만 수렴 가능하며, 이마저도 다양한 난수 시드에 따라 수렴 보장이 되지 않는다.
• FP8-AMAX는 테스트된 모든 모델 크기(67억 및 130억 포함)에서 일관된 수렴을 보이며, 피나이팅 중에 발산 현상이 관찰되지 않았다.
• FP8-AMAX와 FP16의 손실 함수 궤적은 거의 구분되지 않아 안정적인 훈련 동역학을 나타낸다.
• 텐서별 스케일링 분포 분석 결과, FP8-AMAX는 정적 스케일링 전략인 FP8-CSCALE와 달리 훈련 전반에 걸쳐 안정적이고 적응형 스케일링을 유지한다.
• 이 방법은 700억 파라미터의 라마2 모델까지 성공적으로 확장되어 대규모 환경에서의 완전한 FP8 훈련 및 추론의 가능성을 입증한다.

(b) Forward pass for FP8 inference.