This document is relevant for: Trn2, Trn3

Tutorial: Run a Hugging Face model for inference with torch.compile#

This tutorial shows you how to compile a Hugging Face Transformers model with torch.compile and the neuron backend on AWS Trainium. You then generate text with the compiled model. The tutorial uses the steps of the Hugging Face Text generation tutorial. When you complete it, you can compile generate() on a NeuronCore and control its output. You can also shard a model across NeuronCores and compile it.

Overview#

generate() compiles the forward pass of a model when the model uses a static cache and a CompileConfig. With chunked prefill, generate() also runs the prompt through the compiled forward pass. With the neuron backend, torch.compile then makes one Neuron program for the prefill and one for the decoding step. In this tutorial, you:

  • Load Qwen3-8B on a NeuronCore and turn on compilation for generate().

  • Generate text. The first call compiles the model.

  • Change the generation configuration, and save it.

  • Control the output length, the decoding strategy, and a batch of prompts.

  • Use a chat template.

  • Shard a model across two NeuronCores, and compile it.

Before you start#

This tutorial assumes that you have experience in the following areas:

Model details#

This tutorial uses these models from the Hugging Face Hub:

  • Qwen/Qwen3-8B for Steps 1 to 8.

  • Qwen/Qwen2.5-14B-Instruct for the tensor parallelism example in Step 9.

The models load in their checkpoint data type, bfloat16. The Qwen3-8B weights are approximately 16.4 GB. One logical NeuronCore of a Trn2 instance has 24 GB of memory, so Qwen3-8B runs on one NeuronCore. The Qwen2.5-14B-Instruct weights are approximately 29.5 GB, so you must shard this model across two or more NeuronCores.

Prerequisites#

  • A Trn2 instance with torch-neuronx installed and verified: See Quickstart: Install and verify torch-neuronx.

  • Hugging Face Transformers: Tested with transformers 5.19.0.

  • Access to the Hugging Face Hub: The models download on first use. A Hugging Face token is not necessary for these models.

Prepare your environment#

Install or update Transformers in the Python environment that has torch-neuronx:

pip install transformers==5.19.0

You do not have to import torch_neuronx to use the neuron device or the neuron backend. PyTorch loads torch-neuronx automatically when Python imports torch. See TorchNeuron overview.

Step 1: Load the model on a NeuronCore#

In this step, you load Qwen3-8B and its tokenizer.

Pass device_map="neuron" to from_pretrained():

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

Important

Use device_map="neuron", not device_map="auto". Transformers support for neuron with "auto" is not available yet.

Step 2: Turn on compilation for generate()#

In this step, you set the generation configuration of the model. Then each generate() call compiles the model with torch.compile.

Set these values in model.generation_config:

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

  • cache_implementation="static": generate() compiles the model only with a static cache. A static cache has a fixed length, so the shapes of the decoding step do not change.

  • CompileConfig(backend="neuron", dynamic=False, mode="default"): backend="neuron" selects the torch-neuronx backend for torch.compile. The default backend, Inductor, does not support Neuron. dynamic=False is necessary because torch-neuronx does not compile graphs with dynamic shapes. torch-neuronx supports only the default mode. Set mode="default" to prevent a warning for each compiled graph.

  • max_cache_len=128: the length of the static cache. Set it to the longest padded prompt plus max_new_tokens that you expect. Then calls with a different max_new_tokens use the same compiled programs.

  • prefill_chunk_size=64: generate() runs the prompt through the compiled forward pass in chunks of 64 tokens. Pad each prompt to 64 tokens (see Step 3). Then each chunk has the same shape, and the prefill uses one compiled program for all prompts of up to 64 tokens.

  • do_sample=False and no temperature, top_p, or top_k: the Qwen3 checkpoint sets these values for sampling. Greedy decoding gives the same output at each call. Top-p sampling sorts the full vocabulary at each step, outside the compiled program, and is slow on Neuron.

Step 3: Generate text#

In this step, you generate text. The first call compiles the model.

Pad the prompt on the left to 64 tokens, the prefill_chunk_size from Step 2, and move the inputs to model.device. Then pass them to generate() and decode the output tokens:

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])

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

The first call compiles the prefill and the decoding step, so it takes longer than the next calls. Run the same code again: the second call uses the compiled programs. A prompt with a different length also uses them, because the tokenizer pads it to 64 tokens.

torch-neuronx keeps the compiled programs in a cache on the disk. A later call with the same shapes uses the cache, even in a new Python process.

Step 4: Change the generation configuration#

In this step, you read the generation configuration, change it for one call, and save it.

The generation_config attribute shows the values that are different from the defaults, including the compilation settings from Step 2:

print(model.generation_config)

GenerationConfig {
  "bos_token_id": 151643,
  "cache_implementation": "static",
  "compile_config": {
    "backend": "neuron",
    "dynamic": false,
    "fullgraph": false,
    "mode": "default",
    "options": null
  },
  "do_sample": false,
  "eos_token_id": [
    151645,
    151643
  ],
  "max_cache_len": 128,
  "pad_token_id": 151643,
  "prefill_chunk_size": 64
}

To change a setting for one call, pass it to generate(). This example uses top-k sampling:

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

The call uses the same compiled programs as Step 3. The output changes at each call, because the call uses sampling.

To keep a set of settings, save a GenerationConfig to a directory and load it again. save_pretrained() does not save compile_config, so set it again after you load the configuration:

from transformers import GenerationConfig

generation_config = GenerationConfig(
    max_new_tokens=50,
    do_sample=True,
    top_k=50,
    eos_token_id=model.config.eos_token_id,
    pad_token_id=tokenizer.pad_token_id,
    cache_implementation="static",
    max_cache_len=128,
    prefill_chunk_size=64,
)
generation_config.save_pretrained("./my_generation_config")

generation_config = GenerationConfig.from_pretrained("./my_generation_config")
generation_config.compile_config = CompileConfig(
    backend="neuron", dynamic=False, mode="default"
)
generated_ids = model.generate(**model_inputs, generation_config=generation_config)
print(tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0])

Step 5: Control the output length#

In this step, you set the number of new tokens.

If you do not set max_new_tokens, generate() adds a maximum of 20 tokens to the prompt:

model_inputs = tokenizer(
    ["A sequence of numbers: 1, 2"], padding="max_length", max_length=64, return_tensors="pt"
).to(model.device)

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

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

A sequence of numbers: 1, 2, 3, 4, 5, 6, 7, 8,
A sequence of numbers: 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,

Both calls use the programs that Step 3 compiled, because the padded prompt and the new tokens fit in max_cache_len.

Step 6: Select a decoding strategy#

In this step, you compare greedy search and multinomial sampling.

Greedy search selects the most probable token at each step. Set do_sample=True for multinomial sampling:

model_inputs = tokenizer(
    ["I am a cat."], 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])

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

The greedy call prints:

I am a cat. I have a name, but I don't know what it is. I have a collar, but

The sampled output changes at each call. generate() selects the next token outside the compiled programs, so both strategies use the same programs.

Step 7: Generate for a batch of prompts#

In this step, you generate text for two prompts of different lengths in one batch.

Pad the prompts on the left. The model continues from the last token of each row. Padding on the right puts pad tokens between the prompt and the new text.

model_inputs = tokenizer(
    ["1, 2, 3", "A, B, C, D, E"], 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))

['1, 2, 3, 4, 5, 6, 7, 8, 9, ', 'A, B, C, D, E, F, G, H, I, J, K, L, M, N, O']

A batch of 2 runs the prefill and the decoding step for 2 sequences. Thus the first call with this batch size compiles new programs.

Step 8: Use a chat template#

In this step, you format a conversation with the chat template of the model.

Pass the messages to apply_chat_template() to get the prompt text. Set enable_thinking=False to turn off the thinking mode of Qwen3. Then tokenize and pad the prompt as in Step 3:

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])

Arrr, ye be askin’ a question fit for a scallywag! Aye, how many cats be needin’ to change a light bulb? Well, me hearties, it taketh one cat, but only if it

Step 9: Shard a model across NeuronCores and compile it#

In this step, you shard a model that is too large for one NeuronCore with tensor parallelism, and compile it.

On Trn2, one logical NeuronCore has two physical cores that share 24 GB of high-bandwidth memory (HBM). If the weights and the cache of a model are larger than this memory, you must shard the model across NeuronCores. With device_map="neuron", Qwen2.5-14B-Instruct stops with an out-of-memory error.

Transformers can shard a model with the tensor-parallel plan in its configuration. Start one process for each NeuronCore with torchrun. In each process, start the process group before any other Neuron call. Pass distributed_config={"tp_plan": "auto", "tp_size": world_size} to from_pretrained(). Then turn on compilation as in Step 2. Save this code 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()

Stop the Python session from Steps 1 to 8 first. That session uses the NeuronCores, so the torchrun processes cannot get them. Then run the script on two NeuronCores:

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

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

Each process holds half of the sharded weights and compiles its part of the model. The tensor-parallel degree (--nproc-per-node) must divide the number of key-value heads of the model (num_key_value_heads in its config.json). Qwen2.5-14B-Instruct has 8 key-value heads, so you can use 2, 4, or 8 NeuronCores.

Note

If the configuration of a model has no tensor-parallel plan, "tp_plan": "auto" in distributed_config does not shard the model. Then every process loads all of the weights.

To make sure that Transformers shards the weights, check that the parameters of the model are DTensor objects.

Confirmation#

With greedy decoding, the compiled generate() gives the same text as generate() without compilation. To check your setup, run Step 3 two times. The first call compiles the model. Both calls print:

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

Congratulations! You compiled a Hugging Face model with torch.compile on AWS Trainium. If you encountered any issues, see the Common issues section below.

Common issues#

  • The model runs on the CPU: device_map="auto" selects the CPU on a Neuron instance. Use device_map="neuron". 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 the variable, or add import torch_neuronx before you load the model.

  • InductorError ... no scheduling registered for device neuron: The static cache is set, but compile_config is not. Transformers then compiles with the default Inductor backend, which does not support Neuron. Set compile_config=CompileConfig(backend="neuron", dynamic=False, mode="default"). See Step 2: Turn on compilation for generate().

  • RuntimeError: Compilation error occurred on Neuron with [NCC_EMOD023] The StableHLO module could not be converted to HLO: The model compiled again with dynamic shapes, which torch-neuronx does not support. This occurs after a call with a new shape when dynamic is not False. Use CompileConfig(backend="neuron", dynamic=False, mode="default") and set max_cache_len.

  • A call compiles again: Some changes make new shapes for the prefill or the decoding step: a new batch size, a prompt that you do not pad to prefill_chunk_size, or a longer static cache. Then generate() compiles a new program. After one call uses a longer static cache, all later calls on the model also use the longer cache. They do not use the programs from before that call. To prevent this, pad each prompt to prefill_chunk_size. Keep the batch size the same. Set max_cache_len for the longest call that you expect.

  • A prompt is longer than 64 tokens: Set prefill_chunk_size and the max_length of the tokenizer to the length of your longest prompt. Increase max_cache_len by the same number of tokens. Then the prefill stays one chunk with one compiled program.

  • Beam search does not use the compiled program: In Transformers 5.19.0, only greedy search and sampling use the compiled program. A call with num_beams greater than 1 runs without it.

  • Sampling is very slow: The generation configuration has a top_p value. Top-p sampling sorts the full vocabulary at each step, outside the compiled program. Set top_p=None, and use top_k instead.

  • RuntimeError: NRT Execution error occurred on Neuron for operation=neuron::alloc::lazy with a Neuron OOM warning: The model is too large for the memory of one NeuronCore. Shard it across NeuronCores. See Step 9: Shard a model across NeuronCores and compile it.