How to Run Hugging Face Models on AWS Trainium with TorchNeuron: Install, Inference, Fine-Tuning, Errors

Read time: ~13 minutes.

Key facts:

  • TorchNeuron is a new, open-source (Apache 2.0) native PyTorch backend for AWS Trainium. It adds a neuron device to stock PyTorch, with eager execution, torch.compile(backend="neuron"), and the standard distributed APIs (DDP, FSDP2, DTensor, tensor parallelism) (PyTorch blog, AWS Neuron docs).
  • It is a public beta in Neuron SDK 2.33, supported on Trn2 and Trn3 instances only. Starting with Neuron 2.33, the torch-neuronx package for PyTorch 2.13 and 2.14 is TorchNeuron; earlier versions are the legacy PyTorch/XLA library (AWS install guide).
  • Hugging Face Transformers, Diffusers, Accelerate, PEFT and TRL now run on Trainium directly. Hugging Face says 89% of the top 100 text-generation architectures work out of the box, versus 46% with Optimum Neuron v0.4.5, and now recommends using these libraries directly instead of Optimum Neuron (Hugging Face blog).
  • Python 3.12, 3.13 or 3.14 on Linux x86_64. The code change for most scripts is .to("cuda") to .to("neuron"), or device_map="neuron" in Transformers.

What we checked: we read the PyTorch and Hugging Face launch posts, the AWS Neuron quickstart, manual-install and DLAMI pages, the Hugging Face inference tutorial and the TorchNeuron troubleshooting page, and the torch-neuronx GitHub README. We confirmed that torch-neuronx 2.14.3.0 wheels for Python 3.12, 3.13 and 3.14 are published on the Neuron pip repository. We do not have a Trainium instance, so we did not run these commands ourselves; every command, flag and expected output below is copied from AWS or Hugging Face documentation.

If you have been avoiding AWS Trainium because the old toolchain made you “think like a compiler” (lazy tensors, static shapes only, a recompile on every graph break), TorchNeuron is the change that matters. Trainium becomes a regular PyTorch device: you install one package, move your model to neuron, and your existing Transformers or TRL script runs. This guide walks through the shortest path from a fresh Trn2 instance to text generation, a sharded larger model, and a TRL fine-tuning run, plus the errors you are most likely to hit along the way.

For GPU-side serving of open models, see our Qwen3.8-Flash-Next vLLM guide. If you would rather call a hosted model than run your own hardware, the Gemini 3.6 Flash guide covers a cheap API option. More builder tutorials are on the 7minai homepage.

1. What TorchNeuron changes

Until now, the Hugging Face experience on Trainium and Inferentia went through optimum-neuron, which sat on top of the PyTorch/XLA-based torch-neuronx. Hugging Face describes the trade-off plainly: lazy tensors instead of eager execution, static shapes only, and a recompilation for every graph break. Optimum Neuron had to maintain custom modeling classes and trainers to hide that (Hugging Face blog).

TorchNeuron plugs into PyTorch through the standard PrivateUse1 device-backend mechanism and the torch.accelerator APIs (PyTorch blog). In practice that gives you:

FeatureWhat it means for your code
Eager mode (default)Operations run one at a time, like on any PyTorch device. print(), Python control flow and breakpoints work.
torch.compile(backend="neuron")Compiles a model or region into a Neuron executable (NEFF) for production performance.
Distributedtorch.distributed with the neuron backend: DDP, FSDP2, DTensor tensor/sequence parallelism, pipeline and context parallelism.
Custom kernelsNKI (Neuron Kernel Interface) kernels run in eager mode and inside torch.compile.
ProfilingThe same torch.profiler, with Neuron Explorer tracing from framework to hardware.

The AWS docs list the code changes for an existing CUDA script (AWS Neuron docs):

  • Change .to('cuda') to .to('neuron'). If your code already uses torch.accelerator, no change is needed.
  • Mixed precision: use torch.autocast(device_type="neuron").
  • For torch.compile, specify backend="neuron" and remove CUDA-specific settings such as mode="max-autotune-no-cudagraphs".
  • Configure sharding based on NeuronCores rather than GPUs.

You do not even need import torch_neuronx: PyTorch autoloads it on import torch through the torch.backends entry point the package registers. Import it only when you call its own APIs.

Is it production-ready? Not officially. AWS labels it a public beta for Neuron 2.33, and the GitHub README is titled “torch-neuronx (Beta)”. Treat it as the right path for new Trainium work and experiments, and test before moving production traffic.

2. Pick an instance and an image

Instance. TorchNeuron supports only the Trn2 and Trn3 families. The AWS quickstart names trn2.3xlarge and trn2.48xlarge as examples. On a trn2.48xlarge, neuron-ls shows 16 Neuron devices with 4 cores and 96 GB each, and PyTorch reports 64 Neuron devices. One logical NeuronCore on Trn2 has two physical cores sharing 24 GB of HBM (AWS quickstart, AWS HF inference tutorial). That 24 GB figure is the one to plan around: a model whose weights plus KV cache exceed it needs sharding (Section 6).

If you plan to use Trn3, note the AWS warning: per-rank memory, NEURON_RT_NUM_CORES and mesh shapes in the docs are written for trn2.48xlarge and must be recalculated for Trn3.

Image. You have three options:

  1. Neuron DLAMI with torch-neuronx 2.14 (fastest). On Ubuntu 24.04 or 26.04, use the “Deep Learning AMI Neuron torch-neuronx 2.14.0” image. TorchNeuron is preinstalled. AWS notes this DLAMI is not available for Amazon Linux 2023.
  2. Neuron Base DLAMI. Has the driver, EFA and tools; you install the Python packages yourself (Section 3, Step 4).
  3. Plain Ubuntu 24.04/26.04 AMI. You install everything (Section 3). AWS says to size the root EBS volume to at least 512 GB.

To find the DLAMI ID with the AWS CLI (AWS DLAMI guide):

aws ec2 describe-images \
  --owners amazon \
  --filters "Name=name,Values=Deep Learning *Neuron*" \
            "Name=description,Values=*Neuron SDK 2.33.0*" \
  --query 'Images[].{Name:Name, ImageId:ImageId}' \
  --output table \
  --region us-east-1

Pick the torch-neuronx 2.14.0 image for your Ubuntu version, then launch:

aws ec2 run-instances \
  --image-id ami-xxxxxxxxxxxxxxxxx \
  --instance-type trn2.48xlarge \
  --key-name your-key-pair \
  --security-group-ids sg-xxxxxxxxx \
  --subnet-id subnet-xxxxxxxxx

On the single-framework torch-neuronx DLAMI, activate the bundled environment and skip to Section 4:

source /opt/aws_neuronx_venv_torch_neuronx_2_14_0/bin/activate

3. Install TorchNeuron on a plain Ubuntu instance

These steps follow the AWS quickstart for Ubuntu 26.04. If you started from a Neuron Base DLAMI, skip Step 1.

Step 1: add the Neuron apt repository and install the driver and tools.

sudo apt-get update
sudo apt-get install -y curl gnupg
curl -fsSL https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB \
  | sudo gpg --dearmor --yes -o /usr/share/keyrings/neuron-archive-keyring.gpg
. /etc/os-release
printf 'deb [signed-by=/usr/share/keyrings/neuron-archive-keyring.gpg] \
https://apt.repos.neuron.amazonaws.com %s main\n' "${VERSION_CODENAME}" \
  | sudo tee /etc/apt/sources.list.d/neuron.list

sudo apt-get update
sudo apt-get install -y linux-headers-$(uname -r)
sudo apt-get install -y \
  aws-neuronx-dkms=2.* \
  aws-neuronx-tools=2.*

Step 2: install the runtime and collectives.

sudo apt-get update
sudo apt-get install -y \
  python3-venv \
  aws-neuronx-runtime-lib=2.* \
  aws-neuronx-collectives=2.*

If neuron-ls is not found later, add the tools to your path: export PATH=/opt/aws/neuron/bin:$PATH (from the manual install guide).

Step 3: confirm the hardware is visible.

neuron-ls

On a trn2.48xlarge, AWS shows output beginning with instance-type: trn2.48xlarge and logical-neuroncore-config: 2, followed by a table of 16 devices (IDs 0 to 15). If the table is empty, stop here and fix the driver or instance type (Section 8).

Step 4: create a virtual environment and install TorchNeuron.

python3 --version          # must be 3.12, 3.13 or 3.14
python3 -m venv ~/torch-neuronx-venv
source ~/torch-neuronx-venv/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

Both extra index URLs matter. The first serves torch-neuronx, the Neuron compiler (neuronx-cc) and nki; the second serves the CPU PyTorch 2.14 wheel that TorchNeuron builds on. AWS says this release also supports PyTorch 2.13.

Do not pin the old version. torch-neuronx==2.9.* is the legacy PyTorch/XLA library, which last shipped in Neuron 2.31. Only the 2.13/2.14 line is TorchNeuron.

4. Verify the neuron device

Save this as verify_torch_neuronx.py (from the AWS quickstart):

import torch
import torch_neuronx

print(f"PyTorch: {torch.__version__}")
print(f"torch-neuronx: {torch_neuronx.__version__}")

assert torch.neuron.is_available(), "No Neuron device is available."
print(f"PyTorch Neuron devices: {torch.neuron.device_count()}")

torch.manual_seed(0)
cpu_input = torch.randn(2, 3)
expected = cpu_input @ cpu_input.T

neuron_input = cpu_input.to("neuron")
actual = neuron_input @ neuron_input.T

torch.testing.assert_close(actual.cpu(), expected, rtol=1e-4, atol=1e-5)
print("torch-neuronx verification passed")

Run it with python verify_torch_neuronx.py. AWS’s sample output on a trn2.48xlarge:

PyTorch: 2.14.0+cpu
torch-neuronx: 2.14.3.0
PyTorch Neuron devices: 64
torch-neuronx verification passed

The first Neuron operation can include compile time, so a slow first run is normal. Your exact version may differ. The script compares the Neuron matmul result against CPU, so a pass confirms both device execution and the copy back to the host.

The +cpu suffix on PyTorch is expected: TorchNeuron uses the CPU build of PyTorch and adds the neuron device on top.

5. Run a Hugging Face model with compiled generate()

This section follows the AWS tutorial Run a Hugging Face model for inference with torch.compile, which was tested with Transformers 5.19.0 and uses Qwen/Qwen3-8B (about 16.4 GB in bfloat16, so it fits on one 24 GB logical NeuronCore). No Hugging Face token is needed for this model.

pip install transformers==5.19.0

Load the model on a NeuronCore.

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-8B", device_map="neuron")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B", padding_side="left")
print(model.device)   # neuron:0

Use device_map="neuron", not "auto". AWS notes that Transformers does not yet support Neuron with "auto", and on a Neuron instance "auto" silently selects the CPU.

Turn on compilation. generate() only compiles when the model uses a static cache and a CompileConfig:

from transformers import CompileConfig

model.generation_config.cache_implementation = "static"
model.generation_config.compile_config = CompileConfig(
    backend="neuron", dynamic=False, mode="default"
)
model.generation_config.max_cache_len = 128
model.generation_config.prefill_chunk_size = 64
model.generation_config.do_sample = False
model.generation_config.temperature = None
model.generation_config.top_p = None
model.generation_config.top_k = None

Why each setting matters, per AWS:

  • cache_implementation="static": a fixed-length cache keeps the decode step’s shapes constant.
  • backend="neuron": the default Inductor backend does not support Neuron.
  • dynamic=False: torch-neuronx does not compile graphs with dynamic shapes.
  • mode="default": the only supported mode; setting it explicitly avoids a warning per graph.
  • max_cache_len: set it to your longest padded prompt plus max_new_tokens, so different output lengths reuse the same compiled programs.
  • prefill_chunk_size=64: the prompt runs through the compiled forward pass in 64-token chunks.
  • Greedy decoding with top_p=None: Qwen3’s checkpoint enables sampling, and top-p sorts the whole vocabulary outside the compiled program, which is slow on Neuron.

Generate. Pad every prompt on the left to the chunk size so all prompts share one compiled prefill program:

model_inputs = tokenizer(
    ["A list of colors: red, blue"], padding="max_length", max_length=64, return_tensors="pt"
).to(model.device)

generated_ids = model.generate(**model_inputs, max_new_tokens=20)
print(tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0])

AWS’s expected output:

A list of colors: red, blue, green, yellow, purple, orange, black, white, pink, teal, brown, gray

The first call compiles a program for the prefill and one for the decode step, so it is slow. Run it again and it reuses them. Compiled programs are kept in an on-disk cache, so a new Python process with the same shapes skips compilation too. You can point TORCH_NEURONX_NEFF_CACHE_DIR at a shared mount such as Amazon FSx to reuse compiled NEFFs across instances (TorchNeuron overview).

Chat prompts. Use the chat template, then tokenize and pad exactly as above. For Qwen3, enable_thinking=False turns off thinking mode:

messages = [
    {"role": "system", "content": "You are a friendly chatbot who always responds in the style of a pirate"},
    {"role": "user", "content": "How many cats does it take to change a light bulb?"},
]
prompt = tokenizer.apply_chat_template(
    messages, add_generation_prompt=True, enable_thinking=False, tokenize=False
)
model_inputs = tokenizer(
    [prompt], padding="max_length", max_length=64, return_tensors="pt"
).to(model.device)
input_length = model_inputs["input_ids"].shape[1]
generated_ids = model.generate(**model_inputs, max_new_tokens=50)
print(tokenizer.batch_decode(generated_ids[:, input_length:], skip_special_tokens=True)[0])

Things that trigger a recompile. A new batch size, a prompt you did not pad to prefill_chunk_size, or a longer static cache. AWS also warns that once one call uses a longer cache, every later call on that model uses it too and ignores the earlier programs. Keep batch size fixed, pad every prompt, and set max_cache_len once for the longest call you expect.

Sampling and beams. Sampling (do_sample=True, top_k=50) reuses the compiled programs because token selection happens outside them. In Transformers 5.19.0, beam search (num_beams > 1) does not use the compiled program at all.

Saving settings. GenerationConfig.save_pretrained() does not save compile_config, so set it again after loading a saved config.

6. Shard a bigger model across NeuronCores

If a model’s weights and cache exceed 24 GB, device_map="neuron" stops with an out-of-memory error. The AWS example uses Qwen/Qwen2.5-14B-Instruct (about 29.5 GB) and Transformers’ built-in tensor-parallel plan. Save as generate_tp.py:

import torch
import torch.distributed as dist
from transformers import AutoModelForCausalLM, AutoTokenizer, CompileConfig

dist.init_process_group(backend="neuron")
world_size = dist.get_world_size()

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-14B-Instruct",
    dtype=torch.bfloat16,
    distributed_config={"tp_plan": "auto", "tp_size": world_size},
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-14B-Instruct", padding_side="left")

model.generation_config.cache_implementation = "static"
model.generation_config.compile_config = CompileConfig(
    backend="neuron", dynamic=False, mode="default"
)
model.generation_config.max_cache_len = 128
model.generation_config.prefill_chunk_size = 64
model.generation_config.do_sample = False
model.generation_config.temperature = None
model.generation_config.top_p = None
model.generation_config.top_k = None

model_inputs = tokenizer(
    ["A list of colors: red, blue"], padding="max_length", max_length=64, return_tensors="pt"
).to("neuron")
generated_ids = model.generate(**model_inputs, max_new_tokens=20)
if dist.get_rank() == 0:
    print(tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0])
dist.destroy_process_group()

Close any Python session from Section 5 first (it is holding the NeuronCores), then:

torchrun --nproc-per-node 2 generate_tp.py

Two rules from AWS:

  • Start the process group before any other Neuron call in each process.
  • The tensor-parallel degree must divide num_key_value_heads in the model’s config.json. Qwen2.5-14B-Instruct has 8 KV heads, so 2, 4 or 8 NeuronCores work.

If a model’s config has no tensor-parallel plan, "tp_plan": "auto" silently does nothing and every process loads the full weights. Check that the model’s parameters are DTensor objects to confirm sharding happened.

The Hugging Face blog shows the same pattern for a larger model, Qwen/Qwen3.8-27B, using DistributedConfig(tp_size=world_size) with an explicit DeviceMesh("neuron", ...) and a chat-templated prompt padded to a multiple of the prefill chunk (Hugging Face blog).

Serving instead of scripts? For an OpenAI-compatible server with continuous batching and paged KV cache, AWS points to the vLLM Neuron plugin, which compiles models through the TorchNeuron torch.compile backend. Note that vLLM Neuron runs on a separate package, libtorch-neuronx-lite, which the plugin installs for you (Native PyTorch inference on Neuron).

7. Fine-tune with TRL on Trainium

The biggest practical change for training is that you no longer need a Neuron-specific trainer or model fork. Hugging Face’s example is a standard TRL SFTTrainer script with a few Neuron-aware settings (Hugging Face blog). The core of sft_neuron_example.py:

import os
import torch
from datasets import load_dataset
from transformers.distributed import DistributedConfig
from trl import ModelConfig, ScriptArguments, SFTConfig, SFTTrainer, TrlParser

def main(script_args, training_args, model_args):
    # The Neuron backend can't compile Dynamo's automatic dynamic shapes yet
    torch._dynamo.config.automatic_dynamic_shapes = False

    training_args.model_init_kwargs = {
        "revision": model_args.model_revision,
        "trust_remote_code": training_args.trust_remote_code,
        "attn_implementation": model_args.attn_implementation,
        "dtype": model_args.dtype,
        "device_map": None,
    }
    if int(os.environ.get("WORLD_SIZE", "1")) > 1:
        training_args.model_init_kwargs["distributed_config"] = DistributedConfig(
            tp_size=int(os.environ.get("TP_SIZE", "1")),
            fsdp_size=int(os.environ.get("FSDP_SIZE", "1")),
        )

    train_dataset = load_dataset(
        script_args.dataset_name,
        name=script_args.dataset_config,
        split=script_args.dataset_train_split,
    )
    train_dataset = train_dataset.select_columns(["messages"])

    trainer = SFTTrainer(
        model=model_args.model_name_or_path,
        args=training_args,
        train_dataset=train_dataset,
    )
    trainer.train()
    trainer.save_model(training_args.output_dir)
    trainer.end()

if __name__ == "__main__":
    parser = TrlParser((ScriptArguments, SFTConfig, ModelConfig))
    script_args, training_args, model_args = parser.parse_args_and_config()
    main(script_args, training_args, model_args)

Launch it with one process per NeuronCore. Hugging Face ran this on 8 Trainium2 chips (32 NeuronCores, half a trn2.48xlarge):

TP_SIZE=4 FSDP_SIZE=8 HF_DEACTIVATE_ASYNC_LOAD=1 uv run torchrun --nproc_per_node=32 sft_neuron_example.py \
    --model_name_or_path Qwen/Qwen3-8B \
    --dataset_name HuggingFaceH4/ultrachat_200k \
    --dataset_train_split train_sft \
    --output_dir Qwen3-8B-sft-tp4-fsdp8 \
    --dtype bfloat16 \
    --max_length 2048 \
    --packing true \
    --packing_strategy wrapped \
    --learning_rate 2e-5 \
    --warmup_steps 0.03 \
    --num_train_epochs 1 \
    --logging_steps 2 \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 4 \
    --save_only_model true \
    --torch_compile true

What each Neuron-specific piece does:

  • TP_SIZE × FSDP_SIZE must equal --nproc_per_node. Here 4 × 8 = 32. Tensor parallelism splits each layer across the 4 cores of a group; FSDP shards parameters, gradients and optimizer state across the 8 groups.
  • No --fsdp flag or Accelerate FSDP config. The DistributedConfig shards the model as weights load, and the Trainer sees it is already distributed.
  • HF_DEACTIVATE_ASYNC_LOAD=1. Transformers’ asynchronous weight loading is not supported on Neuron yet.
  • Static shapes. Automatic dynamic shapes are off, wrapped packing produces sequences of up to 2048 tokens, and with TRL ≥ 1.16 batches on Neuron are padded to a multiple of 128 tokens to limit the number of distinct shapes.
  • --torch_compile true is all you need; on Trainium, torch.compile detects the device and picks the Neuron backend.

The output is a standard Transformers checkpoint with nothing Neuron-specific in it, so you can push it to the Hub and load it on GPUs or anything else.

For your own PyTorch training loops (not TRL), the PyTorch blog shows that compiling needs only one extra line, torch.compiler.set_default_backend("neuron"), before any model.compile() call.

8. Troubleshooting

All fixes below come from the AWS troubleshooting page, quickstart and HF inference tutorial.

neuron-ls shows no devices / “No Neuron device is available.” Confirm you are on Trn2 or Trn3 (curl http://169.254.169.254/latest/meta-data/instance-type should return trn2.* or trn3.*), check the driver with lsmod | grep neuron, and try sudo systemctl restart neuron-monitor.

ModuleNotFoundError: No module named 'torch_neuronx' Activate the virtual environment, then rerun the install. On a DLAMI, which python should point under /opt/aws_neuronx_venv_*.

Strange import errors or the device is missing after install You probably have the legacy XLA torch-neuronx. Run pip show torch-neuronx; 2.14.x is TorchNeuron, anything earlier is the old library. Use a clean venv, or:

python -m pip uninstall torch-neuronx
python -m pip install "torch-neuronx==2.14.*" \
  --extra-index-url https://pip.repos.neuron.amazonaws.com \
  --extra-index-url https://download.pytorch.org/whl/cpu

Package installation fails Check python --version (3.12 to 3.14 only) and make sure you passed both --extra-index-url options.

The model runs on the CPU You used device_map="auto". Switch to device_map="neuron" and check model.device.

ValueError: When passing device_map as a string, ... but found neuron TORCH_DEVICE_BACKEND_AUTOLOAD=0 is set, so PyTorch did not load torch-neuronx. Unset it, or add import torch_neuronx before loading the model.

InductorError ... no scheduling registered for device neuron You set a static cache but no compile_config, so Transformers fell back to Inductor. Set CompileConfig(backend="neuron", dynamic=False, mode="default").

[NCC_EMOD023] The StableHLO module could not be converted to HLO The model recompiled with dynamic shapes. Use dynamic=False and set max_cache_len.

NRT Execution error ... operation=neuron::alloc::lazy with an OOM warning The model does not fit on one NeuronCore. Shard it (Section 6).

Eager mode is much slower than expected An operator is probably falling back to CPU. List fallbacks with torch_neuronx.get_fallback_ops(), and set TORCH_NEURONX_FALLBACK_ONLY_FOR_UNIMPLEMENTED_OPS=1 so execution errors surface instead of hiding behind a fallback. Optionally try Adaptive Eager Execution (operator fusion) with TORCH_NEURONX_ENABLE_CONCATENATION=1; it is off by default and can cause small numerical differences.

Errors point at the wrong line Neuron ops run asynchronously. Set NEURON_LAUNCH_BLOCKING=1 so each op runs synchronously and errors point at the op that caused them.

Graph break warnings under torch.compile Usually data-dependent control flow such as if tensor.item() > 0. The Neuron backend targets full graphs; use fullgraph=True where you can and restructure code to remove breaks.

9. FAQ

Do I still need Optimum Neuron? For new work on Trn2/Trn3, Hugging Face now recommends using Transformers, Diffusers, Accelerate, PEFT and TRL directly instead of going through Optimum Neuron.

Does it work on Inferentia or Trn1? AWS documents TorchNeuron for Trn2 and Trn3 instances only; Inferentia and Trn1 are not in its supported list.

Does Diffusers work? Yes. Hugging Face says many Diffusers pipelines work in eager and compiled mode, and FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 and H3 also support tensor parallelism. Their Qwen-Image-2.1 example runs with NEURON_LOGICAL_NC_CONFIG=2 torchrun --nproc_per_node=8 run.py.

Can I write custom kernels? Yes, with NKI. Neuron kernels are being published on the Hub under aws-neuron/kernels; Hugging Face cites up to a 2.4x speedup for vision transformers with its ViT encoder kernel.

Is the source available? Yes, at github.com/aws-neuron/torch-neuronx under Apache 2.0, with the public beta going live on October 20, 2026.

Will my checkpoints be locked to Trainium? No. Hugging Face says the saved checkpoint is a standard Transformers checkpoint you can load on any hardware.

Sources