AWS Neuron

AWS Neuron SDK, which supports Inferentia and Trainium, is natively integrated with PyTorch and TensorFlow. This ensures that you can continue using your existing workflows in these popular frameworks and get started with only a few lines of code changes. For distributed model training, the Neuron SDK supports libraries, such as Megatron-LM and PyTorch Fully Sharded Data Parallel (FSDP).

The AWS Neuron SDK enables efficient programming and runtime access to the Trainium and Inferentia chips. It supports a wide range of data types, new rounding modes, control flow, and custom operators to help you choose the optimal configuration for your deep learning workloads. For distributed training, Neuron enables efficient use of Trn1 UltraClusters with tightly coupled support for collective compute operations over Elastic Fabric Adapter (EFA) networking.

Neuron offers just-in-time (JIT) compilation to speed up developer workflows. It offers debugging and profiling tools with the support of the TensorBoard plugin. Neuron supports eager debug mode, which you can use to easily step through the code and evaluate operators one by one. You can use the Neuron helper tools to help you follow best practices for model onboarding and performance optimizations. In addition, Neuron includes tools that provide performance and utilization insights.

AWS Deep Learning AMIs (DLAMI) and AWS Deep Learning Containers come preconfigured with Neuron. If you are using containerized applications, you can deploy Neuron by using Amazon Elastic Container Service (ECS), Amazon Elastic Kubernetes Service (EKS), or your preferred native container engine. Neuron also supports Amazon SageMaker, which data scientists and developers can use to build, train, and deploy ML models.

To increase overall performance, Neuron automatically optimizes neural-net compute to run compute-intensive tasks on Trainium and Inferentia accelerators and other tasks on the CPU.

AWS Neuron supports FP32, TF32, BF16, FP16, and the new configurable FP8. Using the right data types for your workloads helps you optimize performance while meeting accuray goals. 

AWS Neuron takes high precision FP32 trained models and autocasts them to BF16 for high throughput inference at the lower cost and higher speed of 16-bit data type.

Neuron enables hardware-accelerated stochastic rounding. Stochastic rounding enables training at BF16 speeds, with near-FP32 accuracy when autocasting from FP32 to BF16.

NeuronCore Pipeline enables high-throughput model parallelism for latency-sensitive applications, such as natural language processing. The pipeline does this by sharding a compute graph across multiple NeuronCores, caching the model parameters in each core’s on-chip memory, and then streaming training and inference workloads across the cores in a pipelined manner.

AWS Neuron supports various collective communication operations to enable distributed training. Neuron runs these collective communication operators using dedicated hardware. That allows Neuron to run compute and communication in parallel to deliver lower latency and higher overall performance.

AWS Neuron supports custom operators. You can write new custom operators in C++ and Neuron will run those on Trainium inline SIMD cores.

AWS Neuron supports eager debug mode, which you can use to easily step through the code and evaluate operators one by one.

Trainium is an ML training chip that AWS purpose built for high-performance, low-cost deep learning training. Each AWS Trainium chip has two second-generation NeuronCores and supports FP32, TF32, BF16, FP16, and INT8 data types and also configurable FP8 (cFP8), which allows you to achieve the right balance between range and precision. To support efficient data and model parallelism, each Trainium accelerator has 32 GB of high-bandwidth memory, delivers up to 210 TFLOPS of FP16/BF16 compute power, and features NeuronLink, an intra-instance, ultra-high-speed nonblocking interconnect technology.

Learn more »

Inferentia is an ML inference chip that AWS designed and built to deliver high performance at low cost. Each AWS Inferentia chip has four first-generation NeuronCores and supports FP16, BF16, and INT8 data types. AWS Inferentia chips feature a large amount of on-chip memory, which can be used for caching large models. The on-chip memory is especially beneficial for models that require frequent memory access.

Learn more »

Amazon EC2 Trn1 instances, powered by AWS Trainium chips, are purpose built for high-performance ML training. They offer up to 50% cost-to-train savings over equivalent GPU-based EC2 instances. EC2 Trn1 instances feature up to 16 AWS Trainium accelerators and are the first EC2 instances with up to 800 Gbps of Elastic Fabric Adapter (EFA) network bandwidth.

Learn more »

Amazon EC2 Inf1 instances, based on AWS Inferentia chips, deliver up to 2.3x higher throughput and up to 70% lower cost per inference than comparable current generation GPU-based Amazon EC2 instances. Inf1 instances feature up to 16 AWS Inferentia chips, the latest custom second-generation Intel Xeon Scalable processors, and up to 100 Gbps networking to enable high throughput inference.

Learn more »