Tranium is AWS’s machine learning accelerator, and this week at its re:Invent event in Las Vegas, it announced the second generation, the cleverly named Trainium2, purpose-built to enhance the training of large-scale AI models, including foundation models and large language models.

The original Tranium was designed before GenAI began driving the conversation and thus wasn’t competitive against NVIDIA’s training accelerators for those workloads. AWS closes that gap with Trainium2, saying that the new part delivers up to four times faster training performance and up to twice the energy efficiency of Trainium1.

AWS Trainium2

Trainium2 addresses the key shortcomings of its predecessors, Trainium1 and Inferentia2, while introducing new capabilities in hardware and software integration. The new accelerator will improve the relevance of AWS’s in-house silicon in training and inference workloads for LLMs.

Let’s take a look at what AWS disclosed:

1. Hardware Enhancements:
   • Performance:
     • 500W chip with 650 TFLOP/s BF16 compute.
     • 96GB HBM3e memory to support high bandwidth workloads.
   • Networking:
     • AWS introduced a 3D torus topology in the Trainium2 Ultra configuration (64 chips across two racks), enabling improved tensor parallelism and activation sharding.
     • NeuronLink advancements for scale-up networking to better align with Google’s TPU interconnects and Nvidia’s NVLink.
2. Cost Efficiency:
   • Trainium2 Ultra delivers what AWS claims is 64% lower TCO than Nvidia’s H100 in ethernet-based deployments.
   • Effective training costs are estimated to be 45% lower per petaflop-hour than Nvidia H100 deployments.
3. Software Ecosystem:
   • Shift from PyTorch XLA to improved frameworks, including JAX, which are better suited for Trainium’s torus topology.
   • Introduction of the Neuron Kernel Language (NKI), a domain-specific language for high-efficiency kernel programming.
   • Expanded education efforts, including collaboration with Stanford to support NKI adoption.
4. Use Cases:
   • Amazon’s internal workloads and Anthropic’s GenAI applications are primary users of Trainium2.
   • Focused on GenAI inference workloads due to better arithmetic intensity for memory-bound operations.

Competitiveness

AWS’s Trainium2 chip competes directly with Google’s TPUv6e and Nvidia’s H100/Blackwell GPUs.

Key differentiators for Tranium include:

• Arithmetic Intensity:
  • Trainium2: 203 BF16 FLOP/byte.
  • Competitors: TPUv6e/GB200/H100 range from 300 to 560 BF16 FLOP/byte.
• Network Topology:
  • Trainium2: Point-to-point torus topology offers efficient tensor parallelism but lacks the all-to-all connectivity of Nvidia’s NVLink.
  • Nvidia H100: Advanced NVLink domains (72-GPU connectivity) reduce inference costs significantly.
• Software Ecosystem:
  • AWS trails Nvidia’s CUDA ecosystem in maturity but is making strides with NKI and JAX.

Analysis

Trainium2 is the natural evolution of AWS’s AI hardware strategy, addressing prior deficiencies and introducing competitive features for the generative AI workloads that dominate the conversation.

However, Amazon remains heavily reliant on NVIDIA for overall capacity, with Trainium2 being an incremental step rather than a full shift toward silicon independence for the CSP.

• Customer Impact:
  • Trainium2’s cost advantages and improved performance could attract cost-sensitive customers focused on LLM inference.
  • Limited training effectiveness and ecosystem maturity may hinder adoption by high-performance AI customers.
• Competitive Impact:
  • NVIDIA maintains a dominant position with its established ecosystem and performance metrics.
  • Google’s TPU retains an edge in arithmetic intensity and scale-up efficiency.
  • AWS’s investments signal a long-term commitment to competing in custom silicon, likely increasing pressure on Nvidia and Google to innovate further.

AWS’s Trainium2 provides a more viable option for GenAI workloads, particularly in inference. Its success will hinge on the efficiency of its software stack, continued customer adoption, and the realization of cost advantages at scale. While not yet a market disruptor, Trainium2 establishes a foundation for AWS to compete more effectively in custom AI silicon.

Competitive Landscape & Advice to IT Buyers

These sections are only available to NAND Research clients. Please reach out to [email protected] to learn more.