This document is relevant for: Trn2, Trn3

Get started with TorchNeuron

Use these quickstart guides to get your first PyTorch model running on AWS Trainium with TorchNeuron. Install and verify your environment, then build and train your first model.

Install and verify torch-neuronx quickstart

Set up torch-neuronx on a Trainium instance and confirm your environment is ready to run models.

Build your first model with torch-neuronx quickstart

Choose a model, build it, and train it end-to-end on Trainium using TorchNeuron’s native PyTorch backend.

Run an NKI kernel from torch-neuronx quickstart

Call a custom NKI kernel from a PyTorch model on Trainium using @nki_op.

This document is relevant for: Trn2, Trn3