AI
Nota AI가 제안하는 Transformer 모델을 AWS Inferentia/Trainium에 손쉽게 배포하는 방법
두줄요약
Hugging Face Transformer 모델을 AWS Inferentia/Trainium에 포팅하는 과정을 검증했습니다. return_dict=False 설정으로 tracing 오류를 줄여 12개 모델 모두 배포 가능함을 확인했습니다.
문제 상황
- Hugging Face Transformer 모델을 AWS Inferentia/Trainium 환경에 배포할 때 컴파일 실패와 포팅 호환성 검증 필요
- GPU 대비 비용 효율적인 추론 인프라 대안으로 Neuron SDK 활용 가능성 검토
원인 분석
- torch_neuronx.trace() 내부의 정적 tracing 과정에서 복합 딕셔너리 출력 구조 처리 어려움
- Llama-3 계열처럼 Tensor와 past_key_values가 섞인 반환 형식이 타입 추론 실패 유발
해결 방법
- 모델 로드 시 return_dict=False 적용으로 출력 형식을 딕셔너리에서 튜플로 단순화
- 동일 조치 후 torch_neuronx.trace() 재실행으로 12개 Transformer 모델 전체 포팅 성공
적용해볼 점
- Neuron SDK 포팅 전 모델 출력 구조와 tracing 적합성 사전 점검
- 복잡한 출력 반환을 단순화해 AWS Inferentia/Trainium 배포 안정성 확보