什麼是 AWS Neuron？
AWS Neuron 是用於在 AWS Trainium 和 AWS Inferentia 上執行深度學習和生成式 AI 工作負載的開發人員堆疊。Neuron 建立在開放原始碼基礎上，可讓開發人員使用 PyTorch 和 JAX 架構以及 ML 程式庫 (例如 HuggingFace、vLLM、PyTorch Lightning 等) 在本機進行建置、部署和探索，而無需修改您的程式碼。 它包括編譯器、執行時期、訓練和推論程式庫，以及用於監控、分析和偵錯的開發人員工具。Neuron 支援您的端對端機器學習 (ML) 開發生命週期，從建置和部署深度學習與 AI 模型，最佳化以達到最高效能和最低成本，以及深入了解模型行為。
Neuron 可實現快速實驗、前沿模型的生產規模訓練、透過 Neuron Kernel Interface (NKI) 進行自訂核心的低層級效能最佳化、代理式 AI 和強化學習工作負載的成本最佳化推論部署，以及使用 Neuron Explorer 進行全面的分析和偵錯。
專為研究人員打造
Neuron 透過在 Trainium 上原樣執行原生 PyTorch 程式碼來實現快速的 AI 研究。研究人員可以藉助 PyTorch Eager 模式支援嘗試新想法並快速反覆運作。使用 PyTorch 分散式程式庫 (例如 FSDP，DDP 和 DTensor) 輕鬆實現跨晶片分片模型或擴展到多個節點的擴展。Neuron 支援 torch.compile，並且諸如 TorchTitan 和 HuggingFace Transformers 這樣的程式庫現在可以直接在 Trainium 上執行，而無需進行修改。此外，JAX 開發人員可以利用 Neuron 輕鬆開發、最佳化模型並將其部署到 Inferentia 和 Trainium。
專為提高生產力而打造
Neuron 針對代理式 AI 和強化學習工作負載最佳化推論經濟性。標準 vLLM V1 API 可在 Trainium 和 Inferentia 上執行，具有即用的高效能，並具備專家平行性、分解推論和推測性解碼等功能，以及來自 Neuron Kernel Library 的最佳化內核，從而最大限度地提高大規模字符經濟性。 ML 開發人員可以使用 HuggingFace Optimum Neuron、PyTorch Lightning 和 TorchTitan 進行訓練，然後使用標準 vLLM API 部署推論。
專為創新而打造
建置 AI 模型需要快速創新和效能最佳化。雖然像 PyTorch 這樣的標準架構使得擴展實驗變得輕鬆，但要突破效能極限，需要最佳化整個堆疊 (晶片、伺服器和 UltraServer)。Neuron 透過 Neuron Kernel Interface(NKI)、Neuron Explorer 提供的更深入洞察，以及稱為 Neuron Kernel Library (NKILib) 的最佳化內核程式庫，為 ML 效能工程師提供對 AWS AI 晶片的空前存取權。NKI 提供用於記憶體配置和執行排程的 API，並可直接存取 Trainium ISA，從而實現對指令層級程式設計的控制。NKI Compiler 是開放原始碼的工具，在 MLIR 的基礎上建置，可讓開發人員深入了解完整的編譯器管道。開放原始碼的 Neuron Kernel Library 提供最佳化的實作，並附帶原始程式碼、文件和基準測試。Neuron Explorer 提供一套統一的工具，可指引開發人員完成效能最佳化和偵錯過程。效能工程師可以從原始程式碼追蹤執行過程直至硬體操作，分析單節點和分散式應用程式，並取得 AI 驅動的洞察和可行的建議，以此進行內核最佳化和效能改進。
專為開放原始碼打造
AI 創新在開放社群中蓬勃發展，開發人員可以從中進行檢查、修改和做出貢獻。Neuron 致力於打造開放原始碼社群並促進創新。雖然我們正在將堆疊的更多部分移到開放原始碼環境，但 NKI Compiler、Neuron Kernel Driver、Neuron Kernel Library、NxD Inference、Neuron Explorer 以及 PyTorch、JAX 和 vLLM 整合目前已完全開放原始碼。開放原始碼程式庫和工具助力開發人員檢查編譯器實作，為最佳化做出貢獻以及沒有任何阻礙地調整內核程式碼。和我們一起建置。
認識 Neuron
Neuron 提供與 PyTorch 的原生整合，使研究人員和 ML 開發人員能夠在 Trainium 上執行現有程式碼而無需更改。包括 FSDP、DDP 和 DTensor 在內的標準 API 可順暢地用於跨多節點設定的分散式訓練。諸如 TorchTitan、HuggingFace Optimum Neuron、PyTorch Lightning 等熱門的 ML 程式庫，僅需進行少量修改即可直接執行。使用熟悉的工作流程和工具來訓練模型 (從預訓練到使用增強學習的訓練後)，同時利用 Trainium 的效能和成本優勢進行實驗和生產規模訓練。
Neuron 可利用 Trainium 和 Inferentia 上的標準架構和 API 進行生產推論部署。vLLM 與標準 API 的整合利用 Neuron Kernel Library 中最佳化的內核提供高效能服務。包括專家平行度、分解推論和推測性解碼在內的進階功能可最大限度地提高每秒字符數，同時最大限度地降低每個字符的成本。藉助即用的效能最佳化，大規模部署代理式 AI 和強化學習工作負載。
對於尋求最高硬體效率的效能工程師，Neuron 透過 Neuron Kernel Interface (NKI) 提供完全控制，並且直接存取 NeuronISA 指令集、記憶體配置和執行排程。開發人員可以建立標準架構中未提供的新操作，並使用自訂內核最佳化效能關鍵程式碼。在 MLIR 基礎上建置的開放原始碼 NKI Compiler 提供編譯程序的透明度。Neuron Kernel Library 提供生產就緒的最佳化內核，並且附帶完整的原始程式碼、文件和基準測試。
Neuron Explorer 提供一套統一的工具，可指引開發人員完成效能最佳化和偵錯過程。透過將分析、偵錯、實作最佳化以及驗證改進整合到單一環境中，Neuron Explorer 可以消除碎片化工具之間的時間損失。透過程式碼連結對 PyTorch、JAX 和 NKI 進行分層分析，可以追蹤從原始程式碼到硬體操作的執行過程。AI 支援的建議可分析設定檔以識別瓶頸，並為分區策略和核心最佳化提供可行的洞察。該 UI 在 GitHub 上開放原始碼。
Neuron 提供全面的監控和可觀測性功能，助力 ML 開發人員和 MLOps 團隊在生產部署中維持卓越的營運。原生 Amazon CloudWatch 整合可在整個 ML 基礎設施中進行集中監控，並且支援 Kubernetes 和 Amazon EKS 上的容器化應用程式。合作夥伴平台與 Datadog 等工具整合，通過統一的監控、日誌記錄和提供來擴展可觀測性。Neuron 提供各種公用程式，包括用於即時監控的 neuron-top、用於指標收集的 neuron Monitor、用於裝置清單的 neuron-ls 以及用於詳細系統資訊的 neuron Sysfs。
Neuron 透過預先設定的環境和基礎設施工具，簡化 ML 開發人員和 MLOPS 團隊的部署。Neuron 深度學習 AMI (DLAMI) 和深度學習容器 (DLC) 隨附 Neuron 軟體堆疊、熱門的架構和基本程式庫。對於 Kubernetes 部署，Neuron Device Plugin 管理資源配置，Neuron Scheduler Extension 提供智慧工作負載放置，Dynamic Resource Allocation (DRA) 驅動程式透過直觀的大小型資源選取來抽象化硬體拓撲的複雜性。Helm Charts 簡化容器化部署的協同運作。