什麼是 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 架構提供本機分析功能。