配置您的环境

重要

用于单节点任务的 AI 运行时为 公共预览版。 多 GPU 工作负载的分布式训练 API 仍为 Beta 版。

AI 运行时为无服务器 GPU 计算提供两种托管 Python 环境:一种是包含 torchcudatorchvision 的精简标准环境,另一种是预装了 PyTorch、Transformers 以及其他机器学习和深度学习框架的 Databricks AI 环境。 选择标准环境以完全控制依赖项堆栈,或选择 AI 环境,以便进行随时可用的训练设置。 你也可以从任意一种环境开始,并通过 %pip install 添加软件包。

应使用什么环境?

AI 运行时提供两个托管Python环境:标准环境和 Databricks AI 环境。

环境 主要特征 何时使用
标准环境 极简;包括 cuda。 从 环境版本 5 起, torchtorchvision 不再预安装。 你希望完全控制依赖项堆栈,并且希望仅安装所需的内容
Databricks AI 环境 预加载了常用的 ML 框架(PyTorch、转换器等) 你需要一个完整的环境来训练、微调和试验,而无需手动依赖项管理

还可以使用工作区管理员为无服务器 GPU 计算生成的 工作区基础环境 。 请参阅 构建无服务器 GPU 计算(AI 运行时)

标准环境(最小环境)

包含仅供 AI 运行时操作所需的包的最小化稳定环境。 环境包含用于支持 GPU 的 cuda。 从环境版本 5 起, torchtorchvision 不再预安装。 使用 pip install 安装工作负载所需的版本。 有关每个环境版本中安装的包,请参阅以下发行说明。

最适合:希望完全控制其依赖项堆栈的用户,并且希望仅安装所需的内容。

若要选择:在 “环境 ”侧面板中,选择 “标准 v5 ”或“ 标准 v4” 作为基本环境。

有关不同版本中安装的包版本的更多详细信息,请参阅发行说明:

Databricks AI 环境

在环境 4 及更高版本中可用。 AI 环境基于标准环境构建,其中包含特定于 GPU 上的机器学习的常见运行时包和包。 预安装的包包括:

  • PyTorch (支持 CUDA)
  • Transformer(Hugging Face)
  • 以及其他 ML/DL 依赖项

最适合:希望在不需要手动管理依赖项的情况下进行训练、微调和实验的 ML 从业者的完整环境。

若要选择:在 “环境 ”侧面板中,选择 “AI v5 ”或 “AI v4” 作为基础环境。

有关不同版本中安装的包版本的更多详细信息,请参阅发行说明:

工作区基础环境

工作区管理员可以为无服务器 GPU 计算生成 工作区基础环境 ,该环境可通过 “基本环境 ”下拉菜单向工作区中的所有用户提供。 有关详细信息,请参阅构建无服务器 GPU 计算 (AI 运行时)

还可以从提供的基础环境之一(默认或 Databricks AI)开始配置每个项目的深度学习环境,并在笔记本或训练脚本顶部以编程方式 %pip install 安装其他包:

%pip install "trl==1.1.0"
%pip install "peft==0.19.1"
%pip install "transformers==5.5.4"
%pip install "fsspec==2024.9.0"
%pip install "huggingface_hub==1.11.0"
%pip install "datasets==3.2.0"
%pip install "accelerate==1.13.0"

在将@distributed修饰器用于多 GPU 工作负载时,调用.distributed()之前使用%pip install安装的软件包会被自动创建快照并分发到所有分布式进程。 已安装包的总大小不得超过 15 GB。

有关详细信息,请参阅 向笔记本添加依赖项

环境缓存和自定义模块

何时缓存环境?

跨会话缓存环境以加快启动时间。 使用同一环境配置重新连接到 AI 运行时时,以前安装的包可以从缓存获取,从而减少安装时间。

但是,不保证缓存行为,因此请始终确保笔记本包含可重现性所需的 %pip install 命令。

如何导入自定义模块?

可以通过将自定义模块放入/Workspace/Shared并将路径添加到sys.path来实现导入。

import sys
sys.path.append("/Workspace/Shared/my-project/src")
from my_module import my_function

还可以将模块文件作为工作区文件上传,并直接导入它们。 对于多用户协作,请将共享代码 /Workspace/Shared 存储在而不是特定于用户的文件夹中。 对于活动开发,请使用用户特定的文件夹并推送到远程 Git 存储库进行版本控制。

局限性

AI 运行时不提供以下功能:

  • Spark 函数:不能直接导入或使用 PySpark 函数。 AI 运行时是仅限 Python 的环境;Spark 不可用作本地运行时。 但是,Spark Connect 可用于数据加载。 请参阅 在 AI 运行时上加载数据
  • Databricks Runtime ML 库:预安装的包不是 Databricks Runtime ML 的替代项。 Databricks Runtime ML 中提供的某些 ML 库可能不会在 AI 运行时上预安装。
  • 专用项目:在某些情况下,AI 运行时支持专用项目。 有关详细信息,请联系你的帐户团队。