Native PyTorch inference on Neuron#

Native PyTorch inference on AWS Trainium runs through TorchNeuron, which compiles the model to a Neuron executable (NEFF) and runs it on the Neuron Runtime. The vLLM Neuron and vLLM Omni Neuron plugins handle this compilation and provide a complete serving stack. Without a plugin, you compile and run the model yourself and build the serving around it.

If the vLLM Neuron or vLLM Omni Neuron plugin already supports your model, use the vLLM Neuron or vLLM Omni Neuron documentation to deploy it.

This page is for running native PyTorch models for inference on Trainium without a plugin:

  • Implementing a custom model for inference on Trainium directly on native PyTorch.

  • Building a custom inference stack, or adapting a model-serving framework other than vLLM on Trainium.

How vLLM Neuron and vLLM Omni Neuron support native PyTorch#

The vLLM Neuron plugin serves LLMs and brings the full vLLM serving stack to Trainium — continuous batching, paged KV cache, prefix caching, speculative decoding (EAGLE3), structured outputs, multimodal support, parallelism strategies, and disaggregated inference — through the standard vllm serve command and an OpenAI-compatible API. It compiles each model to a NEFF through the TorchNeuron torch.compile backend.

vLLM Neuron supports disaggregated inference, which splits prefill and decode across separate instances. Disaggregated inference uses a serving runtime provided by libtorch-neuronx-lite — the TorchNeuron compile path (PyTorch model to NEFF) packaged with that runtime. vLLM Neuron therefore runs on libtorch-neuronx-lite rather than torch-neuronx.

The vLLM Omni Neuron plugin serves diffusion and multimodal generation models. It supports native PyTorch through the same libtorch-neuronx-lite package that vLLM Neuron uses.

Both plugins install libtorch-neuronx-lite automatically, so you do not install or interact with it directly. For the full feature set and the design of each component, see the vLLM Neuron plugin documentation.

Running custom native PyTorch inference on Trainium#

This covers running native PyTorch inference without the vLLM Neuron or vLLM Omni Neuron plugins. A custom inference stack compiles the model to a NEFF through the TorchNeuron torch.compile backend, runs it, and implements the serving layer. The vLLM Neuron plugin serves as a reference implementation.

Two packages provide the native PyTorch foundation. Which one you install depends on whether you need disaggregated inference, and the package also determines the available execution modes.

Package

When to use

Execution modes

torch-neuronx

Native PyTorch inference without disaggregated inference. The full native PyTorch (TorchNeuron) package, providing the neuron device and the torch.compile backend.

Eager mode and torch.compile.

libtorch-neuronx-lite

Native PyTorch inference that uses disaggregated inference, in a custom stack or through the plugins. Provides the runtime required for disaggregated inference, and vendors the necessary TorchNeuron compile components into its own wheel rather than depending on torch-neuronx, so it installs as a single self-contained package. This is the package the plugins use. For how the plugins build on the TorchNeuron torch.compile backend, see How vLLM is integrated with TorchNeuron.

torch.compile only. No eager execution on the Neuron device: eager operations fall back to the host CPU.

For the full set of TorchNeuron features and supported parameters, see the TorchNeuron reference.

Using libtorch-neuronx-lite in a custom inference stack#

If your custom stack needs disaggregated inference, use libtorch-neuronx-lite. Keep the following in mind when you build on it:

  • Use torch.compile for all model execution. libtorch-neuronx-lite runs models on the Neuron device only through torch.compile; it has no eager execution on the device. Eager operations on neuron tensors fall back to the host CPU, which is slow, and some of them, such as non-contiguous views and dtype casts, raise an error. Move tensors to the device (for example, tensor.to("neuron:0")) only as inputs to, and outputs from, compiled code.

  • Use the gloo backend for torch.distributed. Collectives inside a compiled model are compiled into the NEFF, so the process group is only needed for host-side coordination. Do not create a neuron process group: it belongs to torch-neuronx. vLLM Neuron uses gloo for the same reason.

  • Select the native compile path and register its backend. Set VLLM_NEURON_BACKEND=neuron_native in the environment before the process starts. libtorch-neuronx-lite reads this variable to select its native compile path, also in stacks that do not use vLLM. Importing libtorch_neuronx_lite does not register the torch.compile backend, so register it and pass the returned name:

    import torch
import libtorch_neuronx_lite
from libtorch_neuronx_lite.compile.native_backend import register

compiled = torch.compile(model, backend=register())

  • Run compiled code with gradients disabled. Use torch.no_grad() or torch.inference_mode(), or call model.requires_grad_(False). With gradients enabled and parameters that require them, compilation fails with neuron backend doesn't support events.

  • To reuse compiled graphs across restarts, save them and keep the grad mode fixed. With VLLM_NEURON_BACKEND=neuron_native set, libtorch-neuronx-lite turns on the PyTorch precompile cache. After warmup, call torch.compiler.save_cache_artifacts() to write it. A restart reuses the saved graphs only when it runs them in the same grad mode as warmup: torch.no_grad() and torch.inference_mode() count as different modes. In any other mode, PyTorch traces and compiles the graph again without an error.

Installing both packages in one environment#

torch-neuronx and libtorch-neuronx-lite can be installed in the same Python environment. Both register the neuron device, and only one of them can own it in a given process.

torch-neuronx loads automatically whenever a process runs import torch. In a process that uses libtorch-neuronx-lite, including vllm serve with vLLM Neuron or vLLM Omni Neuron, turn off that automatic loading:

export TORCH_DEVICE_BACKEND_AUTOLOAD=0

Set the variable in the environment before the process starts. PyTorch reads it once, while import torch runs, so setting it later in Python code has no effect. Without it, importing libtorch-neuronx-lite fails with Tried to register multiple backend fallbacks for the same dispatch key PrivateUse1.

Processes that use torch-neuronx need no change: leave the variable unset.

Documentation#

The TorchNeuron documentation covers the native PyTorch workflow on Trainium — setup, quickstarts, the torch.compile pipeline and its supported parameters, operator coverage, and troubleshooting. The vLLM Neuron and vLLM Omni Neuron documentation covers serving with the plugins, which run on the same TorchNeuron compile backend.

overview

TorchNeuron (overview)

TorchNeuron, the PrivateUse1 backend, and the eager and torch.compile execution modes.

quickstart

Get started with TorchNeuron

Running a first model on Trainium with torch.compile.

deep dive

torch.compile on Neuron

The Neuron torch.compile backend and its supported parameters.

reference

TorchNeuron reference

Operator coverage, environment variables, best practices, and troubleshooting.

overview

vLLM Neuron plugin

Serving LLMs on Trainium with the full vLLM serving stack.

overview

vLLM Omni Neuron plugin

Serving diffusion and multimodal generation models on Trainium.