Skip to main content

AWS Trainium 2 (Trn2)

Overview

AWS Trainium 2 reached GA on December 3, 2024 at re:Invent 2024, as the second-generation AWS custom AI training chip. Performance is 4× the first generation, delivering 1,299 TFLOPS FP8 (dense) / 2,563 TFLOPS (sparse), 96GB HBM capacity, and 2.9 TB/s bandwidth. Trn2 instances contain 16 Trainium2 chips. Trn2 UltraServer connects 4 instances via NeuronLink, forming 64 chips.

Core Specifications

ItemSpecification
ArchitectureTrainium2 (NeuronCore-v3)
Process NodeTSMC 4nm (estimated)
NeuronCore8 per chip (NeuronCore-v3)
Memory96 GB HBM (3× 1st gen)
Memory Bandwidth2.9 TB/s (3.6× 1st gen)
FP8 Compute1,299 TFLOPS (dense) / 2,563 (sparse)
BF16/FP16/TF32667 TFLOPS
FP32181 TFLOPS
TDP~700 W
NeuronLinkNeuronLink-v3 (1.28 TB/s bidirectional)
DMA Bandwidth3.5 TB/s

Trainium 1 vs Trainium 2

MetricTrainium 1Trainium 2Improvement
Process7nm4nm
NeuronCore32 per chip8 per chipSimplified
HBM Capacity32 GB96 GB
HBM Bandwidth0.8 TB/s2.9 TB/s3.6×
FP8 Compute191 TFLOPS1,299 TFLOPS6.7×
BF16/FP16191 TFLOPS667 TFLOPS3.4×
FP3248 TFLOPS181 TFLOPS3.7×

Trn2 Instances and UltraServer

Form FactorConfigurationCompute
Trn2.48xlarge16 Trainium220.8 FP8 PFLOPS
Trn2 UltraServer64 Trainium2 (4× Trn2)83.2 FP8 PFLOPS
Project RainierHundreds of thousands of Trainium2 (Anthropic)5× Anthropic current

Vendor Information

ItemDetails
ManufacturerAmazon (Annapurna Labs)
Official Websitehttps://aws.amazon.com/machine-learning/trainium/
Product Pagehttps://aws.amazon.com/ec2/instance-types/trn2/
SDKhttps://awsdocs-neuron.readthedocs-hosted.com/
ReleaseDecember 3, 2024 GA

Key Features

  • Logical NeuronCore Configuration (LNC): Multiple physical cores merged into logical cores
  • NeuronLink 2D Torus: Low latency, high bandwidth
  • 30-40% price-performance advantage vs EC2 P5e/P5en (H100)
  • EFAv3 Interconnect: 3.2 Tbps
  • UltraServer 64 chips: Trillion-parameter training

Use Cases

  • LLM training (100B to trillion parameters)
  • Large-scale inference
  • Training + inference fungible