[KR] OwLite와 함께하는 FP8 Quantization
지난 몇 년간 AI 분야는 엄청난 발전을 이루었고, ChatGPT와 같은 대규모 언어 모델(LLM) 기반의 어플리케이션은 이미 다양한 영역에서 우리 삶 속에 스며들어 큰 도움을 주고 있습니다. 딥러닝 알고리즘의 발전, 폭발적으로 증가하는 데이터, 그리고 컴퓨팅 파워의 혁신적인 발전이 이러한 성과를 이끌었다고 볼 수 있습니다. 이러한 배경 속에서 효율적인 데이터 처리와 모델 배포를 가능하게 하는 새로운 데이터 표현 방식들이 등장하고 있습니다. 그중에서도 FP8 (Floating-point 8-bit) format은 AI 하드웨어 아키텍처에서 주목할 만한 혁신 중 하나로 떠오르고 있습니다.
NVIDIA의 최신 GPU 아키텍처 (ex, Hopper, Ada Lovelace, Blackwell) 및 Intel의 Gaudi-2 &3와 같은 최신 하드웨어는 FP8 format의 연산을 지원하고 있습니다. 이런 하드웨어들의 spec sheet를 보면 BF16 TFLOPS (Tera-Floating-Point Operations per Second)와 함께 FP8 TFLOPS를 함께 report하고 있고, 최근에는 TensorRT-LLM 프레임워크에서 Transformer 모델 (주로 LLM)의 inference 효율을 높이기 위해 FP8 format을 적극적으로 활용하려는 움직임을 보이고 있습니다. 하지만 최근까지도 Transformer 모델들만 지원하는 등 이를 뒷받침하는 SW가 충분하지 않아 기존의 Computer Vision과 같은 분야에서는 FP8 format이 전혀 활용되지 못하고 있었습니다.
스퀴즈비츠에서는 NVIDIA의 TensorRT가 공식적으로 지원하기 전부터 FP8 Convolution을 지원하는 등 발 빠르게 업데이트를 해오고 있었습니다. 그러던 중, 드디어 최근 TensorRT 업데이트에서 FP8 Convolution을 지원하기 시작했고, OwLite와 TensorRT가 지원하는 FP8 quantized model의 성능을 비교해 볼 수 있게 되었습니다.
이번 포스트에서는 FP8 format에 대해 간단히 알아보고, OwLite를 통해 FP8 format으로 양자화한 TensorRT 엔진의 성능을 살펴보겠습니다.
FP8 (Floating-point 8-bit) format이란?
FP8 format은 8-bit의 데이터 길이를 갖는 부동소수점 데이터 타입입니다. 부동소수점은 (1)부호 비트(sign bit), (2)지수부(exponent), 그리고 (3)가수부(mantissa)로 구성됩니다. IEEE-754 표준에 따르면 BF16, FP32, 그리고 FP16의 부호 비트, 지수부, 가수부는 다음과 같습니다.
FP8 format은 표준이 없기 때문에 다양한 크기의 지수부와 가수부 조합으로 연구가 이루어지고 있었으며, 2022년에 NVIDIA/ARM/Intel에서 IEEE-754 표준을 따르는 두 가지 8비트 부동소수점 타입인 FP8 E5M2와 FP8 E4M3를 제안했습니다.
NVIDIA의 Hopper 및 Ada Lovelace 아키텍처에서는 위의 두 가지 FP8 format을 지원합니다.
추론에서는 주로 E4M3을 사용합니다 [1]. 별 다른 언급이 없다면 아래에서 언급되는 FP8 format은 E4M3을 의미합니다.
TensorRT에서의 FP8 Format 지원
TensorRT는 8.6 버전부터 FP8 데이터 타입(E4M3)을 지원했고, 9.0 버전부터 explicit quantization을 통한 FP8 연산을 지원했습니다. 하지만, FP8 Gemm 연산만 지원했기 때문에 Convolution 레이어를 포함하는 모델에서는 사용할 수 없었고, 주로 TensorRT-LLM을 통한 대규모 언어 모델(LLM)에서 사용되고 있었습니다.
그러나 최근 업데이트된 10.2 버전에서는 normal FP8 convolution이 추가되었습니다. 이로 인해 Convolution 레이어를 포함하는 모델을 FP8 format으로 양자화하여 돌려볼 수 있게 되었습니다. 하지만 여전히 아래의 제약 사항이 존재하여 모든 모델을 지원하지는 않습니다.
- 입출력 채널이 16배수가 아닌 Convolution에 대한 FP8 연산의 구현이 없습니다.
- Group Convolution 및 Depthwise Convolution 구현이 없습니다.
테스트는 ResNet18 모델을 사용하였는데, 입출력 채널에 대한 제한 때문에 첫 번째 Convolution 레이어를 제외한 나머지 레이어에 대해 Post-Training Quantization (PTQ)를 적용하였습니다.
PTQ는 OwLite를 사용하여 적용할 수 있습니다.
테스트 환경과 단위 시간 당 이미지 처리량을 측정한 결과는 다음과 같습니다.
- GPU: NVIDIA L40S (Ada Lovelace 아키텍처)
- Data Types: FP16(baseline), FP8
- PTQ Calibration Method: AbsMax, Per-Tensor Quantization for Activations and Weights
- Datasets: ImageNet
위 결과를 통해 FP8로 빌드한 엔진의 처리량이 FP16로 빌드한 엔진의 처리량보다 작다는 것을 확인할 수 있습니다. 이러한 결과를 보여주는 이유는 TensorRT에서 아직 완벽하게 최적화된 커널이 업데이트되지 않은 것으로 추정됩니다.
TensorRT 10.2에서 FP8 데이터 타입으로 엔진을 빌드할 때,
--stronglyTyped플래그를 추가하지 않으면 FP8 implementation tactics가 활성화되지 않는 버그가 있습니다. 해당 이슈는 NVIDIA-TensorRT 팀과 공유되었으며 TensorRT 10.3 버전에서 수정될 것으로 예상됩니다.
OwLite의 FP8 Format 지원
스퀴즈비츠의 OwLite에서는 TensorRT 10.2 버전이 업데이트되기 전부터 FP8 format을 지원하기 시작했습니다. NVIDIA CUTLASS 라이브러리를 활용하여 FP8 GeMM 및 Convolution을 구현하여 TensorRT의 커스텀 플러그인으로 지원할 수 있도록 개발하였습니다. 이를 통해, 보다 편리하게 FP8 format으로의 양자화를 적용할 수 있고, 보다 빠른 TensorRT 엔진을 빌드할 수 있습니다.
방법은 간단합니다. 먼저, owlite 패키지를 통해 ResNet18 baseline을 업로드합니다. 그리고, 웹 GUI를 통해 experiment를 생성하고 양자화를 적용하고자 하는 연산에 적용할 데이터 타입을 아래와 같이 fp8_e4m3로 선택해 줍니다.
적용하고자 하는 연산에 대해 설정이 완료되었다면, 저장 버튼을 누르고 다시 코드로 돌아와서 calibration 및 benchmark를 수행하면 모든 단계가 마무리됩니다.
자세한 방법은 ResNet18 Tutorial을 확인해주세요.
이렇게 추출된 ONNX 파일은 아래와 같습니다.
이 ONNX 파일을 사용하여 TensorRT 엔진을 빌드하면, TensorRT에서 네이티브로 지원하는 FP8 구현으로 빌드됩니다. 하지만, 위에서 보셨다시피 FP16으로 빌드된 엔진보다 성능이 떨어지는 것을 확인할 수 있었습니다. OwLite에서는 조금 더 빠른 엔진을 빌드하기 위해 커스텀 플러그인으로 구현한 Gemm 및 Convolution 연산을 제공합니다.
ONNX를 통해 커스텀 플러그인을 적용하기 위해서는 네이티브 연산이 아닌 커스텀 플러그인 연산으로 인식하도록 ONNX 변환이 필요합니다. OwLite backend에서는 엔진을 빌드하기 전에 필요한 ONNX 그래프 변환을 먼저 수행합니다.
커스텀 플러그인을 사용할 수 있도록 변환된 ONNX는 다음과 같습니다.
변환된 ONNX 그래프를 보면 기존 그래프에 비해 너무 간결해졌다는 생각이 들 수 있을 것 같습니다. 앞서 언급했듯이 TensorRT의 커스텀 플러그인을 사용하기 위해선 커스텀 플러그인으로 인식할 수 있는 커스텀 노드로 변환해 주어야 합니다. 하지만, 커스텀 노드를 사용하게 되면 TensorRT가 수행하던 그래프 최적화를 커스텀 노드에 적용할 수가 없게 됩니다. OwLite에서는 최적의 연산 효율을 달성하기 위해서 ONNX 변환 이후, TensorRT의 그래프 최적화와 유사한 자체 그래프 최적화를 수행하게 됩니다. 이 모든 과정을 거치고 나면 위의 ONNX 그래프를 얻을 수 있습니다.
이렇게 변환된 ONNX는 스퀴즈비츠의 커스텀 플러그인 라이브러리와 함께 TensorRT 엔진으로 빌드됩니다. 커스텀 플러그인 연산을 사용하여 빌드된 엔진의 성능 결과는 아래와 같습니다.
위 결과를 통해 OwLite의 커스텀 플러그인을 사용하면 네이티브 FP8 연산을 사용하도록 빌드한 엔진보다 상당한 성능 향상을 얻을 수 있음을 확인할 수 있습니다.
—
지금까지 PTQ를 통해 빌드한 엔진의 latency에 대해서 살펴봤습니다. PTQ에서는 모델을 최적화하여 latency를 줄이는 것과 동시에 기존 모델의 accuracy를 최대한 유지하는 것이 중요합니다.
OwLite를 사용하면 웹 GUI를 통해 각 experiment로부터 측정한 다양한 메트릭을 한 눈에 비교해 볼 수 있는 기능을 제공합니다. 이를 이용하여 배치 사이즈가 128인 경우의 accuracy를 아래와 같이 비교 해볼 수 있습니다.
커스텀 플러그인을 사용한 모델(
fp8_plugin)의 배포 사이즈가 큰 이유는 TensorRT 엔진에 커스텀 플러그인 라이브러리가 포함되어 있기 때문입니다.
위 결과를 통해 FP8 PTQ를 통해 모델을 경량화하더라도 baseline(69.788)과 비슷한 수준의 accuracy(69.569)를 갖는 모델을 얻을 수 있다는 것을 알 수 있습니다.
마치며
INT8 format과 마찬가지로 FP8 format은 AI 하드웨어 아키텍처에서 메모리 효율성과 연산 속도를 크게 향상시킬 수 있는 데이터 타입입니다. 최근 AI 하드웨어 아키텍처에서는 FP8 format을 지원하며 자체 연산 가속기도 제공합니다. 이를 통해 AI 모델의 성능을 극대화하고 다양한 가능성을 제공하는데 기여하고 있습니다.
FP8 format이 더 좋을지, INT8 format이 더 좋을지는 응용 분야, 모델 특성, 하드웨어 특성에 따라 달라지겠지만 OwLite는 최대한 다양한 경량화 옵션을 제공함으로써 사용자의 모델에 최적화된 경량화 방법을 찾을 수 있게 지원하고 있습니다. 혹시 INT8 quantization의 결과물이 만족스럽지 않으셨나요? 그렇다면 이 기회에 FP8 quantization을 시도해보시는 것은 어떨까요?
—
지금 바로 OwLite를 통해 FP8 format을 활용해보세요!
References
[1] Micikevicius, P., Stosic, D., Judd, P., Kamalu, J., Oberman, S., Shoeybi, M., Siu, M., Wu, H., Burgess, N., Ha, S., Grisenthwaite, R., Mellempudi, N., Cornea, M., Heinecke, A., and Dubey, P., FP8 Formats for Deep Learning. arxiv:2209.05433, 2022.
