AWS Trainium3: NeuronCore-v4 기반 분산 LLM 학습 인프라
AWS Trainium3의 NeuronCore-v4, HBM3e, NeuronLink-v4 구조와 Trn3 기반 분산 LLM 학습 설계 및 GPU 전환 조건을 정리한다.
2026-08-14 · 최초 발행 2026-04-27
AWS는 2025년 12월 re:Invent 2025에서 발표한 Trainium3를 정식 출시(GA)했다. 3nm 공정 기반의 이 학습 전용 가속기는 NeuronCore-v4와 NeuronLink-v4를 중심으로 설계됐으며, Trainium2와 비교해 칩당 성능은 4.4배, 전력 효율은 4배 향상됐다.
LLM 학습은 수천 개의 가속기를 몇 주에서 몇 달 동안 연속 운용하는 워크로드다. 이 환경에서는 단일 칩 연산 성능 외에도 메모리 대역폭, 칩 간 인터커넥트, 수치 정밀도를 선택하는 유연성이 학습 속도에 직접 영향을 준다. AWS는 2022년 Trainium1과 2024년 Trainium2를 선보였고, Trainium3에서는 메모리 기술과 수치 형식, 인터커넥트 구조까지 다시 설계했다. Amazon이 2026년 AI 인프라에 약 2,000억 달러를 투자하겠다고 밝힌 방향과도 연결된다.
칩 안에서 분리된 연산과 로컬 메모리
Trainium3 한 장에는 NeuronCore-v4 8개가 들어간다. 각 코어는 정밀도와 처리량이 다른 연산 경로를 분리해 운용한다.
BF16 계산은 128×128 시스톨릭 어레이가 맡고, 처리량 중심의 학습 단계에는 512×128 MXFP8 시스톨릭 어레이를 사용한다. 칩 전체 FP8 컴퓨팅 성능은 2.52 PFLOPs다.
코어마다 32 MiB SRAM도 배치된다. 빈번하게 다시 쓰는 파라미터와 중간 활성화 값을 HBM으로 되돌아가지 않고 로컬에 캐싱할 수 있어, 어텐션 연산이나 레이어 정규화처럼 데이터 접근이 반복되는 패턴에서 활용할 수 있다.
지원하는 수치 형식은 FP32, BF16, MXFP8, MXFP4다. MXFP(Microscaling Floating Point)는 스케일 팩터를 그룹 단위로 공유해 FP8보다 낮은 비트 폭에서도 양호한 수렴 특성을 유지하도록 구성된다.
HBM3e와 NeuronLink-v4가 맡는 통신 계층
Trainium3의 메모리와 인터커넥트는 대규모 병렬 학습을 전제로 한다.
| 항목 | 사양 |
|---|---|
| 칩당 HBM3e 용량 | 144GB, Trainium2 대비 1.5배 |
| 메모리 대역폭 | 4.9 TB/s, Trainium2 대비 1.7배·Trn2 전체 대비 3.9배 |
| NeuronLink-v4 양방향 대역폭 | 2 TB/s |
NeuronLink-v4는 CPU나 NIC를 거치지 않고 칩끼리 직접 통신하도록 한다. AllReduce와 AllGather 같은 집합 통신을 이 계층에서 처리하므로, 텐서 병렬화(Tensor Parallelism)나 파이프라인 병렬화(Pipeline Parallelism)에서 통신 오버헤드가 병목이 되는 상황을 완화할 수 있다.
단일 서버 경계에서 확장하는 Trn3 UltraServer
Trn3 UltraServer는 Trainium3를 최대 규모로 묶은 단일 서버 구성이다.
| 구분 | 사양 |
|---|---|
| Trainium3 칩 수 | 최대 144개 |
| FP8 컴퓨팅 | 362 PFLOPs |
| 총 HBM3e 용량 | 20.7 TB |
| 집계 메모리 대역폭 | 706 TB/s |
수백억~수천억 파라미터 규모의 모델을 단일 서버 경계에서 학습할 수 있고, EFA(Elastic Fabric Adapter) 네트워킹을 사용하면 멀티 서버 클러스터로도 확장할 수 있다. GPT-OSS 벤치마크에서는 Trainium2 대비 칩당 처리량이 3배, 응답 시간은 4배 단축된 것으로 보고됐다.
PyTorch 코드와 NeuronSDK 사이의 접점
하드웨어 성능은 개발 환경에 흡수될 수 있을 때 의미가 있다. AWS는 Trainium3 접근 계층으로 NeuronSDK를 제공한다.
- Neuron Compiler는 그래프 최적화, 커널 퓨전, 데이터 레이아웃 변환을 자동 수행한다.
- Neuron Runtime은 칩 간 메모리 관리, 실행 스케줄링, 오류 복구를 담당한다.
- Neuron Training Library는 분산 학습 집합 통신과 자동 혼합 정밀도(AMP)를 지원한다.
- NKI(Neuron Kernel Interface)는 Trainium ISA에 저수준으로 접근하는 커스텀 커널 개발 인터페이스다.
AWS는 torch.nn 모듈 단위로 작성한 기존 학습 코드에서 디바이스 지정만 바꾸면 Trainium3에서 실행할 수 있다고 밝히고 있다. Hugging Face Transformers와 FSDP(Fully Sharded Data Parallel) 등 주요 학습 프레임워크와의 호환성도 계속 확장되고 있다.
NKI는 범용 컴파일러가 최적화하기 어려운 연산을 직접 구현하는 경로다. 플래시어텐션(FlashAttention) 변형이나 커스텀 희소 어텐션처럼 특수한 연산을 다룰 수 있어, 새로운 어텐션 메커니즘 또는 활성화 함수를 실험하는 연구 팀에 하드웨어 수준의 제어권을 제공한다.
서버 내부와 클러스터 외부의 병렬화
Trainium3 분산 학습은 서버 내부의 NeuronLink-v4와 서버 간 EFA를 다른 계층으로 사용한다.
서버 안에서는 NeuronLink-v4가 칩 간 통신을 맡고, 서버를 넘는 집합 통신은 EFA로 처리한다. NeuronSDK 컴파일러는 두 통신 계층의 토폴로지를 인식해 병렬화 전략과 통신 스케줄링을 최적화한다.
학습 구성에서는 다음 전략을 함께 조합하는 경우가 많다.
- 데이터 병렬화(DP): 배치를 칩 단위로 나누고 AllReduce로 그래디언트를 동기화한다.
- 텐서 병렬화(TP): 가중치 행렬을 칩 단위로 분할하고 NeuronLink-v4로 중간 결과를 교환한다.
- 파이프라인 병렬화(PP): 레이어를 서버 단위로 나누며, 마이크로배치 파이프라이닝으로 버블을 축소한다.
CUDA 의존성이 전환 비용을 가른다
Trainium3는 SageMaker, EC2, EKS를 포함한 AWS 에코시스템과 네이티브로 통합된다. 전력 효율을 우선하는 설계와 AWS 공식 벤치마크에서 NVIDIA H100 대비 유리하다고 제시한 가격 성능비도 도입 검토의 근거가 된다.
반대로 NVIDIA GPU는 CUDA 생태계의 수십만 개 기존 커널, 라이브러리, 도구를 활용할 수 있다. 커뮤니티에 쌓인 사례와 트러블슈팅 자료, 다른 클라우드 및 온프레미스로의 이식성도 NVIDIA 환경의 강점이다.
CUDA 종속성이 강한 커스텀 연산이 없고 PyTorch 표준 연산 위주로 구성한 학습 파이프라인이라면 NeuronSDK의 PyTorch 통합만으로도 마이그레이션을 검토할 수 있다. TensorRT-LLM이나 FlashInfer처럼 CUDA에 깊이 의존하는 추론 최적화 도구가 핵심인 시스템은 NKI 수준의 커스터마이징이 필요하거나, 단기적으로 NVIDIA 환경을 유지하는 편이 현실적이다.
TCO에 포함할 비용 항목
Amazon은 2026년 AI 인프라에 약 2,000억 달러를 투자할 계획을 밝혔고, Trainium3는 그 투자에서 핵심 자산으로 제시된다. 클라우드 사용자는 Trn3 인스턴스를 직접 쓰는 방식과 관리형 서비스를 연계하는 방식을 선택할 수 있다.
EC2 Trn3 인스턴스 또는 SageMaker 학습 잡으로 특정 워크로드를 실행하면 동등한 H100 구성과 학습 비용을 직접 비교할 수 있다. AWS Bedrock이나 SageMaker JumpStart를 활용하는 경우에는 Trainium3 기반 모델 서빙을 관리형 서비스로 사용해 인프라 운영 부담 없이 비용 효율을 확보하는 접근이 가능하다.
어느 방식을 택하든 초기 TCO 계산에는 NeuronSDK 학습 곡선과 마이그레이션 비용을 포함해야 한다. Trainium3는 NeuronCore-v4의 이중 시스톨릭 어레이, HBM3e 메모리 대역폭, NeuronLink-v4 인터커넥트를 결합해 LLM 학습에 맞춘 하드웨어 스택을 구성한다. 표준 PyTorch 연산 기반의 LLM 학습 파이프라인을 운영한다면 Trn3 인스턴스 파일럿으로 비용을 검증할 수 있다.
Sources
- https://aws.amazon.com/machine-learning/trainium/
- https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/arch/neuron-hardware/trainium3.html
- https://aws.amazon.com/blogs/machine-learning/aws-trainium3-is-now-generally-available/
- https://aws.amazon.com/ec2/instance-types/trn2/
- https://awsdocs-neuron.readthedocs-hosted.com/en/latest/neuron-guide/nki/index.html
- https://www.aboutamazon.com/news/aws/amazon-aws-investment-ai-data-centers-2025