Hugging Face Transformers와 AWS Inferentia를 결합하여 BERT 추론 레이턴시를 5-6ms로 단축하고 GPU 대비 80% 비용 절감

Context

BERT 및 Transformer 기반 모델은 NLP, Computer Vision, Speech 등에서 최첨단 성능을 달성하고 있으나, 기본적으로 전통 머신러닝 알고리즘 대비 느리고 크며 복잡한 구조를 가지고 있다. 프로덕션 환경에서 대규모 워크로드를 처리하기 위해 Transformer 가속화가 필수적인 과제로 대두되었다.

Technical Solution

• AWS Neuron SDK를 사용하여 PyTorch 기반 Hugging Face Transformer 모델을 AWS Inferentia 호환 Neuron 모델로 컴파일: torch_neuron의 trace 메서드로 동적 형태 대신 정적 입력 크기(배치 크기 1, 시퀀스 길이 128)로 고정
• distilbert-base-uncased-finetuned-sst-2-english 모델을 로드하고 AutoTokenizer로 토큰화한 뒤 더미 입력(시퀀스 길이 128)으로 모델 추적 수행
• 변환된 Neuron 모델과 토크나이저, 모델 설정을 파일로 저장
• Amazon SageMaker에 커스텀 inference.py 스크립트 작성: model_fn에서 Neuron 모델 로드, predict_fn에서 텍스트 분류 파이프라인 구성
• NEURON_RT_NUM_CORES=1 환경 변수로 각 HTTP 워커가 1개의 Neuron Core 사용하도록 구성하여 처리량 극대화
• SageMaker Endpoint로 배포 후 CloudWatch로 모델 레이턴시 모니터링

Impact

• 시퀀스 길이 128에서 단일 Neuron Core당 추론 레이턴시 5-6ms 달성
• AWS Inferentia는 GPU 기반 Amazon EC2 인스턴스 대비 추론당 비용 최대 80% 절감, 처리량 최대 2.3배 향상
• 4개의 Neuron Core를 활용하여 4개 모델을 병렬 실행 가능

Key Takeaway

BERT와 같은 인코더 기반 Transformer 모델(텍스트 분류, 토큰 분류, 질의응답)을 프로덕션 환경에서 운영할 때 AWS Inferentia 같은 커스텀 추론 칩을 도입하면 CPU 기준 레이턴시 개선과 GPU 기준 처리량 향상을 동시에 달성할 수 있다. 단, 동적 입력 형태를 지원하지 않는 제약이 있으므로 고정 시퀀스 길이 워크로드에 적합하다.