已启用 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 支持以下实例类型:

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

GPU 调度

GPU 计划可在大量 GPU 之间高效地分配 Spark 任务。

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

备注

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

用户定义的 GPU 计划仅适用于 Databricks Runtime 7.1 及更高版本。 对于版本低于 7.1 的 Databricks Runtime,Databricks 会自动配置 GPU 计算,让每个节点最多有一个正在运行的任务。 如此,该任务可以使用节点上的所有 GPU,不会与其他任务发生冲突。

AI 和 ML 适用的 GPU 计划

spark.task.resource.gpu.amount 是你可能需要配置的与 GPU 感知型计划相关的唯一 Spark 配置。 默认配置为每个任务使用一个 GPU。如果你使用所有 GPU 节点,这对分布式推理工作负载和分布式训练来说是不错的基线。

要减少分布式训练期间的通信开销,Databricks 建议在计算 Spark 配置中将 spark.task.resource.gpu.amount 设置为每个工作器节点的 GPU 数。 此设置只会为每个 Spark 工作器创建一个 Spark 任务,并将该工作器节点中的所有 GPU 分配给同一任务。

若要增加分布式深度学习推理的并行度,可将 spark.task.resource.gpu.amount 设置为分数值,例如 1/2、1/3、1/4…… 1/N。 此设置会创建比 GPU 更多的 Spark 任务,允许更多的并发任务并行处理推理请求。 例如,若将 spark.task.resource.gpu.amount 设置为 0.50.330.25,则可用 GPU 数将拆分为任务数的两倍、三倍或四倍。

GPU 索引

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

如果你使用 Scala,则可从 TaskContext.resources().get("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 驱动程序版本为 535.154.05

有关所包含的库的版本,请参阅你使用的特定 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 驱动程序版本,因为该版本必须与主机上的驱动程序版本相匹配。

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