跳至主要內容

AWS AI 晶片

Build on Trainium

一項 1.1 億 USD 的投資計劃，透過 AWS Trainium 加速 AI 研究和教育

與我們合作
提交研究建議

什麼是 Build on Trainium？

Build on Trainium 是一項 1.1 億 USD 的投資計畫，著重在 AI 研究和大學教育，以支援下一代 AWS Trainium 的創新和開發。AWS Trainium 是一種 AI 脈動陣列晶片，專為推進最先進的 AI 構想和應用程式而設計。Build on Trainium 資助 Trainium 的新型 AI 研究，投資領先的學術團隊，在關鍵領域建立創新，包括新模型架構、ML 程式庫、最佳化、大規模分散式系統等。這項多年的計劃透過激勵學術界利用、投資並為圍繞在 Trainium 的開放原始碼社群做出貢獻，為 AI 的未來奠定基礎。將這些優勢與 Neuron 軟體開發套件 (SDK) 和最近推出的 Neuron Kernel Interface (NKI) 相結合，Trainium 客戶現在可以在雲端中大規模創新。
An abstract illustration depicting a stylized hand supporting a neural network or data sphere, represented by interconnected blue nodes and lines. The design features geometric shapes and a gradient blue and purple color palette, symbolizing artificial intelligence and machine learning concepts.

AWS Trainium 研究叢集

我們建立了一個專用的 Trainium 研究叢集，其中包含多達 40,000 個 Trainium 晶片，可透過使用 Amazon EC2 UltraCluster 在單一非阻攔式 PB 規模網路上連線的 Amazon EC2 Trn1 執行個體提供。研究團隊和學生可以透過使用適用於 ML 的 Amazon EC2 容量區塊的自我管理容量區塊存取這些晶片。
An abstract close-up image featuring a vibrant pattern of red and violet hues, resembling a geometric or honeycomb texture, with bright highlights and a modern, digital art aesthetic.

Amazon 研究獎

我們正在向廣泛的研究社群進行多輪的 Amazon 研究獎 (ARA) 徵求提案 (CFP)，選取的提案獲得 AWS Trainium 積分，並可存取 Trainium 研究叢集。Build on Trainium 歡迎研究提案，利用流行的開放原始碼 ML 程式庫和架構，並為開放原始碼做出貢獻，以增強 ML 開發人員社群的資源。
Illustration of two blue silhouettes drawing or collaborating on a complex digital workflow or process diagram against an orange background. Represents teamwork, planning, or designing digital infrastructure or compute processes.

Neuron Kernel Interface

Neuron Kernel Interface (NKI) 是適用於 AWS AI 晶片、Trainium 和 Inferentia 的新程式設計介面。NKI 可直接存取 AWS Trainium 和 Inferentia 上提供的硬體基本類型和指示，使研究人員能夠建置和調整運算核心以獲得最佳效能。它是以 Python 為基礎的程式設計環境，採用常用的類 Triton 語法和圖磚層級語義。研究人員可以使用 NKI 以新功能、最佳化和科學創新來增強深度學習模型。如需進一步了解，請參閱 NKI 文件頁面。
Abstract illustration featuring various geometric shapes, such as triangles, circles, rectangles, and grids, in gradient pastel yellow and purple tones on a lavender background. Arrows and dotted lines add motion and connectivity, suitable for themes of application integration and creative design.

優勢

存取專用 AWS Trainium 研究叢集，並使用世界級的 AI 硬體和可擴展雲端基礎設施，為您最雄心勃勃的研究專案提供支援。

建置創新且最佳化的運算核心，效能優於現有架構和技術，突破生成式 AI 研究和開放原始碼創新的界限。建置高度最佳化的核心，以最佳化模型中最關鍵或差異化的部分。

使用與 PyTorch 和 JAX 無縫整合的 Neuron SDK 輕鬆開始使用。Neuron Kernel Interface 以 Python 為基礎的程式設計環境採用常用的類 Triton 語法來協助您快速提升。

與 AWS 專家和更廣大的研究社群合作，以擴大您的工作實際影響。

參與大學

以下是領先的大學如何從 Build on Trainium 計畫中受益。

麻省理工學院

麻省理工學院的 Device Realization Lab 使用 AWS Trainium 來挑戰醫學 AI 研究的極限。我們使用比以往更快、更有效率的方式訓練 3D 超音波分段和聲速估算模型，從而在縮短一半以上實驗時間的同時，達成頂尖的準確度。AWS Trainium 讓我們得以透過傳統 GPU 系統無法達成的方式擴展研究。透過在 AWS Trainium (trn.32xlarge) 上訓練我們的 3D 完全卷積神經網路，我們得以實現最先進的效能，與 NVIDIA A100 執行個體相比，輸送量提高 50%，而且成本更為低廉。透過 32 節點的 Trainium 叢集，我們有效率地進行了 180 多項消融實驗，將總訓練時間從數個月縮短為數週，並加速了麻省理工學院的醫療 AI 創新。未來，我們計劃使用 Trainium 訓練 AI 代理程式模型，以操作和自動化數位超音波工作流程，從而節省大量臨床醫生時間，並為患者提供更好的照護。」


Missing alt text value

卡內基梅隆大學

「CMU Catalyst 研究小組致力於最佳化 ML 系統。我們的專案旨在讓使用者能夠於不同的 ML 系統中更輕鬆地進行最佳化。Trainium 能夠以獨一無二的方式，透過 Neuron Kernel Interface (NKI) 提供低層級控制和可存取的程式設計介面。

在 AWS 透過 Build on Trainium 計畫的支援下，我們的研究人員得以探索 FlashAttention 這項關鍵核心的進階最佳化作業。最令我們驚訝的是我們能夠進行迭代的速度：我們僅花費一周時間，就透過公開可用的 NKI、Neuron Profiler 和架構文件，在先前最先進的成果上實現了有意義的改善。強大工具和清晰硬體見解的結合，讓我們的團隊能夠輕鬆進行複雜的低層級最佳化。

AWS Trainium 和 Neuron Kernel Interface (NKI) 使我們這類研究人員能夠更快地進行創新，消除通常會減慢硬體特定最佳化工作的障礙。」

 

Missing alt text value

加州大學柏克萊分校

「透過 Build on Trainium 計畫，他的團隊獲得了 AWS Neuron 全新 NKI 開放原始碼編譯器堆疊的完整存取權限，包括針對 Trainium ISA 和 API 的直接可見性，以便進行精確的排程和記憶體配置。這種可見性和控制程度，讓他的學生得以更輕鬆地分析最佳化機會，並更有效地探索高效能的實作方式。」

加州大學柏克萊分校資訊系副教授 Christopher Fletcher
 

Missing alt text value

伊利諾大學厄巴納/香檳分校

「AWS Trainium 和 Inferentia 的存取權對於推動我們對大規模、高效 AI 系統的研究和教育至關重要。我們使用這些平台進行混合專家訓練和推論最佳化，並為新的執行時期和排程技術設計原型，以提高新興加速器架構的可擴展性、效率和可攜性。UIUC 研究人員正在利用 Neuron 開發人員堆疊開發新的執行時期和排程技術，以提高 AI 工作負載的效率和可攜性。團隊對 Neuron 開發人員堆疊的開放性留下了深刻印象，這項特性使這些平台在執行時期研究中格外有價值，並在稀疏度、記憶體階層和通訊效率方面實現了超越傳統 GPU 架構的創新。」

Missing alt text value

加州大學洛杉磯分校

「透過利用 AWS Trainium 和 Build on Trainium 計畫，我和我的學生得以大幅加速量子電路模擬作業。這項專案聚集了一群優秀的學生，他們合作構置了一個高效能模擬器，從而大規模完成先前不可能實現的更深入實驗和實踐學習。」

Missing alt text value

雪梨科技大學

「我們 UTS 的研究團隊正在透過開發自訂 Neuron NKI 核心來探索樹輪水印演算法的整合。透過 Build on Trainium 計畫存取開放原始碼 Neuron 堆疊，帶來了變革性的結果。這為我們提供前所未有的 Trainium 架構可見性，以及直接在硬體層級作業的能力。Trainium 的存取權讓我們的團隊能夠大幅加速我們的水印工作負載、減少迭代週期，並使我們得以探索更複雜的模型和技術。這種深入的存取權使我們的研究人員能夠為新想法建立原型、實驗低層級最佳化，並挑戰水印系統在現代 AI 加速器上的限制。」

Missing alt text value