Comprehensive reference for the AWS Neuron SDK, covering installation, compilation, runtime configuration, and troubleshooting.
- SDK Overview and Components
- Installation and Setup
- Compilation Workflow
- Runtime Configuration
- Performance Tuning
- Troubleshooting SDK Issues
- API Reference
- Advanced Topics
┌─────────────────────────────────────────────────────────────┐
│ Neuron SDK Stack │
├─────────────────────────────────────────────────────────────┤
│ │
│ ┌─────────────────────────────────────────────────────┐ │
│ │ ML Frameworks (User Code) │ │
│ │ • PyTorch │ │
│ │ • TensorFlow │ │
│ │ • JAX │ │
│ │ • MXNet │ │
│ └─────────────────┬───────────────────────────────────┘ │
│ │ │
│ ┌─────────────────▼───────────────────────────────────┐ │
│ │ Framework Integrations │ │
│ │ • torch-neuronx (PyTorch + XLA) │ │
│ │ • tensorflow-neuronx │ │
│ │ • jax-neuronx │ │
│ │ • transformers-neuronx (HuggingFace) │ │
│ └─────────────────┬───────────────────────────────────┘ │
│ │ │
│ ┌─────────────────▼───────────────────────────────────┐ │
│ │ Neuron Compiler (neuronx-cc) │ │
│ │ • XLA → Neuron IR → Binary │ │
│ │ • Graph optimization │ │
│ │ • Operator fusion │ │
│ │ • Memory layout optimization │ │
│ └─────────────────┬───────────────────────────────────┘ │
│ │ │
│ ┌─────────────────▼───────────────────────────────────┐ │
│ │ Neuron Runtime (libnrt.so) │ │
│ │ • Model loading & execution │ │
│ │ • Memory management │ │
│ │ • Collective operations (NCCL-like) │ │
│ └─────────────────┬───────────────────────────────────┘ │
│ │ │
│ ┌─────────────────▼───────────────────────────────────┐ │
│ │ Neuron Driver (aws-neuronx-dkms) │ │
│ │ • Kernel module (/dev/neuronN) │ │
│ │ • DMA engine │ │
│ │ • Hardware access │ │
│ └─────────────────────────────────────────────────────┘ │
│ │
│ ┌─────────────────────────────────────────────────────┐ │
│ │ Tools & Utilities │ │
│ │ • neuron-top (monitoring) │ │
│ │ • neuron-ls (device listing) │ │
│ │ • neuron-monitor (metrics collection) │ │
│ │ • neuron-profile (profiling) │ │
│ └─────────────────────────────────────────────────────┘ │
└─────────────────────────────────────────────────────────────┘
| Component | Version | Purpose |
|---|---|---|
| Neuron Driver | 2.16.1 | Kernel module for hardware access |
| Neuron Runtime | 2.18.1 | Model execution engine |
| Neuron Compiler | 1.13.1 | Graph compilation (XLA → Neuron) |
| torch-neuronx | 2.1.2.1.0 | PyTorch integration |
| transformers-neuronx | 0.9.474 | Optimized transformers |
| neuronx-distributed | 0.6.0 | Distributed training/inference |
| Optimum Neuron | 1.16.0 | HuggingFace integration |
| ML Framework | Neuron SDK Version | Python Version | Notes |
|---|---|---|---|
| PyTorch 2.1 | 2.18+ | 3.8, 3.9, 3.10 | Recommended for training |
| PyTorch 1.13 | 2.15-2.17 | 3.7, 3.8, 3.9 | Legacy, inference only |
| TensorFlow 2.10 | 2.18+ | 3.8, 3.9, 3.10 | Training + inference |
| JAX 0.4 | 2.18+ | 3.9, 3.10 | Experimental |
Operating Systems:
- Ubuntu 20.04 LTS (recommended)
- Ubuntu 22.04 LTS
- Amazon Linux 2
- RHEL 8+
Hardware:
- Trainium: trn1.2xlarge, trn1.32xlarge, trn1n.32xlarge
- Inferentia2: inf2.xlarge, inf2.8xlarge, inf2.24xlarge, inf2.48xlarge
# 1. Update system
sudo apt-get update -y
sudo apt-get upgrade -y
# 2. Install kernel headers (required for driver)
sudo apt-get install linux-headers-$(uname -r) -y
# 3. Configure Neuron repository
. /etc/os-release
sudo tee /etc/apt/sources.list.d/neuron.list > /dev/null <<EOF
deb https://apt.repos.neuron.amazonaws.com ${VERSION_CODENAME} main
EOF
# 4. Add repository key
wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | sudo apt-key add -
# 5. Update package list
sudo apt-get update -y
# 6. Install Neuron driver (kernel module)
sudo apt-get install aws-neuronx-dkms -y
# 7. Install Neuron runtime libraries
sudo apt-get install aws-neuronx-collectives -y
sudo apt-get install aws-neuronx-runtime-lib -y
# 8. Install Neuron tools
sudo apt-get install aws-neuronx-tools -y
# 9. Verify driver installation
sudo modprobe neuron
ls /dev/neuron*
# Expected output: /dev/neuron0 /dev/neuron1 ... /dev/neuron15 (on trn1.32xlarge)
# 10. Check device listing
neuron-ls
# Expected output:
# +--------+--------+--------+-----------+
# | DEVICE | TYPE | CORES | STATUS |
# +--------+--------+--------+-----------+
# | 0 | Trainium | 2 | Available |
# | 1 | Trainium | 2 | Available |
# | ... | ... | ... | ... |
# +--------+--------+--------+-----------+# Create virtual environment
python3.10 -m venv neuron_env
source neuron_env/bin/activate
# Upgrade pip
pip install --upgrade pip
# Install PyTorch Neuron (for training)
pip install torch-neuronx==2.1.2.1.0 \
--extra-index-url=https://pip.repos.neuron.amazonaws.com
# Install neuronx-distributed (for distributed training)
pip install neuronx-distributed==0.6.0
# Install transformers-neuronx (for inference)
pip install transformers-neuronx==0.9.474
# Install optimum[neuronx] (HuggingFace integration)
pip install optimum[neuronx]==1.16.0
# Install additional ML libraries
pip install transformers==4.36.0
pip install datasets==2.16.0
pip install accelerate==0.25.0
# Verify installation
python -c "import torch; import torch_xla.core.xla_model as xm; print(xm.xla_device())"
# Expected output: xla:0# Dockerfile.neuron
FROM public.ecr.aws/neuron/pytorch-training-neuronx:2.1.2-neuronx-py310-sdk2.18.1-ubuntu20.04
# Install additional dependencies
RUN pip install transformers==4.36.0 datasets==2.16.0
# Copy application code
WORKDIR /app
COPY train.py /app/
COPY requirements.txt /app/
RUN pip install -r requirements.txt
# Set default command
CMD ["python", "train.py"]Build and run:
# Build image
docker build -t my-neuron-training:latest -f Dockerfile.neuron .
# Run container with Neuron device access
docker run --device=/dev/neuron0 --device=/dev/neuron1 \
-v /opt/aws/neuron:/opt/aws/neuron \
my-neuron-training:latest# 1. Configure Neuron repository
sudo tee /etc/yum.repos.d/neuron.repo > /dev/null <<EOF
[neuron]
name=Neuron YUM Repository
baseurl=https://yum.repos.neuron.amazonaws.com
enabled=1
metadata_expire=0
EOF
# 2. Update and install
sudo yum update -y
# 3. Install Neuron driver
sudo yum install aws-neuronx-dkms -y
# 4. Install Neuron runtime
sudo yum install aws-neuronx-collectives -y
sudo yum install aws-neuronx-runtime-lib -y
# 5. Install tools
sudo yum install aws-neuronx-tools -y
# 6. Load kernel module
sudo modprobe neuron
# 7. Verify
neuron-ls┌─────────────────────────────────────────────────────────────┐
│ Neuron Compilation Pipeline │
├─────────────────────────────────────────────────────────────┤
│ │
│ Step 1: Framework Graph Capture │
│ ┌─────────────────────────────────────────────────────┐ │
│ │ PyTorch Model → torch.jit.trace() → TorchScript │ │
│ │ TorchScript → torch-xla → XLA HLO Graph │ │
│ └─────────────────┬───────────────────────────────────┘ │
│ │ │
│ Step 2: XLA Optimization │
│ ┌─────────────────▼───────────────────────────────────┐ │
│ │ XLA HLO Graph → Algebraic Simplification │ │
│ │ → Constant Folding │ │
│ │ → Dead Code Elimination │ │
│ └─────────────────┬───────────────────────────────────┘ │
│ │ │
│ Step 3: Neuron Compiler (neuronx-cc) │
│ ┌─────────────────▼───────────────────────────────────┐ │
│ │ HLO Graph → Neuron IR │ │
│ │ → Operator Fusion │ │
│ │ (e.g., matmul + bias + gelu → 1 kernel)│ │
│ │ → Memory Layout Optimization │ │
│ │ (HBM → SRAM → Tensor Engine) │ │
│ │ → Collective Op Scheduling │ │
│ │ (overlap compute + communication) │ │
│ └─────────────────┬───────────────────────────────────┘ │
│ │ │
│ Step 4: Code Generation │
│ ┌─────────────────▼───────────────────────────────────┐ │
│ │ Neuron IR → NeuronCore Assembly (.neff) │ │
│ │ → Store in cache (~/.cache/neuron/) │ │
│ └─────────────────────────────────────────────────────┘ │
└─────────────────────────────────────────────────────────────┘
# compile_example.py
import torch
import torch_xla.core.xla_model as xm
from transformers import AutoModelForCausalLM, AutoTokenizer
# 1. Load model
model_id = "meta-llama/Llama-2-7b-hf"
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16)
tokenizer = AutoTokenizer.from_pretrained(model_id)
# 2. Move to Neuron device (triggers lazy compilation)
device = xm.xla_device()
model = model.to(device)
# 3. Prepare input (fixed shapes for compilation)
input_text = "Hello, how are you?"
inputs = tokenizer(input_text, return_tensors="pt", padding="max_length", max_length=128)
input_ids = inputs["input_ids"].to(device)
print("Starting compilation... (5-10 minutes)")
# 4. First forward pass = compilation
import time
start = time.time()
with torch.no_grad():
outputs = model(input_ids)
xm.mark_step() # Synchronize XLA operations
compile_time = time.time() - start
print(f"Compilation complete: {compile_time:.2f}s")
# 5. Subsequent runs use cached binary (fast)
print("Running compiled model...")
start = time.time()
for i in range(10):
with torch.no_grad():
outputs = model(input_ids)
xm.mark_step()
inference_time = (time.time() - start) / 10
print(f"Average inference time: {inference_time:.3f}s")
# 6. Check compiled cache
import os
cache_dir = os.path.expanduser("~/.cache/neuron/")
print(f"Compiled binaries stored in: {cache_dir}")
print(f"Cache size: {os.path.getsize(cache_dir) / 1e6:.1f} MB")# Set via environment variable
export NEURON_CC_FLAGS="--flag1=value1 --flag2=value2"
# Or in Python
import os
os.environ['NEURON_CC_FLAGS'] = '--model-type=transformer --auto-cast=bf16'Key Flags:
| Flag | Values | Description | Use Case |
|---|---|---|---|
--model-type |
transformer, cnn, generic |
Model architecture hint | Enables transformer-specific optimizations |
--auto-cast |
bf16, fp16, none |
Automatic precision casting | Convert FP32 → BF16 where safe |
--distribution-strategy |
llm-training, llm-inference, none |
Distributed workload type | Enable collective op optimizations |
--enable-mixed-precision-accumulation |
true, false |
Use FP32 for gradient accumulation | Improves training stability |
--internal-hlo-simplification |
true, false |
Enable aggressive HLO optimization | May improve performance but longer compile time |
--verbose |
error, warning, info, debug |
Logging level | Debugging compilation issues |
Example configurations:
# Training (Llama 2 7B)
export NEURON_CC_FLAGS="
--model-type=transformer
--auto-cast=bf16
--distribution-strategy=llm-training
--enable-mixed-precision-accumulation=true
"
# Inference (vLLM)
export NEURON_CC_FLAGS="
--model-type=transformer
--auto-cast=bf16
--distribution-strategy=llm-inference
"
# Debugging
export NEURON_CC_FLAGS="
--verbose=debug
--dump-hlo-snapshot=/tmp/hlo_snapshot
"# Check cache location
echo $NEURON_COMPILE_CACHE_URL
# Default: ~/.cache/neuron/
# Set custom cache location
export NEURON_COMPILE_CACHE_URL="/mnt/efs/neuron_cache"
# Clear cache (force recompilation)
rm -rf ~/.cache/neuron/*
# Share cache across team (S3)
export NEURON_COMPILE_CACHE_URL="s3://my-bucket/neuron-cache"
# Requires AWS credentials with S3 access# Number of NeuronCores to use
export NEURON_RT_NUM_CORES=16 # Use all 16 cores on trn1.32xlarge
# Visible cores (for multi-model serving)
export NEURON_RT_VISIBLE_CORES="0,1,2,3" # Use only cores 0-3
# Execution timeout
export NEURON_RT_EXEC_TIMEOUT=30 # 30 seconds
# Async execution (overlap host/device)
export NEURON_RT_ASYNC_EXEC_MAX_INFLIGHT_REQUESTS=5
# Logging
export NEURON_RT_LOG_LEVEL=INFO # ERROR, WARNING, INFO, DEBUG
export NEURON_RT_LOG_LOCATION="/var/log/neuron/"
# Profiling
export NEURON_PROFILE=profile.json # Enable profiling
# Memory allocation
export NEURON_RT_ALLOCATOR_MODE=2 # 0=legacy, 1=hybrid, 2=new (recommended)
# Collective operations (distributed training)
export NEURON_RT_ROOT_COMM_ID=localhost:54321 # Master node for collectives{
"neuron_runtime": {
"num_cores": 16,
"exec_timeout": 30,
"async_exec": {
"enabled": true,
"max_inflight_requests": 5
},
"logging": {
"level": "INFO",
"destination": "/var/log/neuron/runtime.log"
},
"memory": {
"allocator_mode": 2,
"max_hbm_per_core": "30GB"
}
}
}Load config:
import os
os.environ['NEURON_RT_CONFIG_FILE'] = '/path/to/config.json'import neuronxcc.nki as nki
import neuronperf
# Get runtime info
runtime_info = neuronperf.get_runtime_info()
print(f"Runtime version: {runtime_info['version']}")
print(f"Number of devices: {runtime_info['num_devices']}")
# Get per-core metrics
for core_id in range(16):
metrics = neuronperf.get_core_metrics(core_id)
print(f"Core {core_id}:")
print(f" Utilization: {metrics['utilization_pct']:.1f}%")
print(f" HBM used: {metrics['hbm_used_bytes'] / 1e9:.2f} GB")
print(f" Execution latency: {metrics['exec_latency_ms']:.2f} ms")
# Reset metrics
neuronperf.reset_metrics()# Find optimal batch size
import torch
import torch_xla.core.xla_model as xm
model = load_model()
device = xm.xla_device()
batch_sizes = [1, 2, 4, 8, 16, 32]
results = []
for bs in batch_sizes:
try:
# Test if batch size fits in memory
inputs = torch.randint(0, 32000, (bs, 2048)).to(device)
# Warmup
for _ in range(5):
outputs = model(inputs)
xm.mark_step()
# Measure throughput
start = time.time()
for _ in range(100):
outputs = model(inputs)
xm.mark_step()
elapsed = time.time() - start
throughput = (bs * 100 * 2048) / elapsed
results.append((bs, throughput))
print(f"Batch size {bs}: {throughput:,.0f} tokens/s")
except RuntimeError as e:
print(f"Batch size {bs}: OOM")
break
# Optimal batch size
optimal_bs = max(results, key=lambda x: x[1])
print(f"\nOptimal batch size: {optimal_bs[0]} (throughput: {optimal_bs[1]:,.0f} tok/s)")# Profile memory usage
import neuronperf
def profile_memory(model, inputs):
# Get baseline memory
baseline = neuronperf.get_memory_stats()
print(f"Baseline HBM: {baseline['hbm_used_bytes'] / 1e9:.2f} GB")
# Load model
model = model.to(device)
after_load = neuronperf.get_memory_stats()
print(f"After model load: {after_load['hbm_used_bytes'] / 1e9:.2f} GB")
print(f"Model weights: {(after_load['hbm_used_bytes'] - baseline['hbm_used_bytes']) / 1e9:.2f} GB")
# Forward pass
outputs = model(inputs)
after_forward = neuronperf.get_memory_stats()
print(f"After forward: {after_forward['hbm_used_bytes'] / 1e9:.2f} GB")
print(f"Activations: {(after_forward['hbm_used_bytes'] - after_load['hbm_used_bytes']) / 1e9:.2f} GB")
# Backward pass (training only)
if model.training:
loss = outputs.loss
loss.backward()
after_backward = neuronperf.get_memory_stats()
print(f"After backward: {after_backward['hbm_used_bytes'] / 1e9:.2f} GB")
print(f"Gradients: {(after_backward['hbm_used_bytes'] - after_forward['hbm_used_bytes']) / 1e9:.2f} GB")# Enable aggressive fusion
export NEURON_CC_FLAGS="--internal-hlo-simplification=True"
# Disable fusion (debugging)
export NEURON_CC_FLAGS="--disable-fusion"
# Profile operator execution
export NEURON_PROFILE=profile.json
python train.py
# Analyze profile
neuron-profile view profile.json --show-fusion-opportunities
# Output:
# Fusion Opportunity: matmul + bias + gelu
# Current: 3 kernels, 2.5ms
# Fused: 1 kernel, 1.2ms (52% faster)Symptom:
RuntimeError: Compilation failed for model layer_0
neuronx-cc error: Unsupported operation 'aten::foo'
Diagnosis:
# Enable verbose compilation logging
export NEURON_CC_FLAGS="--verbose=debug"
python compile_model.py 2>&1 | tee compile_log.txt
# Check for unsupported ops
grep "Unsupported" compile_log.txtSolutions:
- Fallback to CPU for unsupported ops:
import torch_xla.core.xla_model as xm
# Mark specific ops for CPU execution
xm.set_replication(False) # Disable XLA for unsupported layers
class HybridModel(nn.Module):
def __init__(self):
super().__init__()
self.neuron_layers = NeuronLayers() # Runs on Neuron
self.cpu_layers = CPULayers() # Runs on CPU
def forward(self, x):
# Neuron execution
x = self.neuron_layers(x)
# Move to CPU for unsupported op
x = x.cpu()
x = self.cpu_layers(x)
# Move back to Neuron
x = x.to(xm.xla_device())
return x- Replace unsupported op:
# Example: Replace torch.where (sometimes unsupported)
# Instead of:
output = torch.where(condition, x, y)
# Use:
output = condition.float() * x + (1 - condition.float()) * y- Update SDK (may have added support):
pip install --upgrade torch-neuronx \
--extra-index-url=https://pip.repos.neuron.amazonaws.comSymptom:
RuntimeError: No Neuron devices found
Failed to open /dev/neuron0: No such file or directory
Diagnosis:
# Check if driver is loaded
lsmod | grep neuron
# Check devices
ls -l /dev/neuron*
# Check dmesg for errors
sudo dmesg | grep neuronSolutions:
- Load driver:
sudo modprobe neuron
# Verify
ls /dev/neuron*- Reinstall driver:
sudo apt-get remove aws-neuronx-dkms -y
sudo apt-get install aws-neuronx-dkms -y
sudo modprobe neuron- Check instance type:
# Verify running on Trainium/Inferentia2
curl http://169.254.169.254/latest/meta-data/instance-type
# Should output: trn1.32xlarge or inf2.xlarge, etc.Diagnosis:
# Monitor NeuronCore utilization
neuron-top
# Expected: >85% utilization
# If <60%: CPU bottleneck or small batch sizeSolutions:
- Increase batch size:
# Before
batch_size = 4
# After
batch_size = 16 # Increase until HBM ~80% full- Enable async execution:
export NEURON_RT_ASYNC_EXEC_MAX_INFLIGHT_REQUESTS=5- Optimize data loading:
# Use multiple workers
dataloader = DataLoader(
dataset,
batch_size=16,
num_workers=4, # Parallel preprocessing
pin_memory=True, # Faster CPU→Neuron transfer
prefetch_factor=2 # Prefetch next batches
)- Profile execution:
export NEURON_PROFILE=profile.json
python train.py
neuron-profile view profile.json
# Look for:
# - High data transfer time: Increase prefetch
# - Low compute time: Increase batch size
# - Frequent synchronization: Reduce xm.mark_step() callsSymptom:
RuntimeError: [Neuron] Out of HBM memory
Allocation failed: requested 8.5 GB, available 4.2 GB
Solutions:
- Reduce batch size:
batch_size = 8 # Was 16- Enable gradient checkpointing:
model.gradient_checkpointing_enable()- Use ZeRO optimizer:
from neuronx_distributed.optimizer import NeuronZero1Optimizer
optimizer = NeuronZero1Optimizer(
model.parameters(),
torch.optim.AdamW,
lr=2e-5
)- Reduce sequence length:
max_length = 2048 # Was 4096import torch_xla.core.xla_model as xm
# Device management
device = xm.xla_device() # Get Neuron device
devices = xm.get_xla_supported_devices() # List all devices
# Execution control
xm.mark_step() # Synchronize XLA ops
xm.wait_device_ops() # Wait for all ops to complete
# Distributed training
xm.xrt_world_size() # Number of devices
xm.get_ordinal() # Device rank (0-indexed)
# Profiling
xm.save_compiled_binary(model, "model.pt") # Save compiled model
xm.load_compiled_binary("model.pt") # Load compiled model
# Optimization
xm.optimizer_step(optimizer, barrier=True) # Sync gradients across devicesfrom neuronx_distributed.parallel_layers import parallel_state
# Initialize parallelism
parallel_state.initialize_model_parallel(
tensor_model_parallel_size=8,
pipeline_model_parallel_size=1,
data_parallel_size=1
)
# Parallel layers
from neuronx_distributed.parallel_layers import (
ColumnParallelLinear,
RowParallelLinear,
parallel_state
)
# Column-wise parallelism (split output dim)
layer = ColumnParallelLinear(
in_features=4096,
out_features=4096,
gather_output=False # Don't gather, keep distributed
)
# Row-wise parallelism (split input dim)
layer = RowParallelLinear(
in_features=4096,
out_features=4096,
input_is_parallel=True # Input already distributed
)
# Collective operations
from neuronx_distributed.parallel_layers import mappings
output = mappings.all_reduce(tensor)
output = mappings.all_gather(tensor)
output = mappings.reduce_scatter(tensor)from transformers_neuronx.llama.model import LlamaForSampling
# Load model for inference
model = LlamaForSampling.from_pretrained(
"meta-llama/Llama-2-7b-hf",
tp_degree=2, # Tensor parallelism
amp='bf16', # Mixed precision
batch_size=1,
n_positions=2048,
context_length_estimate=512
)
# Compile for Neuron
model.to_neuron()
# Inference
input_ids = tokenizer.encode("Hello", return_tensors="pt")
output_ids = model.sample(input_ids, sequence_length=100)from optimum.neuron import NeuronModelForCausalLM
# Automatic compilation and optimization
model = NeuronModelForCausalLM.from_pretrained(
"meta-llama/Llama-2-7b-hf",
export=True, # Compile for Neuron
batch_size=1,
sequence_length=2048,
num_cores=2,
auto_cast_type="bf16"
)
# Save compiled model
model.save_pretrained("llama2-7b-neuron")
# Load pre-compiled model
model = NeuronModelForCausalLM.from_pretrained(
"llama2-7b-neuron",
export=False # Already compiled
)Neuron Kernel Interface (NKI) allows writing custom kernels for NeuronCore:
import neuronxcc.nki as nki
import neuronxcc.nki.language as nl
@nki.jit
def custom_attention_kernel(q, k, v, output):
"""
Custom attention kernel optimized for NeuronCore
"""
# Allocate on-chip SRAM
sram_q = nl.ndarray((128, 64), dtype=nl.bfloat16, buffer=nl.sbuf)
sram_k = nl.ndarray((128, 64), dtype=nl.bfloat16, buffer=nl.sbuf)
# Load from HBM to SRAM
nl.load(sram_q, q)
nl.load(sram_k, k)
# Compute attention scores (Tensor Engine)
scores = nl.matmul(sram_q, nl.transpose(sram_k))
# Softmax (Vector Engine)
attn = nl.softmax(scores, axis=-1)
# Apply to values
output_sram = nl.matmul(attn, v)
# Store back to HBM
nl.store(output, output_sram)
# Use custom kernel
output = custom_attention_kernel(q_tensor, k_tensor, v_tensor, output_tensor)# Serve multiple models on single instance
import os
class MultiModelServer:
def __init__(self):
self.models = {}
def load_model(self, name, model_path, cores):
# Assign specific cores to each model
os.environ['NEURON_RT_VISIBLE_CORES'] = cores
model = NeuronModelForCausalLM.from_pretrained(model_path)
self.models[name] = model
print(f"Loaded {name} on cores {cores}")
def infer(self, model_name, input_text):
model = self.models[model_name]
return model.generate(input_text)
# Load 3 models on inf2.48xlarge (12 cores)
server = MultiModelServer()
server.load_model("llama3", "llama3-neuron", cores="0,1")
server.load_model("mistral", "mistral-neuron", cores="2,3")
server.load_model("codellama", "codellama-neuron", cores="4,5")
# Inference
output = server.infer("llama3", "Hello, world!")SDK Version: 2.18.1 (as of 2026-04-04)
Next Steps:
- Install Neuron SDK following the setup guide
- Try example compilation scripts
- Profile your workload with neuron-top
- Join the AWS Neuron community for support