AWS Trainium은 Amazon Web Services(AWS)가 클라우드에서 머신러닝 훈련 및 추론 작업을 가속화하기 위해 개발한 맞춤형 주문형 집적 회로(ASIC) 제품군입니다. 2020년 11월에 발표되었고 2022년에 처음 사용 가능해진 Trainium 칩은 대규모 AI 워크로드, 특히 딥러닝 모델(예: 대규모 언어 모델 및 트랜스포머 아키텍처)에 대해 범용 그래픽 처리 장치(GPU)보다 비용 효율적이고 에너지 효율적인 대안을 제공하도록 설계되었습니다. 이 칩은 AWS의 Elastic Compute Cloud(EC2) 인스턴스에 통합되어 고객에게 신경망 연산의 특정 계산 패턴에 최적화된 전용 하드웨어를 제공합니다.
Trainium의 개발은 AWS가 AI 하드웨어 제공을 다양화하고 Nvidia와 같은 외부 칩 공급업체에 대한 의존도를 줄이기 위한 광범위한 전략의 일부입니다. 자체 실리콘을 설계함으로써 AWS는 플랫폼에서 실행되는 주요 AI 워크로드에 대해 성능 대비 달러 및 와트당 성능을 최적화하는 것을 목표로 합니다. Trainium은 추론에 특화된 별도의 칩 제품군인 AWS Inferentia로 보완되며, 이 둘은 함께 AWS의 맞춤형 AI 가속화 포트폴리오의 핵심을 형성합니다. Trainium 제품군은 이후 세대를 거쳐 진화했으며, 2023년에 발표되고 2024년에 더 광범위하게 배포되기 시작한 2세대 Trainium2와 2024년 말에 발표된 3세대 Trainium3가 포함됩니다.
하드웨어 아키텍처
Trainium 칩은 신경망 훈련의 핵심 연산인 고처리량 행렬 곱셈과 데이터 이동을 강조하는 목적 기반 아키텍처로 구축되었습니다. 각 Trainium 칩에는 FP32, FP16, BF16 및 FP8과 같은 데이터 유형을 지원하는 혼합 정밀도 계산을 수행할 수 있는 여러 텐서 코어가 포함되어 있습니다. TSMC가 7나노미터 공정으로 제조한 1세대 Trainium 칩은 FP16 및 BF16 연산에 대해 최대 200테라플롭스의 성능을 제공합니다. 또한 대용량 온칩 메모리와 외부 메모리로의 고대역폭 인터페이스를 포함하여 대규모 모델 매개변수와 활성화를 효율적으로 처리할 수 있습니다.
2세대인 Trainium2는 성능과 메모리 용량을 크게 향상시킵니다. 5나노미터 공정으로 제조된 각 Trainium2 칩은 최대 400테라플롭스의 FP16/BF16 성능을 제공하며 32기가바이트의 고대역폭 메모리(HBM)를 갖추고 있습니다. 랙 수준 시스템인 단일 Trainium2 UltraServer에는 고속 패브릭으로 상호 연결된 64개의 Trainium2 칩이 포함되어 총 25.6페타플롭스의 컴퓨팅 성능과 6테라바이트의 메모리를 제공합니다. 이 설계는 생성형 AI 애플리케이션에서 사용되는 것과 같은 수천억 개의 매개변수를 가진 모델 훈련을 가능하게 합니다.
2024년 12월에 발표된 Trainium3는 3나노미터 공정으로 제조될 것으로 예상되며 훈련 및 추론 효율성에 중점을 두고 Trainium2의 약 두 배 성능을 제공할 것으로 예상됩니다. AWS는 Trainium3가 2025년에 제공될 것이며 일부 고객에게 조기 액세스가 제공될 것이라고 밝혔습니다.
EC2 인스턴스 및 가용성
AWS는 Trn1 및 Trn2 인스턴스 제품군을 통해 EC2 서비스에서 Trainium 기반 컴퓨팅 용량을 제공합니다. 2022년 10월에 출시된 1세대 Trn1 인스턴스는 Trainium 칩 1개가 있는 trn1.2xlarge와 Trainium 칩 16개가 있는 trn1.32xlarge의 두 가지 크기로 제공됩니다. trn1.32xlarge는 최대 3.2페타플롭스의 FP16/BF16 성능과 512기가바이트의 메모리를 제공하며, 여러 인스턴스에 걸친 분산 훈련을 위해 Elastic Fabric Adapter(EFA)를 통한 고속 네트워킹을 지원합니다.
2024년 12월에 출시된 2세대 Trn2 인스턴스는 Trainium2 칩을 기반으로 합니다. trn2.48xlarge 인스턴스에는 16개의 Trainium2 칩이 포함되어 6.4페타플롭스의 성능과 1.5테라바이트의 메모리를 제공합니다. AWS는 또한 4개의 trn2.48xlarge 인스턴스를 64개의 Trainium2 칩이 있는 단일 논리 시스템으로 결합하여 최대 1조 개의 매개변수를 가진 모델 훈련을 가능하게 하는 Trn2 UltraServer를 도입했습니다. 이러한 인스턴스는 미국 동부(버지니아 북부), 미국 서부(오레곤), 유럽(아일랜드)을 포함한 일부 AWS 리전에서 제공되며 확장이 계획되어 있습니다.
소프트웨어 스택 및 통합
Trainium 하드웨어는 컴파일러, 런타임 및 프로파일링 도구를 포함하는 소프트웨어 개발 키트인 AWS Neuron SDK로 지원됩니다. Neuron SDK는 PyTorch 및 TensorFlow와 같은 널리 사용되는 프레임워크로 작성된 모델을 Trainium에서 실행되는 최적화된 실행 파일로 변환합니다. 또한 자동 모델 병렬 처리를 지원하여 대규모 모델을 최소한의 사용자 개입으로 여러 칩과 인스턴스에 분할할 수 있습니다. AWS는 또한 Trainium을 SageMaker 플랫폼과 통합하여 고객이 GPU 기반 훈련에 사용되는 것과 동일한 관리 인터페이스를 통해 Trainium 인스턴스에서 훈련 작업을 시작할 수 있게 했습니다.
또한 AWS는 OpenAI 및 Anthropic과 협력하여 각각의 모델을 Trainium에 최적화했습니다. Amazon으로부터 상당한 투자를 받은 Anthropic은 Trainium2 UltraServer를 사용하여 Claude 모델 제품군을 훈련하고 배포하는 주요 초기 채택자였습니다. OpenAI는 또한 2024년에 일부 프론티어 모델 훈련에 Trainium2를 사용할 것이라고 발표하여 맞춤형 칩의 성능에 대한 주목할 만한 지지를 나타냈습니다.
성능 및 비용 효율성
AWS는 Trainium이 Nvidia A100 또는 H100 GPU로 구동되는 것과 같은 비교 가능한 GPU 기반 EC2 인스턴스에 비해 훈련 작업당 최대 50% 낮은 비용을 제공한다고 마케팅합니다. MLPerf와 같은 독립적인 벤치마크는 Trainium2가 특히 트랜스포머 기반 모델의 표준 훈련 작업에서 경쟁력 있는 성능을 달성한다는 것을 보여주었습니다. 추론의 경우 Trainium2는 낮은 지연 시간으로 높은 처리량을 제공하도록 설계되어 챗봇 및 코드 생성과 같은 실시간 애플리케이션에 적합합니다.
그러나 Trainium 주변의 생태계는 Nvidia의 CUDA 플랫폼보다 덜 성숙합니다. 일부 머신러닝 프레임워크와 라이브러리는 제한된 지원을 제공하며, 희소 어텐션 또는 사용자 지정 커널과 같은 특정 고급 기능은 추가 엔지니어링 노력이 필요할 수 있습니다. AWS는 Neuron SDK의 기능을 확장하는 데 적극적으로 투자하고 있으며 고객이 거의 최대 활용률을 달성할 수 있도록 일련의 성능 최적화 가이드를 출시했습니다.
채택 및 생태계
Anthropic 및 OpenAI 외에도 여러 다른 조직이 AI 워크로드에 Trainium을 채택했습니다. GPU 인프라 전문 클라우드 서비스 제공업체인 CoreWeave는 2024년에 고객에게 Trainium 기반 인스턴스를 제공할 것이라고 발표하여 AWS 자체 클라우드를 넘어 칩의 가용성을 확장했습니다. AI 칩 분야의 경쟁사인 Cerebras Systems도 Trainium과 함께 자체 하드웨어를 제공하기 위해 AWS와 파트너십을 맺었지만 두 제품은 서로 다른 시장 부문을 대상으로 합니다.
연구 커뮤니티에서 Trainium은 학술 환경에서 대규모 모델 훈련에 사용되었습니다. 예를 들어, 스탠포드 AI 연구소 및 버클리 AI 연구의 연구자들은 모델 확장 및 효율성 실험에 Trainium2를 사용한 연구를 발표했습니다. AWS는 또한 오픈 소스 AI 프로젝트를 진행하는 학술 기관에 무료 컴퓨팅 크레딧을 제공하는 Trainium 연구 프로그램을 설립했습니다.
경쟁 구도
Trainium은 Google Cloud의 Tensor Processing Units(TPU), Microsoft Azure의 Maia 칩 및 Groq의 언어 처리 장치를 포함한 다른 맞춤형 AI 가속기와 직접 경쟁합니다. TPU는 몇 년 동안 사용 가능했고 성숙한 소프트웨어 스택을 갖추고 있지만, Trainium의 장점은 AWS의 광범위한 클라우드 서비스와의 긴밀한 통합 및 공격적인 가격 책정에 있습니다. Trainium3의 등장은 특히 AI 컴퓨팅 수요가 공급을 계속 초과함에 따라 경쟁을 심화시킬 것으로 예상됩니다.
AWS의 맞춤형 실리콘 투자는 또한 대체 GPU 및 가속기 옵션을 제공하는 AMD 및 Intel과 데이터 센터에서 주목을 받고 있는 Arm Holdings 기반 설계에 대항하여 자리 잡고 있습니다. Trainium의 성공은 AWS가 강력한 소프트웨어 생태계를 유지하고 장기적으로 명확한 성능 및 비용 이점을 입증하는 능력에 달려 있습니다.
향후 방향
앞으로 AWS는 Trainium이 Anthropic과 협력하여 구축된 대규모 슈퍼컴퓨터인 Project Rainier에서 핵심 역할을 할 것이라고 밝혔습니다. 2024년에 발표된 Project Rainier는 수십만 개의 Trainium2 칩으로 구성될 예정이며 2025년 운영을 목표로 세계에서 가장 큰 AI 훈련 클러스터 중 하나가 될 것으로 예상됩니다. 이 프로젝트는 성장하는 AI 인프라 시장을 포착하기 위한 전략적 자산으로서 맞춤형 실리콘에 대한 AWS의 약속을 강조합니다.
또한 AWS는 엣지 및 온프레미스 배포를 위한 Trainium 사용을 탐색하고 있지만 구체적인 제품은 발표되지 않았습니다. 회사는 또한 현재 UltraServer 폼 팩터를 넘어 Trainium 시스템을 확장하기 위해 고급 패키징 및 상호 연결 기술에 투자하고 있으며, 잠재적으로 수조 개의 매개변수를 가진 모델 훈련을 가능하게 할 수 있습니다.
결론
AWS Trainium은 AI를 위한 맞춤형 실리콘에 대한 Amazon의 중요한 도박을 나타냅니다. 클라우드 플랫폼과 긴밀하게 통합된 목적 기반 칩을 제공함으로써 AWS는 고객에게 성능, 비용 및 확장성의 매력적인 조합을 제공하는 것을 목표로 합니다. 소프트웨어 성숙도와 생태계 채택에는 여전히 과제가 있지만, Trainium 제품군의 빠른 진화와 Anthropic 및 OpenAI와 같은 주요 AI 연구소의 지원은 Trainium이 앞으로 몇 년 동안 AI 하드웨어 환경에서 주요 플레이어가 될 것임을 시사합니다.