AWS Neuron

在 AWS Trainium 和 AWS Inferentia 上最佳化 AI 和深度學習的 SDK

什麼是 AWS Neuron?

AWS Neuron 是用於在由 AWS Inferentia 和 AWS Trainium 提供支援的 Amazon Elastic Compute Cloud (Amazon EC2) 執行個體上執行深度學習和生成式 AI 工作負載的軟體開發套件 (SDK)。它包括編譯器、執行階段、訓練和推論程式庫,以及用於監控、分析和偵錯的開發人員工具。Neuron 支援您的端對端機器學習 (ML) 開發生命週期,包括建置和部署深度學習與 AI 模型,最佳化以達到最高效能和最低成本,以及深入了解模型行為。

背景圖案

與熱門的 ML 架構和程式庫原生整合

Neuron 與 PyTorch 和 JAX 以及重要的 ML 程式庫 (例如 Hugging Face Optimum Neuron、PyTorch Lightning 和 AXLearn) 原生整合。Neuron 還支援 OpenXLA,包括 StableHLO 和 GSPMD,助力 PyTorch、XLA 和 JAX 開發人員利用 Neuron 針對 Inferentia 和 Trainium 的編譯器最佳化。Neuron 可讓您將 Trainium 和 Inferentia 型執行個體與各種服務 (例如 Amazon SageMaker、Amazon EKS、Amazon ECS、AWS ParallelCluster 和 AWS Batch) 以及 Ray (Anyscale)、Domino Data Lab、Datadog 和 Weights & Biases 等第三方服務搭配使用。

背景圖案

分散式訓練和推論程式庫

Neuron 使用開放原始碼 PyTorch 程式庫 NxD Training 和 NxD Inference 對分散式訓練和推論進行立即可用的最佳化。NxD Training 簡化並最佳化了大規模分散式訓練,並支援各種模型架構、並行策略和訓練工作流程。NxD Inference 為最佳化模型推論提供全面的解決方案,其主要功能包括裝置上抽樣、QKV 權重融合、連續批次處理、推測解碼、動態儲存貯體和分散式推論等。NxD Inference 還與服務解決方案整合,例如 vLLM 和 Hugging Face TGI。它們均包括用於不同模型架構的模型中心。

背景圖案

進階的應用科學能力

Neuron 擁有多項應用科學能力,賦能科學家和研究人員在 Trainium 和 Inferentia 上突破開放原始碼 AI 研究和創新的界限。Neuron Kernel Interface (NKI) 可直接存取 Trainium 和 Inferentia 上提供的硬體基本類型和指示,使研究人員能夠建置和調整運算核心以獲得最佳效能。它是以 Python 為基礎的程式設計環境,採用常用的類 Triton 語法和圖磚層級語義。研究人員可以使用 NKI 以新功能、最佳化和科學創新來增強深度學習模型。Neuron 的自訂 C++ 運算子使開發人員能夠透過建立針對 Inferentia 和 Trainium 最佳化的運算子來擴展 SDK 的功能。

背景圖案

強大的開發人員工具

AWS Neuron SDK 提供了一套全面的工具集,可深入了解 AWS Inferentia 和 Trainium 所支援 EC2 執行個體上的深度學習模型的監控、管理和最佳化。它提供 neuron-top、neuron-monitor 和 Neuron Sysfs 等公用程式來監控硬體資源、模型執行和系統詳細資訊。對於 Kubernetes 和 EKS 上的容器化應用程式,Neuron 透過 Amazon CloudWatch 整合和其他熱門的可觀測性工具 (如 Data Dog 和 Weights & Biases) 簡化了監控。此外,neuron-profile 工具有助於識別和解決單節點和分散式應用程式中的效能瓶頸,並為熱門的 ML 架構提供本機分析功能。

背景圖案

入門

Neuron 深度學習 Amazon Machine Images (Neuron DLAMIs) 已預先設定有 Neuron SDK、流行架構和有用的程式庫,讓您可以快速開始在 AWS Inferentia 上進行訓練和執行推論。Neuron DLAMIs 簡化您的工作流程並最佳化效能,同時消除設定複雜性,因此您可以專注於建置和部署 AI 模型。開始使用 Neuron DLAMIs

使用預先設定的 AWS Neuron 深度學習容器 (Neuron DLC) 以及針對 Trainium 和 Inferentia 最佳化的架構快速部署模型。對於自訂解決方案,可建置自己的容器並利用 Kubernetes 功能,例如 Neuron 裝置外掛程式、Neuron 排程器擴充和 Helm Chart。與 Amazon EKS、AWS Batch 和 Amazon ECS 等 AWS 服務無縫整合,以實現可擴展的部署。開始使用 Neuron DLC

Optimum Neuron 將 Hugging Face 轉換器與 AWS Neuron SDK 連接起來,為 Trainium 和 Inferentia 提供標準 Hugging Face API。它提供用於訓練和推論的解決方案,包括支援大規模模型訓練和 AI 工作流程部署。Optimum Neuron 支援 Amazon SageMaker 和預先建置的深度學習容器,從而簡化適用於 ML 的 Trainium 和 Inferentia 的使用。此整合可讓開發人員使用熟悉的 Hugging Face 介面,同時利用 Trainium 和 Inferentia 進行其以轉換器為基礎的專案。開始使用 Hugging Face Optimum Neuron

您可以使用 Amazon SageMaker JumpStart,藉助 Neuron 訓練和部署模型。JumpStart 提供對微調和部署流行模型的支援,例如 Meta 的 Llama 模型系列。開始使用 SageMaker JumpStart