Nota AI가 제안하는 Transformer 모델을 AWS Inferentia/Trainium에 손쉽게 배포하는 방법

Nota AI는 인공지능(AI, Artificial Intelligence) 모델의 경량화 및 최적화 기술을 전문적으로 연구·개발하는 기업입니다. 인공지능 모델을 분석해 특정 하드웨어(특히 NPU, Neural Processing Unit)의 호환성을 지원하고 하드웨어의 특성에 맞게 모델을 변경하여 추론 성능 및 메모리 효율성을 극대화하는 자사 플랫폼인 NetsPresso를 기반으로 모바일, 자동차, 로보틱스, 스마트시티 등 자원이 제한된 다양한 산업군에 고성능 AI 솔루션을 제공합니다.

최근에는 Vision-Language Model(VLM)과 대규모 언어 모델(LLM)의 경량화 연구를 통해, 더욱 다양한 도메인에 특화된 AI를 확산해 나가고 있습니다.

그리고 지금, Nota AI가 주목하고 있는 또 하나의 핵심 과제는 LLM 추론의 비용 효율화입니다.

고성능 LLM 추론, 비용이 고민이라면?

LLM을 실제 서비스에 도입하려는 수요는 빠르게 증가하고 있지만, 고가의 GPU 인프라는 여전히 큰 장벽으로 작용하고 있습니다. 이러한 한계를 극복하기 위해, Nota AI는 GPU를 대체할 수 있는 고성능 NPU 환경에서 LLM을 최적화하고 안정적으로 배포하는 기술을 개발하고 있습니다.

이 과정에서 주목한 플랫폼이 바로 AWS Inferentia/Trainium 입니다. 단순히 연산 성능이 뛰어난 저비용 칩을 넘어서, AWS 클라우드 상에서 손쉽게 활용 가능한 추론 전용 인프라를 제공합니다. 특히, Neuron SDK를 통해 복잡한 하드웨어 설정 없이도 LLM을 컴파일하고 배포할 수 있어, LLM 인프라 구축 비용과 운영 복잡도를 모두 줄일 수 있는 대안으로 주목받고 있습니다.

Transformer 모델의 Neuron SDK 기반 포팅 가능성 검토

본격적인 최적화 작업에 앞서, Hugging Face Transformers 라이브러리 기반의 주요 사전학습 모델들이 Neuron SDK로 AWS Inferentia에 포팅 가능한지를 사전 검토했습니다. 이는 추후 최적화 기법 적용을 위한 기반 확보는 물론, 실제 NPU 환경에서의 실행 가능성 검증을 목표로 합니다.

테스트 대상 모델은 실제 서비스에서 많이 활용되는 주요 Transformer 계열 모델 12종을 선정했습니다. 다양한 모델 구조(GPT, BERT, OPT 등)를 포함하여 Neuron SDK의 호환성 범위를 폭넓게 검증하고자 하였습니다.

1차 포팅 테스트 및 결과 분석

총 12개의 Transformer 기반 모델을 대상으로 torch_neuronx.trace()를 활용하여 Neuron SDK 포팅을 시도하였습니다. 이 함수는 PyTorch 모델을 AWS Neuron 하드웨어에 최적화된 형식으로 변환하며, 모델이 AWS Inferentia/Trainium에서 실행될 수 있도록 컴파일하는 역할을 합니다.

아래는 초기 테스트 결과입니다.

성공적으로 변환된 모델은 2개뿐이었고, 나머지는 모두 RuntimeError로 인해 컴파일 단계에서 실패했습니다.

포팅 실패 사례 분석: Llama-3.2 모델

대표적인 실패 사례는 meta-llama/Llama-3-2-1B 모델입니다. 해당 모델은 torch_neuronx.trace()를 통해 모델 컴파일을 시도하던 중 다음과 같은 오류가 발생했습니다:

사용한 코드:

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3-2-1B"
).eval()
# 오류 발생 지점
model_neuron = torch_neuronx.trace(
    model,input_sample,compiler_args="--model-type=transformer"
)

에러 로그:

RuntimeError: Tracer cannot infer type of
{'last_hidden_state': tensor(...), 'past_key_values': ((tensor(...)), ...)}
Dictionary inputs to traced functions must have consistent type.

원인 분석

이 오류는 torch_neuronx.trace()가 내부적으로 사용하는 torch.jit.trace() 함수가 모델을 정적 TorchScript 그래프로 변환하는 과정에서 발생합니다.
torch.jit.trace()는 샘플 입력을 기준으로 모델의 실행 경로를 추적(tracing)하며, 이 과정에서 출력값의 구조를 정적으로 명확히 파악할 수 있어야 합니다.

하지만 meta-llama/Llama-3-2와 같은 모델은 forward() 함수에서 다음과 같이 다양한 자료형이 혼합된 딕셔너리를 반환합니다:

{
  'last_hidden_state': Tensor,
  'past_key_values': Tuple[Tuple[Tensor, Tensor], ...]
  ...
  ...
} 

이러한 복합적인 출력 구조는 tracing 기반 TorchScript 변환 방식에 적합하지 않으며, 출력 구조를 정적으로 예측할 수 없기 때문에 컴파일 과정에서 오류가 발생합니다.

따라서 Llama-3와 같이 다양한 출력 타입을 포함하는 Transformer 계열 모델을 Neuron SDK로 컴파일하기 위해서는, tracing 전에 출력 구조를 단순화하거나 필요한 출력만 선택적으로 반환하도록 처리하는 작업이 필요합니다.

해결 방법

이 문제는 모델을 로드할 때 return_dict=False 옵션을 사용하면 간단하게 해결할 수 있습니다.이 설정을 적용하면 모델의 출력이 딕셔너리 대신 튜플 형태로 반환되어, tracing 과정에서 발생할 수 있는 오류를 방지할 수 있습니다. return_dict 옵션에 대한 설명은 여기에서 참고하실 수 있습니다.

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3-2-1B",
    return_dict=False # 해결 방법
).eval()

이후 torch_neuronx.trace()를 실행하면, 정상적으로 모델 컴파일이 수행됩니다.

model_neuron = torch_neuronx.trace(
    model,
    input_sample,
    compiler_args="--model-type=transformer"
)

변환된 모델이 원본 모델과 동일한 출력을 생성하는지 확인하기 위해, 다음과 같은 예시 문장을 입력으로 테스트했습니다.

예시 입력:

What do you feel like having for dinner tonight?

출력:

I’m having a little trouble deciding what to order tonight. I’m not a picky eater, but I’m picky about my food.

변환 전후 동일한 응답을 생성함으로써, 출력 구조 변경이 모델의 성능에 영향을 주지 않았음을 확인하였으며, NPU 환경에서도 성능 저하 없이 안정적으로 추론하는 것을 검증할 수 있었습니다.

최종 포팅 결과: 모든 모델 포팅 성공

위와 같은 방식으로 나머지 모델에도 동일한 조치를 적용한 결과, 12개 모델 모두 포팅에 성공하였습니다.

결론

이번 실험을 통해 Hugging Face 기반 Transformer 모델들이 AWS Inferentia 및 Trainium 환경에서도 안정적으로 추론이 가능하다는 점을 확인했습니다. 특히 return_dict=False와 같은 간단한 설정만으로도 다양한 포팅 오류를 해결할 수 있었으며, 이를 통해 복잡한 출력 구조를 가진 Transformer 모델도 Neuron SDK를 통해 손쉽게 컴파일하고 배포할 수 있음을 검증했습니다.

최종적으로 12개 모델 전부를 성공적으로 변환하였고, 출력 결과도 원본 모델과 일관성을 유지함을 확인했습니다.

마무리

Nota AI는AI 모델 최적화 플랫폼인 NetsPresso를 활용하여 기존 NVIDIA V100 GPU를 T4 GPU로 대체하고도 채널 수를 2배 확장하고, 추론 속도(FPS)를 14% 향상시키고, 정확도는 97.14%에서 97.03%로 단 0.1% 이내로 유지하면서도 월 운영 비용을 약 85% 절감($2,203 → $391)하는 성과를 달성하였습니다.
Nota AI는 이러한 정량적 성과를 바탕으로, AWS Inferentia 및 Trainium 환경에서 다양한 모델의 메모리 효율성과 추론 속도를 더욱 향상시킬 계획입니다. 이를 통해 더 많은 기업이 GPU 외 대안들을 활용하여 고성능 AI 모델을 효율적으로 배포하고, 운영 비용을 절감할 수 있도록 기술 확장을 이어나가고자 합니다.

AWS Trainium과 Inferentia는 단순히 모델 학습과, 추론 가속을 넘어, 대규모 언어 모델을 위한 강력하고 안정적인 인프라를 제공합니다. 특히 Neuron SDK를 통해 복잡한 하드웨어 설정 없이도 모델을 손쉽게 컴파일하고 배포할 수 있어, 실제 서비스 환경에서 매우 실용적으로 활용될 수 있다는 점이 인상적이었습니다. 이처럼 AWS의 고도화된 AI 추론 플랫폼은 성능, 비용, 그리고 접근성 측면에서 뛰어난 균형을 갖춘 솔루션임을 확인할 수 있었습니다.

앞으로도 Nota AI는 Transformer를 비롯한 다양한 AI 모델을 지속적으로 AWS Neuron 환경에 최적화해 나가며, 더 많은 기업이 AWS Inferentia 및 Trainium 기반의 AI 서비스를 효율적으로 운영할 수 있도록 적극 지원하겠습니다.