README
AWS Neuron EKS Samples
This repository contains samples for Amazon Elastic Kubernetes Service (EKS) and AWS Neuron, the software development kit (SDK) that enables machine learning (ML) inference and training workloads on the AWS ML accelerator chips Inferentia and Trainium.
The samples in this repository demonstrate the types of patterns that can be used to deliver inference and distributed training on EKS using Inferentia and Trainium. The samples can be used as-is, or easily modified to support additional models and use cases.
Samples are organized by use case below:
Training
| Link | Description | Instance Type |
|---|---|---|
| BERT pretraining | End-end workflow for creating an EKS cluster with 2 trn1.32xl nodes and running BERT phase1 pretraining (64-worker DataParallel) | Trn1 |
| MLP training | Introductory workflow for creating an EKS cluster with 1 node and running a simple MLP training job | Trn1 |
| Llama 3.1 8B finetuning with Ray+PTL | End-end workflow for creating a Ray cluster with 2 trn1.32xlarge nodes on EKS and running Llama 3.1 8B finetuning | Trn1 |
Inference
| Link | Description | Instance Type |
|---|---|---|
| SD inference | SD Inference workflow for creating an inference endpoint forwarded by ALB LoadBalancer powered by Karpenter's NodePool | Inf2 |
| Flux inference | FLUX.1 dev Inference workflow for creating an inference endpoint forwarded by ALB LoadBalancer powered by Karpenter's NodePool and S3 mountpoints | Trn1/Inf2 |
| Optimal TP/DP for LLM serving | Demonstrates optimal tensor parallelism configuration for LLM serving with vLLM, comparing TP1, TP2, and TP4 performance on Qwen models | Trn2 |
| Speculative decoding | Accelerate LLM inference using speculative decoding with vLLM, comparing baseline vs draft model performance with Neuron DRA and S3 persistence | Trn2 |
| Disaggregated inference | Prefill/decode disaggregated serving with vLLM, using DRA for Neuron + EFA allocation and NIXL/LIBFABRIC KV transfer; dynamic xPyD scaling routed by vLLM production-stack or AIBrix | Trn2/Trn3 |
Getting Help
If you encounter issues with any of the samples in this repository, please open an issue via the GitHub Issues feature.
Contributing
Please refer to the CONTRIBUTING document for details on contributing additional samples to this repository.
Release Notes
Please refer to the Change Log.
Similar repos
Languages
Python
95.5%
Shell
2.7%
README
AWS Neuron EKS Samples
This repository contains samples for Amazon Elastic Kubernetes Service (EKS) and AWS Neuron, the software development kit (SDK) that enables machine learning (ML) inference and training workloads on the AWS ML accelerator chips Inferentia and Trainium.
The samples in this repository demonstrate the types of patterns that can be used to deliver inference and distributed training on EKS using Inferentia and Trainium. The samples can be used as-is, or easily modified to support additional models and use cases.
Samples are organized by use case below:
Training
| Link | Description | Instance Type |
|---|---|---|
| BERT pretraining | End-end workflow for creating an EKS cluster with 2 trn1.32xl nodes and running BERT phase1 pretraining (64-worker DataParallel) | Trn1 |
| MLP training | Introductory workflow for creating an EKS cluster with 1 node and running a simple MLP training job | Trn1 |
| Llama 3.1 8B finetuning with Ray+PTL | End-end workflow for creating a Ray cluster with 2 trn1.32xlarge nodes on EKS and running Llama 3.1 8B finetuning | Trn1 |
Inference
| Link | Description | Instance Type |
|---|---|---|
| SD inference | SD Inference workflow for creating an inference endpoint forwarded by ALB LoadBalancer powered by Karpenter's NodePool | Inf2 |
| Flux inference | FLUX.1 dev Inference workflow for creating an inference endpoint forwarded by ALB LoadBalancer powered by Karpenter's NodePool and S3 mountpoints | Trn1/Inf2 |
| Optimal TP/DP for LLM serving | Demonstrates optimal tensor parallelism configuration for LLM serving with vLLM, comparing TP1, TP2, and TP4 performance on Qwen models | Trn2 |
| Speculative decoding | Accelerate LLM inference using speculative decoding with vLLM, comparing baseline vs draft model performance with Neuron DRA and S3 persistence | Trn2 |
| Disaggregated inference | Prefill/decode disaggregated serving with vLLM, using DRA for Neuron + EFA allocation and NIXL/LIBFABRIC KV transfer; dynamic xPyD scaling routed by vLLM production-stack or AIBrix | Trn2/Trn3 |
Getting Help
If you encounter issues with any of the samples in this repository, please open an issue via the GitHub Issues feature.
Contributing
Please refer to the CONTRIBUTING document for details on contributing additional samples to this repository.
Release Notes
Please refer to the Change Log.
Similar repos
Languages
Python
95.5%
Shell
2.7%