Nota AI, AWS Inferentia/Trainium에서 LLM 모델 양자화로 성능 UP!
Nota AI는 AWS Inferentia/Trainium에서 LLM 모델의 추론 성능 향상을 위해 양자화 기법을 활용
NxD Inference를 사용하여 int8, fp8, 선택적 양자화 등 다양한 양자화 옵션을 실험
NetsPresso를 활용한 민감도 분석으로 양자화에 민감한 모듈을 제외하여 정확도 손실 최소화
Autoround 기반 Fake Quantization을 통해 NxD 환경 호환성을 유지하며 정확도 손실 완화
향후 Activation 양자화, KV Cache 최적화, Cross-backend 호환성 강화 계획 발표
AWS Inferentia/Trainium 환경과 NxD Inference
Nota AI는 GPU를 대체하기 위해 AWS Inferentia와 Trainium을 활용하여 LLM 모델을 최적화하고 있다. Inferentia/Trainium은 고성능 처리와 빠른 응답 시간을 제공하도록 설계되었으며, AWS는 PyTorch 기반의 NxD Inference를 제공하여 모델 배포를 지원한다.
NxD Inference: 최소한의 코드 수정으로 PyTorch 모델 배포 지원
HuggingFace 모델 및 vLLM 서빙 엔진과의 통합
이러한 환경은 LLM 모델의 추론 성능 극대화를 위한 기반을 제공하며, Nota AI는 NxD Inference의 양자화 옵션을 활용하여 성능을 더욱 향상시켰다.
NxD Inference의 양자화 옵션 분석
NxD Inference는 컴파일 시 다양한 양자화 옵션을 제공하며, weight 텐서에 대해 int8, fp8을 지원한다. 특히, fp8은 fp8e4m3, fp8e5m2 두 가지 형태로 정밀도와 범위 간의 트레이드 오프를 제공한다.
Weight 텐서: per_tensor_symmetric, per_channel_symmetric 옵션 지원
KVcache: fp8 형태 지원
Nota AI는 Mistral-12B 모델의 weight 텐서를 int8, KV cache를 fp8로 양자화하는 예제를 제시하며, 모델 크기를 24.1GB에서 13.2GB로 줄이는 결과를 보였다. 이는 메모리 사용량 감소와 추론 속도 향상에 기여한다.
Low-sensitive mixed-precision weight 양자화
LLM의 weight 텐서를 낮은 정밀도로 양자화할 때, 양자화 비율이 높을수록 모델 크기는 줄어들지만, 정확도 손실이 발생한다. 특히, LayerNorm, Embedding, LM head weight 텐서는 양자화에 민감하다.
NetsPresso: 모듈별 양자화 민감도 계산 기능 제공
민감도 분석 결과: layernorm weight 텐서가 양자화에 절대적으로 민감
Nota AI는 민감도 분석 결과를 바탕으로, 민감도가 높은 모듈을 제외하는 선택적 양자화(Selective Quantization)를 수행하여 정확도 손실을 최소화했다. 그 결과, 모든 weight 텐서를 int8로 양자화했을 때보다 정확도 손실을 줄이면서, 모델 크기 및 메모리 사용량 증가를 최소화했다.
Fake-quantization 기법과 Autoround
Nota AI는 NxD Inference 환경에서 호환성을 유지하기 위해 Fake-quantization 기법을 활용했다. Fake-quantization은 실제 양자화된 데이터 타입을 사용하는 대신, float32 기반 연산을 통해 동일한 출력을 내는 방식이다.
Autoround: 가중치마다 양자화 레벨을 학습하는 알고리즘
Fake-quantization 적용: NVIDIA GPU에서 수행, bf16 형태의 weight 텐서를 huggingface .safetensors 형태로 저장
Nota AI는 Autoround를 적용한 fake-quantization을 통해 GSM8k 벤치마크에서 정확도 손실을 완화했다. 이는 모델 정확도 유지와 추론 성능 향상을 동시에 달성하는 효과적인 방법이다.
향후 연구 방향
Nota AI는 Activation 양자화, KV Cache 효율 최적화, Cross-backend 호환성 강화, 다양한 벤치마크 기반 평가 확장을 통해 LLM 모델 최적화를 지속할 계획이다.
Activation Quantization 및 Mixed Precision 연구: 동적 양자화(dynamic quantization) 및 per-token 정밀도 조절 방식 적용
KV Cache 효율 최적화: 긴 시퀀스 입력(Long-context inference) 상황에서의 성능 및 지연(latency) 분석
Cross-backend 호환성 강화: GPU 학습 모델을 Inferentia 환경에서 원활히 컴파일 및 추론
다양한 벤치마크 기반 평가 확장: GSM8k 외 MMLU, ARC, HellaSwag 등 활용
이러한 연구를 통해 LLM 모델의 성능과 효율성을 지속적으로 개선할 것으로 기대된다.
