Hugging Face Transformers와 AWS Inferentia를 사용하여 BERT 추론 가속화하기

TL;DR

Hugging Face는 AWS Neuron SDK를 사용하여 Transformers 모델을 컴파일하고 Amazon SageMaker를 통해 AWS Inferentia 칩에 배포함으로써 BERT 추론을 가속화하는 방법을 시연했습니다. 이 접근 방식은 시퀀스 길이가 128일 때 지연 시간을 5-6ms로 줄이며, 기존 GPU 기반 EC2 인스턴스에 비해 더 높은 처리량과 더 낮은 비용을 제공합니다.

Transformer 가속화를 위한 AWS Inferentia

AWS Inferentia는 최적화된 추론 워크로드를 위해 특별히 설계된 맞춤형 머신러닝 칩입니다. AWS에 따르면, Inferentia는 유사한 현재 세대의 GPU 기반 Amazon EC2 인스턴스보다 추론당 비용을 최대 80% 낮추고 처리량을 최대 2.3배 높일 수 있습니다.

이 아키텍처는 칩 내의 맞춤형 가속기인 "Neuron Cores"를 활용합니다. 각 Inferentia 칩에는 4개의 Neuron Core가 포함되어 있어 사용자는 다음 중 하나를 선택할 수 있습니다:

처리량을 최대화하려면 코어당 하나의 모델을 로드합니다.
지연 시간을 최소화하려면 모든 코어에 걸쳐 하나의 모델을 로드합니다.

BERT 배포를 위한 기술적 워크플로우

BERT와 유사한 모델(예: distilbert-base-uncased-finetuned-sst-2-english)을 AWS Inferentia에 배포하려면 다단계 컴파일 및 배포 프로세스가 필요합니다.

1. AWS Neuron SDK를 통한 모델 변환

AWS Neuron SDK는 PyTorch 및 TensorFlow 모델을 EC2 Inf1 인스턴스용 neuron-compatible 포맷으로 변환하기 위한 딥러닝 컴파일러 및 런타임을 제공합니다.

AWS Neuron SDK는 동적 형상(dynamic shapes)을 지원하지 않기 때문에, 컴파일 및 추론 과정에서 입력 크기가 정적이어야 합니다. 예를 들어, 모델이 배치 크기 1과 시퀀스 길이 128로 컴파일된 경우, 해당 정확한 형상의 입력만 처리할 수 있습니다.

2. 커스텀 추론 스크립트 작성

Hugging Face Inference Toolkit이 많은 모델에 대해 코드 없는 배포를 지원하지만, AWS Inferentia는 현재 커스텀 inference.py 스크립트가 필요합니다. 이 스크립트는 다음을 정의해야 합니다:

model_fn: 토크나이저, neuron 모델, 모델 설정을 로드하기 위해 사용됩니다.
predict_fn: 입력 임베딩을 처리하고, 패딩 및 절단을 통해 정적 시퀀스 길이에 맞는지 확인하며, 예측을 실행합니다.

처리량을 최대화하기 위해, 각 HTTP 워커가 단일 Neuron core를 사용하도록 보하는 환경 변수 NEURON_RT_NUM_CORES=1이 사용됩니다.

3. SageMaker 배포

모델이 컴파일되고 추론 스크립트가 생성되면, 아티팩트는 model.tar.gz 파일로 압축되어 Amazon S3에 업로드됩니다. 모델은 Amazon SageMaker의 HuggingFaceModel 클래스를 사용하여 ml.inf1.xlarge 인스턴스를 대상으로 하는 실시간 추론 엔드포인트로 배포됩니다.

성능 결과

10,000개의 동기식 요청으로 구성된 부하 테스트에서, BERT 모델은 시퀀스 길이 128에 대해 평균 5-6ms의 지연 시간을 달성했습니다.

Hugging Face는 이 설정이 CPU 기반 추론보다 빠르며, Neuron Cores를 통해 4개의 모델을 병렬로 실행함으로써 GPU보다 높은 처리량을 제공한다고 결론지었습니다. 이는 AWS Inferentia를 텍스트 분류, 토큰 분류, 질의응답과 같은 인코더 작업용 BERT와 유사한 Transformer 모델을 사용하는 기업들에게 실행 가능한 옵션이 됩니다.