Skip to main content
Ctrl+K
Neuron 2.33.0 is released! Check the What's New and Release Notes for more details.
Logo image
Ctrl+K
Search Engine: Default Google
  • About Neuron
    • Contributing Guidelines FAQ
    • Neuron FAQ
    • Neuron Features
      • Data types
      • Logical NeuronCore configuration
      • Neuron persistent cache
      • NeuronCore batching
      • NeuronCore pipeline
      • Rounding modes
    • Open Source
    • SDK Maintenance Policy
    • Security
    • Term Glossary
    • Troubleshooting
    • What is AWS Neuron?
  • Neuron architecture
    • AWS Inferentia
    • AWS Inferentia2
    • AWS Trainium
    • AWS Trainium2
    • AWS Trainium3
    • NeuronCore v1
    • NeuronCore v2
    • NeuronCore v3
    • NeuronCore v4
    • Inf1 Architecture
    • Inf2 Architecture
    • Trn1 Architecture
    • Trn2 Architecture
    • Trn3 Architecture
  • What's new
    • Release Notes
      • Neuron 2.33.0
      • Previous Neuron versions
        • Neuron 2.32.0
        • Neuron 2.31.0
        • Neuron 2.30.0
        • Neuron 2.29.1
        • Neuron 2.29.0
        • Neuron 2.28.1
        • Neuron 2.28.0
        • Neuron 2.27.1
        • Neuron 2.27.0
        • Neuron 2.26.1
        • Neuron 2.26.0
        • Neuron 2.25.0
      • Component release notes
        • TorchNeuron
        • TorchTitan Neuron
        • PyTorch NeuronX
        • JAX NeuronX
        • vLLM Neuron
        • vLLM Omni Neuron
        • Neuron Kernel Interface
        • NKI Library
        • Neuron Compiler
        • Neuron Runtime
        • Neuron Developer Tools
        • Neuron Observer
        • Neuron Agentic Development
        • Neuron DLAMI
        • Neuron Containers
        • Older components and features
      • Release artifacts
  • Announcements
  • News & blogs
  • Contribute

Get started

  • Quickstarts
  • Set up environments
    • PyTorch
      • New DLAMI
      • Update Existing DLAMI
      • New DLC
      • Update Existing DLC
      • New Manual Configuration
      • Update Manual Configuration
    • JAX
      • Install JAX via Deep Learning AMI
      • Install JAX via Deep Learning Container
      • Install JAX Manually
    • Multi-framework
    • Troubleshooting
  • Deploy model workloads
    • Deployment Overview
    • Pre-configured Environments
      • Deep Learning AMIs
      • DLC Images
      • Customize DLC
      • Docker Setup
      • Quickstart: Deploy DLC with vLLM
    • Use Amazon EC2
      • Train on EC2
      • Inference (Inf2, Trn1)
      • Inference with DLC
      • Inference (Inf1)
    • Use Amazon EKS
      • Set up EKS for Neuron
      • Components Included
      • Inference on EKS
      • Training on EKS
      • Dynamic Resource Allocation
        • Support Files
      • Device Taints
      • Neuron UltraServer Operator
      • UltraServer MPI Jobs
      • EKS Prerequisites
    • Use Amazon ECS
      • Inference on ECS
      • Training on ECS
      • Node Problem Detector
    • Use AWS Batch
      • Train on AWS Batch
    • Use AWS ParallelCluster
      • Train on ParallelCluster
    • Use Amazon SageMaker
      • Train on SageMaker
    • Infrastructure Components
      • Plugins Overview
      • Scheduler Extension
      • Scheduler Flow
      • Neuron Monitor
      • Node Problem Detector
      • NPD Permissions (IRSA)
    • Tutorials
      • Docker Environment Setup
      • Build and Run Containers
      • OCI Hook Setup
      • Inference Tutorials
        • Run Inference in PyTorch Neuron Container
        • Deploy a TensorFlow Resnet50 model as a Kubernetes service
      • Training Tutorials
        • Run Training in PyTorch Neuron Container
        • Deploy a simple mlp training script as a Kubernetes job
    • Docker Examples
      • Inference DLC Dockerfile
      • Libmode Dockerfile
      • TF Serving Dockerfile
      • TorchServe Config
      • Libmode Entrypoint
      • TorchServe Neuron
      • Training DLC Dockerfile
      • MLP Training Example
      • v1 App+RT Different Containers
      • v1 App+RT Same Container
      • v1 Neuron Runtime
      • v1 Torch Neuron
      • v1 Entrypoint
    • Third-party Solutions
    • FAQ
    • Troubleshooting
    • Container Release Notes
    • DLAMI Release Notes

PyTorch support

  • TorchNeuron
    • About TorchNeuron
    • Get started
      • Install and verify torch-neuronx
      • Build your first model
      • Run an NKI kernel from torch-neuronx
    • Tutorials
      • Simple training loop on Neuron
      • Distributed training with torch.distributed
      • Hugging Face inference with torch.compile
    • Guides
      • Work with NKI kernels from torch-neuronx
      • Use torch.distributed on Neuron
      • Profile memory usage in torch-neuronx
      • torch.compile on Neuron
        • How torch.compile works
        • Optimize a model
        • Backend reference
    • Reference
      • API reference
        • Operator compatibility table
      • Environment variables
      • Best practices
      • Guidelines for memory management
      • TorchNeuron resources
    • Troubleshooting

Training

  • Training on Neuron
    • Training paths (native and legacy XLA)
    • Best practices for training
  • TorchTitan Neuron
    • About TorchTitan Neuron
    • Get started
      • Set up TorchTitan Neuron
    • Tutorials
      • Basic training
      • SimpleFSDP training with Llama and DeepSeek
      • GPT-OSS training
    • Best practices
    • Troubleshooting

Model inference

  • Inference on Neuron
  • Native PyTorch inference
  • vLLM Neuron
    • Get started
      • Setting up
      • Online serving quickstart
      • Offline serving quickstart
      • Migrating from NxDI
    • Deploy & serve
      • Features guide
      • Configuration reference
      • Profiling workloads
      • Compilation route
      • Precompile cache
      • Feature–model compatibility
    • Model Recipes
      • Llama 3
      • GPT-OSS
      • Qwen3-32B
      • Qwen3-VL
      • Qwen3-Embedding
      • Qwen3-ASR
    • Tutorials
      • Deploying Llama 3.3 70B
      • EAGLE3 speculative decoding (Llama 3.1)
      • EAGLE3 speculative decoding (GPT-OSS)
      • DFlash speculative decoding
      • Disaggregated inference (1P1D and xPyD)
      • Deploying gpt-oss
      • High-performance GPT-OSS 120B 8K/1K deployment
      • Benchmarking prefix caching (GPT-OSS)
      • Deploying dense Qwen3-32B
      • Deploying Qwen3-VL-32B
      • Disaggregated encoder (1E1PD and xEyPD)
      • Deploying Qwen3-Embedding-8B
      • Deploying Qwen3-ASR-1.7B
    • Model development
      • Onboarding a model
      • Onboarding a vision-language model
      • Optimizing a vision-language model
      • CPU development workflow
      • NKI CPU simulator
      • Debugging model code
      • Debugging accuracy issues
      • Accuracy debugger tools
    • Concepts & architecture
      • Parallelism
        • Attention DP
        • Component-Level DP Sharding
        • Data Parallelism (DP)
        • DCP (Decode Context Parallelism)
        • Encoder Parallelism
        • Expert Parallelism (EP)
        • Tensor Parallelism (TP)
      • Speculation
        • Speculative Decoding (EAGLE3) in the vLLM Neuron Framework
        • DFlash Speculative Decoding in the vLLM Neuron Framework
      • Multimodal
        • Block Packing Vision Attention
        • On-Device Encoder Cache
        • M-RoPE (Multimodal Rotary Position Embedding)
        • Encoder Cache Metrics
      • vLLM integration
        • vLLM Integration Design Reference
        • KV Caching with vLLM
        • Additional Configuration Options
        • Async Scheduling and Async Execution
        • Decode Context-Length Bucketing
        • determine_available_memory design
        • KV Cache Quantization
        • vLLM Metrics Integration
        • Runtime NRT Profiling via vLLM’s Profiler Interface
        • Scheduler Design
        • Padding and Batching with Block KV
        • Prefill Processing and Prefix Caching
        • Pooling models on Trainium
        • Prompt embeddings on Neuron
        • Structured Outputs and Tool Calling
        • Disaggregated Inference — Design Document
        • GPT-OSS-120B giga-kernel decode
      • Framework
        • Asynchronous Execution Framework Capability
        • Checkpoint Loading Framework Capability
        • Scalable Model Registry Design
        • Per-Graph Compiler Arguments
      • Compilation
        • How vLLM is integrated with TorchNeuron
        • Compilation Cache
        • CPU Compilation
        • FX Passes Architecture Design
        • Aliasing Output Rewrite Pass
        • FX Device Rewriting for XLA Compilation
        • In-Place to Out-of-Place Rewrite Pass
        • SPMD Compilation
        • Compile-Artifact Manifest
      • Accuracy
        • Accuracy Debugging Design
        • Dataset Evaluation Design
        • KV Cache Analysis Design
        • vLLM Neuron Logit Validation Design
        • Module-Level Accuracy Test Guidelines
        • Tensor Capture Design
        • Tensor Compare Design
        • Tensor Replacement — Design
        • Input Snapshot Design
  • vLLM Omni Neuron
    • Get started
      • Setup guide
      • Online serving quickstart
      • Offline serving quickstart
    • Guides
      • Features guide
    • Model recipes
      • WAN 2.2 T2V
      • WAN 2.2 I2V
    • Tutorials
      • Deploy Wan2.2-A14B
    • Model development
      • Onboarding a model
      • Accuracy evaluation and debugging
      • Kernel implementations
        • Const-max ring attention
        • Fused adaptive LayerNorm and FP8 quantization
        • QKV projection with QK Distributed RMSNorm and RoPE fusion
      • Optimizing offline video generation
    • Concepts & architecture
      • vLLM Omni Neuron overview
      • Context parallelism

Other frameworks (reference)

  • JAX
    • JAX NeuronX Known Issues
    • API Reference Guide for JAX Neuronx
      • JAX NeuronX Environment Variables
    • Release Notes
  • Legacy inference frameworks
    • NxD Inference
      • Overview
      • Setup
      • vLLM
        • Quickstart: Offline Model Serving
        • Quickstart: Online Model Serving
        • vLLM on Neuron User Guide
        • Model Recipes
        • Deploy Llama4 with vLLM
      • Tutorials
        • Disaggregated Inference (1P1D)
        • Disaggregated Inference
        • Flux Inference
        • Flux Inpainting
        • Benchmark using Performance CLI
        • GPT-OSS 120B
        • Llama3.1 405B on Trn2
        • Llama3.1 405B with Speculative Decoding
        • Llama3.1 70B Instruct Accuracy Evaluation
        • Llama3.1 8B with Multi-LoRA
        • Llama3.3 70B FP8
        • Llama3.3 70B with APC
        • Llama3.3 70B with Data Parallelism
        • Llama3.3 70B with Speculative Decoding
        • Llama4
        • Llama4 Legacy
        • Pixtral
        • Qwen3 MoE 235B
        • Qwen2 VL 7B
        • Speculative Decoding
        • Qwen3-VL 8B
      • Developer Guides
        • Accuracy Evaluation
        • Custom Quantization
        • Disaggregated Inference
        • Feature Guide
        • Using Pipeline Execution Mode
        • LLM Benchmarking
        • Migrate from TNX
        • Model Reference
        • MoE Architecture
        • Migrate from NxD Core
        • Onboarding Models
        • Performance Benchmarking CLI
        • vLLM Guide (Legacy)
        • vLLM Guide v1
        • Weights Sharding
        • Writing Tests
      • API Reference Guide
        • NxD Inference API Reference
      • App Notes
        • Parallelism Techniques for LLM Inference
      • Release Notes
      • Misc
        • This page has moved
        • Troubleshooting Guide for NxD Inference
    • NxD Core (inference)
      • Setup
      • App Notes
        • Tensor Parallelism Overview
        • Pipeline Parallelism Overview
        • Activation Memory Reduction
        • Context Parallelism Overview
      • API Reference Guide
        • Distributed Strategies APIs
        • Training APIs
        • Inference APIs
        • ModelBuilderV2 API Reference
      • Developer Guide
        • About NeuronX-Distributed (NxD) Inference
      • LoRA Guide
      • Tutorials
        • Training Tutorials
        • Inference Tutorials
      • Misc
        • This page has moved
  • Legacy training frameworks

Use custom kernels

  • Neuron Kernel Interface (NKI)
    • Get started
      • Environment Setup
      • First NKI Kernel
      • NKI Language Guide
      • Concepts
        • Memory Hierarchy
        • Data Representation
        • Indexing
        • Tiling
        • Direct Memory Access
        • Logical Neuron Cores
    • Guides
      • Tutorials
        • Matrix Multiplication
        • AveragePool2D
        • Transpose2D
        • Fused Mamba
        • Introduction to NKI Kernel Optimization
      • Architecture
        • Trainium/Inferentia2 Guide
        • Trainium2 Guide
        • Trainium3 Guide
      • NKI CPU Simulator
      • Insert NKI Kernels into Models
      • Use NKI Scheduling APIs
      • Profile a NKI Kernel
    • Deep dives
      • Performance Optimizations
      • MXFP8/4 Matrix Multiplication
      • NKI Access Patterns
      • NKI Loops
      • Descriptor Generation Engine (DGE)
      • DMA Bandwidth Guide
      • About the NKI Compiler
    • Migration
    • NKI FAQ
  • NKI API reference
    • nki
      • nki.jit
      • nki.simulate
    • nki.isa
      • nki.isa.nc_matmul
      • nki.isa.nc_matmul_mx
      • nki.isa.nc_transpose
      • nki.isa.activation
      • nki.isa.activation_reduce
      • nki.isa.activate2
      • nki.isa.tensor_reduce
      • nki.isa.tensor_partition_reduce
      • nki.isa.tensor_tensor
      • nki.isa.tensor_tensor_scan
      • nki.isa.scalar_tensor_tensor
      • nki.isa.tensor_scalar
      • nki.isa.tensor_scalar_reduce
      • nki.isa.tensor_scalar_cumulative
      • nki.isa.tensor_copy
      • nki.isa.tensor_copy_predicated
      • nki.isa.exponential
      • nki.isa.reciprocal
      • nki.isa.quantize_mx
      • nki.isa.iota
      • nki.isa.dropout
      • nki.isa.affine_select
      • nki.isa.range_select
      • nki.isa.select_reduce
      • nki.isa.sequence_bounds
      • nki.isa.memset
      • nki.isa.bn_stats
      • nki.isa.bn_aggr
      • nki.isa.local_gather
      • nki.isa.nc_n_gather
      • nki.isa.dma_copy
      • nki.isa.dma_transpose
      • nki.isa.dma_compute
      • nki.isa.max8
      • nki.isa.nc_find_index8
      • nki.isa.nc_match_replace8
      • nki.isa.nc_stream_shuffle
      • nki.isa.register_alloc
      • nki.isa.register_load
      • nki.isa.register_move
      • nki.isa.register_store
      • nki.isa.core_barrier
      • nki.isa.sendrecv
      • nki.isa.rng
      • nki.isa.rand2
      • nki.isa.rand_set_state
      • nki.isa.rand_get_state
      • nki.isa.set_rng_seed
      • nki.isa.nonzero_with_count
      • nki.isa.topk
      • nki.isa.engine
      • nki.isa.dma_engine
      • nki.isa.reduce_cmd
      • nki.isa.dge_mode
      • nki.isa.oob_mode
      • nki.isa.matmul_perf_mode
      • nki.isa.nc_version
      • nki.isa.get_nc_version
      • nki.isa.VirtualRegister
      • nki.isa.NkiInstruction
    • nki.language
      • nki.language.ndarray
      • nki.language.zeros
      • nki.language.ones
      • nki.language.full
      • nki.language.zeros_like
      • nki.language.empty_like
      • nki.language.shared_constant
      • nki.language.shared_identity_matrix
      • nki.language.rand
      • nki.language.random_seed
      • nki.language.NkiTensor
      • nki.language.load
      • nki.language.load_transpose2d
      • nki.language.store
      • nki.language.copy
      • nki.language.matmul
      • nki.language.transpose
      • nki.language.abs
      • nki.language.abs_max
      • nki.language.abs_min
      • nki.language.add
      • nki.language.arctan
      • nki.language.ceil
      • nki.language.cos
      • nki.language.exp
      • nki.language.floor
      • nki.language.log
      • nki.language.maximum
      • nki.language.minimum
      • nki.language.multiply
      • nki.language.negative
      • nki.language.power
      • nki.language.reciprocal
      • nki.language.rsqrt
      • nki.language.sign
      • nki.language.sin
      • nki.language.sqrt
      • nki.language.square
      • nki.language.subtract
      • nki.language.tan
      • nki.language.tanh
      • nki.language.trunc
      • nki.language.relu
      • nki.language.sigmoid
      • nki.language.silu
      • nki.language.silu_dx
      • nki.language.gelu
      • nki.language.gelu_dx
      • nki.language.gelu_apprx_sigmoid
      • nki.language.gelu_apprx_sigmoid_dx
      • nki.language.gelu_apprx_tanh
      • nki.language.mish
      • nki.language.softplus
      • nki.language.softmax
      • nki.language.erf
      • nki.language.erf_dx
      • nki.language.dropout
      • nki.language.rms_norm
      • nki.language.all
      • nki.language.max
      • nki.language.mean
      • nki.language.min
      • nki.language.prod
      • nki.language.sum
      • nki.language.var
      • nki.language.equal
      • nki.language.not_equal
      • nki.language.less
      • nki.language.less_equal
      • nki.language.greater
      • nki.language.greater_equal
      • nki.language.logical_and
      • nki.language.logical_or
      • nki.language.logical_xor
      • nki.language.logical_not
      • nki.language.bitwise_and
      • nki.language.bitwise_or
      • nki.language.bitwise_xor
      • nki.language.invert
      • nki.language.left_shift
      • nki.language.right_shift
      • nki.language.broadcast_to
      • nki.language.ds
      • nki.language.expand_dims
      • nki.language.where
      • nki.language.gather_flattened
      • nki.language.affine_range
      • nki.language.dynamic_range
      • nki.language.sequential_range
      • nki.language.static_range
      • nki.language.fori_loop
      • nki.language.while_loop
      • nki.language.psum
      • nki.language.sbuf
      • nki.language.hbm
      • nki.language.private_hbm
      • nki.language.shared_hbm
      • nki.language.is_psum
      • nki.language.is_sbuf
      • nki.language.is_hbm
      • nki.language.is_on_chip
      • nki.language.device_print
      • nki.language.no_reorder
      • nki.language.program_id
      • nki.language.num_programs
      • nki.language.program_ndim
      • nki.language.bool_
      • nki.language.int8
      • nki.language.int16
      • nki.language.int32
      • nki.language.uint8
      • nki.language.uint16
      • nki.language.uint32
      • nki.language.float16
      • nki.language.float32
      • nki.language.bfloat16
      • nki.language.tfloat32
      • nki.language.float8_e4m3
      • nki.language.float8_e5m2
      • nki.language.float8_e4m3fn
      • nki.language.float8_e8m0fnu
      • nki.language.float8_e5m2_x4
      • nki.language.float8_e4m3fn_x4
      • nki.language.float4_e2m1fn_x4
      • nki.language.tile_size
      • nki.language.prelu
      • nki.language.bypass
    • nki.collectives
      • nki.collectives.all_reduce
      • nki.collectives.all_gather
      • nki.collectives.all_gather_v
      • nki.collectives.reduce_scatter
      • nki.collectives.all_to_all
      • nki.collectives.all_to_all_v
      • nki.collectives.collective_permute
      • nki.collectives.collective_permute_implicit
      • nki.collectives.collective_permute_implicit_reduce
      • nki.collectives.collective_permute_implicit_current_processing_rank_id
      • nki.collectives.rank_id
      • nki.collectives.ReplicaGroup
    • NKI ISA API Common Fields
  • NKI Library
    • Get started
    • Overview
    • Kernel Design Specs
      • RMSNorm-Quant
    • Kernel API Reference
      • Argsort Unstable
      • Attention Block TKG
      • Attention Const-Max
      • Attention CTE
      • Attention KV-Parallel Segmented CTE
      • Attention MXFP8 TKG
      • Attention Segmented CTE
      • Attention TKG
      • Average Pooling 3D
      • Average Pooling 3D Backward
      • Batch Norm Backward
      • Blockwise MM Backward
      • Blockwise MM Backward MXFP8
      • Blockwise MM Block MX
      • Blockwise MM Forward MXFP8
      • Blockwise MM Shard-on-Block
      • Build All-to-All-V Metadata
      • Build All2All Dispatch Metadata
      • Build Packed Routing
      • Chunked Top-K
      • Collective Communication Kernels
      • Conv1D
      • Conv3D
      • Conv3D Backward
      • Conv3D Temporal Unroll
      • Conv3D Transpose
      • Cross Entropy
      • Cumsum
      • DeepSeek-V3 Router
      • DeepSeek-V3.2 Sparse Indexer
      • Depthwise Conv1D
      • Dynamic Elementwise Add
      • FGCC
      • Find Nonzero Indices
      • Fine-Grained All-Gather
      • Flash Attention 2
      • Foreach Elementwise
      • Foreach Norm
      • FP8 Quantize
      • Fused Adam
      • Gather
      • GDN Block TKG
      • GDN Conv1D
      • GDN CTE
      • GDN TKG
      • GPSIMD Top-K
      • GPT-OSS Tail Chunked Megakernel
      • KDA Block TKG
      • KDA TKG
      • Linear Scan
      • Matmul MXFP8
      • Matmul MXFP8 Backward
      • Max Pooling 3D
      • Max Pooling 3D Backward
      • MLA Attention Block Gather Decode
      • MLA Decode Attention
      • MLA Decode QKV
      • MLA QKV CTE
      • MLA Sparse Attention CTE
      • MLA V-Up O-Proj CTE
      • MLP
      • MLP Backward MXFP8
      • MLP DeepSeek MX
      • MLP Forward MXFP8
      • MoE CTE
      • MoE TKG
      • MS Deformable Attention
      • MS Deformable Attention Backward
      • MX MoE Block TKG Wrapper
      • Output Projection CTE
      • Output Projection TKG
      • Pad
      • Permute A2AV
      • Permute Routed Tokens
      • QKV
      • QKV Batch Shard
      • QKV CTE MLA
      • Quantize MXFP8
      • Ring Attention Backward
      • Ring Attention Forward
      • Ring Attention Unpermute
      • RMSNorm MX Prefill
      • RMSNorm Residual Prefill
      • RMSNorm Router Top-K A2AV
      • RMSNorm Router Top-K TKG
      • RMSNorm-Quant
      • RNG
      • RoPE
      • Router Top-K
      • SBUF-to-SBUF All-Gather
      • Scatter-Add
      • Selective Scan
      • Sparse Attention Indexer MX BF16Score
      • SSD
      • SWA Fused CTE
      • Top-K Reduce
      • Transformer TKG
      • Unpermute A2AV
    • Kernel Utilities
      • SbufManager (Allocator)
      • stream_shuffle_broadcast
    • Release Notes

Developer tools

  • Neuron tools
    • System Tools
      • Neuron-Monitor User Guide
      • Neuron-Top User Guide
      • Neuron-LS User Guide
      • Neuron-Sysfs User Guide
      • NCCOM-TEST User Guide
    • Third-party Tools
    • Tutorials
      • Profiling a vLLM Inference Workload on AWS Trainium
      • Track System Resource Utilization during Training with neuron-monitor using PyTorch Neuron
    • Release Notes
  • Neuron Agentic Development
    • Getting Started
    • Developer Guides
      • Autoport NxDI Deep Dive
      • Autoport vLLM Neuron Deep Dive
      • Equivalence Validation Deep Dive
      • Profiling vLLM Neuron Deep Dive
      • Profile Analysis vLLM Neuron Deep Dive
      • Parallelism Planner Deep Dive
      • Hardware Entitlement Deep Dive
    • Tutorials
      • Port a model with Autoport
      • Port a model to vLLM Neuron
      • Validate a port with Equivalence
      • Capture profiles with Profiling
      • Analyze a profile with Profile Analysis
      • Plan a serving config with the Parallelism Planner
      • Find the binding operator with Hardware Entitlement
  • Neuron Explorer
    • Get Started
    • Capture Profiles
    • UI Overview
    • Device Trace Viewer
    • System Trace Viewer
    • Hierarchy Viewer
    • Source Code Viewer
    • Summary Viewer
    • Tensor Viewer
    • Memory Viewer
    • Region Highlighter
    • AI Recommendation Viewer
    • Profile Parquet Schema Reference
    • Analyze Profile Output
    • View Profiles with Perfetto
    • Neuron Profiler to Neuron Explorer Migration Guide
    • Troubleshooting & FAQs
    • Environment Variables
    • Glossary
  • Neuron Observer
    • Get Started
    • Install and Configure
    • Configuration Reference
    • Metrics and Faults
    • Troubleshooting and FAQs

Compiler, runtime, & collectives

  • Graph compiler
    • About neuronx-cc
    • NeuronX Compiler for Trn* & Inf2
      • API Reference Guide
      • How-to: Generate a NEFF file
      • How-to: Convolution
        • * `UNet training sample <https://github.com/aws-neuron/aws-neuron-samples/tree/master/torch-neuronx/training/unet_image_segmentation>`_ - Sample UNet training implementation
      • FAQ
    • Error Codes
      • NCC_EARG001
      • NCC_EBIR023
      • NCC_EBVF030
      • NCC_EHCA005
      • NCC_EMOD025
      • NCC_EOOM001
      • NCC_EOOM002
      • NCC_ESFH002
      • NCC_ESPP004
      • NCC_ESPP047
      • NCC_EUOC002
      • NCC_EVRF001
      • NCC_EVRF004
      • NCC_EVRF005
      • NCC_EVRF006
      • NCC_EVRF007
      • NCC_EVRF009
      • NCC_EVRF010
      • NCC_EVRF011
      • NCC_EVRF013
      • NCC_EVRF015
      • NCC_EVRF016
      • NCC_EVRF017
      • NCC_EVRF018
      • NCC_EVRF019
      • NCC_EVRF022
      • NCC_EVRF031
      • NCC_EVRF035
      • NCC_EVRF036
      • NCC_EVRF037
      • NCC_EVRF038
      • NCC_EVRF039
      • NCC_EVRF040
      • NCC_EVRF041
      • NCC_EVRF042
      • NCC_EVRF043
      • NCC_EVRF044
      • NCC_EVRF045
      • NCC_EVRF046
      • NCC_EVRF047
      • NCC_EVRF048
      • NCC_EVRF049
      • NCC_EVRF050
      • NCC_EVRF051
      • NCC_EVRF052
      • NCC_EVRF053
      • NCC_EVRF054
      • NCC_EVRF055
      • NCC_EVRF056
      • NCC_EVRF057
      • NCC_EVRF058
      • NCC_EVRF059
      • NCC_EVRF060
      • NCC_EVRF061
      • NCC_EVRF062
      • NCC_EVRF063
      • NCC_EVRF064
      • NCC_EVRF065
      • NCC_EXSP001
      • NCC_EXTP004
    • Release Notes
    • neuron-cc (Legacy)
      • API Reference Guide
        • Neuron compiler CLI Reference Guide (neuron-cc)
      • CLI Reference
      • Developer Guide
      • FAQ
  • Neuron Runtime
    • Overview
      • Quickstart: Generating a core dump
    • Deep Dives
      • Compute-Communication Overlap
      • Neuron Device Memory
      • Direct HBM Tensor Allocation
      • Runtime Performance Tips
      • Neuron Runtime Core Dumps
      • Inter-node Collectives
      • Intra-node Collectives
    • How-To Guides
      • Developer guide
      • Migrate to the explicit async APIs
      • Configuration guide
    • API Reference
      • NRT API
      • NRT Status
      • NRT Version
      • Debug Stream APIs
      • NRT Async APIs Overview
      • NRT Async APIs Core Concepts
      • NRT Async APIs Best Practices
      • NRT Async APIs Examples
      • NRT Async
      • NRT Async Send/Recv
      • NRT Profile
      • NRT System Trace
      • Debug Stream
      • NEC API
      • NDL API
      • Neuron Driver Shared
      • Tensor Batch Operations
      • Neuron Datastore
      • NRT Experimental
      • Environment Variables
    • nrtpy (Neuron Runtime Python)
      • Getting started with nrtpy
      • Test and debug a NKI kernel end-to-end
      • Validate, benchmark, and trace a kernel
      • Troubleshooting
      • Model API
      • Tensor API
      • Configuration
      • Error handling
    • NRT Debug Stream
    • Compatibility
    • Troubleshooting on Inf1 and Trn1
    • Release Notes
    • FAQ
  • Collectives

Archive

  • Archived content
    • Fine-tune T5 model on Trn1
    • Running SSD300 with AWS Neuron
    • Megatron GPT Pretraining
    • Launch a GPT-3 pretraining job using neuronx-nemo-megatron
    • Training GPT-NeoX 20B with Tensor Parallelism and ZeRO-1 Optimizer
    • Fine-tuning Llama2 7B with tensor parallelism and ZeRO-1 optimizer using Neuron PyTorch-Lightning
    • Training Llama-2-7B/13B/70B using Tensor Parallelism and Pipeline Parallelism with Neuron PyTorch-Lightning
    • Training CodeGen2.5 7B with Tensor Parallelism and ZeRO-1 Optimizer
    • Profiling Multi-Node Training Jobs with Neuron Explorer
    • Training GPT-NeoX 6.9B with Tensor Parallelism and ZeRO-1 Optimizer
    • Neuron application notes
      • Neuron Runtime Library
      • Mixed precision (neuron-cc)
      • Mixed precision (neuronx-cc)
      • Introducing NxD Inference
      • Introducing NxD Training
      • Performance
      • Parallel execution
      • PyTorch for Neuron
        • Running inference on variable input shapes with bucketing
        • Running R-CNNs on Inf1
        • Data Parallel Inference on Torch Neuron
      • PyTorch for NeuronX
        • Introducing PyTorch 2.6 Support
        • Introducing PyTorch 2.7 Support
        • Introducing PyTorch 2.8 Support
        • Introducing PyTorch 2.9 Support
        • Introducing PyTorch 2.5 Support
        • Migration From XLA_USE_BF16/XLA_DOWNCAST_BF16
        • Data Parallel Inference on torch_neuronx
        • Graph Partitioner on torch_neuronx
    • Neuron Plugin for TensorBoard (Inf1)
    • TensorBoard plugin (archived)
      • NeuronX Plugin for TensorBoard (Trn1) (archived)
      • Neuron Plugin for TensorBoard (Inf1)
      • Profiling PyTorch NeuronX with TensorBoard (archived)
      • Track Training Progress in TensorBoard using PyTorch Neuron (archived)
    • NeuronPerf (Beta)
      • NeuronPerf Install
      • Overview
      • Terminology
      • Examples
      • Benchmark Guide
      • Evaluate Guide
      • Compile Guide
      • Model Index Guide
      • API
      • Framework Notes
      • FAQ
      • Troubleshooting
      • What’s New
        • Component Release Notes for Neuron Developer Tools
    • Neuron Profiler (archived)
      • Neuron Profiler User Guide (archived)
      • Neuron Profiler 2.0 (Beta) User Guide (archived)
    • Helper Tools
      • Check Model
      • GatherInfo
    • Transformers NeuronX (transformers-neuronx)
      • Setup
      • Developer Guide
        • Transformers NeuronX (transformers-neuronx) Developer Guide
        • Transformers NeuronX (transformers-neuronx) Developer Guide for Continuous Batching
      • Tutorials
        • Hugging Face meta-llama/Llama-2-13b autoregressive sampling on Inf2 & Trn1
        • Hugging Face facebook/opt-13b autoregressive sampling on Inf2 & Trn1
        • Hugging Face facebook/opt-30b autoregressive sampling on Inf2 & Trn1
        • Hugging Face facebook/opt-66b autoregressive sampling on Inf2
      • Misc
      • API Reference Guide
      • Generative LLM Inference with Neuron
      • Neuron Calculator
    • AWS Neuron Reference for NeMo Megatron
    • Apache MXNet Neuron — Archived
      • Apache MXNet Neuron Release Notes
      • MXNet Neuron Setup
      • Install MXNet Neuron
      • MXNet Neuron (“mxnet-neuron”) Setup on Amazon Linux 2
      • MXNet Neuron (“mxnet-neuron”) Setup on Amazon Linux 2
      • Update to latest MXNet Neuron
      • Install MXNet Neuron (Neuron 1.14.2)
      • Install MXNet Neuron (Neuron 1.15.0)
      • Install MXNet Neuron (Neuron 1.15.1)
      • Install MXNet Neuron (Neuron 1.15.2)
      • Install MXNet Neuron
      • Install MXNet Neuron
      • Install MXNet Neuron
      • Install MXNet Neuron
      • Inference (mxnet-neuron) (maintenance)
        • Tutorials
        • API Reference Guide
        • Developer Guide
        • Misc
      • Compile with Framework API and Deploy on EC2 Inf1 — Archived (MXNet)
      • Compile with SageMaker Neo and Deploy on SageMaker Hosting (Inf1) — Archived (MXNet)
      • MXNet Tutorial Setup
      • Tutorial: Neuron Apache MXNet Model Serving
    • TensorFlow Neuron
      • Tensorflow Neuron Setup
      • Get Started with TensorFlow Neuron (Archived)
      • Inference (Inf2 & Trn1)
        • Tutorials
        • API Reference Guide
        • Misc
      • Inference (Inf1)
        • Tutorials
        • Additional Examples
        • API Reference Guide
        • Misc
      • TensorFlow Neuron (Inf1)
        • TensorFlow 1.x (tensorflow-neuron) Compilation API
        • Install TensorFlow Neuron
        • Install previous TensorFlow Neuron releases
        • Install Previous Tensorflow Neuron Releases for Ubuntu (tensorflow-neuron)
        • Install Previous Tensorflow Neuron Releases for Ubuntu (tensorflow-neuron)
        • Install Previous Tensorflow Neuron Releases for Ubuntu (tensorflow-neuron)
        • Update to latest TensorFlow Neuron
        • Update to latest TensorFlow Neuron (tensorflow-neuron)
        • Update to latest TensorFlow Neuron (tensorflow-neuron)
        • Compile with Framework API and Deploy on EC2 Inf1 — Archived (TensorFlow)
        • Deploy Neuron Container on EC2 — Archived (TensorFlow)
        • Deploy Neuron Container on ECS — Archived (TensorFlow)
        • Deploy Neuron Container on EKS — Archived (TensorFlow)
        • Compile with SageMaker Neo and Deploy on SageMaker Hosting (Inf1) — Archived (TensorFlow)
        • TensorFlow 2.x FAQ
        • TensorFlow Tutorials
      • TensorFlow NeuronX (Inf2 & Trn1)
        • TensorFlow Setup Guide for Inf2 & Trn1
      • Setup (legacy Inf1)
    • PyTorch Neuron (torch-neuron) — Archived
      • API Reference Guide (torch-neuron)
        • PyTorch Neuron trace Python API
        • torch.neuron.DataParallel API
        • PyTorch Neuron (torch-neuron) Core Placement API
      • PyTorch-Neuron trace python API
      • PyTorch Neuron (torch-neuron) Core Placement API
      • torch.neuron.DataParallel API
      • Developer Guide (torch-neuron)
        • Running Inference on Variable Input Shapes with Bucketing
        • Data Parallel Inference on PyTorch Neuron
        • Developer Guide - PyTorch Neuron (torch-neuron) LSTM Support
        • PyTorch Neuron (torch-neuron) Core Placement
      • Troubleshooting Guide for PyTorch Neuron (torch-neuron)
      • Tutorials for Inference with torch-neuron (Inf1)
        • Computer Vision Tutorials
        • Natural Language Processing (NLP) Tutorials
        • Utilizing Neuron Capabilities Tutorials
      • Using NeuronCore Pipeline with PyTorch Tutorial
        • Using NeuronCore Pipeline with PyTorch
        • Using NeuronCore Pipeline with PyTorch
      • PyTorch Tutorial Setup
      • Setup Guide for Inf1
        • Fresh install
        • Update to latest release
        • Install previous releases
        • Install with support for cxx11 ABI
        • Install Previous PyTorch Neuron Releases for Amazon Linux (torch-neuron)
        • Update to latest PyTorch Neuron (torch-neuron)
        • Update to latest PyTorch Neuron (torch-neuron)
      • Install PyTorch Neuron (torch-neuron)
      • Update to latest PyTorch Neuron (torch-neuron)
      • Inference with torch-neuron (Inf1)
        • Tutorials
        • Additional Examples
        • API Reference Guide
        • Developer Guide
        • Misc
      • Additional Examples (torch-neuron)
        • AWS Neuron Samples GitHub Repository
      • Misc (torch-neuron)
        • PyTorch Neuron (torch-neuron) Supported operators
        • Troubleshooting Guide for PyTorch Neuron (torch-neuron)
        • Component Release Notes for Neuron PyTorch Framework
    • Bring Your Own Neuron Container to SageMaker Hosting (Inf1) — Archived
    • Bring Your Own Neuron Container to SageMaker Hosting (Inf2 or Trn1) — Archived
    • Deploy on Sagemaker Hosting
    • Compile with Sagemaker Neo and Deploy on Sagemaker Hosting (inf1)
    • Deploy Neuron Container on Kubernetes
    • Tutorial: Docker environment setup for Neuron Runtime 1.x
    • Archived setup guides
      • PyTorch Neuron setup for Rocky Linux 9 — Archived
      • Inf1 installation (legacy) — Archived
        • PyTorch on Inf1 (legacy) — Archived
      • Install previous PyTorch NeuronX releases for Ubuntu 20.04 — Archived
      • Install previous PyTorch NeuronX releases for Amazon Linux 2 — Archived
      • Install PyTorch NeuronX (Neuron 2.9.0) — Archived
      • Install PyTorch NeuronX (Neuron 2.8.0) — Archived
      • Install PyTorch NeuronX (Neuron 2.7.0) — Archived
    • NxD Training (archived)
      • Setup
      • YAML Configuration Settings
      • Neuronx Distributed Training Library Features
      • Developer Guide for Training with PyTorch NeuronX
      • How to debug models in PyTorch NeuronX
      • PyTorch NeuronX Environment Variables
      • PyTorch Neuron (torch-neuronx) for Training Troubleshooting Guide
      • Known Issues and Workarounds
      • CPU Mode Overview
      • Integrating a New Model
      • Integrating a new dataset/dataloader
      • Registering an optimizer and LR scheduler
      • NxD Training Compatibility with NeMo
      • Migrating from Neuron-NeMo-Megatron to Neuronx Distributed Training
      • Tensor Parallelism Overview
      • Pipeline Parallelism Overview
      • Context Parallelism Overview
      • Activation Memory Reduction
      • Analyze for Training Tutorial
      • Hugging Face BERT Pretraining Tutorial (Data-Parallel)
      • Multi-Layer Perceptron Training Tutorial
      • PyTorch Neuron for Trainium Hugging Face BERT MRPC task finetuning using Hugging Face Trainer API
      • ZeRO-1 Tutorial
      • Checkpoint Conversion
      • HuggingFace Llama3.1/Llama3-8B Pretraining
      • HuggingFace Llama3.1/Llama3-8B Supervised Fine-tuning
      • HuggingFace Llama3.1/Llama3-8B Efficient Supervised Fine-tuning with LoRA (Beta)
      • HF Llama3.1/Llama3-8B Direct Preference Optimization (DPO) and Odds Ratio Preference Optimization (ORPO) based Fine-tuning (Beta)
      • HuggingFace Llama3.1/Llama3-70B Pretraining
    • ML framework support on AWS Neuron SDK
    • PyTorch Support on Neuron
      • About PyTorch on Neuron
      • PyTorch Setup
      • Inference
        • Tutorials
        • Additional Examples
        • API Reference Guide
        • Developer Guide
        • Misc
      • torch-neuron v. torch-neuronx
      • Release Notes
      • Supported operators
      • Profiling API
      • Profiling developer guide
      • neuron_parallel_compile CLI
      • Multi-node EFA setup
    • Work with training and inference libraries
      • HF Transformers
        • Setup
        • Developer Guide
        • Tutorials
        • Misc
        • API Reference Guide
        • Generative LLM Inference with Neuron
        • Neuron Calculator
      • NxD Core Release Notes
    • NxD Core for Training
      • Setup
      • App Notes
        • Tensor Parallelism Overview
        • Pipeline Parallelism Overview
        • Activation Memory Reduction
        • Context Parallelism Overview
      • API Reference Guide
        • Distributed Strategies APIs
        • Training APIs
        • Inference APIs
        • ModelBuilderV2 API Reference
      • Developer Guide
        • Developer guide for Tensor Parallelism
        • Developer guide for Pipeline Parallelism
        • Developer guide for Activation Memory reduction
        • Developer guide for save/load checkpoint
        • Developer guide for Neuron-PT-Lightning
        • Developer guide for model and optimizer wrapper
        • Developer guide for LoRA finetuning
      • Tutorials
        • Training Tutorials
        • Inference Tutorials
      • Misc
        • This page has moved
    • Quickstart: Train a Model on Trainium (archived)
    • Quickstart: Run Inference on Inferentia (archived)
    • Ask Amazon AI helper tools (archived)
    • Model samples and tutorials (archived)
      • Training on Trn1
      • Inference on Inf2/Trn1/Trn2
      • Inference on Inf1
    • Neuron performance benchmarks (archived)
      • Inf2 Inference Performance (archived)
      • Trn1/Trn1n Inference Performance (archived)
      • Inf1 Inference Performance (archived)
      • Trn1/Trn1n Training Performance (archived)
    • Other Neuron FAQs (archived)
      • ONNX FAQ
      • Inference with Neuron FAQ
      • Troubleshooting for Inf1 FAQ
      • Training with Neuron FAQ
    • Neuron Custom C++ Operators (archived)
    • Neuron Custom C++ Operators [Beta] (archived)
      • API Reference Guide (archived)
        • Custom Operators API Reference Guide [Beta] (archived)
      • Developer Guide (archived)
        • Neuron Custom C++ Operators Developer Guide [Beta] (archived)
      • Tutorials (archived)
        • Neuron Custom C++ Operators in MLP Training (archived)
        • Neuron Custom C++ Operators Performance Optimization (archived)
      • Misc (Neuron Custom C++ Operators) (archived)
        • Neuron Custom C++ Tools Release Notes
        • Neuron Custom C++ Library Release Notes
  • Repository
  • Suggest edit
  • Open issue
  • .rst

API Reference Guide for JAX Neuronx

API Reference Guide for JAX Neuronx#

  • JAX NeuronX Environment Variables

previous

JAX NeuronX Known Issues

next

JAX NeuronX Environment Variables

By AWS

© Copyright 2026, Amazon.com.