How to Self-Host AI Inference on AWS Trainium2 With Native PyTorch (Instead of NVIDIA GPUs)

Key points:

• Trainium2 inference runs on the AWS Neuron SDK, not CUDA. You need the Neuron driver, runtime and compiler (neuronx-cc) plus TorchNeuron, the native PyTorch backend for Trainium. Then you serve the model with the vLLM Neuron plugin.
• TorchNeuron is real, and it installs as torch-neuronx. AWS calls it "an entirely new native PyTorch backend" that keeps its predecessor's package name. Versions 2.13 and 2.14 are TorchNeuron (neuron device, torch.compile, no XLA); 2.9 is the last PyTorch/XLA release. torch-neuron (no x) is the archived Inferentia1 package.
• TorchNeuron is in public beta as of Neuron SDK 2.33.0 (October 19, 2026) and targets Trn2 and Trn3 only. Trn1 and Inf2 stay on the XLA-based torch-neuronx.
• The entry point is trn2.48xlarge: 16 Trainium2 chips, 128 physical NeuronCores (shown as 64 logical cores by default), and 1.5 TB of HBM.
• On EKS, pods request aws.amazon.com/neuron or aws.amazon.com/neuroncore instead of nvidia.com/gpu. The old neurondevice name was dropped in Neuron 2.20.
• The real tradeoff is ecosystem coverage, not throughput. You get a compile step, a narrower model list, and no drop-in path for custom CUDA kernels. Benchmark your exact model before you commit capacity.
• You do not need a platform team to run this. Qovery operates the cluster layer inside your own AWS account, so Trn2 Capacity Blocks and reservations stay on your bill. The same workflow runs on GCP, Azure, Scaleway, or an existing Kubernetes cluster.

Most teams stay on NVIDIA because of the toolchain, not the silicon. CUDA works with every model on day one, and porting to anything else used to mean rewriting your PyTorch code around PyTorch/XLA.

That changed at the end of 2025, when AWS announced TorchNeuron, a native PyTorch backend for Trainium. It reached public beta in Neuron SDK 2.33.0 on October 19, 2026. This guide covers what it takes to self-host inference on Trainium2 today, what it costs next to H100s, and when you should not bother. I checked every version, command and price against AWS documentation on October 22, 2026.

What do you need to run inference on AWS Trainium2 instead of an NVIDIA GPU?

To run inference on Trainium2 you need a Trn2 instance, the Neuron driver (aws-neuronx-dkms), the Neuron runtime and collectives, the neuronx-cc compiler, and TorchNeuron (torch-neuronx 2.13 or 2.14) as the PyTorch backend. Every CUDA layer you know has a Neuron equivalent, and the vLLM Neuron plugin gives you an OpenAI-compatible server on top.

Start with the naming, because that is where most people get confused:

• TorchNeuron is the open-source native PyTorch backend for AWS Trainium. It registers as a native PyTorch device through PrivateUse1, so your code uses device="neuron" and torch.compile(model, backend="neuron").
• torch-neuronx is the package name. AWS's own wording: TorchNeuron "retains the same Python package name as its predecessor (torch-neuronx)" but is "an entirely new native PyTorch backend". Per the 2.33 release notes, torch-neuronx for PyTorch 2.13 and 2.14 is TorchNeuron, and PyTorch 2.9 is the last version on PyTorch/XLA.
• torch-neuron (no x) is the Inferentia1 package. It is archived and no longer actively developed.

Three things behave differently from CUDA. First, everything goes through the compiler: in eager mode each operation "compiles on first use and is cached for later calls", and torch.compile requires static shapes (dynamic=False). Second, the scheduling unit is the NeuronCore. A Trainium2 chip has eight NeuronCore-v3 cores, and the default Logical NeuronCore setting of 2 pairs them, so a trn2.48xlarge "presents 64 available NeuronCores". Third, the hardware generation decides your software path: TorchNeuron and vLLM Neuron run on Trn2 and Trn3, while Inf2 and Trn1 stay on the XLA lineage.

The shortcut I recommend is the Neuron Deep Learning AMI or the Neuron Deep Learning Containers. They ship the full toolchain pinned and tested together, which beats chasing package versions by hand. Run neuron-ls and neuron-top on a fresh instance to confirm the chips and cores are visible before you touch Python.

Which models and frameworks actually work on Trainium2 today?

Mainstream decoder-only LLMs such as Llama 3, Qwen3 and GPT-OSS run on Trainium2 through the vLLM Neuron plugin, and a wider list including Mixtral, Llama 4 and DBRX is supported by the older NxD Inference library. Custom CUDA kernels and brand-new architectures do not run without porting work, and the supported list changes with every Neuron SDK release.

As of SDK 2.33, the vLLM Neuron plugin (version 0.29.0.1.2.0, beta) lists validated models including Llama 3 (1B, 8B, 70B), GPT-OSS (20B, 120B), Qwen3-32B, Qwen3-VL-32B and Qwen3-Embedding-8B. The NxD Inference model reference covers Llama 2 through 3.3, Llama 4 Scout and Maverick, Mixtral, DBRX, Qwen2.5, Qwen3 and Qwen3 MoE, with FLUX.1 in beta.

You have three realistic paths:

• vLLM Neuron plugin for LLM serving behind an OpenAI-compatible API. This is where AWS is putting its effort, and since 2.33 it "uses native PyTorch compilation by default".
• NxD Inference if your model is only on its list. AWS has put it in maintenance status with the guidance "Migrate to vLLM Neuron", so treat it as a bridge.
• TorchNeuron directly (or optimum-neuron for Hugging Face models) for encoders, embeddings, vision and anything you serve with your own code.

The friction is real, and I would rather you hear it from me than from your on-call engineer:

• Custom CUDA kernels have no drop-in equivalent. NKI lets you write your own kernels for Trainium, but somebody on your team writes them.
• TorchNeuron has documented gaps. It does not support torch.sparse, torch.func, torch.nested or torch.distributions, has no TorchAO integration, and runs float64 as float32.
• Tensor parallel degree is baked into the compiled artifact. Change it, or the SDK version, and you recompile.
• Both TorchNeuron and vLLM Neuron are labelled beta. Pin versions and read release notes before every upgrade.

The spike that saves you a quarter: compile your exact model with your exact sequence lengths, then measure tokens per second, time-to-first-token and p99 latency. Write down the SDK release you tested, because the answer may differ next release.

How do you compile and serve a model with torch-neuronx, step by step?

Launch a Trn2 instance from the Neuron DLAMI, install TorchNeuron from the AWS Neuron pip repository, compile with torch.compile(backend="neuron") or let vLLM Neuron compile for you, serve with vllm serve, then containerize from the Neuron Deep Learning Container. The first compile is slow and gets cached, so plan for it in your deploy pipeline.

Step 1 - Get capacity and verify the hardware. In the US, trn2.48xlarge is available through EC2 Capacity Blocks in US East (Ohio). Launch from the Neuron DLAMI, then run neuron-ls.

Step 2 - Install TorchNeuron. The DLAMI already has it. For a manual install, the AWS setup guide gives these commands for PyTorch 2.14 (copied verbatim; system packages for the driver and runtime come from the Neuron apt repo on the same page):

python3.12 -m venv aws_neuron_venv_pytorch_2_14
source aws_neuron_venv_pytorch_2_14/bin/activate
python -m pip install --upgrade pip
python -m pip install "torch-neuronx==2.14.*" \
"neuronx-cc==2.*" nki \
--extra-index-url https://pip.repos.neuron.amazonaws.com \
--extra-index-url https://download.pytorch.org/whl/cpu

The quickstart notes the standalone wheels in this release support Ubuntu 26.04 only. On 22.04 or 24.04, use the DLAMI. Keep the version pins: earlier torch-neuronx versions are the XLA-based product.

Step 3 - Compile. For encoders and custom models, move the model to the neuron device and wrap it with torch.compile(model, backend="neuron"). No import torch_neuronx is required, since PyTorch autoloads the backend. For LLMs, vLLM Neuron runs the compilation itself. Large graphs can take "several minutes" to compile, according to the vLLM Neuron design docs, and the online-serving quickstart raises VLLM_NEURON_COMPILATION_TIMEOUT to 1200 seconds for GPT-OSS.

Step 4 - Serve. Install the plugin from the vllm-neuron repo and start the server. This example is from the plugin README:

vllm serve openai/gpt-oss-20b \
    --tensor-parallel-size 8 \
    --max-model-len 8192 \
    --max-num-batched-tokens 4096 \
    --max-num-seqs 4

Your clients keep calling /v1/chat/completions, so nothing changes on the application side.

Step 5 - Containerize. Start from the Neuron Deep Learning Containers. For SDK 2.33 the vLLM image is public.ecr.aws/neuron/pytorch-inference-vllm-neuronx:0.29.0.1.2.0-neuronx-py314-sdk2.33.0-ubuntu26.04. Point the compile cache at a shared volume. TorchNeuron uses TORCH_NEURONX_NEFF_CACHE_DIR, and vLLM Neuron documents a precompile cache under VLLM_CACHE_ROOT. NEURON_COMPILE_CACHE_URL with s3:// support belongs to the legacy XLA stack.

Pitfalls I would plan for:

• Compiled artifacts are specific to the Neuron SDK version, tensor parallel degree and shapes. Key your cache on all three.
• The driver (aws-neuronx-dkms) lives on the host, not in the container. A container on a node without it sees no devices.
• Neuron DLC images are large. Pre-pull them on new nodes, or your first pod waits on the registry.
• Once the cache is warm, cold start is dominated by loading weights, so keep them close (S3 in-region or FSx for Lustre).

How do you run Trainium2 inference on Amazon EKS in production?

On EKS, you add Trn2 nodes with the Neuron driver, install the Neuron Helm chart (which deploys the device plugin), and have pods request aws.amazon.com/neuron for whole devices or aws.amazon.com/neuroncore for individual cores. Karpenter, HPA and rolling updates then work the way they do for any other workload.

The official install is one Helm command:

helm upgrade --install neuron-helm-chart oci://public.ecr.aws/neuron/neuron-helm-chart \
--set "npd.enabled=false"

Add --set "scheduler.enabled=true" to deploy the Neuron scheduler extension, which places pods on contiguous cores. Ignore any guide that still uses aws.amazon.com/neurondevice: AWS dropped that resource name in Neuron 2.20. Newer clusters can use the Neuron DRA driver instead, though AWS says it cannot run alongside the device plugin on the same cluster's nodes.

What matters in practice:

• Match requests to your compile. A model compiled with tensor parallel degree 8 needs exactly that many cores in the pod. Request whole devices when in doubt.
• Plan capacity explicitly. Capacity Blocks cover Trn2 and Trn1, for one to 64 instances, reserved up to eight weeks ahead, for up to six months.
• Scale idle capacity away. An idle trn2.48xlarge costs the same per hour as a busy one. Karpenter consolidation and scheduled scale-down matter far more here than on CPU nodes.
• Export neuron-monitor metrics to Prometheus. Alert on NeuronCore utilization, device memory, runtime errors per pod, and failed compiles.
• Gate readiness on model load. The readiness probe should pass only once vLLM answers, never when the process merely starts.

Is Trainium2 actually cheaper than NVIDIA H100 GPUs for inference?

On AWS's published Capacity Block rates, a trn2.48xlarge costs $35.76 per hour versus $47.76 for a p5.48xlarge with 8 H100s, and AWS claims Trn2 offers "30-40% better price performance" than P5e and P5en instances. That claim comes from the vendor. Your real number is cost per million tokens at your latency target, and only your own benchmark gives you that.

Prices below were checked on October 22, 2026, from AWS's on-demand pricing data for us-east-1 and the Capacity Blocks pricing page. AWS publishes no us-east-1 On-Demand price for trn2.48xlarge; in the US it is sold through Capacity Blocks in Ohio. P5 prices reflect AWS's June 2025 cut of up to 45%.

The Trn2 product page lists 46 TBps of memory bandwidth and 3.2 Tbps of EFAv3 networking for the instance. Trn2 UltraServers connect 64 chips across four instances with 6 TB of shared accelerator memory. Trn3 UltraServers went GA in December 2025, and TorchNeuron supports them as well.

The hourly rate is rarely what decides your bill:

• Utilization. A model compiled for 16 chips pins a whole instance. At 20% utilization, the cheaper instance loses to a right-sized GPU.
• Porting time. Engineer-weeks spent on unsupported ops cost more than a month of the price gap.
• Recompiles. Every model update or SDK bump means a compile pass and a fresh benchmark.
• Discount mechanics. AWS states that "Savings Plans and Reserved Instance discounts don't apply to Capacity Blocks". Trn1 and Inf2 do have Reserved Instance rates.

Scale is no longer the open question. Amazon reports that Project Rainier, the cluster built for Anthropic, became fully operational in late October 2025 with "nearly half a million Trainium2 chips". I could not find a broad, independent benchmark set for Trn2 inference, which is one more reason to run your own.

My rule: Trainium2 wins on steady, high-volume inference of one or two supported models. GPUs win on breadth, exotic kernels, and spiky low-volume traffic.

When should you NOT use Trainium2, and what are the real alternatives?

Skip Trainium2 if your model relies on custom CUDA kernels, you need new architectures the week they launch, your traffic is low and bursty, or your team will not accept a compile step in the deploy pipeline. In those cases, NVIDIA GPUs in your own account, a managed inference API, or another accelerator family is the better fit.

The alternatives worth considering:

• NVIDIA GPUs in your own account: p5, p5en and g6e on AWS, A3 and G2 on GCP, the ND and NC series on Azure, or H100 instances on Scaleway. You get the widest model coverage and zero porting.
• Managed inference APIs such as Amazon Bedrock, OpenAI, Anthropic or Together AI, when volume is low, spiky, or unproven.
• Amazon SageMaker endpoints if you want AWS to manage the serving layer and accept SageMaker's pricing and constraints.
• Modal or RunPod for GPU serving without owning a cluster, accepting that the workload runs outside your account.
• Google Cloud TPU or AMD MI300X if you already run multi-cloud and your model is on their supported lists.

The pattern I see work: GPUs for experiments and new models, Trainium2 for the one or two models that carry most of the traffic. Make the decision per model, not per company, and revisit it at every Neuron SDK release.

How do you ship a Trainium2 inference service without a dedicated platform team?

The hard part of self-hosting inference is the EKS cluster, IAM, autoscaling, rollouts and on-call around the chip, and Qovery runs that layer inside your own AWS account. A small team can ship a Trn2-backed endpoint without hiring a platform engineer first.

What that looks like in practice:

• Your account, your bill. Qovery deploys into your own cloud account (BYOC), so Trn2 instances, Capacity Blocks and reservations stay in your name.
• Managed cluster lifecycle. Qovery handles Kubernetes upgrades, so your accelerated node groups and the Neuron Helm chart do not rot when the control plane version moves.
• A developer workflow. Git-push deploys of the vLLM Neuron container, preview environments per pull request for prompt and model changes, and environment auto-stop so idle non-production capacity shuts down overnight.
• Per-environment RBAC. Research, staging and production endpoints stay separated, with different access for each.
• The same workflow on every cloud. Qovery runs on AWS, GCP, Azure, Scaleway, or an existing Kubernetes cluster, which matters when your GPUs live on one provider and Trainium lives on AWS.

The boundary is clear: Qovery does not compile your model or tune Neuron. It runs the platform around it, so your engineers spend their time on the benchmark rather than on node groups.

How do I set this up from Claude Code, Cursor, or another coding agent?

Install the Qovery Agent Skills once, then ask your agent in plain language to deploy the vLLM Neuron container into your own AWS account. The skills drive the Qovery side and ask you to approve each action before it runs.

Install globally for every compatible agent (Claude Code, Cursor, Codex, OpenCode, Gemini CLI):

curl -fsSL https://skill.qovery.com/install.sh | bash

Or scope it to the current project:

curl -fsSL https://skill.qovery.com/install.sh | bash -s -- --project

To let Claude Code inspect what is already running, add the Qovery MCP Server. It uses OAuth and is read-only by default:

claude mcp add --transport http qovery https://mcp.qovery.com/mcp --callback-port 4242

Enable write access only when you want it. It also has to be enabled in the Qovery Console:

claude mcp add --transport http qovery "https://mcp.qovery.com/mcp?read_write=true" --callback-port 4242

The skills deploy (qovery-onboard, qovery-deploy, qovery-troubleshoot, qovery-optimize, qovery-speedup, qovery-preview, qovery-terraform), and the MCP Server operates what already exists. Then ask: "deploy this vLLM Neuron container to my AWS account with Qovery". The AI agent quickstart and the Claude Code integration page cover the details.

What is the difference between torch-neuron and torch-neuronx?

torch-neuron is the archived Inferentia1 package. torch-neuronx is the current package for Trainium and Inferentia2, and it now covers two products: for PyTorch 2.13 and 2.14 it is TorchNeuron, the native PyTorch backend for Trn2 and Trn3, while PyTorch 2.9 is the last PyTorch/XLA release used on Trn1 and Inf2.

What is TorchNeuron?

TorchNeuron is AWS's open-source native PyTorch backend for Trainium. It registers a neuron device, supports eager mode and torch.compile(backend="neuron"), and handles DDP, FSDP2 and DTensor. It shipped first in Neuron 2.27 and entered public beta in Neuron SDK 2.33.0 on October 19, 2026. It installs as torch-neuronx.

Does vLLM support AWS Trainium2?

Yes, through the vLLM Neuron plugin, currently in beta. Version 0.29.0.1.2.0 targets vLLM 0.29.0 and Neuron SDK 2.33 on Trn2 and Trn3, uses native PyTorch compilation by default, and exposes the standard OpenAI-compatible API through vllm serve.

Which EC2 instance types have Trainium2 chips, and how many NeuronCores do they expose?

trn2.48xlarge has 16 Trainium2 chips with 8 NeuronCore-v3 cores each, 128 physical cores in total. With the default Logical NeuronCore setting of 2, it presents 64 NeuronCores. Trn2 UltraServers link 64 Trainium2 chips across four instances. A smaller trn2.3xlarge size is offered through Capacity Blocks in some regions.

Can I run Llama or Mistral on Trainium2 without rewriting my PyTorch code?

For supported models, mostly yes. Llama 3 runs on vLLM Neuron with a standard vllm serve command, and the NxD Inference reference lists Llama 2 through 3.3, Mistral and Mixtral. Custom PyTorch code needs TorchNeuron, static shapes for torch.compile, and no ops from its unsupported list.

How do I schedule Trainium2 workloads on Amazon EKS?

Install the Neuron Helm chart (oci://public.ecr.aws/neuron/neuron-helm-chart) to deploy the device plugin, then have pods request aws.amazon.com/neuron for whole devices or aws.amazon.com/neuroncore for individual cores. Enable the scheduler extension for contiguous core placement. Do not use aws.amazon.com/neurondevice, which was dropped in Neuron 2.20.

Is Trainium2 inference cheaper than running NVIDIA H100 GPUs?

Per hour, usually. Capacity Block rates on October 22, 2026 were $35.76 for trn2.48xlarge and $47.76 for p5.48xlarge, and AWS claims 30-40% better price performance than P5e and P5en. Per token, it depends on your model, your utilization and your porting cost, so benchmark before you commit.

Do I need to recompile my model when the Neuron SDK or the tensor parallel degree changes?

Yes. Compiled artifacts are tied to the SDK version, the tensor parallel degree and the input shapes. Keep a persistent compile cache on shared storage, key it on all three, and budget a recompile and a fresh benchmark for every SDK upgrade.

Run a one-afternoon spike on your exact model with TorchNeuron and vLLM Neuron before you reserve any long-term Trn2 capacity.