Comprehensive technical guide to AWS Trainium (Trn1/Trn1n) for high-performance, cost-efficient LLM training on Amazon EKS.
- Introduction to AWS Trainium
- Trainium Chip Architecture
- Training Optimizations
- Distributed Training Patterns
- Neuron SDK for Training
- EKS Deployment Patterns
- Performance Tuning Guide
- Best Practices for LLM Training
- Real-World Use Cases
- Troubleshooting
AWS Trainium is Amazon's second-generation custom ML training chip, purpose-built for training large-scale neural networks. Launched in 2022, Trainium is designed to deliver:
- 50% cost savings compared to GPU-based training
- Linear scaling from single instances to 1000+ NeuronCores
- High-performance training for transformer models (LLMs, vision transformers)
- Native PyTorch, TensorFlow, and JAX support via Neuron SDK
| Challenge | GPU Reality | Trainium Solution |
|---|---|---|
| Training cost | $500K-$2M per 70B LLM run | 50% lower with comparable performance |
| GPU availability | Weeks to months wait time | Available immediately in 6 AWS regions |
| Scaling efficiency | 85-90% scaling beyond 256 GPUs | 95%+ scaling to 1000+ NeuronCores |
| Energy consumption | 300W per A100 GPU | 225W per Trn1 instance (32 cores) |
| Framework support | CUDA lock-in | Open ecosystem (PyTorch, JAX, TensorFlow) |
| Instance | NeuronCores | HBM Memory | vCPUs | System RAM | Network | Cost/Hour | Best For |
|---|---|---|---|---|---|---|---|
| trn1.2xlarge | 1 core | 32 GB | 8 | 32 GB | 12.5 Gbps | $1.34 | Development, testing |
| trn1.32xlarge | 16 cores | 512 GB | 128 | 512 GB | 800 Gbps EFA | $21.50 | Single-node training |
| trn1n.32xlarge | 16 cores | 512 GB | 128 | 512 GB | 1600 Gbps EFA | $24.78 | Multi-node, distributed training |
Key Differences:
- trn1.32xlarge: Standard network (800 Gbps EFA) for single-node or small-scale distributed training
- trn1n.32xlarge: 2× faster network (1600 Gbps EFA) for large-scale distributed training (64+ instances)
Regional Availability (as of 2026):
- US East (N. Virginia)
- US West (Oregon)
- Asia Pacific (Singapore, Tokyo)
- Europe (Frankfurt, Ireland)
Each Trainium chip contains 2 NeuronCore-v2 processors. A trn1.32xlarge instance contains 16 NeuronCores (8 Trainium chips).
┌─────────────────────────────────────────────────────────────┐
│ NeuronCore-v2 Processor │
├─────────────────────────────────────────────────────────────┤
│ ┌───────────────┐ ┌───────────────┐ ┌──────────────┐ │
│ │ Tensor Engine │ │ Vector Engine │ │ Scalar Engine│ │
│ │ (Matrix Ops) │ │ (Element-wise)│ │ (Control) │ │
│ │ │ │ │ │ │ │
│ │ • BF16/FP16 │ │ • BF16/FP16 │ │ • FP32 │ │
│ │ • FP8 │ │ • FP32 │ │ • INT32 │ │
│ │ • INT8 │ │ • Custom ops │ │ │ │
│ └───────┬───────┘ └───────┬───────┘ └──────┬───────┘ │
│ └──────────────────┼────────────────┘ │
│ │ │
│ ┌──────────────▼──────────────┐ │
│ │ Local Memory (SRAM) │ │
│ │ - 96 MB on-chip cache │ │
│ │ - 3 TB/s bandwidth │ │
│ └──────────────┬──────────────┘ │
│ │ │
│ ┌──────────────▼──────────────┐ │
│ │ HBM2E Memory Interface │ │
│ │ - 32 GB HBM per core │ │
│ │ - 820 GB/s bandwidth │ │
│ └─────────────────────────────┘ │
└─────────────────────────────────────────────────────────────┘
1. Tensor Engine (Matrix Math Unit)
- Purpose: High-throughput matrix multiplication for transformer layers
- Operations:
Y = XW(dense matrix multiplication)Q @ K^T(attention scores)- Supports BF16 (native), FP16, FP8, INT8
- Performance:
- 190 TFLOPS BF16 per NeuronCore
- 380 TFLOPS FP8 per NeuronCore (2× speedup)
2. Vector Engine (Element-wise Operations)
- Purpose: Non-matrix operations (activations, normalization)
- Operations:
- GELU, SiLU, Softmax
- LayerNorm, RMSNorm
- Element-wise addition, multiplication
- Performance:
- 95 TFLOPS BF16 per NeuronCore
- Optimized for memory-bound ops
3. Scalar Engine (Control Logic)
- Purpose: Branching, indexing, scheduling
- Operations:
- Loop control
- Memory addressing
- Inter-core synchronization
4. On-Chip Memory Hierarchy
- L1 Cache: 96 MB SRAM per NeuronCore (3 TB/s)
- HBM2E: 32 GB per NeuronCore (820 GB/s)
- Local Scratch: 128 MB for intermediate results
┌─────────────────────────────────────────────────────────────┐
│ trn1.32xlarge Memory Layout │
├─────────────────────────────────────────────────────────────┤
│ │
│ NeuronCore 0 → HBM: 32 GB │ NeuronCore 8 → HBM: 32 GB │
│ NeuronCore 1 → HBM: 32 GB │ NeuronCore 9 → HBM: 32 GB │
│ NeuronCore 2 → HBM: 32 GB │ NeuronCore 10 → HBM: 32 GB │
│ NeuronCore 3 → HBM: 32 GB │ NeuronCore 11 → HBM: 32 GB │
│ NeuronCore 4 → HBM: 32 GB │ NeuronCore 12 → HBM: 32 GB │
│ NeuronCore 5 → HBM: 32 GB │ NeuronCore 13 → HBM: 32 GB │
│ NeuronCore 6 → HBM: 32 GB │ NeuronCore 14 → HBM: 32 GB │
│ NeuronCore 7 → HBM: 32 GB │ NeuronCore 15 → HBM: 32 GB │
│ │
│ Total HBM: 512 GB (32 GB × 16 cores) │
│ │
│ ┌─────────────────────────────────────┐ │
│ │ System Memory (DDR4): 512 GB │ │
│ │ - Model loading │ │
│ │ - Dataset preprocessing │ │
│ │ - Checkpointing │ │
│ └─────────────────────────────────────┘ │
└─────────────────────────────────────────────────────────────┘
Memory Characteristics:
- HBM bandwidth: 820 GB/s per NeuronCore = 13 TB/s aggregate
- HBM latency: 120 ns (vs 200 ns for GDDR6)
- System memory: Used for data loading, not training compute
Intra-Instance Communication (between NeuronCores on same trn1.32xlarge):
┌─────────────────────────────────────────────────────────────┐
│ NeuronLink-v2 Topology │
├─────────────────────────────────────────────────────────────┤
│ │
│ [NC0]───[NC1]───[NC2]───[NC3] [NC8]───[NC9]──... │
│ │ │ │ │ │ │ │
│ └───────┼───────┼───────┤ └───────┼────... │
│ │ │ │ │ │
│ [NC4]───[NC5]───[NC6]───[NC7] [NC12]─... │
│ │
│ Bandwidth: 384 GB/s bidirectional between each pair │
│ Latency: <5 μs for RDMA operations │
│ Collective ops: All-Reduce, All-Gather, Reduce-Scatter │
└─────────────────────────────────────────────────────────────┘
Inter-Instance Communication (EFA for multi-node):
| Network Type | Bandwidth | Use Case |
|---|---|---|
| trn1.32xlarge | 800 Gbps (100 GB/s) EFA | Single-node or 2-8 node training |
| trn1n.32xlarge | 1600 Gbps (200 GB/s) EFA | Large-scale distributed (16-128 nodes) |
Trainium natively supports multiple precision formats optimized for different training stages:
Format:
- 1 sign bit, 8 exponent bits, 7 mantissa bits
- Same dynamic range as FP32 (10^-38 to 10^38)
- Lower precision than FP16 (7 vs 10 mantissa bits)
Advantages on Trainium:
- Native hardware support in Tensor Engine
- No loss scaling required (unlike FP16)
- Same numerical stability as FP32 for most models
- 2× memory savings vs FP32
Usage:
import torch
import torch_xla.core.xla_model as xm
# Enable BF16 training
model = MyLlamaModel().to(xm.xla_device())
model = model.bfloat16() # Convert to BF16
# Training step (automatic mixed precision)
with torch.autocast(device_type='xla', dtype=torch.bfloat16):
outputs = model(input_ids)
loss = outputs.loss
loss.backward()Format (E4M3):
- 1 sign bit, 4 exponent bits, 3 mantissa bits
- Range: 10^-9 to ~448
Advantages:
- 2× throughput vs BF16 (380 TFLOPS vs 190 TFLOPS per core)
- 2× memory savings (critical for large models)
- Lower power consumption
Limitations:
- Requires careful quantization-aware training
- Not all layers benefit (embeddings, norms typically stay BF16)
- Still experimental for LLMs (stable for vision models)
Usage:
# FP8 training (requires Neuron SDK 2.18+)
from neuronx_distributed.quantization import quantize_model
model = quantize_model(
model,
dtype=torch.float8_e4m3fn,
layers_to_quantize=['attention.qkv', 'mlp.dense']
)Trainium's Tensor Engine implements stochastic rounding for BF16 accumulation:
Traditional rounding: 1.501 + 1.501 = 3.0 (rounded down twice)
Stochastic rounding: 1.501 + 1.501 = 3.0 (50% chance) or 4.0 (50% chance)
Expected value: 3.002 (unbiased)
Impact: Reduces gradient accumulation errors in long training runs.
Trainium's asynchronous gradient accumulation overlaps compute and communication:
# Traditional (sequential)
for micro_batch in range(gradient_accumulation_steps):
loss = model(micro_batch)
loss.backward() # Wait for backward pass
xm.optimizer_step(optimizer) # Then update weights
# Neuron-optimized (overlapped)
from neuronx_distributed.pipeline import GradientAccumulationConfig
config = GradientAccumulationConfig(
num_microbatches=8,
overlap_grad_reduce=True # Overlap all-reduce with backward
)
# Gradients are reduced asynchronously during backward pass
# No blocking wait before optimizer stepPerformance Gain: 15-20% faster training by hiding all-reduce latency.
Problem: Training large models requires storing all activations for backpropagation.
- Llama 2 70B: ~150 GB activations per batch (doesn't fit in 32 GB HBM)
Solution: Recompute activations during backward pass instead of storing them.
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
checkpoint_wrapper,
CheckpointImpl,
apply_activation_checkpointing
)
import torch_xla.core.xla_model as xm
# Checkpoint transformer blocks
def checkpoint_policy(submodule):
return isinstance(submodule, TransformerBlock)
model = apply_activation_checkpointing(
model,
checkpoint_wrapper_fn=lambda m: checkpoint_wrapper(
m, checkpoint_impl=CheckpointImpl.NO_REENTRANT
),
check_fn=checkpoint_policy
)
# Now can train Llama 70B on 16 NeuronCoresTrade-off:
- Memory savings: 80% reduction (150 GB → 30 GB)
- Compute overhead: 30% slower (recomputation cost)
- Net benefit: Can train models that otherwise wouldn't fit
Trainium excels at distributed training via Neuron Distributed, which supports three parallelism strategies:
Concept: Replicate model on each NeuronCore, split data across cores.
┌─────────────────────────────────────────────────────────────┐
│ Data Parallelism │
├─────────────────────────────────────────────────────────────┤
│ │
│ NeuronCore 0: [Full Model] ← Batch 0 │
│ NeuronCore 1: [Full Model] ← Batch 1 │
│ NeuronCore 2: [Full Model] ← Batch 2 │
│ ... │
│ NeuronCore 15: [Full Model] ← Batch 15 │
│ │
│ After forward/backward: All-Reduce gradients across cores │
└─────────────────────────────────────────────────────────────┘
When to Use:
- Model fits in single NeuronCore memory (< 32 GB)
- Examples: Llama 2 7B, Mistral 7B, GPT-2
Code Example:
import torch_xla.core.xla_model as xm
import torch_xla.distributed.parallel_loader as pl
from torch.utils.data import DataLoader
# Initialize distributed training
device = xm.xla_device()
model = MyModel().to(device)
# Wrap dataloader for data parallelism
dataloader = DataLoader(dataset, batch_size=8)
parallel_loader = pl.MpDeviceLoader(dataloader, device)
# Training loop
for batch in parallel_loader:
outputs = model(batch)
loss = outputs.loss
loss.backward()
# Synchronize gradients across all NeuronCores
xm.optimizer_step(optimizer, barrier=True)Performance:
- Scaling efficiency: 95%+ on trn1.32xlarge (16 cores)
- Communication overhead: ~5% (all-reduce every step)
Concept: Split model layers across NeuronCores (column/row-wise).
┌─────────────────────────────────────────────────────────────┐
│ Tensor Parallelism (TP=4) │
├─────────────────────────────────────────────────────────────┤
│ │
│ Layer: Attention Q/K/V projection (4096 → 4096) │
│ │
│ NeuronCore 0: W[0:1024, :] ← Process columns 0-1023 │
│ NeuronCore 1: W[1024:2048, :] ← Process columns 1024-2047 │
│ NeuronCore 2: W[2048:3072, :] ← Process columns 2048-3071 │
│ NeuronCore 3: W[3072:4096, :] ← Process columns 3072-4095 │
│ │
│ After layer: All-Gather to combine results │
└─────────────────────────────────────────────────────────────┘
When to Use:
- Model too large for single NeuronCore (Llama 2 70B, GPT-3)
- Need to maximize per-step throughput
- Have fast interconnect (NeuronLink)
Code Example:
from neuronx_distributed.parallel_layers import (
ColumnParallelLinear,
RowParallelLinear,
parallel_state
)
# Initialize tensor parallelism (TP=8 for 70B model)
parallel_state.initialize_model_parallel(
tensor_model_parallel_size=8
)
# Replace standard layers with tensor-parallel versions
class ParallelAttention(nn.Module):
def __init__(self, config):
super().__init__()
# Split Q/K/V across 8 NeuronCores (columns)
self.qkv = ColumnParallelLinear(
config.hidden_size,
3 * config.hidden_size,
gather_output=False # Don't gather yet
)
# Split output projection across 8 NeuronCores (rows)
self.out_proj = RowParallelLinear(
config.hidden_size,
config.hidden_size,
input_is_parallel=True # Input already split
)Performance:
- Scaling efficiency: 90-93% on trn1.32xlarge (TP=16)
- Communication overhead: ~7-10% (all-gather/reduce-scatter per layer)
- Memory savings: 8× (model split across 8 cores)
Concept: Split model layers into stages, each stage on different NeuronCores.
┌─────────────────────────────────────────────────────────────┐
│ Pipeline Parallelism (PP=4) │
├─────────────────────────────────────────────────────────────┤
│ │
│ Stage 0 (NC 0-3): Layers 0-7 ← Embedding + Early layers│
│ ↓ (pass activations) │
│ Stage 1 (NC 4-7): Layers 8-15 ← Middle layers │
│ ↓ │
│ Stage 2 (NC 8-11): Layers 16-23 ← Late layers │
│ ↓ │
│ Stage 3 (NC 12-15): Layers 24-31 ← Final layers + Head │
│ │
│ Microbatching: Process multiple batches in parallel │
└─────────────────────────────────────────────────────────────┘
When to Use:
- Very large models (GPT-3, Llama 2 70B)
- Multi-node training (minimize inter-node communication)
- Combined with TP for massive models
Code Example:
from neuronx_distributed.pipeline import NeuronPipelineModule
# Define model stages
stages = [
nn.Sequential(embedding, *layers[0:8]), # Stage 0
nn.Sequential(*layers[8:16]), # Stage 1
nn.Sequential(*layers[16:24]), # Stage 2
nn.Sequential(*layers[24:32], lm_head) # Stage 3
]
# Initialize pipeline parallelism
parallel_state.initialize_model_parallel(
pipeline_model_parallel_size=4,
tensor_model_parallel_size=4 # Also use TP within each stage
)
# Wrap model for pipeline execution
model = NeuronPipelineModule(
stages,
num_microbatches=16, # Split batch into 16 microbatches
loss_fn=cross_entropy_loss
)Performance:
- Scaling efficiency: 85-88% (pipeline bubbles reduce efficiency)
- Communication overhead: ~12-15% (stage boundaries)
- Memory savings: 4× (model split across 4 stages)
For massive scale (Llama 2 70B on 64+ instances), combine all three:
┌─────────────────────────────────────────────────────────────┐
│ Llama 2 70B: 64× trn1n.32xlarge (1024 cores) │
├─────────────────────────────────────────────────────────────┤
│ │
│ Data Parallel (DP=8): 8 model replicas │
│ ├─ Replica 0 (128 cores) │
│ │ ├─ Pipeline Parallel (PP=4): 4 stages │
│ │ │ ├─ Stage 0 (32 cores): Tensor Parallel (TP=32) │
│ │ │ ├─ Stage 1 (32 cores): Tensor Parallel (TP=32) │
│ │ │ ├─ Stage 2 (32 cores): Tensor Parallel (TP=32) │
│ │ │ └─ Stage 3 (32 cores): Tensor Parallel (TP=32) │
│ ├─ Replica 1 (128 cores) ... same structure │
│ └─ ... Replica 7 │
│ │
│ Effective throughput: 165K tokens/s │
│ Cost: $1,584/hour (vs $2,048/hour for 64× A100) │
└─────────────────────────────────────────────────────────────┘
Configuration:
parallel_state.initialize_model_parallel(
data_parallel_size=8, # 8 model replicas
tensor_model_parallel_size=32, # 32-way TP within each stage
pipeline_model_parallel_size=4 # 4 pipeline stages
)
# Total: 8 × 32 × 4 = 1024 NeuronCores┌─────────────────────────────────────────────────────────────┐
│ Neuron SDK Stack │
├─────────────────────────────────────────────────────────────┤
│ │
│ ┌─────────────────────────────────────────────────────┐ │
│ │ ML Frameworks (User Code) │ │
│ │ • PyTorch 2.1+ │ │
│ │ • TensorFlow 2.10+ │ │
│ │ • JAX 0.4+ │ │
│ └─────────────────┬───────────────────────────────────┘ │
│ │ │
│ ┌─────────────────▼───────────────────────────────────┐ │
│ │ Framework Integration │ │
│ │ • torch-xla (PyTorch) │ │
│ │ • tensorflow-neuronx (TensorFlow) │ │
│ │ • jax-neuronx (JAX) │ │
│ └─────────────────┬───────────────────────────────────┘ │
│ │ │
│ ┌─────────────────▼───────────────────────────────────┐ │
│ │ Neuron Compiler (neuronx-cc) │ │
│ │ • XLA → Neuron IR → Binary │ │
│ │ • Operator fusion, memory optimization │ │
│ └─────────────────┬───────────────────────────────────┘ │
│ │ │
│ ┌─────────────────▼───────────────────────────────────┐ │
│ │ Neuron Runtime (neuronx-runtime) │ │
│ │ • Kernel execution, memory management │ │
│ │ • NeuronLink collective operations │ │
│ └─────────────────┬───────────────────────────────────┘ │
│ │ │
│ ┌─────────────────▼───────────────────────────────────┐ │
│ │ Neuron Driver (aws-neuronx-dkms) │ │
│ │ • Hardware access, DMA, interrupts │ │
│ └─────────────────────────────────────────────────────┘ │
└─────────────────────────────────────────────────────────────┘
# On trn1 instance (Amazon Linux 2 or Ubuntu 20.04)
# 1. Install Neuron driver (kernel module)
sudo apt-get update -y
sudo apt-get install linux-headers-$(uname -r) -y
sudo apt-get install aws-neuronx-dkms -y
# 2. Install Neuron runtime
sudo apt-get install aws-neuronx-collectives -y
sudo apt-get install aws-neuronx-runtime-lib -y
# 3. Install Neuron tools
sudo apt-get install aws-neuronx-tools -y
# 4. Verify installation
neuron-ls # Should list 16 NeuronCores on trn1.32xlarge
# 5. Install PyTorch Neuron (in virtual environment)
pip install torch-neuronx==2.1.2.1.0 --extra-index-url=https://pip.repos.neuron.amazonaws.com
pip install neuronx-distributed==0.6.0
pip install transformers-neuronx==0.9.474Neuron uses ahead-of-time (AOT) compilation via XLA:
import torch
import torch_xla.core.xla_model as xm
# 1. Define model (standard PyTorch)
model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
model = model.to(xm.xla_device()) # Move to Neuron device
model = model.bfloat16()
# 2. First forward pass triggers compilation
input_ids = torch.randint(0, 32000, (1, 2048)).to(xm.xla_device())
# This will take 5-10 minutes (compiling all layers)
print("Compiling model... (first run only)")
outputs = model(input_ids)
# 3. Subsequent runs use cached compiled binary
print("Running compiled model... (fast)")
outputs = model(input_ids) # <1sCompilation Stages:
- Graph Capture (torch-xla): Convert PyTorch ops to XLA HLO
- Optimization (neuronx-cc):
- Operator fusion (e.g., matmul + bias + GELU → single kernel)
- Memory layout optimization (HBM → SRAM → Tensor Engine)
- Collective operation scheduling (overlap compute + comms)
- Code Generation: Emit NeuronCore-v2 assembly
- Caching: Store binary in
~/.cache/neuron(reuse across runs)
Compilation Flags:
# Set via environment variable
export NEURON_CC_FLAGS="--model-type=transformer --auto-cast=bf16 --distribution-strategy=llm-training --internal-hlo-simplification=True"
# Or in Python
import os
os.environ['NEURON_CC_FLAGS'] = '--model-type=transformer --auto-cast=bf16'Common Flags:
--model-type=transformer: Optimize for transformer models--auto-cast=bf16: Automatic BF16 casting where safe--distribution-strategy=llm-training: Enable distributed training optimizations--enable-mixed-precision-accumulation: Use FP32 for gradient accumulation
┌─────────────────────────────────────────────────────────────┐
│ EKS Control Plane │
│ (Kubernetes API Server) │
└─────────────────────────────┬───────────────────────────────┘
│
┌───────────────┼───────────────┐
│ │ │
┌─────────▼─────────┐ ┌──▼───────┐ ┌────▼──────────┐
│ Karpenter │ │ Neuron │ │ Monitoring │
│ (Autoscaler) │ │ Device │ │ (Prometheus) │
│ │ │ Plugin │ │ │
└─────────┬─────────┘ └──┬───────┘ └────┬──────────┘
│ │ │
┌─────────▼───────────────▼───────────────▼──────────┐
│ Trainium Node Pool │
│ ┌──────────────┐ ┌──────────────┐ │
│ │ trn1.32xlarge│ │ trn1.32xlarge│ ... │
│ │ 16 cores │ │ 16 cores │ │
│ └──────────────┘ └──────────────┘ │
│ │
│ Taints: aws.amazon.com/neuron=trainium:NoSchedule │
│ Labels: instance-type=trn1.32xlarge │
└─────────────────────────────────────────────────────┘
Deploy Neuron Device Plugin to expose NeuronCores as Kubernetes resources:
# neuron-device-plugin.yaml
apiVersion: apps/v1
kind: DaemonSet
metadata:
name: neuron-device-plugin
namespace: kube-system
spec:
selector:
matchLabels:
name: neuron-device-plugin
template:
metadata:
labels:
name: neuron-device-plugin
spec:
# Only run on Trainium nodes
nodeSelector:
node.kubernetes.io/instance-type: trn1.32xlarge
tolerations:
- key: aws.amazon.com/neuron
operator: Exists
effect: NoSchedule
containers:
- name: neuron-device-plugin
image: public.ecr.aws/neuron/neuron-device-plugin:2.18.1
env:
- name: KUBECONFIG
value: /etc/kubernetes/kubelet.conf
securityContext:
privileged: true
volumeMounts:
- name: device-plugin
mountPath: /var/lib/kubelet/device-plugins
- name: neuron-devices
mountPath: /dev
volumes:
- name: device-plugin
hostPath:
path: /var/lib/kubelet/device-plugins
- name: neuron-devices
hostPath:
path: /devDeploy:
kubectl apply -f neuron-device-plugin.yaml
# Verify NeuronCores are advertised
kubectl describe node <trn1-node> | grep neuroncore
# Should show: aws.amazon.com/neuroncore: 16# karpenter-trainium-nodepool.yaml
apiVersion: karpenter.sh/v1beta1
kind: NodePool
metadata:
name: trainium-training
spec:
template:
metadata:
labels:
workload-type: ml-training
accelerator: trainium
spec:
# Taint to prevent non-Neuron workloads
taints:
- key: aws.amazon.com/neuron
value: trainium
effect: NoSchedule
nodeClassRef:
name: trainium-nodeclass
requirements:
- key: node.kubernetes.io/instance-type
operator: In
values:
- trn1.32xlarge
- trn1n.32xlarge # For multi-node training
- key: kubernetes.io/arch
operator: In
values:
- amd64
- key: karpenter.sh/capacity-type
operator: In
values:
- on-demand # Training jobs need reliability
# Autoscaling limits
limits:
cpu: "2048" # Max 16 nodes (128 vCPU × 16)
memory: 8192Gi
# Node lifecycle
disruption:
consolidationPolicy: WhenUnderutilized
consolidateAfter: 30m # Don't scale down too quickly
expireAfter: 720h # 30 days
---
apiVersion: karpenter.k8s.aws/v1beta1
kind: EC2NodeClass
metadata:
name: trainium-nodeclass
spec:
amiFamily: AL2
role: KarpenterNodeRole-eks-cluster
# Use Deep Learning AMI (has Neuron drivers pre-installed)
amiSelectorTerms:
- name: "Deep Learning AMI Neuron PyTorch*"
owner: "amazon"
subnetSelectorTerms:
- tags:
karpenter.sh/discovery: eks-cluster
securityGroupSelectorTerms:
- tags:
karpenter.sh/discovery: eks-cluster
# Enable EFA for multi-node training
networkInterfaces:
- deviceIndex: 0
interfaceType: efa
userData: |
#!/bin/bash
# Additional Neuron setup
sudo modprobe neuron
sudo systemctl enable neuron-monitor
sudo systemctl start neuron-monitor
# Increase shared memory for large models
sudo mount -o remount,size=64G /dev/shm
# Configure EFA (for multi-node)
sudo sed -i 's/^# *export FI_EFA_USE_DEVICE_RDMA=1/export FI_EFA_USE_DEVICE_RDMA=1/' /etc/profile.d/efa.sh
blockDeviceMappings:
- deviceName: /dev/xvda
ebs:
volumeSize: 500Gi # Large disk for model checkpoints
volumeType: gp3
iops: 16000
throughput: 1000
deleteOnTermination: true# llama2-7b-training.yaml
apiVersion: batch/v1
kind: Job
metadata:
name: llama2-7b-training
spec:
backoffLimit: 2
template:
metadata:
labels:
app: llama2-training
spec:
# Tolerate Trainium taint
tolerations:
- key: aws.amazon.com/neuron
value: trainium
effect: NoSchedule
# Request trn1.32xlarge
affinity:
nodeAffinity:
requiredDuringSchedulingIgnoredDuringExecution:
nodeSelectorTerms:
- matchExpressions:
- key: node.kubernetes.io/instance-type
operator: In
values:
- trn1.32xlarge
containers:
- name: trainer
image: <account>.dkr.ecr.us-east-1.amazonaws.com/llama-training:latest
command:
- python3
- train.py
- --model_id=meta-llama/Llama-2-7b-hf
- --output_dir=/models/llama2-7b-finetuned
- --tensor_parallel_degree=8
- --num_train_epochs=3
- --bf16
resources:
requests:
aws.amazon.com/neuroncore: "16" # All 16 cores
memory: 400Gi
cpu: "96"
limits:
aws.amazon.com/neuroncore: "16"
memory: 512Gi
cpu: "128"
env:
- name: NEURON_RT_NUM_CORES
value: "16"
- name: NEURON_CC_FLAGS
value: "--model-type=transformer --auto-cast=bf16"
- name: XLA_USE_BF16
value: "1"
volumeMounts:
- name: model-output
mountPath: /models
- name: shm
mountPath: /dev/shm
volumes:
- name: model-output
persistentVolumeClaim:
claimName: efs-models-pvc
- name: shm
emptyDir:
medium: Memory
sizeLimit: 64GiFor large models (Llama 2 70B), use MPIJob with torchrun:
# llama2-70b-multinode.yaml
apiVersion: kubeflow.org/v2beta1
kind: MPIJob
metadata:
name: llama2-70b-training
spec:
slotsPerWorker: 16 # 16 NeuronCores per worker
runPolicy:
cleanPodPolicy: Running
mpiReplicaSpecs:
Launcher:
replicas: 1
template:
spec:
containers:
- name: launcher
image: <account>.dkr.ecr.us-east-1.amazonaws.com/llama-training:latest
command:
- mpirun
- -np
- "128" # 8 nodes × 16 cores
- --host
- <worker-0>:16,<worker-1>:16,...,<worker-7>:16
- --mca
- btl_tcp_if_include
- eth0
- --mca
- oob_tcp_if_include
- eth0
- torchrun
- --nnodes=8
- --nproc_per_node=16
- --node_rank=$OMPI_COMM_WORLD_RANK
- train_llama_70b.py
- --tensor_parallel_degree=32
- --pipeline_parallel_degree=4
Worker:
replicas: 8 # 8× trn1n.32xlarge
template:
spec:
tolerations:
- key: aws.amazon.com/neuron
value: trainium
effect: NoSchedule
affinity:
nodeAffinity:
requiredDuringSchedulingIgnoredDuringExecution:
nodeSelectorTerms:
- matchExpressions:
- key: node.kubernetes.io/instance-type
operator: In
values:
- trn1n.32xlarge # Need high-bandwidth network
containers:
- name: worker
image: <account>.dkr.ecr.us-east-1.amazonaws.com/llama-training:latest
resources:
requests:
aws.amazon.com/neuroncore: "16"
memory: 400Gi
cpu: "96"
limits:
aws.amazon.com/neuroncore: "16"
memory: 512Gi
cpu: "128"
env:
- name: NEURON_RT_NUM_CORES
value: "16"
- name: FI_EFA_USE_DEVICE_RDMA
value: "1"
- name: FI_PROVIDER
value: "efa"1. Model Sharding (ZeRO-1/2/3):
from neuronx_distributed.optimizer import NeuronZero1Optimizer
# ZeRO-1: Shard optimizer states only (4× memory savings for AdamW)
optimizer = NeuronZero1Optimizer(
model.parameters(),
torch.optim.AdamW,
lr=2e-5,
pin_layout=False # Let Neuron choose optimal layout
)
# ZeRO-2: Also shard gradients (additional 2× savings)
from neuronx_distributed.parallel_layers import DistributedDataParallel
model = DistributedDataParallel(
model,
gradient_as_bucket_view=True,
zero_optimization_stage=2
)2. Activation Checkpointing:
# Checkpoint every N layers (trade compute for memory)
from torch.utils.checkpoint import checkpoint
class CheckpointedTransformerBlock(nn.Module):
def forward(self, x):
# Recompute activations during backward pass
return checkpoint(self._forward, x, use_reentrant=False)
def _forward(self, x):
# Standard transformer block
return self.attention(self.ln1(x)) + xMemory Budget (trn1.32xlarge with 512 GB HBM):
| Model | Parameters | Weights (BF16) | Gradients (BF16) | Optimizer States (FP32) | Activations (batch=8) | Total | Fits? |
|---|---|---|---|---|---|---|---|
| Llama 2 7B | 7B | 14 GB | 14 GB | 28 GB | 12 GB | 68 GB | ✅ Single core with ZeRO-1 |
| Llama 2 13B | 13B | 26 GB | 26 GB | 52 GB | 18 GB | 122 GB | ✅ TP=4 (30 GB per core) |
| Llama 2 70B | 70B | 140 GB | 140 GB | 280 GB | 45 GB | 605 GB | ✅ TP=32 (19 GB per core) |
1. Increase Batch Size:
# Larger batches improve hardware utilization
# Start with batch_size that fits in memory, then increase
# Too small: 60% utilization
batch_size = 1
gradient_accumulation_steps = 128
# Optimal: 95% utilization
batch_size = 8
gradient_accumulation_steps = 162. Tune Gradient Accumulation:
# Balance between memory and communication frequency
# More frequent updates (better convergence, more communication)
gradient_accumulation_steps = 4
# Less frequent updates (better throughput, delayed convergence)
gradient_accumulation_steps = 32
# Optimal for most LLMs: 8-16 steps3. Enable Asynchronous Execution:
# Overlap data loading with training
import torch_xla.distributed.parallel_loader as pl
dataloader = DataLoader(dataset, batch_size=8, num_workers=4, pin_memory=True)
parallel_loader = pl.MpDeviceLoader(dataloader, device, loader_prefetch_size=2)
# Overlap gradient reduction with backward pass
from neuronx_distributed.parallel_layers import reduce_grads_async
def training_step(model, batch, optimizer):
loss = model(batch).loss
loss.backward()
# Start gradient reduction asynchronously
reduce_grads_async(model)
# Do other work while reduction happens
log_metrics(loss.item())
# Wait for reduction to complete before optimizer step
optimizer.step()1. Precompile Models:
# Compile once, use many times
import torch_xla.core.xla_model as xm
def precompile_model(model, input_shape):
device = xm.xla_device()
model = model.to(device)
# Trigger compilation with dummy input
dummy_input = torch.randint(0, 32000, input_shape).to(device)
print("Compiling model (10-15 minutes)...")
with torch.no_grad():
model(dummy_input)
# Save compiled binary
xm.save_compiled_binary(model, "model_compiled.pt")
print("Compilation complete!")
# Later: Load precompiled binary (instant)
model = xm.load_compiled_binary("model_compiled.pt")2. Reduce Compilation Time:
# Use fewer cores for compilation (compiles faster, same runtime performance)
NEURON_RT_NUM_CORES=1 python train.py # First run (compilation)
# Then switch to full cores for training
NEURON_RT_NUM_CORES=16 python train.py # Subsequent runs# Development workflow:
# 1. Test on small model locally
model = "meta-llama/Llama-2-7b-hf"
tensor_parallel_degree = 1
num_epochs = 1
max_steps = 100
# 2. Verify on single trn1.32xlarge
tensor_parallel_degree = 8
max_steps = 1000
# 3. Scale to multi-node
tensor_parallel_degree = 32
pipeline_parallel_degree = 4
num_nodes = 8# Use neuron-top to monitor NeuronCore usage
neuron-top
# Expected metrics:
# NeuronCore utilization: 90-95% (good)
# HBM usage: 70-85% (optimal)
# Tensor Engine TFLOPS: 170-185 per core (near theoretical 190)
# Low utilization (<70%)? Increase batch size or check for CPU bottlenecksfrom transformers import TrainingArguments
args = TrainingArguments(
output_dir="/models/checkpoints",
# Save every N steps (balance between safety and storage)
save_steps=500,
save_total_limit=3, # Keep only 3 most recent checkpoints
# Save to EFS/S3 (not local disk)
save_on_each_node=False, # Only rank 0 saves
# Resume from checkpoint
resume_from_checkpoint="/models/checkpoints/checkpoint-1500"
)
# For S3 checkpointing:
import boto3
def save_checkpoint_to_s3(checkpoint_path, step):
s3 = boto3.client('s3')
s3.upload_file(
checkpoint_path,
"my-training-bucket",
f"checkpoints/step-{step}/pytorch_model.bin"
)import torch
import random
import numpy as np
import torch_xla.core.xla_model as xm
def set_seed(seed):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
xm.set_rng_state(seed) # XLA random state
set_seed(42)
# Also pin dataset order
dataset = dataset.shuffle(seed=42)import math
class LossMonitor:
def __init__(self, window_size=100):
self.losses = []
self.window_size = window_size
def check_divergence(self, loss):
self.losses.append(loss)
if len(self.losses) < self.window_size:
return False
recent = self.losses[-self.window_size:]
# Check for NaN
if math.isnan(loss):
print(f"❌ NaN detected at step {len(self.losses)}")
return True
# Check for explosion (10× increase in 100 steps)
if loss > 10 * min(recent):
print(f"❌ Loss exploded: {loss:.3f} vs min {min(recent):.3f}")
return True
return False
monitor = LossMonitor()
for step, batch in enumerate(dataloader):
loss = training_step(batch)
if monitor.check_divergence(loss.item()):
print("Stopping training due to divergence")
breakScenario: Fine-tune Llama 2 7B on domain-specific data (20K samples)
Configuration:
Instance: 1× trn1.32xlarge
NeuronCores: 16 (tensor parallel = 8)
Batch size: 8 per core (64 global batch)
Training time: 4.2 hours
Cost: $90 (vs $137 on 8× A100)Results:
- Training throughput: 1,050K tokens/s
- Final loss: 1.72 (converged)
- Model quality: Same as GPU baseline
- Savings: 34% vs A100
Scenario: Continue pre-training GPT-NeoX 20B (100B tokens)
Configuration:
Instances: 8× trn1n.32xlarge
NeuronCores: 128 (TP=16, PP=2, DP=4)
Batch size: 4 per core (512 global batch)
Training time: 72 hours
Cost: $14,240 (vs $18,545 on 64× A100)Results:
- Throughput: 4.2M tokens/s
- Scaling efficiency: 94%
- Savings: 23% vs A100
Scenario: Fine-tune SDXL on custom images (10K images)
Configuration:
Instance: 1× trn1.32xlarge
NeuronCores: 8 (TP=8)
Batch size: 16
Training time: 6 hours
Cost: $129Results:
- Images/s: 42
- FID score: 12.3 (same as GPU)
- 50% cost reduction vs A100
Symptoms:
RuntimeError: [Neuron] Out of memory on NeuronCore 0
HBM usage: 32.0 GB / 32.0 GB (100%)
Solutions:
- Reduce batch size:
per_device_train_batch_size = 8 # Try 4
gradient_accumulation_steps = 8 # Increase to 16- Enable activation checkpointing:
model.gradient_checkpointing_enable()- Use ZeRO-1 optimizer:
from neuronx_distributed.optimizer import NeuronZero1Optimizer
optimizer = NeuronZero1Optimizer(model.parameters(), torch.optim.AdamW, lr=2e-5)- Increase tensor parallelism:
# Split model across more cores
tensor_parallel_degree = 16 # Was 8Symptoms:
Compiling layer 0/32... (elapsed: 15 minutes)
Solutions:
- Use cached compilation:
# Check cache
ls -lh ~/.cache/neuron/
# If empty, compilation is running
# Compilation cache persists across runs
# Second run will be instant- Compile with fewer cores:
NEURON_RT_NUM_CORES=1 python train.py # Faster compilation- Use pre-compiled models:
# Download from Hugging Face Hub
from optimum.neuron import NeuronModelForCausalLM
model = NeuronModelForCausalLM.from_pretrained(
"aws-neuron/Llama-2-7b-hf-neuron-latency",
export=False # Already compiled
)Diagnosis:
# Check gradient norms
total_norm = 0
for p in model.parameters():
if p.grad is not None:
param_norm = p.grad.data.norm(2)
total_norm += param_norm.item() ** 2
total_norm = total_norm ** 0.5
print(f"Gradient norm: {total_norm}")
# Should be 0.1-10 (if >100 or <0.001, something is wrong)Solutions:
- Adjust learning rate:
# Too high: Loss oscillates or diverges
learning_rate = 5e-5 # Try 1e-5
# Too low: Loss decreases very slowly
learning_rate = 1e-6 # Try 5e-5- Check data quality:
# Verify tokenization
batch = next(iter(dataloader))
print(tokenizer.decode(batch['input_ids'][0]))
# Should be readable text, not gibberish- Enable gradient clipping:
from torch.nn.utils import clip_grad_norm_
# After loss.backward()
clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()# Collect detailed performance metrics
NEURON_PROFILE=profile.json python train.py
# Analyze profile
neuron-profile view profile.json
# Look for:
# - Low NeuronCore utilization (<70%): Increase batch size
# - High CPU usage: Optimize data preprocessing
# - Frequent compilation: Cache miss, check cache directoryNext Steps:
- Explore Inferentia2 Architecture for inference optimization
- Learn vLLM-Neuron Deployment for production serving
- Set up Benchmarking to measure your workload performance