操作系统映像和 Docker 映像

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

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

了解图片类别

图片分为以下几类：

  • AI 和机器学习框架与库：预配置了机器学习框架和库的二进制文件的 Docker 映像，可简化机器学习模型的创建、训练和使用。在 AI Hypercomputer 上，您可以使用深度学习软件层 (DLSL) Docker 映像在 Google Kubernetes Engine (GKE) 集群上运行 NeMO 和 MaxText 等机器学习模型。
  • 集群部署和编排：用于部署和管理运行 AI 工作负载的性能优化基础架构的操作系统 (OS) 映像。您应使用的映像取决于您使用的 GPU 基础设施（集群化 GPU 或通用 GPU）以及您计划如何部署集群：

AI 和机器学习框架和库

Google Cloud 提供打包了热门 AI 和机器学习框架及库的 Docker 映像。这些映像提供了在 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 映像图表

硬件层

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 深度学习映像，该映像用作 GPU 训练映像的基础映像。GPU 映像还包含 Transformer Engine，这是一个自定义 NVIDIA 库，用于在 NVIDIA GPU 上加速 Transformer 模型。

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

JAX AI 映像中的库：

TPU 映像

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

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

PyGrain array-record

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

机器学习吞吐量测量 机器学习集合

云专用工具
Google Cloud

Google Cloud SDK Google Cloud storage

诊断

Cloud Accelerator Diagnostics Cloud TPU Diagnostics

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 storage

诊断 Cloud Accelerator Diagnostics

当前 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工具，例如 gcsfuse 和 gcloud CLI，用于将工作负载部署到 Google Kubernetes Engine。

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 工具，例如 gcsfuse 和 gcloud CLI，用于将工作负载部署到 Google Kubernetes Engine。

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 等工具，例如 gcsfuse 和 gcloud CLI，用于将工作负载部署到 Google Kubernetes Engine。

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。此映像还包含用于在 Google Kubernetes Engine 上运行工作负载的依赖项，例如 dnsutils。

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 的 Container-Optimized OS (cos_containerd) 节点映像包含支持 AI 和机器学习工作负载的优化。

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

Slurm 操作系统映像

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

如需预配 AI 优化型 Slurm 集群，您必须使用 Cluster Toolkit。默认情况下，A4X、A4 和 A3 Ultra 集群的集群蓝图会部署使用预构建 Advanced Compute Images 的计算实例。对于其他机器系列，或在构建 Slurm 自定义映像时，集群蓝图会在部署期间自动构建自定义操作系统映像，该映像会在 Slurm 节点上安装用于集群和 Workload Manager 的必需系统软件。您可以在部署默认蓝图之前对其进行修改，以自定义映像包含的软件。

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

A4X Max

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

  • Ubuntu 24.04 LTS
  • Slurm：版本 25.05.2
  • 以下 Slurm 依赖项：
    • munge
    • mariadb
    • libjwt
    • lmod
  • Open MPI：版本 5.0.x
  • NFS 客户端和服务器
  • NVIDIA 580 系列驱动程序
  • NVIDIA Enroot
  • NVIDIA Pyxis
  • 以下 NVIDIA 工具：
    • 数据中心 GPU 管理器 (dcgmi)
    • nvidia-container-toolkit：版本 1.17.7
    • nvidia-imex
    • nvidia-fabricmanager
  • CUDA 工具包：版本 13.0
  • 适用于 A4X Max 的 NCCL 插件 (nccl-gib-a4x-max-arm64)
  • DOCA-OFED 驱动程序
  • Ops Agent
  • Cloud Storage FUSE

A4X

GitHub 上提供的 A4X 蓝图默认使用预构建的 Advanced Compute Images。aci-gpu-u2404-slurm-2511-cuda-130-nvidia-580-arm64 映像系列默认包含以下软件和配置：

  • Ubuntu 24.04 LTS (Arm64)
  • Slurm：版本 25.11 及其所有依赖软件包
  • NVIDIA 580 系列驱动程序
  • CUDA 工具包：版本 13.0
  • NVIDIA Container Toolkit、Enroot 和 Pyxis
  • 数据中心 GPU Manager (dcgmi) 和 Fabric Manager
  • NCCL/gIB 插件
  • Ops Agent
  • Cloud Storage FUSE

如需查看预安装工具和清单详情的完整列表，请参阅 Advanced Compute Images 概览。

A4

GitHub 上提供的 A4 蓝图默认使用预构建的 Advanced Compute Images。默认情况下，aci-gpu-u2404-slurm-2511-cuda-130-nvidia-580-amd64 映像系列包含以下软件和配置：

  • Ubuntu 24.04 LTS (x86_64)
  • Slurm：版本 25.11 及所有依赖软件包
  • NVIDIA 580 系列驱动程序
  • CUDA 工具包：版本 13.0
  • NVIDIA Container Toolkit、Enroot 和 Pyxis
  • 数据中心 GPU Manager (dcgmi) 和 Fabric Manager
  • NCCL/gIB 插件
  • Infiniband 和多队列网络支持
  • Ops Agent
  • Cloud Storage FUSE

如需查看预安装工具和清单详情的完整列表，请参阅 Advanced Compute Images 概览。

A3 Ultra

GitHub 上提供的 A3 Ultra 蓝图默认使用预构建的 Advanced Compute Images。默认情况下，aci-gpu-u2404-slurm-2511-cuda-130-nvidia-580-amd64 映像系列包含以下软件和配置：

  • Ubuntu 24.04 LTS (x86_64)
  • Slurm：版本 25.11 及所有依赖软件包
  • NVIDIA 580 系列驱动程序
  • CUDA 工具包：版本 13.0
  • NVIDIA Container Toolkit、Enroot 和 Pyxis
  • 数据中心 GPU Manager (dcgmi) 和 Fabric Manager
  • NCCL/gIB 插件
  • Infiniband 和多队列网络支持
  • Ops Agent
  • Cloud Storage FUSE

如需查看预安装工具和清单详情的完整列表，请参阅 Advanced Compute Images 概览。

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
  • TCPXO for A3 Mega：NCCL TCPXO 和 1.6 Tbps 网络栈的专用安装程序
  • Ops Agent
  • Cloud Storage FUSE

A3 High（8 个 GPU）

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 加速器操作系统映像：

基础设施 映像系列 架构 机器系列 映像项目
集群化 GPU rocky-linux-10-optimized-gcp-nvidia-580 x86 A4、A3 Ultra rocky-linux-accelerator-cloud
rocky-linux-10-optimized-gcp-nvidia-580-arm64 Arm A4X rocky-linux-accelerator-cloud
rocky-linux-9-optimized-gcp-nvidia-580 x86 A4、A3 Ultra rocky-linux-accelerator-cloud
rocky-linux-9-optimized-gcp-nvidia-580-arm64 Arm A4X rocky-linux-accelerator-cloud
rocky-linux-8-optimized-gcp-nvidia-580 x86 A4、A3 Ultra rocky-linux-accelerator-cloud
rocky-linux-8-optimized-gcp-nvidia-580-arm64 Arm A4X rocky-linux-accelerator-cloud

Ubuntu LTS 加速器

每个机器系列都提供以下 Ubuntu LTS 加速器操作系统映像：

基础设施 映像系列 架构 机器系列 映像项目
集群化 GPU ubuntu-accelerator-2404-arm64-with-nvidia-580 Arm A4X ubuntu-os-accelerator-images
ubuntu-accelerator-2404-amd64-with-nvidia-580 x86 A4、A3 Ultra ubuntu-os-accelerator-images
ubuntu-accelerator-2204-arm64-with-nvidia-580 Arm A4X ubuntu-os-accelerator-images
ubuntu-accelerator-2204-amd64-with-nvidia-580 x86 A4、A3 Ultra ubuntu-os-accelerator-images

后续步骤