Cloud Architecture

Custom AI Silicon in Production: AWS Trainium 3, Google Ironwood, and When to Break Your NVIDIA Dependency

A practitioner's guide to AWS Trainium 3, Google Ironwood TPU, and Microsoft Maia: when hyperscaler custom AI silicon beats NVIDIA GPUs on cost and performance, and what it actually takes to migrate.

Side-by-side comparison of AWS Trainium 3, Google Ironwood TPU, and NVIDIA B200 chip architectures with performance and memory specifications

For most of my career I told clients the same thing about AI hardware: if you need to train or run large models, use NVIDIA. Period. The CUDA ecosystem, the tooling, the documentation, the community knowledge base built up over fifteen years - it was so far ahead of anything else that choosing non-NVIDIA was choosing pain. AMD kept promising, Intel kept trying, and custom silicon from hyperscalers was either not publicly available or required so much framework surgery that the savings evaporated in engineering hours.

That calculus has shifted. Not completely, not painlessly, but enough that in 2026 I am having genuinely different conversations with teams about AI infrastructure. AWS Trainium 3 is now a credible option for large-scale training workloads. Google’s Ironwood TPU has closed the performance gap with NVIDIA Blackwell on specific workloads while dramatically cutting power consumption. Microsoft Maia 200 is quietly powering massive chunks of Azure OpenAI infrastructure. The hyperscalers have been building custom silicon for years, and those investments are maturing at exactly the moment when NVIDIA GPU pricing and availability constraints are making people look harder at alternatives.

This is the article I wish I had when I started evaluating these chips in production. I will tell you what they are good at, where they still fall short, and how to think about making the switch without gambling your team’s time on a dead end.

Why Every Hyperscaler Wants Out of NVIDIA’s Ecosystem

Before getting into the chips, it is worth understanding what is driving this. The hyperscalers have an obvious incentive to develop their own AI hardware: buying NVIDIA GPUs at scale costs tens of billions of dollars per year, and that money goes straight to a competitor they have no influence over.

But there is more to it than margin. When you train a large model on NVIDIA hardware, your software stack becomes deeply dependent on CUDA, cuDNN, and NCCL. Those are NVIDIA’s proprietary libraries. Swapping hardware means swapping everything below your model code. The hyperscalers want a world where they control the full stack from silicon to software, which gives them the ability to co-optimize in ways NVIDIA never will for a customer’s specific workloads.

The economic signal is hard to ignore. Custom ASIC shipments are projected to grow 44% year-over-year in 2026, compared to 16% for merchant GPUs. That growth is not speculative. It is hyperscalers deploying chips they already designed and validated internally before making them available externally.

For teams running GPU cloud infrastructure at scale, this represents the most significant shift in the hardware landscape since the A100 launch. Understanding the alternatives is no longer optional.

AWS Trainium and Inferentia: The Neuron SDK Story

AWS has the most mature story for external customers, and it starts with understanding the two-chip strategy: Trainium for training and Inferentia for inference. This is not arbitrary product differentiation. Training and inference have fundamentally different computational profiles.

Training involves large gradient computations, frequent model updates, and heavy all-reduce operations across nodes. Inference is about throughput and latency for forward passes only. Building separate silicon optimized for each workload makes sense when you are deploying at the scale AWS operates.

Trainium 3: The Current Flagship

Trainium 3 launched in 2025 and represents a substantial leap over the previous generation. AWS claims up to 4.4x more compute performance than Trainium 2, with a rack-scale system (the Trainium 3 UltraServer) that competes directly with NVIDIA’s NVL72 configurations.

The instance types you care about for production are the trn3.32xlarge with 16 Trainium 3 accelerators and 512 GB of HBM3e memory, and the newer UltraServer configurations that scale across an entire rack. The memory bandwidth is competitive with NVIDIA Blackwell for most transformer workloads.

The real story is cost. AWS’s published data suggests 30-50% cost savings for training workloads compared to equivalent GPU instances. In my experience talking to teams running production training pipelines, 30-35% is more realistic for general workloads, with higher savings possible for workloads that align well with the Neuron architecture.

The Neuron SDK: What You Actually Have to Deal With

This is where I need to be honest with you. The Neuron SDK has improved dramatically, but migration from CUDA is not a weekend project.

Historically, running PyTorch on Trainium required using torch_neuronx, a replacement API that wrapped standard PyTorch but handled compilation and graph optimization for the Neuron hardware. This created a two-code-path problem: you maintained one version for GPU/CUDA and another for Neuron.

Starting with Neuron 2.28, AWS has been transitioning to a native PyTorch backend. The TorchNeuron Native interface provides eager execution, torch.compile integration, and standard distributed APIs (FSDP, DTensor, DDP, tensor parallelism) through the same PyTorch interfaces you already use. This is a meaningful improvement. The goal is that your training code runs on Trainium with minimal modification beyond changing your distributed training configuration.

The catch is that “minimal modification” still means modification. The XLA compilation that underlies Neuron execution has quirks. Dynamic shapes in your model can cause recompilation at runtime, which destroys your throughput. Certain PyTorch operations are not supported in the NeuronX graph execution path and fall back to CPU, which you will discover through profiling, not through upfront documentation.

The practical migration path is: start with a stable, batch-oriented training workload where you control the input shapes. Get it running on a trn3.48xlarge spot instance. Measure actual cost versus your current GPU setup. Do not attempt to migrate your most dynamic, research-oriented training code first.

Neuron also integrates with HuggingFace, vLLM, and PyTorch Lightning, which matters for LLM inference workloads where you might want to serve models on Inferentia after training on Trainium.

AWS Trainium 3 Neuron SDK software stack showing TorchNeuron Native, XLA compilation layer, and hardware abstraction across trn3 instance families

Inferentia 3 for Production Inference

Inferentia 3 is worth a separate mention because the cost case here is often stronger than for training. For high-throughput, latency-tolerant inference on well-quantized models, Inferentia instances consistently deliver better tokens-per-dollar than equivalent GPU instances. I have seen teams run LLaMA 3.1 70B inference on inf3 instances at 40-50% the cost of equivalent H100 inference instances with acceptable latency for batch and asynchronous workloads.

Real-time, low-latency inference is a different story. If you need sub-100ms time-to-first-token for interactive applications, you will probably stay on NVIDIA H200 or Blackwell. Inferentia’s compilation model creates overhead that GPU instances avoid. But for async inference pipelines, document processing, embedding generation at scale - Inferentia is genuinely compelling.

Google TPU: The XLA Heritage

Google has been running TPUs internally since 2016. The external version has been available since v2, and the ecosystem has matured through multiple generations. The current externally available chips are Trillium (v6e) and Ironwood, with Ironwood being the newest general-purpose generation.

TPU Architecture and Why It Is Different

TPUs are systolic array accelerators. Where a GPU uses many general-purpose compute cores (CUDA cores plus tensor cores), a TPU routes data through a fixed matrix of multiply-accumulate units in a wave pattern. This is spectacularly efficient for the forward and backward passes of transformer training, where the data flow patterns are predictable. It is much less efficient for workloads with irregular memory access patterns or heavy branching.

The memory situation is more nuanced than it appears. Trillium v6e has less HBM than H100/H200, which limits it for very large models without model parallelism. Ironwood addresses this directly with 192 GB of HBM3e per chip - matching the NVIDIA B200 - and 7.4 TB/s of memory bandwidth that exceeds what Blackwell offers. The power efficiency is particularly striking: Trillium runs at roughly 300W per accelerator versus 700W for H100, which has real implications for data center thermal design and electricity costs.

The XLA Ecosystem and PyTorch Support

Here is the honest reality of TPU software compatibility: JAX is a first-class citizen. PyTorch is a second-class citizen. TensorFlow is technically supported but increasingly legacy even at Google.

If your team is building with JAX (which a lot of ML research teams now are, particularly for custom training loops and research-oriented work), TPUs are an excellent choice. The XLA compilation that JAX uses is exactly what TPUs were designed to run, and the optimization passes that XLA applies to JAX programs produce excellent results on TPU hardware.

PyTorch on TPUs runs through torch_xla, which compiles PyTorch operations to XLA and runs them on the TPU. This works, but you need to be thoughtful about it. The lazy tensor execution model in torch_xla means that PyTorch code runs in a compiled graph mode rather than eager mode. Operations that seem benign in standard PyTorch (like Python-side control flow that depends on tensor values) can cause the XLA graph to break and recompile, resulting in massive performance degradation.

The operational model is also different. TPUs are accessed through Google Cloud TPU nodes, which are separate VM types from standard Compute Engine instances. Your driver VM connects to TPU slice over the network. This network hop introduces coordination overhead that GPU clusters (where the GPU is directly attached to the host) do not have. Google has improved this, but it is a real difference in the operational model.

For teams running workloads on enterprise AI platforms and evaluating GCP’s AI infrastructure, Vertex AI and TPU integration is the natural path. Google makes it reasonably straightforward to provision TPU slices through Vertex AI Training rather than managing raw TPU nodes.

Cost and Performance Reality

Google’s numbers for Ironwood are impressive on paper: 4,614 TFLOPS at FP8, 192 GB HBM3e, 2x better performance per watt than Trillium. For JAX-based transformer training, Trillium v6e has already demonstrated roughly 45% training cost reduction versus comparable NVIDIA configurations in published benchmarks.

The catch is that benchmark conditions are always optimized for the hardware being benchmarked. For production workloads with all the messiness that entails (variable sequence lengths, mixed precision issues, checkpoint/restore overhead), the real-world advantage is meaningful but smaller than the headline numbers suggest.

My rule of thumb: if you are running JAX-based training at scale, evaluate TPUs seriously. If you are running PyTorch at scale and have any significant use of dynamic shapes, custom CUDA kernels, or third-party libraries with GPU-specific paths, stay on NVIDIA unless you have budget to run a careful evaluation.

Google Ironwood TPU versus NVIDIA B200 comparison showing FP8 TFLOPS, memory capacity, and power consumption across workload types

Microsoft Maia 200: The Mostly-Internal Option

Microsoft Maia 200 is worth understanding even though it is not broadly available as an external compute option. It is the chip powering a significant portion of Azure OpenAI’s infrastructure and increasingly powering Microsoft’s own Copilot products.

What Maia represents is Microsoft’s bet that training and running large language models internally (rather than on NVIDIA hardware) makes economic sense at their scale. Microsoft has not published Maia’s detailed specs, and external availability through Azure is limited. When I have seen Maia available through Azure AI Infrastructure, it has been for specific partnerships or early access programs rather than general availability.

The strategic implication is more interesting than the operational one. Microsoft is demonstrating that you can build hyperscale AI infrastructure without being completely dependent on NVIDIA. As Maia matures and potentially becomes more broadly available, it could become a serious option for Azure-native workloads. For now, treat it as infrastructure intelligence rather than an immediate migration target.

Meta’s MTIA (Meta Training and Inference Accelerator) is in a similar category. Meta uses it at massive scale internally but does not offer it externally. It matters for understanding where the industry is heading, not for your next infrastructure decision.

When Custom Silicon Beats NVIDIA: A Practical Framework

After running evaluations of all three platforms, here is how I think about the decision:

Stay on NVIDIA if:

Your training code relies on CUDA-specific extensions or custom CUDA kernels. Migrating custom CUDA kernels to Neuron or XLA is genuinely difficult. The ecosystem of PyTorch extensions (FlashAttention, Triton, specialized quantization kernels) has grown assuming CUDA as the target. Much of it does not port cleanly.

Your workload involves dynamic shapes that change frequently during training. Both Trainium and TPU depend on XLA-style ahead-of-time compilation, which is efficient when the computational graph is stable but expensive when it recompiles constantly.

Your team needs to iterate quickly on model architecture research. NVIDIA’s debugging tools (Nsight, the CUDA profiler) are mature and well-understood. Debugging custom silicon is harder because the tooling is less mature and the internal architecture is less documented.

You need to run arbitrary ML frameworks beyond PyTorch and JAX. ONNX support, deployment with TensorRT, custom deployment pipelines built around CUDA assumptions - all of these work reliably on NVIDIA and require work on custom silicon.

Seriously evaluate Trainium if:

You are running stable, large-scale training jobs that run for hours or days with consistent input shapes. This is the sweet spot. Large language model pre-training, continued pre-training on domain-specific data, and FSDP-based training with fixed sequence lengths are exactly where Trainium delivers.

You are already heavily invested in the AWS ecosystem. If your data is in S3, your orchestration is on EKS or SageMaker, and your team knows AWS operations, the integration overhead of Trainium is much lower than if you were starting from scratch.

Your current GPU training costs are significant enough that a 30% reduction is strategically meaningful. Below a certain threshold (rough rule: less than $50K/month in GPU costs), the engineering investment to evaluate and migrate probably does not pay back quickly enough.

Seriously evaluate TPUs if:

Your training is JAX-based, which increasingly describes cutting-edge research teams. If you are writing custom training loops in JAX, model parallelism with JAX sharding APIs, or using frameworks like MaxText or Praxis, TPUs are architecturally aligned with your stack.

You are primarily a GCP shop and want native integration with Vertex AI, BigQuery ML, and Google’s data infrastructure. The ecosystem integration story is real and reduces operational overhead.

Your workloads are transformer-heavy with stable shapes. The systolic array architecture that TPUs use is genuinely efficient for standard transformer attention and feedforward operations. You pay a price for diversity; you gain efficiency for conformity.

The Hidden Costs of Migration

I want to be direct about something that often gets undersold in the vendor materials: migrating to custom silicon has real engineering costs that do not appear in the GPU comparison spreadsheet.

The first is framework porting time. Even with Neuron’s improved PyTorch native support, you should budget two to four weeks of engineering time per training job to evaluate, port, and validate. Multiply that by your number of distinct training workflows.

The second is debugging complexity. When something goes wrong on NVIDIA, you have fifteen years of Stack Overflow answers, blog posts, and community knowledge. When something goes wrong on Trainium or TPU, you are working with much thinner documentation and a smaller community. Support tickets to AWS and Google can be slow for unusual issues.

The third is performance validation overhead. You cannot simply assume that a model running correctly produces the same outputs on different hardware. Numerical precision differences between CUDA and XLA implementations can produce subtle divergence in model behavior that is hard to catch without careful output comparison. I have seen teams spend weeks tracing a quality degradation in a fine-tuned model back to a subtle difference in how mixed precision was handled on the porting path.

For LLM fine-tuning workloads, this validation overhead is particularly important. LoRA and QLoRA implementations depend on specific numerical behavior that may not translate cleanly.

The fourth is opportunity cost. Every engineer-week spent on a hardware migration is an engineer-week not spent on model quality, product features, or other infrastructure work. The ROI calculation needs to be honest about this.

Total cost of ownership comparison for AI training workloads: NVIDIA H100 clusters vs AWS Trainium 3 vs Google TPU Trillium v6e over 12 months at 100 GPU-equivalent scale

A Migration Strategy That Does Not Blow Up in Your Face

If you have decided the economics make sense, here is a migration approach that has worked for teams I have advised:

Start with a shadow evaluation, not a production migration. Pick one training job that runs regularly, is not on the critical path, and produces outputs you can validate automatically. Run the exact same job on your current NVIDIA setup and on Trainium or TPU in parallel for two to four weeks. Compare cost, runtime, and most importantly, output quality. If the outputs diverge in ways you cannot explain, stop and investigate before going further.

Pick workloads that are already containerized and have clean dependency management. Training jobs that depend on a soup of ad-hoc CUDA packages installed on a custom AMI are much harder to migrate than jobs running in standardized containers with pinned dependencies.

Use the official integrations where they exist. Both AWS (through SageMaker) and Google (through Vertex AI Training) provide managed training infrastructure on custom silicon. These integrations handle a lot of the operational complexity around job scheduling, checkpointing, and monitoring. Going straight to raw trn3 or TPU node management without the managed layer means taking on more operational burden than the savings justify for most teams.

Budget explicitly for the performance validation phase. Define in advance what constitutes acceptable output equivalence for your workloads. For pre-training jobs, loss curve convergence behavior is the right signal. For fine-tuning and alignment, you need eval benchmark scores. For inference, you need output quality metrics specific to your use case.

Consider AI FinOps practices as part of the migration. Custom silicon changes your pricing model (some workloads now benefit from reserved capacity on Trainium rather than on-demand NVIDIA), and setting up proper cost allocation and tracking before the migration gives you clean data to measure ROI.

The Trajectory to Watch

The reason this conversation is worth having in 2026 but was not worth having in 2022 is trajectory. NVIDIA’s custom silicon competitors have been on a steep improvement curve while NVIDIA’s own advantage has been moderating.

Trainium 3 is 4.4x faster than Trainium 2, which was itself a major step from Trainium 1. Google has shipped seven generations of TPUs, with each generation closing the gap further. Microsoft is deploying Maia at scale for its most cost-sensitive AI workloads. Meta MTIA is reportedly running the majority of Meta’s recommendation model inference.

The software story is also improving. Neuron’s native PyTorch support is removing the biggest porting barrier for AWS. Google’s continued investment in PyTorch/XLA and the increasing prevalence of JAX in research pipelines creates a larger talent pool with TPU experience. The ecosystem is catching up to the hardware.

The counterargument - which is legitimate - is that NVIDIA is not standing still. Blackwell is a significant architecture, and Rubin is coming. NVIDIA’s NIM microservices and deployment tooling are getting better. The CUDA ecosystem continues to attract investment from every major framework team. The NVIDIA advantage in software maturity may persist longer than the custom silicon enthusiasts predict.

My view is that over a three-to-five year horizon, teams will not choose between NVIDIA and custom silicon. They will have heterogeneous AI compute infrastructure and route workloads to the most cost-effective hardware based on workload characteristics - similar to how mature cloud teams route general compute workloads based on instance type and spot availability. Understanding serverless GPU alternatives is part of the same picture: the market for AI compute is diversifying, and that is good for everyone running AI workloads at scale.

The work I am doing with teams right now is building the tooling and organizational knowledge to be hardware-flexible rather than hardware-locked. That investment pays off regardless of which chip wins the next generation.

What to Do Next

If you are spending more than $100K per month on GPU compute and have not done a serious evaluation of Trainium or TPU options, you have probably left real money on the table. The economics are different enough now to justify the evaluation time.

Start by identifying two or three training jobs that are good candidates based on the framework I described above: stable shapes, PyTorch or JAX codebase, no custom CUDA kernels. Get access to a trn3.48xlarge or a TPU v5p slice and run your benchmark. Do not start with your most complex workload. Do not try to migrate everything at once. The teams that have succeeded with this have treated it as an infrastructure engineering project with a proper discovery, evaluation, and migration phase - not as a weekend hack.

The Kubernetes infrastructure for AI workloads your team is already running will need to accommodate Trainium and TPU node types as you expand, which means understanding how your job scheduling, resource quotas, and cost allocation frameworks handle heterogeneous accelerator types. That is a solvable problem, but it is worth planning for before you are trying to do it under deadline pressure.

The NVIDIA monopoly on serious AI compute is ending. Not completely, not quickly, but definitively. Understanding the alternatives is increasingly part of what it means to design AI infrastructure that stays economically competitive over the next five years.