This document is relevant for: Trn2, Trn3

Get started with AWS Neuron#

Use these quickstarts to run your first workload with the AWS Neuron SDK. Each group matches a task. Pick the task you want to do, then start with the first quickstart in that group.

Note

First time using AWS Neuron? These quickstarts assume you have:

  • An active AWS account with EC2 access

  • Basic familiarity with PyTorch

  • SSH access to launch and connect to EC2 instances

Most of these quickstarts run on Trn2 or Trn3 instances. For installation instructions, see the Setup Guide.

Train a model#

Use TorchNeuron, the native PyTorch backend for Trainium, to train with eager mode and torch.compile. Use TorchTitan on Neuron for large-scale pre-training.

Install and verify TorchNeuron

Install torch-neuronx and confirm that PyTorch can use the Neuron device.

Time: ~10–15 minutes

Trn2 Trn3

Build your first model with TorchNeuron

Run a PyTorch model in eager mode, with torch.compile, and in a short training loop.

Time: ~10 minutes

Trn2 Trn3

Set up TorchTitan Neuron

Install and configure TorchTitan for pre-training on Trainium with TorchNeuron.

Trn2 Trn3

Run a basic TorchTitan training

Complete a first pre-training run with TorchTitan on Trainium.

Trn2 Trn3

Serve LLMs#

Use the vLLM Neuron plugin to serve large language models with the standard vllm serve command and an OpenAI-compatible API.

Online serving with vLLM Neuron

Launch an OpenAI-compatible API server and send your first request.

Trn2 Trn3

Offline batch inference with vLLM Neuron

Run high-throughput batch inference without a server.

Trn2 Trn3

Deploy a vLLM server with a Neuron DLC

Start a vLLM server from a prebuilt Neuron Deep Learning Container.

Generate images and video#

Use the vLLM Omni Neuron plugin to run diffusion models, such as Wan2.2 text-to-video, on Trainium.

Offline video generation with Wan2.2

Generate a video clip from a text prompt with Wan2.2-T2V-A14B.

Trn2 Trn3

Online video generation with Wan2.2

Serve Wan2.2-T2V-A14B with the vLLM Omni OpenAI-compatible video API.

Trn2 Trn3

Write custom kernels#

Use the Neuron Kernel Interface (NKI) to write kernels with direct access to the NeuronCore, or use pre-optimized kernels from the NKI Library.

Implement and run your first NKI kernel

Write a simple NKI kernel and run it on a NeuronCore.

Time: ~10 minutes

Run an NKI kernel from TorchNeuron

Call an NKI kernel from a PyTorch model with nki_op, in eager mode and with torch.compile.

Time: ~15 minutes

Trn2 Trn3

Use the NKI Library

Add pre-optimized NKI Library kernels to your model code.

Time: ~10 minutes

Profile and optimize#

Use Neuron Explorer to capture profiles and find performance bottlenecks.

Get started with Neuron Explorer

Set up Neuron Explorer, capture a profile, and view it in the web UI.

Monitor cluster health#

Use Neuron Observer to detect hardware faults and find the node and device that cause them.

Get started with Neuron Observer

Install Neuron Observer on Amazon EKS, verify it with the nobs CLI, and view metrics.

Time: ~25 minutes

Trn2 Trn3

More quickstarts#

Framework setup guides#

PyTorch

Set up TorchNeuron with a DLAMI, a DLC, or a manual installation.

JAX

Set up JAX 0.7+ with a DLAMI, a DLC, or a manual installation.