已启用 GPU 的计算

注意

一些启用了 GPU 的实例类型为 Beta 版本,并且当你在计算创建过程中选择驱动程序和工作器类型时,它们将在下拉列表中如此标记。

概述

Azure Databricks 支持通过图形处理单元 (GPU) 加速的计算。 本文介绍了如何创建包含启用了 GPU 的实例的计算,并介绍了在这些实例上安装的 GPU 驱动程序和库。

要详细了解启用了 GPU 的计算上的深度学习,请参阅深度学习

创建 GPU 计算

创建 GPU 计算类似于创建任何计算。 应当记住以下事项:

  • Databricks Runtime 版本必须是启用了 GPU 的版本,例如 Runtime 13.3 LTS ML(GPU、Scala 2.12.15、Spark 3.4.1)
  • 工作器类型驱动程序类型必须为 GPU 实例类型。

支持的实例类型

Azure Databricks 支持以下实例类型:

  • NC 实例类型系列:Standard_NC12、Standard_NC24
  • NC v3 实例类型系列:Standard_NC6s_v3、Standard_NC12s_v3、Standard_NC24s_v3
  • NC T4 v3 实例类型系列:Standard_NC4as_T4_v3、Standard_NC8as_T4_v3、Standard_NC16as_T4_v3、Standard_NC64as_T4_v3
  • NC A100 v4 实例类型系列:Standard_NC24ads_A100_v4、Standard_NC48ads_A100_v4、Standard_NC96ads_A100_v4
  • ND A100 v4 实例类型系列:Standard_ND96asr_v4
  • NV A10 v5 实例类型系列:Standard_NV36ads_A10_v5、Standard_NV36adms_A10_v5、Standard_NV72ads_A10_v5
    • Standard_NV72ads_A10_v5 与 PyTorch 2.0+ 不兼容,并且在运行多 GPU 工作负荷时会引发 CUDA 错误。 由于 Databricks Runtime ML 14.0 及更高版本预安装了 PyTorch 2.0+,因此如果你在多 GPU 工作负荷上运行 PyTorch,我们建议使用 Databricks Runtime ML 13.x 或手动运行 pip install torch==1.13.1

有关支持的 GPU 实例类型及其可用性区域的最新列表,请参阅 Azure Databricks 定价。 你的 Azure Databricks 部署必须位于受支持的区域内,才能启动启用了 GPU 的计算。

GPU 调度

Databricks Runtime 支持 Apache Spark 3.0 中的 GPU 感知型计划。 Azure Databricks 会在 GPU 计算上预先配置它。

单节点计算上不会启用 GPU 计划。

spark.task.resource.gpu.amount 是你可能需要更改的与 GPU 感知型计划相关的唯一 Spark 配置。 默认配置为每个任务使用一个 GPU。如果你使用所有 GPU 节点,这对于分布式推理工作负荷和分布式训练非常理想。 要在部分节点上进行分布式训练(这有助于减少分布式训练期间的通信开销),Databricks 建议在计算 Spark 配置中将 spark.task.resource.gpu.amount 设置为每个工作器节点的 GPU 数。

对于 PySpark 任务,Azure Databricks 自动将分配的 GPU 重新映射到索引 0、1…。 在为每个任务使用一个 GPU 的默认配置下,你的代码可以直接使用默认 GPU,不需检查为任务分配了哪个 GPU。 如果为每个任务设置了多个 GPU(例如 4 个),则你的代码可以假设已分配 GPU 的索引始终为 0、1、2、3。 如果你需要已分配 GPU 的物理索引,则可从 CUDA_VISIBLE_DEVICES 环境变量获取它们。

如果你使用 Scala,则可从 TaskContext.resources().get("gpu") 获取分配给任务的 GPU 的索引。

对于低于 7.0 的 Databricks Runtime 版本,为了避免多个尝试使用同一 GPU 的 Spark 任务之间发生冲突,Azure Databricks 会自动配置 GPU 计算,以便每个节点最多有一个正在运行的任务。 在这种情况下,任务可以使用节点上的所有 GPU,而不会与其他任务发生冲突。

NVIDIA GPU 驱动程序、CUDA 和 cuDNN

Azure Databricks 在 Spark 驱动程序和工作器实例上安装使用 GPU 所需的 NVIDIA 驱动程序和库:

  • CUDA 工具包,安装在 /usr/local/cuda下。
  • cuDNN:NVIDIA CUDA 深度神经网络库。
  • NCCL:NVIDIA 集体通信库。

包含的 NVIDIA 驱动程序的版本为 535.54.03,支持 CUDA 11.0。 对于 NV A10 v5 实例类型系列,包含的 NVIDIA 驱动程序版本为 525.105.17

有关所包含的库的版本,请参阅你使用的特定 Databricks Runtime 版本的发行说明

注意

此软件包含 NVIDIA Corporation 提供的源代码。 具体而言,为支持 GPU,Azure Databricks 包括了 CUDA 示例中的代码。

NVIDIA 最终用户许可协议 (EULA)

当你在 Azure Databricks 中选择启用了 GPU 的“Databricks Runtime 版本”时,你默示同意 NVIDIA EULA 中列出的有关 CUDA、cuDNN 和 Tesla 库的条款和条件,以及 NCCL 库的 NVIDIA 最终用户许可协议(包含 NCCL 补充条款)

GPU 计算上的 Databricks 容器服务

重要

此功能目前以公共预览版提供。

可以在具有 GPU 的计算上使用 Databricks 容器服务,以通过自定义库创建可移植的深度学习环境。 有关说明,请参阅使用 Databricks 容器服务自定义容器

要为 GPU 计算创建自定义映像,必须选择标准运行时版本,而不是 Databricks Runtime ML for GPU。 如果选择“使用自有 Docker 容器”,则可以选择具有标准运行时版本的 GPU 计算。 GPU 的自定义映像基于正式版 CUDA 容器,这与 Databricks Runtime ML for GPU 不同。

为 GPU 计算创建自定义映像时,无法更改 NVIDIA 驱动程序版本,因为它必须与主机上的驱动程序版本相匹配。

databricksruntimedatabricksruntime 包含具有 GPU 功能的示例基础映像。 用于生成这些映像的 Dockerfile 位于示例容器 GitHub 存储库中,其中还详细介绍了示例映像提供的内容以及如何对其进行自定义。