AWS Trainium2와 i14로 효율적인 인공지능 개발

AWS Trainium2와 i14의 Laplace 3B 모델: 효율적인 인공지능 개발의 새로운 기반

AWS Trainium2와 i14의 파트너십은 인공지능 모델의 효율성을 극대화하는 데 주력합니다. 이 글에서는 i14가 AWS Trainium2를 어떻게 활용하여 Laplace 3B 모델을 개선했는지를 살펴보겠습니다.

긴 도입부: AWS와 i14의 협력 배경

AWS Trainium2는 인공지능 모델의 학습과 추론을 지원하는 AWS의 혁신적인 머신러닝 가속기 제품군입니다. 모델 학습과 추론의 용이한 구현을 돕기 위해 i14는 AWS와의 협력을 통해 Laplace 3B 모델을 개선했습니다. Laplace 3B는 키-값 캐시를 줄여가며 긴 문맥을 유지할 수 있는 Linear Attention 모델로 설계되었습니다.

본문: 핵심 내용과 사례

i14는 GPU 최적화 Triton 커널을 통해 구현한 Laplace 3B의 GLKA(게이티드 선형 주의) 커널을 표준 PyTorch 연산자를 사용하여 AWS Trainium2에 맞게 변형하였습니다. 이로 인해 최적화된 모델은 120,000 토큰/초로 학습하고, 한 번에 1,280개의 동시 시퀀스를 처리할 수 있었습니다.

AWS Trainium2의 주요 특징은 고속의 칩 대 칩 연결과 정적 컴파일러입니다. 각각의 NeuronCore는 소프트웨어 관리 디바이스 온칩 메모리를 제공하며, 이는 데이터 전송 시 지연 없이 필요한 곳에 정확히 도착하도록 합니다. 이로 인해 Laplace 3B 모델은 메모리 사용을 예측 가능하게 유지합니다.

결론: AWS Trainium2와 i14의 혁신적인 접근

이번 프로젝트는 AWS Trainium2가 i14의 고유한 모델 아키텍처를 지원할 수 있음을 증명했습니다. TorchNeuron Native를 통한 네이티브 PyTorch를 활용해, GLKA 기반 모델을 고정된 재귀 상태를 유지하며 학습 및 제공할 수 있었습니다. 결국 이러한 접근은 긴 문맥을 확보하면서도 메모리 사용을 일정하게 유지할 수 있는 기능을 제공하여, 초대형 언어 모델의 새로운 가능성을 열었습니다.

참고 자료

[1] 원문 URL: https://aws.amazon.com/blogs/machine-learning/beyond-the-kv-cache-training-and-serving-i14s-laplace-3b-linear-attention-model-on-aws-trainium2-with-native-pytorch/