새로운Amazon Elastic Compute Cloud(Amazon EC2) Trn2 인스턴스와 Trn2 UltraServer는 ML 학습과 추론을 위한 가장 강력한 EC2 컴퓨팅 옵션입니다. 2세대 AWS Trainium 칩(AWS Trainium2)을 적용한 Trn2 인스턴스는 1세대 Trn1 인스턴스에 비해 4배의 속도와 4배의 메모리 대역폭, 3배의 메모리 용량을 제공합니다. Trn2 인스턴스는 기존의 GPU 기반 EC2 P5e 인스턴스와 P5en 인스턴스에 비해 가격 효율이 30~40% 더 좋습니다.
16개의 Trainium2 칩 외에도 각각의 Trn2 인스턴스에는 192 vCPUs, 메모리 2 TiB, Elastic Fabric Adapter(EFA) v3 네트워크 대역폭 3.2 Tbps를 갖추고 있어 이전 세대에 비해 지연 시간이 최대 35%까지 줄었습니다.
완전히 새로운 컴퓨팅 옵션인 Trn2 UltraServer는 64x Trainium2 칩을 대역폭은 높고 지연 시간은 짧은 NeuronLink 인터커넥트에 연결하여 프론티어 기반 모델에서 최고의 추론 및 학습 성능을 제공합니다.
이미 수만 개의 Trainium 칩이 Amazon과 AWS 서비스를 지원하고 있습니다. 예를 들어 최근 프라임 데이에서는 8만 개가 넘는 AWS Inferentia 칩과 Trainium1 칩이 Rufus 쇼핑 어시스턴트를 지원했습니다. Trainium2 칩은 이미 Amazon Bedrock에서 Llama 3.1 405B와 Claude 3.5 Haiku 모델의 지연 속도 최적화 버전을 지원하고 있습니다.
양적 확장과 질적 확장
혁신적인 아키텍처 형태가 마찬가지로 혁신적인 컴퓨팅 성능과 결합되어 프론티어 모델의 규모와 복잡성을 지속적으로 성장시킵니다. 복잡성이 낮았던 시기에는 확장성의 아키텍처 설계에 대해 더 큰 컴퓨터를 사용한 질적 확장(스케일업)과 더 많은 컴퓨터를 사용한 양적 확장(스케일아웃)으로 나누어 이야기할 수 있었습니다. 잠깐 동안에도 Trainium2 칩, Trn2 인스턴스는 물론 그보다 대규모의 컴퓨팅 제품들에 대해 이야기할 수 있는 요즘의 현실에서는 두 모델을 동시에 적용할 수 있는 것 같이 보이면서도 전체적인 계층 구조에서는 수준 차이가 있는 것 같습니다. NeuronCore부터 UltraCluster까지 확장되는 Trn2 구성 요소에 대해 다음과 같이 살펴봅시다.
NeuronCores는 Trainium2 칩의 핵심적인 부분입니다. 각각의 3세대 NeuronCore에는 스칼라 엔진(입력 1개에서 출력 1개) 하나, 벡터 엔진(여러 입력에서 여러 줄력) 하나, 텐서 엔진(시스톨릭 배열 곱셈, 컨볼루션, 전치) 하나, GPSIMD (범용 단일 명령 다중 데이터) 코어 하나가 들어있습니다.
각각의 Trainium2 칩에는 NeuronCore 8개와 고대역폭 메모리(HBM) 96GiB가 탑재되어 있으며 HBM 대역폭을 초당 2.9TB로 지원합니다. 코어를 개별적으로 취급하여 사용하거나 물리적 코어 쌍을 단일 논리적 코어로 그룹화할 수 있습니다. 단일 Trainium2 칩은 최대 1.3페타플롭의 고밀도 FP8 컴퓨팅과 최대 5.2페타플롭의 스파스 FP8 컴퓨팅을 제공하며 HBM 대기열의 자동 재정렬 덕분에 메모리 대역폭을 95% 활용할 수 있습니다.
각각의 Trn2 인스턴스에는 차례대로 16개의 Trainum2 칩이 있습니다. 총 128개의 뉴런 코어, 1.5TiB의 HBM, 초당 46Tb의 HBM 대역폭에 해당합니다. 이를 모두 합하면 고밀도 FP8 컴퓨팅의 경우 최대 20.8페타플롭으로, 스파스 FP8 컴퓨팅의 경우 최대 83.2페타플롭으로 증가합니다. Trainium2 칩은 NeuronLink를 통해 2D 토러스로 연결되어 1GB/초의 속도로 대역폭은 높고, 지연 시간은 짧은 칩-투-칩 통신을 제공합니다.
UltraServer에는 지연 시간이 짧고 대역폭이 높은 NeuronLink에 연결된 Trn2 인스턴스 4개가 있습니다. 이는 NeuronCore 512개, Trainium2 칩 64개, HBM의 6TiB, 초당 185TB의 HBM 대역폭에 해당됩니다. 계산해 보면 고밀도 FP 컴퓨팅의 경우 최대 83페타플롭으로, 스파스 FP8 컴퓨팅의 경우 최대 332페타플롭으로 증가합니다. 인스턴스 내 NeuronCore를 연결하는 2D 토러스 외에도, 인스턴스 4개가 각각 해당 XY 위치에 있는 코어에서 고리 모양으로 연결됩니다. 추론을 위해 UltraServer는 업계 최고 수준의 응답 시간을 제공하여 최상의 실시간 경험을 제공합니다. 학습을 위해 UltraServer는 독립형 인스턴스에 비해 모델 병렬 처리를 위한 더 빠른 집단 통신을 통해 모델 학습 속도와 효율성을 높입니다. UltraServer는 1조 파라미터 수준 이상의 학습 및 추론을 지원하도록 설계되며 미리 보기 형식으로 제공되는데, 미리 보기에 참여하려면 당사에 문의하면 됩니다.
Trn2 인스턴스와 UltraServer는 EC2 UltraCluster에서 배포되어 단일 페타비트 규모의 논블로킹 네트워크에서 Amazon FSx for Lustre 고성능 스토리지에 액세스 가능한 상태로 수만 개의 Trainium 칩에서 스케일아웃 분산 학습을 가능하게 합니다.
Trn2 인스턴스 사용
현재 Trn2 인스턴스는 미국 동부(오하이오) AWS 리전에서 프로덕션 용도로 사용할 수 있으며 Amazon EC2 Capacity Blocks for ML을 사용하여 예약할 수 있습니다. 최대 6개월 간 인스턴스 64개를 예약할 수 있는데 필요한 경우 즉시 시작 시간과 예약 연장 기능을 갖춘 상태로 최대 8주 미리 예약 승인을 받을 수 있습니다. 자세한 내용은 기계 학습 워크로드를 위한 GPU 용량 확보용 Amazon EC2 Capacity Blocks for ML 발표를 참조하세요.
소프트웨어 측면에서는 AWS Deep Learning AMI를 통해 시작할 수 있습니다. 이러한 이미지는 이미 알고 있으며 사용 중이실 PyTorch, JAX 등의 프레임워크, 도구를 통해 미리 구성되어 있습니다.
AWS Neuron SDK를 사용하여 앱을 구축한 경우, 해당 앱을 가져와 Trn2 인스턴스용으로 다시 컴파일할 수 있습니다. 이 SDK는 기본적으로 JAX, PyTorch는 물론, 필수 라이브러리인 Hugging Face, PyTorch Lightning, NeMo 등과 통합됩니다. Neuron에는 오픈 소스 PyTorch 라이브러리인 NxD Training 및 NxD Inference를 사용하여 분산 학습 및 추론을 위한 기본 최적화 기능이 포함되어 있으며 프로파일링 및 디버깅에 대한 심층적인 인사이트를 제공합니다. 또한, Neuron은 안정적인 HLO 및 GSPMD를 비롯한 OpenXLA를 지원하므로 PyTorch/XLA 개발자, JAX 개발자가 Trainium2용 Neuron 컴파일러 최적화를 활용할 수 있습니다.
— Jeff;