GPU capacity is the single biggest line item on most SageMaker AI bills, and it is also the hardest capacity to get. AWS's own accelerators — Inferentia2 (ml.inf2.*) and Trainium (ml.trn1.*/trn2) — are the standing answer to both problems: lower price per hour, generally better availability, and, for steady high-throughput serving, a materially lower cost per million tokens.
The catch is that Neuron is not a drop-in swap for CUDA. Models must be compiled ahead of time for a fixed tensor shape and tensor-parallel degree, and that compile step is where most teams give up. This tutorial walks the whole path for a 7B-class chat model: compile, cache, deploy behind a SageMaker AI real-time endpoint, benchmark, and decide honestly whether to keep it.
Assumes the SageMaker Python SDK v2 (pip install "sagemaker>=2.200"), an execution role with S3 and ECR access, and an account quota for at least one ml.inf2.* inference instance (check Service Quotas first — this is the step that silently blocks a Friday deploy).
1. Decide before you compile
Neuron pays off when the workload is steady and throughput-bound. Run through this before spending an afternoon on compilation:
| Signal | Inferentia2 is a good fit | Stay on GPU |
|---|---|---|
| Traffic shape | Continuous, predictable RPS | Spiky, mostly idle, scale-to-zero |
| Model | Llama, Mistral, Qwen, Mixtral, BERT-family, Stable Diffusion — well-trodden architectures | Brand-new architecture or custom CUDA kernels |
| Iteration speed | Weights change monthly | Weights change daily; every change means a recompile |
| Sequence lengths | A small, known set of buckets | Wildly variable, long-context experiments |
| Ops appetite | Team can own a compile step in CI | Team wants model.deploy() and nothing else |
If your endpoint is idle most of the day, inference components with scale-to-zero on a GPU will beat a cheaper-per-hour accelerator that is also mostly idle. Cost per hour is not the metric; cost per million tokens at your real utilisation is.
2. Compile with Optimum Neuron
The least painful compiler front-end for transformer models is Hugging Face Optimum Neuron, which wraps the Neuron SDK (neuronx-cc, transformers-neuronx/NxD). Run the compile on the same instance family you will serve on — an ml.inf2.xlarge notebook or a SageMaker training job on ml.inf2.8xlarge — because the artifact is specific to the Neuron core topology.
Do it as a SageMaker training job so the artifact lands in S3 and the whole thing is repeatable in CI:
import sagemaker
from sagemaker.pytorch import PyTorch
role = sagemaker.get_execution_role()
compile_job = PyTorch(
entry_point="compile.py",
source_dir="src",
role=role,
instance_type="ml.inf2.8xlarge",
instance_count=1,
framework_version="2.6",
py_version="py310",
output_path="s3://my-bucket/neuron-artifacts/",
environment={"HF_MODEL_ID": "mistralai/Mistral-7B-Instruct-v0.3"},
max_run=7200,
)
compile_job.fit()
And src/compile.py:
import os
from optimum.neuron import NeuronModelForCausalLM
model_id = os.environ["HF_MODEL_ID"]
out_dir = "/opt/ml/model"
# These three numbers ARE the contract. Changing any of them means recompiling.
compiler_args = {"auto_cast_type": "bf16", "num_cores": 8}
input_shapes = {"batch_size": 4, "sequence_length": 4096}
model = NeuronModelForCausalLM.from_pretrained(
model_id,
export=True,
**compiler_args,
**input_shapes,
)
model.save_pretrained(out_dir)
from transformers import AutoTokenizer
AutoTokenizer.from_pretrained(model_id).save_pretrained(out_dir)
Three things bite people here:
num_coresmust match the instance. Aninf2.xlargehas 2 NeuronCores,inf2.8xlargehas 2,inf2.24xlargehas 12,inf2.48xlargehas 24. Compile for 8 cores and deploy to a 2-core instance and the endpoint will fail at model load, not at deploy time.sequence_lengthandbatch_sizeare baked in. Requests longer than the compiled sequence length are rejected or truncated by the serving stack, not silently handled.- Compilation is slow — 15–60 minutes for a 7B model is normal. Budget for it, and never do it inside a container start-up path.
Gated models need HF_TOKEN passed through the job environment (store it in Secrets Manager, not in the bundle).
3. Cache the artifact — this is the whole trick
The compiled model is the reusable asset. Once it is in S3, deploys are fast and reproducible, and CI can gate on "did the compile hash change?" rather than recompiling on every push.
s3://my-bucket/neuron-artifacts/<job-name>/output/model.tar.gz
Tag the artifact with the model ID, Neuron SDK version, num_cores, batch_size, and sequence_length. When someone six months later asks why the endpoint won't start on a new instance type, that tag set is the answer. Register the artifact in the SageMaker Model Registry as a separate package group from your GPU variant (mistral-7b-neuron vs mistral-7b-gpu) so promotion pipelines never mix them.
4. Deploy with the LMI-Neuron container
AWS publishes Large Model Inference containers built against the Neuron SDK. Point one at the compiled artifact:
from sagemaker.model import Model
from sagemaker import image_uris
image_uri = image_uris.retrieve(
framework="djl-neuronx",
region="us-east-1",
version="0.32.0",
)
model = Model(
image_uri=image_uri,
model_data="s3://my-bucket/neuron-artifacts/<job-name>/output/model.tar.gz",
role=role,
env={
"OPTION_ROLLING_BATCH": "vllm",
"OPTION_TENSOR_PARALLEL_DEGREE": "8",
"OPTION_MAX_ROLLING_BATCH_SIZE": "4",
"OPTION_MAX_MODEL_LEN": "4096",
"OPTION_MODEL_LOADING_TIMEOUT": "1200",
},
)
predictor = model.deploy(
initial_instance_count=1,
instance_type="ml.inf2.24xlarge",
endpoint_name="mistral-7b-neuron",
container_startup_health_check_timeout=1200,
)
Raise the health-check timeout. Neuron model load involves mapping the compiled graph onto the cores and is much slower than a GPU load; the default health check will kill a perfectly healthy container mid-load, and the CloudWatch log line that explains it is easy to miss.
Smoke test:
import json
print(predictor.predict({
"inputs": "Summarise the SageMaker AI inference options in one sentence.",
"parameters": {"max_new_tokens": 128, "temperature": 0.2},
}))
5. Benchmark like a finance person
Do not compare hourly prices. Measure tokens per second at your target concurrency and p95 latency, then divide.
import time, threading, boto3, json
rt = boto3.client("sagemaker-runtime")
PROMPT = "Explain vector databases to a CFO." * 8
results = []
def worker(n=20):
for _ in range(n):
t0 = time.time()
r = rt.invoke_endpoint(
EndpointName="mistral-7b-neuron",
ContentType="application/json",
Body=json.dumps({"inputs": PROMPT,
"parameters": {"max_new_tokens": 256}}),
)
body = json.loads(r["Body"].read())
results.append((time.time() - t0, 256))
threads = [threading.Thread(target=worker) for _ in range(8)]
t0 = time.time()
[t.start() for t in threads]
[t.join() for t in threads]
wall = time.time() - t0
tok = sum(t[1] for t in results)
lat = sorted(t[0] for t in results)
print(f"throughput: {tok / wall:,.0f} tok/s")
print(f"p50 {lat[len(lat)//2]:.2f}s p95 {lat[int(len(lat)*0.95)]:.2f}s")
Then, with the on-demand hourly price for the instance:
cost per 1M output tokens = (hourly_price / 3600) / tokens_per_second * 1_000_000
Run the identical script against your existing ml.g5.12xlarge or ml.g6.12xlarge endpoint with the same prompt, concurrency, and max_new_tokens. Anything else is not a comparison. In our engagements the honest outcome is usually: Inferentia2 wins clearly on throughput-bound batch and chat workloads at sustained concurrency, ties on lightly loaded endpoints, and loses on anything that needs the newest kernels or frequent recompiles.
Also record time to first token separately. Streaming chat UIs live or die on TTFT, and a configuration that wins on aggregate throughput can lose on perceived responsiveness.
6. Watch the right metrics in production
Neuron endpoints emit their own CloudWatch metrics alongside the standard endpoint ones. Alarm on these, not on GPUUtilization (which will be flat zero and mislead every dashboard you inherit):
NeuronCoreUtilization— your real saturation signalNeuronDeviceMemoryUsage— a compiled graph that barely fits will OOM under longer promptsModelLatencyandOverheadLatency— separates the model from the serving stackInvocation4XXErrors— usually requests exceeding the compiled sequence length
Autoscale on SageMakerVariantInvocationsPerInstance or a concurrency target rather than a hardware metric, and remember that scale-out on Neuron is slow because each new instance repeats the model load. Keep MinCapacity at least at your steady-state need.
7. Keep a GPU escape hatch
Maintain both model packages in the registry and keep the GPU endpoint config on file. When a base model changes and the recompile fails on a new Neuron SDK version — it happens — you want a one-command fallback, not an incident. Practise the fallback once so it is documented.
Teardown
predictor.delete_endpoint()
predictor.delete_model()
Compiled artifacts in S3 cost nothing worth mentioning; a forgotten ml.inf2.24xlarge costs a great deal. Tag every endpoint with an owner and a purge date.
Where teams get stuck
The compile step is the whole difficulty, and it is a build engineering problem more than an ML problem: pinning Neuron SDK versions, caching artifacts, wiring the recompile into CI, and keeping the shape contract documented so nobody deploys an 8-core graph onto a 2-core instance at 6pm on a Friday. Teams that treat it that way get 30–50% off their serving bill and better capacity availability. Teams that treat it as a one-off notebook experiment abandon it after the second failed model load.
If you want a second pair of hands on a Neuron migration — benchmarking the real cost per million tokens for your workload before committing, or building the compile-and-cache pipeline — get in touch. We do SageMaker AI, exclusively.