OS 和 Docker 映像檔

Google Cloud 提供內含常見作業系統、架構、程式庫和驅動程式的映像檔。 Google Cloud 將這些預先設定的映像檔最佳化，支援人工智慧 (AI) 和機器學習 (ML) 工作負載。

本文概述用於部署、管理及執行 AI Hypercomputer 環境中工作負載的映像檔。

瞭解圖片類別

圖片會歸入下列類別：

• AI 和機器學習架構與程式庫：預先設定機器學習架構和程式庫二進位檔的 Docker 映像檔，可簡化機器學習模型的建立、訓練和使用程序。在 AI Hypercomputer 上，您可以使用深度學習軟體層 (DLSL) Docker 映像檔，在 Google Kubernetes Engine (GKE) 叢集執行 NeMO 和 MaxText 等機器學習模型。
• 叢集部署和自動化調度管理：您用來部署及管理作業系統 (OS) 映像檔，這些映像檔會用於執行 AI 工作負載的效能最佳化基礎架構。您應使用的映像檔取決於所用的 GPU 基礎架構 (叢集 GPU 或一般 GPU)，以及叢集的部署方式：
  • GKE 節點映像檔：您可以使用這些映像檔部署 GKE 叢集。
  • Slurm OS 映像檔：Cluster Toolkit 會建構及部署這些映像檔，並為 Slurm 節點安裝必要的系統軟體。
  • 加速器 OS 映像檔：您可以使用這些映像檔建立個別或群組的 Compute Engine 執行個體。

AI 和機器學習架構與程式庫

Google Cloud 提供 Docker 映像檔，其中包含熱門的 AI 和 ML 架構與程式庫。這些映像檔提供所需的軟體，可簡化在 AI Hypercomputer 上執行的 AI 最佳化叢集模型開發、訓練和部署作業。

JAX AI 圖像

適用於Google Cloud TPU 和 GPU 的 JAX AI 映像檔 (JAII，舊稱 JAX Stable Stack 映像檔) 提供可立即使用的 Docker 映像檔，其中包含 JAX 架構、精選的相容程式庫，以及 Google Cloud 基礎架構的設定。JAX AI TPU 映像檔已預先設定 JAX 程式庫和 TPU 程式庫。JAX AI GPU 映像檔已預先設定 JAX 程式庫和相關的 CUDA/NVIDIA 程式庫。

硬體層

JAX AI 圖片位於硬體層之上，硬體層由加速器 (TPU 或 GPU) 和相關聯的 VM 組成。如要使用 JAX AI 圖片，請佈建 TPU 或 GPU VM。您可以使用 TPU API、Compute Engine API 或 GKE API 執行這項操作。

架構層

架構層提供建構機器學習工作負載的工具和程式庫。JAX AI 映像檔為以 JAX 為基礎的機器學習工作負載提供預先設定的基礎，包括核心 JAX 程式庫和其他必要依附元件，確保開發體驗一致且高效能。

JAX AI 映像檔中的 LibTPU 層是專門建構而成，並與對應的 JAX 版本一併封裝。使用其他 JAX 版本可能會導致非預期的行為或錯誤。

JAX AI 映像檔中的 CUDA 層包含 NVIDIA 管理的元件，例如 NGC CUDA 深度學習映像檔，這個映像檔會做為 GPU 訓練映像檔的基本映像檔。GPU 映像檔也包含 Transformer Engine，這是 NVIDIA 的自訂程式庫，可加速 NVIDIA GPU 上的 Transformer 模型。

除了 JAX AI 映像檔提供的套件外，您可能還需要其他應用程式專用套件，才能執行特定的機器學習工作負載。

JAX AI 圖像中的程式庫：

目前 JAX AI 圖像