操作系统映像和 Docker 映像

Google Cloud 提供包含常见操作系统、框架、库和驱动程序的映像。 Google Cloud 优化这些预配置的映像，以支持人工智能 (AI) 和机器学习 (ML) 工作负载。

本文档简要介绍了用于在 AI Hypercomputer 环境中部署、管理和运行工作负载的映像。

了解图片类别

图片分为以下几类：

  • AI 和机器学习框架与库：预配置了机器学习框架和库的二进制文件的 Docker 映像，可简化机器学习模型的创建、训练和使用。在 AI 超级计算机上，您可以使用 Deep Learning Software Layer (DLSL) Docker 映像在 Google Kubernetes Engine (GKE) 集群上运行 NeMO、MaxText 等机器学习模型。
  • 集群部署和编排：用于部署和管理经过性能优化的基础设施（AI 工作负载在其中运行）的操作系统 (OS) 映像。您可以将 AI 工作负载部署至 GKE 集群、Slurm 集群或 Compute Engine 实例。如需了解详情，请参阅虚拟机和集群创建概览。以下操作系统映像可用于集群或实例部署：

AI 和机器学习框架和库

Google Cloud 提供打包了热门 AI 和 ML 框架及库的 Docker 映像。这些映像提供了所需的软件，可简化在 AI Hypercomputer 上运行的 AI 优化型集群中开发、训练和部署模型的过程。

JAX AI 映像

适用于Google Cloud TPU 和 GPU 的 JAX AI 映像（JAII，之前称为 JAX 稳定堆栈映像）提供即用型 Docker 映像，其中包含 JAX 框架、精选的兼容库以及 Google Cloud 基础架构的设置。JAX AI TPU 映像预配置了 JAX 库和 TPU 库。JAX AI GPU 映像预配置了 JAX 库和相关的 CUDA/NVIDIA 库。

JAX AI 映像的示意图

硬件层

JAX AI 映像位于硬件层之上，硬件层由加速器（TPU 或 GPU）及其关联的虚拟机组成。如需使用 JAX AI 映像，您需要预配 TPU 或 GPU 虚拟机。您可以使用 TPU API、Compute Engine API 或 GKE API 来执行此操作。

框架层

框架层提供用于构建机器学习工作负载的工具和库。JAX AI 映像为基于 JAX 的机器学习工作负载提供预配置的基础，包括核心 JAX 库和其他必要的依赖项，以确保一致且高性能的开发体验。

JAX AI 映像中的 LibTPU 层是专门构建的，并与相应的 JAX 版本捆绑在一起。使用其他 JAX 版本可能会导致意外行为或错误。

JAX AI 映像中的 CUDA 层包含由 NVIDIA 管理的组件，例如 NGC CUDA Deep Learning Image，该映像用作 GPU 训练映像的基础映像。GPU 映像还包含 Transformer Engine，这是一个自定义 NVIDIA 库，用于在 NVIDIA GPU 上加速 Transformer 模型。

除了 JAX AI 映像中提供的软件包之外，您的特定机器学习工作负载可能还需要其他应用专用软件包。

JAX AI 映像中的库：

TPU 图片

功能 软件包名称
核心库或组件
机器学习框架 JAX 和 JAX 库
LibTPU Cloud LibTPU 稳定版
层/模型库 Flax
检查点 Orbax
量化 Qwix
优化器 Optax

aqtp
自定义内核 Tokamax
配置 Fiddle
输入流水线 tf.data

PyGrain array-record

性能分析、调试 Tensorboard
Utils 常见循环实用程序

机器学习有效吞吐量衡量 机器学习集合

云专用工具
Google Cloud

Google Cloud SDK Google Cloud 存储空间

诊断

云加速器诊断 Cloud TPU 诊断

GPU 映像

功能 软件包名称
核心库或组件
机器学习框架 JAX 和 JAX 库
NVIDIA CUDA 库 CUDA DL 映像
层/模型库 Flax
检查点 Orbax
优化器

Optax aqtp TransformerEngine

配置 Fiddle
输入流水线

tf.data PyGrain array-record

性能分析、调试 Tensorboard
Utils

常见循环实用程序 机器学习有效吞吐量测量 机器学习集合

云专用工具
Google Cloud

Google Cloud SDK Google Cloud 存储空间

诊断 Cloud Accelerator 诊断

当前 JAX AI 映像

TPU 图片

JAX AI 映像 发布日期
JAX 0.9.0 修订版 1 2026-02-03
JAX 0.8.2 修订版 1 2026-01-14
JAX 0.8.1 修订版 1 2025-11-21
JAX 0.8.0 修订版 1 2025-10-28
JAX 0.7.2 修订版 1 2025-09-30
JAX 0.7.0 修订版 1 2025-07-29
JAX 0.6.1 修订版 1 2025-06-05
JAX 0.5.2 修订版 2 2025-04-25
JAX 0.5.2 修订版 1 2025-03-17
JAX 0.4.37 修订版 1 2024-12-12
JAX 0.4.35 修订版 1 2024-10-30

GPU 映像

JAX AI 映像 发布日期
JAX 0.7.2，搭配 CUDA DL 25.06 修订版 1 2025-09-30
JAX 0.6.1，搭配 CUDA DL 25.03 修订版 1 2025-06-05
JAX 0.5.1，搭配 CUDA DL 25.02 修订版 1 2025-03-17

应用层

您可以在应用层（位于框架层之上）实现特定的机器学习工作负载。应用层包含应用专属的代码、模型和逻辑，所有这些都是使用框架层提供的工具和库构建的。

虽然此映像为基于 JAX 的 AI 工作负载提供了可靠且经过充分测试的基础，但您可能需要添加特定于应用的依赖项。这样做时，我们建议您尽量减少对预配置的基础层（包括 JAX 及其核心依赖项）的干扰。引入会替换或冲突现有依赖项的应用级依赖项可能会导致以下副作用：

  • 意外行为：与向 JAX AI 映像添加其他依赖项之前相比，您的机器学习工作负载可能会表现出不同的行为。
  • 性能回归：替换经过优化的 JAX 相关库可能会对 JAX AI 映像提供的性能优势产生负面影响
  • 稳定性问题：您添加的依赖项与核心 JAX 依赖项之间的冲突可能会导致应用出现不稳定性和运行时错误。

发布频率

最初，JAX AI 映像每季度提供一次，近期的目标是与每个 JAX 版本同步发布。这样，您就可以在新功能和改进推出后的第一时间利用。

支持

每个 JAX AI 映像版本都遵循有限时间的支持生命周期。在此时间范围内，我们会处理对现有 JAX AI 映像的修改请求，具体类别如下：

  • 安全漏洞：我们会优先处理在 JAX 稳定版栈 Docker 映像的基础映像或依赖项中发现的安全漏洞。我们将发布更新后的映像，以缓解潜在风险。
  • 重大变更：如果 JAX AI 图片使用的底层库或框架发生重大变更， Google Cloud会评估并实现必要的更新，以保持兼容性。这可能需要使用更新后的依赖项重新构建 Docker 映像。

如果在 JAII 中的某个库中发现安全漏洞或 bug，我们会将更新后的库纳入 JAII 中，并固定所有其他库版本，以保持整体稳定性。这会生成新的 JAII 修订版本。

修订的最小更改：

如果我们在 JAX-0.4.30-rev1 中的软件包“X”中发现 bug，我们会将“X”更新到其下一个版本（例如 v2.0），同时尽量保持所有其他软件包不变。这会生成一个新修订版本：JAX-0.4.30-rev2，我们会尽快发布该版本。

深度学习软件层 (DLSL) Docker 映像

这些映像打包了 NVIDIA CUDA、NCCL、机器学习框架和模型。它们可为深度学习工作负载提供即用型环境。这些预构建的 DLSL Docker 映像可与您的 GKE 集群无缝协作，因为我们在内部可重现性测试和回归测试期间对这些映像进行了测试和验证。

DLSL Docker 映像具有以下优势：

  • 预配置的软件：DLSL Docker 映像复制了内部重现性和回归测试所用的设置。这些映像提供预配置、经过测试和优化的环境，可大幅节省安装和配置时间。
  • 版本管理：DLSL Docker 映像会经常更新。这些版本更新提供了最新的稳定版框架和驱动程序，并且还解决了安全补丁问题。
  • 基础架构兼容性：DLSL Docker 映像经过构建和测试，可与 AI Hypercomputer 上提供的 GPU 机器类型无缝协作。
  • 快速入门说明：某些 DLSL Docker 映像附带示例配方，可向您展示如何启动使用预配置映像的工作负载。

NeMo + PyTorch + NCCL gIB 插件

这些 Docker 映像基于 NVIDIA NeMo NGC 映像。它们包含 Google 的 NCCL gIB 插件，并捆绑了在每个受支持的加速器机器上运行工作负载所需的所有 NCCL 二进制文件。这些映像还包括 Google Cloud等工具，例如用于将工作负载部署到 Google Kubernetes Engine 的 gcsfuse 和 gcloud CLI。

DLSL 映像版本 依赖项版本 机器系列 发布日期 支持终止日期 DLSL 映像名称
nemo25.04-gib1.0.6-A4
  • NeMo NGC:25.04.01
  • NCCL giB plugin: 1.0.6
 A4 2025 年 7 月 3 日 2026 年 7 月 3 日 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/pytorch-gpu-nemo-nccl:nemo25.04-gib1.0.6-A4
nemo25.04-gib1.0.6-A3U
  • NeMo NGC:25.04.01
  • NCCL giB plugin: 1.0.6
 A3 Ultra 2025 年 7 月 3 日 2026 年 7 月 3 日 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/pytorch-gpu-nemo-nccl:nemo25.04-gib1.0.6-A3U
nemo25.02-gib1.0.5-A4
  • NeMo NGC:25.02
  • NCCL giB plugin: 1.0.5
 A4 2025 年 3 月 14 日 2026 年 3 月 14 日 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/pytorch-gpu-nemo-nccl:nemo25.02-gib1.0.5-A4
nemo24.07-gib1.0.2-A3U
  • NeMo NGC:24.07
  • NCCL giB plugin: 1.0.2
 A3 Ultra 2025 年 2 月 2 日 2026 年 2 月 2 日 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/pytorch-gpu-nemo-nccl:nemo24.07-gib1.0.2-A3U
nemo24.07-gib1.0.3-A3U
  • NeMo NGC:24.07
  • NCCL giB plugin: 1.0.3
 A3 Ultra 2025 年 2 月 2 日 2026 年 2 月 2 日 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/pytorch-gpu-nemo-nccl:nemo24.07-gib1.0.3-A3U
nemo24.12-gib1.0.3-A3U
  • NeMo NGC:24.12
  • NCCL giB plugin: 1.0.3
 A3 Ultra 2025 年 2 月 7 日 2026 年 2 月 7 日 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/pytorch-gpu-nemo-nccl:nemo24.12-gib1.0.3-A3U
nemo24.07-tcpx1.0.5-A3Mega
  • NeMo NGC:24.07
  • GPUDirect-TCPX: 1.0.5
 A3 Mega 2025 年 3 月 12 日 2026 年 3 月 12 日 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/pytorch-gpu-nemo-nccl:nemo24.07-tcpx1.0.5-A3Mega
nemo24.07-tcpx1.0.5-A3High
  • NeMo NGC:24.07
  • GPUDirect-TCPX: 1.0.5
 A3 High（8 个 GPU） 2025 年 3 月 12 日 2026 年 3 月 12 日 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/pytorch-gpu-nemo-nccl:nemo24.07-tcpx1.0.5-A3High

NeMo + PyTorch

此 Docker 映像基于 NVIDIA NeMo NGC 映像，并包含 Google Cloud 工具，例如用于将工作负载部署到 Google Kubernetes Engine 的 gcsfuse 和 gcloud CLI。

DLSL 映像版本 依赖项版本 机器系列 发布日期 支持终止日期 DLSL 映像名称
nemo24.07--A3U NeMo NGC:24.07 A3 Ultra 2024 年 12 月 19 日 2025 年 12 月 19 日 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/pytorch-gpu-nemo:nemo24.07-A3U
nemo24.07-tcpx1.0.5-A3Mega
  • NeMo NGC:24.07
  • GPUDirect-TCPX: 1.0.5
 A3 Mega 2025 年 3 月 12 日 2026 年 3 月 12 日 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/pytorch-gpu-nemo:nemo24.07-tcpx1.0.5-A3Mega
nemo24.07-tcpx1.0.5-A3High
  • NeMo NGC:24.07
  • GPUDirect-TCPX: 1.0.5
 A3 High（8 个 GPU） 2025 年 3 月 12 日 2026 年 3 月 12 日 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/pytorch-gpu-nemo:nemo24.07-tcpx1.0.5-A3High

MaxText + JAX 工具箱

此 Docker 映像基于 NVIDIA JAX 工具箱映像，包含Google Cloud 等工具，例如用于将工作负载部署到 Google Kubernetes Engine 的 gcsfuse 和 gcloud CLI。

DLSL 映像版本 依赖项版本 机器系列 发布日期 支持终止日期 DLSL 映像名称
toolbox-maxtext-2025-01-10-A3U JAX toolbox: maxtext-2025-01-10 A3 Ultra 2025 年 3 月 11 日 2026 年 3 月 11 日 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/jax-maxtext-gpu:toolbox-maxtext-2025-01-10-A3U
jax0.5.1-cu12-A3Mega
  • JAX: 0.5.1
  • CUDA: 12.x
  • MaxText: Latest
 A3 Mega 2025 年 3 月 17 日 2026 年 3 月 17 日 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/jax-gpu-maxtext:jax0.5.1-cu12-A3Mega
jax0.5.1-cu12-A3High
  • JAX: 0.5.1
  • CUDA: 12.x
  • MaxText: Latest
 A3 High（8 个 GPU） 2025 年 3 月 17 日 2026 年 3 月 17 日 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/jax-gpu-maxtext:jax0.5.1-cu12-A3High

MaxText + JAX 稳定栈

此 Docker 映像基于 JAX 稳定版栈和 MaxText。此映像还包含 dnsutils 等依赖项，以便在 Google Kubernetes Engine 上运行工作负载。

DLSL 映像版本 依赖项版本 机器系列 发布日期 支持终止日期 DLSL 映像名称
jax-maxtext-gpu:jax0.5.1-cuda_dl25.02-rev1-maxtext-20150317
  • JAX Stable stacks:jax0.5.1-cuda_dl25.02-rev1
  • maxtext 提交：54e98c9e62caa426cf5902be068533ddb4fb79f5
 A4 2025 年 3 月 17 日 2026 年 3 月 17 日 us-central1-docker.pkg.dev/deeplearning-images/reproducibility/jax-maxtext-gpu:jax0.5.1-cuda_dl25.02-rev1-maxtext-20150317

集群部署和编排

操作系统映像包含在计算实例或 GKE 节点上部署操作系统所需的所有软件组件。操作系统管理底层硬件资源，例如加速器和网络。 这可为 AI 工作负载提供计算资源。

GKE 节点映像

GKE 使用节点映像部署集群。这些节点映像适用于各种操作系统，例如 Container-Optimized OS、Ubuntu 和 Windows Server。您需要部署 GKE Autopilot 集群的预装了 containerd (cos_containerd) 的 Container-Optimized OS 节点映像包含多项优化，可支持您的 AI 和 ML 工作负载。

如需详细了解这些节点映像，请参阅节点映像。

Slurm 操作系统映像

Slurm 集群将计算节点和控制器节点部署为 Compute Engine 上的计算实例。

如需配置 AI 优化型 Slurm 集群，您必须使用 Cluster Toolkit。在 Slurm 集群部署期间，集群蓝图会自动构建自定义操作系统映像，该映像会在 Slurm 节点上安装集群和工作负载管理所需的系统软件。您可以在部署默认蓝图之前对其进行修改，以自定义映像中包含的某些软件。

以下部分总结了集群蓝图在 A4、A3 Ultra、A3 Mega 和 A3 High（8 个 GPU）Slurm 节点上安装的软件。集群蓝图扩展了 Ubuntu LTS 加速器操作系统映像。

A4

GitHub 上提供的 A4 蓝图默认包含以下软件：

  • Ubuntu 22.04 LTS
  • Slurm：版本 24.11.2
  • 以下 Slurm 依赖项：
    • munge
    • mariadb
    • libjwt
    • lmod
  • Open MPI：4.1.x 的最新版本
  • PMIx：版本 4.2.9
  • NFS 客户端和服务器
  • NVIDIA 570 系列驱动程序
  • NVIDIA enroot 容器运行时：版本 3.5.0，包含发布后 bug 修复
  • NVIDIA pyxis
  • 以下 NVIDIA 工具：
    • 数据中心 GPU 管理器 (dcgmi)
    • nvidia-utils-570
    • nvidia-container-toolkit
    • libnvidia-nscq-570
  • CUDA 工具包：版本 12.8
  • Infiniband 支持，包括 ibverbs-utils
  • Ops Agent
  • Cloud Storage FUSE

A3 Ultra

GitHub 上提供的 A3 Ultra 蓝图默认包含以下软件：

  • Ubuntu 22.04 LTS
  • Slurm：版本 24.11.2
  • 以下 Slurm 依赖项：
    • munge
    • mariadb
    • libjwt
    • lmod
  • Open MPI：4.1.x 的最新版本
  • PMIx：版本 4.2.9
  • NFS 客户端和服务器
  • NVIDIA 570 系列驱动程序
  • NVIDIA enroot 容器运行时：版本 3.5.0，包含发布后 bug 修复
  • NVIDIA pyxis
  • 以下 NVIDIA 工具：
    • 数据中心 GPU 管理器 (dcgmi)
    • libnvidia-cfg1-570-server
    • libnvidia-nscq-570
    • nvidia-compute-utils-570-server
    • nsight-compute
    • nsight-systems
  • CUDA 工具包：版本 12.8
  • Infiniband 支持，包括 ibverbs-utils
  • Ops Agent
  • Cloud Storage FUSE

A3 Mega

GitHub 上提供的 A3 Mega 蓝图默认包含以下软件：

  • Ubuntu 22.04 LTS
  • Slurm：版本 24.11.2
  • 以下 Slurm 依赖项：
    • munge
    • mariadb
    • libjwt
    • lmod
  • Open MPI：4.1.x 的最新版本
  • PMIx：版本 4.2.9
  • NFS 客户端和服务器
  • NVIDIA 550 系列驱动程序
  • NVIDIA enroot 容器运行时：版本 3.4.1
  • NVIDIA pyxis
  • 以下 NVIDIA 工具：
    • 数据中心 GPU 管理器 (dcgmi)
    • libnvidia-cfg1-550-server
    • libnvidia-nscq-550
    • nvidia-compute-utils-550-server
    • nsight-compute
    • nsight-systems
  • CUDA 工具包：版本 12.4
  • Infiniband 支持，包括 ibverbs-utils
  • 适用于 A3 Mega 的 TCPXO：适用于 NCCL TCPXO 和 1.6 Tbps 网络堆栈的专用安装程序
  • Ops Agent
  • Cloud Storage FUSE

A3 High

GitHub 上提供的 A3 High（8 个 GPU）蓝图默认包含以下软件：

  • Ubuntu 22.04 LTS
  • Slurm：版本 24.11.2
  • 以下 Slurm 依赖项：
    • munge
    • mariadb
    • libjwt
    • lmod
  • Open MPI：4.1.x 的最新版本
  • PMIx：版本 4.2.9
  • NFS 客户端和服务器
  • NVIDIA 535 或 550 系列驱动程序
  • NVIDIA enroot 容器运行时：版本 3.4.1
  • NVIDIA pyxis
  • 以下 NVIDIA 工具：
    • 数据中心 GPU 管理器 (dcgmi)
    • libnvidia-cfg1-535/550-server
    • libnvidia-nscq-535/550
    • nvidia-compute-utils-535/550-server
    • nsight-compute
    • nsight-systems
  • CUDA 工具包：版本 12.2 或 12.4
  • Infiniband 支持，包括 ibverbs-utils
  • 适用于 A3 High（8 个 GPU）的 TCPX：适用于 NCCL TCPX (GPUDirect-TCPX) 和 800 Gbps 网络堆栈（4 个 NIC）的专用安装程序
  • Ops Agent
  • Cloud Storage FUSE

加速器操作系统映像

借助 AI Hypercomputer，您可以预配单个计算实例或一组计算实例。 如果您想创建这些实例，则必须在创建实例期间指定操作系统映像。

Google Cloud 提供了一套用于创建实例的操作系统映像。 Google Cloud 还提供了一套专为 AI 优化型实例定制的加速器操作系统映像。这些操作系统映像包含 GPU 和网络功能的核心驱动程序，例如 NVIDIA 驱动程序、Mellanox 驱动程序及其依赖项。

如需详细了解每个操作系统，请参阅 Compute Engine 文档中的操作系统详细信息页面。

加速器操作系统映像适用于 Rocky Linux 和 Ubuntu LTS 操作系统。

Rocky Linux 加速器

以下 Rocky Linux 加速器操作系统映像适用于每个机器系列：

操作系统版本 映像系列 架构 机器系列 映像项目
Rocky Linux 9 加速器 rocky-linux-9-optimized-gcp-nvidia-580-arm64 Arm A4X Max、A4X rocky-linux-accelerator-cloud
rocky-linux-9-optimized-gcp-nvidia-580 x86 A4、A3 Ultra、A3 Mega、A3 High（8 个 GPU） rocky-linux-accelerator-cloud
rocky-linux-9-optimized-gcp-nvidia-570-arm64 Arm A4X rocky-linux-accelerator-cloud
rocky-linux-9-optimized-gcp-nvidia-570 x86 A4、A3 Ultra、A3 Mega、A3 High（8 个 GPU） rocky-linux-accelerator-cloud
Rocky Linux 8 加速器 rocky-linux-8-optimized-gcp-nvidia-580-arm64 Arm A4X Max、A4X rocky-linux-accelerator-cloud
rocky-linux-8-optimized-gcp-nvidia-580 x86 A4、A3 Ultra、A3 Mega、A3 High（8 个 GPU） rocky-linux-accelerator-cloud
rocky-linux-8-optimized-gcp-nvidia-570 x86 A4、A3 Ultra、A3 Mega、A3 High（8 个 GPU） rocky-linux-accelerator-cloud

Ubuntu LTS 加速器

以下 Ubuntu LTS 加速器操作系统映像适用于各个机器系列：

操作系统版本 映像系列 架构 机器系列 映像项目
Ubuntu 24.04 LTS 加速器 ubuntu-accelerator-2404-arm64-with-nvidia-580 Arm A4X Max、A4X ubuntu-os-accelerator-images
ubuntu-accelerator-2404-amd64-with-nvidia-580 x86 A4、A3 Ultra、A3 Mega、A3 High（8 个 GPU） ubuntu-os-accelerator-images
ubuntu-accelerator-2404-arm64-with-nvidia-570 Arm A4X ubuntu-os-accelerator-images
ubuntu-accelerator-2404-amd64-with-nvidia-570 x86 A4、A3 Ultra、A3 Mega、A3 High（8 个 GPU） ubuntu-os-accelerator-images
Ubuntu 22.04 LTS 加速器 ubuntu-accelerator-2204-arm64-with-nvidia-580 Arm A4X Max、A4X ubuntu-os-accelerator-images
ubuntu-accelerator-2204-amd64-with-nvidia-580 x86 A4、A3 Ultra、A3 Mega、A3 High（8 个 GPU） ubuntu-os-accelerator-images
ubuntu-accelerator-2204-arm64-with-nvidia-570 Arm A4X ubuntu-os-accelerator-images
ubuntu-accelerator-2204-amd64-with-nvidia-570 x86 A4、A3 Ultra、A3 Mega、A3 High（8 个 GPU） ubuntu-os-accelerator-images

后续步骤