Amazon SageMaker AI에서 LLM을 사용한 Strands 에이전트용 커스텀 모델 제공자 구축하기

소개

인공지능 및 기계 학습의 발달 속에서 여러 조직이 그들의 배포 환경에 맞춘 커스텀 대형 언어 모델(LLM)을 활용하고자 Amazon SageMaker에서 실시간 AI 엔드포인트를 사용하고 있습니다. 이러한 배포 방식은 SGLang, vLLM, 그리고 TorchServe와 같은 프레임워크를 통해 이루어지며 이는 더욱 효율적이고 비용 최적화된 운영을 가능케 합니다. 하지만 이러한 유연성에는 기술적 도전 과제가 따르는데, Strands 에이전트와의 응답 형식 불일치가 바로 그것입니다.

본문

커스텀 모델 파서 구축의 필요성

Amazon SageMaker AI 실시간 엔드포인트에서 다양한 모델을 배포할 수 있는 유연성은 많은 이점을 가지지만, Strands 에이전트와의 응답 형식 불일치 문제를 야기할 수 있습니다. 대부분의 커스텀 서빙 프레임워크는 OpenAI-호환 형식의 응답을 반환하지만, Strands 에이전트는 Bedrock 메시지 API 형식에 맞춘 응답을 기대합니다. 이러한 차이로 인해 파싱 오류가 발생하게 됩니다.

구현 개요

이 글에서는 SageMaker에서 호스팅된 LLM, 특히 Bedrock 메시지 API 형식을 기본적으로 지원하지 않는 모델과 함께 사용할 Strands 에이전트용 커스텀 모델 파서를 구축하는 방법을 살펴봅니다. 우리는 SGLang으로 Llama 3.1을 SageMaker에서 배포하고 Strands 에이전트와 통합할 커스텀 파서를 구현하는 과정을 단계별로 진행합니다.

Strands용 커스텀 모델 파서 구축하기

1. 모델 배포: Llama 3.1을 SGLang으로 SageMaker에 배포하여 OpenAI-호환 응답을 반환.
2. 파서 레이어: LlamaModelProvider 클래스를 확장하여 Llama 3.1의 응답 형식을 처리.
3. 에이전트 레이어: 대화형 AI에 적합한 Strands 에이전트를 사용, 모델의 응답을 적절히 파싱.

Llama 3.1 모델 배포 및 파서 구현

• ml-container-creator 설치 및 설정: awslabs/ml-container-creator를 사용하여 SageMaker 이미지 생성.
• 배포 프로젝트 생성: yeoman을 사용하여 다양한 설정과 환경 구축.
• 서빙 컨테이너 빌드 및 배포: AWS CodeBuild를 활용해 SageMaker에 모델 배포.

npm install -g yo
git clone https://github.com/awslabs/ml-container-creator
cd ml-container-creator
npm install && npm link
yo ml-container-creator

테스트 및 확장

Strands 에이전트를 통해 구현된 커스텀 모델 파서의 호환성을 검증하며, 필요한 경우 추가적인 커스터마이징을 통해 더욱 정교한 AI 응답 체계를 마련할 수 있습니다.

결론

Strands 에이전트를 위한 커스텀 모델 파서는 AWS SageMaker에서 LLM 배포를 최적화하는 데 중요한 역할을 합니다. SageMakerAIModel을 확장하고 stream() 메서드를 구현하여, 다양한 배포 형식의 모델을 통합하면서도 깔끔한 에이전트 인터페이스를 유지할 수 있습니다.

[1] https://aws.amazon.com/blogs/machine-learning/building-custom-model-provider-for-strands-agents-with-llms-hosted-on-sagemaker-ai-endpoints/