在 AI 运行时上加载数据

重要

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

数据和模型资产对于深度学习和大型语言模型（LLM）和视觉语言模型（VLM）的训练后工作负载至关重要。通过 AI 运行时，所有数据和模型资产都可通过 Unity 目录访问：

Unity 目录卷：主要用于大型数据集和非结构化文件，包括图片、音频和文本。
Unity 目录表：用于结构化和表格化数据，通过 Spark Connect 访问。

你的卷和表必须在 Unity Catalog 中注册，并且你的用户或服务主体可以访问它们。

Unity目录卷用于非结构化数据

Unity 目录卷提供对任何格式的非表格数据的受控访问，包括结构化、半结构化和非结构化数据。在 AI 运行时中，卷是访问大型数据集、文本、模型资产和模型检查点的主要机制。

用户可以使用熟悉的文件系统操作，在 Unity 目录卷中列出、读取和写入文件，类似于处理本地磁盘上的文件：

import os

dir_path = "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir"
file_path = os.path.join(dir_path, "test_file")

os.makedirs(dir_path, exist_ok=True)

# Write to the file
with open(file_path, "w") as file:
    file.write("Hello, World!")

类似地，Shell 操作的工作方式也是一样的：

%sh ls -l /Volumes/<catalog-name>/<schema-name>/<volume-name>
%sh mkdir -p /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir
%sh touch /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/test_file

Unity 目录卷的几个特性使其非常适合机器学习工作负载：

分布式存储：Unity Catalog 支持分布式存储，允许 AI 运行时工作负载从笔记本和基于 CLI 的工作负载中读写数据和建模资产。
针对机器学习访问模式优化：底层存储和访问路径针对常见机器学习工作负载进行了优化，尤其是具有顺序读写的大文件。这使得Unity Catalog非常适合训练数据加载、模型资产加载和模型检查点编写。
类似文件系统的访问：用户可以通过熟悉的文件系统操作在 Unity 目录卷中列出、读取和写入文件，类似于处理本地磁盘上的文件。

由于自动后台提交，用户可以期待对 Unity 目录卷数据的稳定访问：

写入：AI 运行时会自动提交写入操作，从而使访问同一 Unity Catalog 卷的其他应用程序和工作负载能够看到这些更改。
读取：AI 运行时会自动检测到卷中的变化，无需用户显式执行任何刷新或同步操作。

调校音量性能

如前所述，Unity 目录卷由分布式存储支持，并针对具有顺序读写的大文件进行了优化。

有几个建议可以帮助你从AI运行时获得最佳性能：

将数据串接到更大的文件中：如果可能，将数据合并到较少且较大的文件中，每个文件大约1到10 GiB。这使得 AI 运行时能够积极预取数据，自动实现接近最优的顺序读取性能。

对于小文件工作负载，使用本地磁盘：如果你的工作负载包含许多小文件，考虑在处理文件前使用并行复制到本地磁盘。这可以减少通过卷反复访问大量小文件所产生的开销。

# Recommended using parallel copy (256 concurrency in this example, you can tune)
#
# This takes only 22 seconds to copy 15,375 150KiB small image files.
%sh cd /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ && find . -type f -print0 | xargs -0 -P 256 -I {} cp --parents "{}" /tmp/


# !!! Avoid doing this !!!
#
# Because the files are copied in serial, this copies the same 15,375 150KiB small image files much more slowly.
# %sh cp -r /Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/* /tmp

你可以用 UCVolumeDataset 它来处理你的机器学习工作负载。它结合了上述优化，以实现从 Unity 目录卷中高效的数据访问和加载。有关更多详细信息，请参阅以下部分。

使用 `UCVolumeDataset` 加载非结构化数据

对于存储在 Unity 目录卷中的非结构化数据，如图片、音频和文本文件，请使用 databricks.air.data 模块中的 UCVolumeDataset。 UCVolumeDataset 是一个 PyTorch IterableDataset 函数，它会在首次访问时将卷中的每个文件复制到快速的本地缓存中，并返回缓存的本地文件路径。它处理了原本需要你手动实现的性能和分发相关问题：

本地缓存。 文件会在首次访问时从 FUSE 挂载点复制到本地缓存目录，此后直接从缓存中提供，因此多 epoch 训练不会重复读取卷。
自动分区。 初始化 torch.distributed 时，文件会按 rank 进行分区，然后进一步分配给 DataLoader 个 worker，因此每个 (rank, worker) 对都会收到一个互不重叠的切片，无需额外配置。

注释

UCVolumeDataset 和 databricks.air.data.DataLoader 来自 databricks-sdk-air 软件包。用额外组件 data 安装，它还会安装兼容的 torch 版本：

%pip install "databricks-sdk-air[data]"

UCVolumeDataset 生成原始本地文件路径。要将这些文件解码为张量，请将其包装在第二个 IterableDataset 中，该装饰器会消费路径流并应用你的解析逻辑。这样可以将 I/O 和解析相关的关注点分开。

from databricks.air.data import UCVolumeDataset
from torch.utils.data import IterableDataset
from PIL import Image
import torchvision.transforms.functional as TF

class ImageDataset(IterableDataset):
    """Decodes each cached file path from UCVolumeDataset into a tensor."""

    def __init__(self, path_dataset: UCVolumeDataset):
        self._path_dataset = path_dataset

    def __iter__(self):
        for local_path in self._path_dataset:
            image = Image.open(local_path).convert("RGB")
            yield TF.to_tensor(image)

path_dataset = UCVolumeDataset("/Volumes/catalog/schema/my_volume/images")
dataset = ImageDataset(path_dataset)

封装器接收已缓存的本地路径，因此解析步骤完全不会访问该卷。可以链接其他包装器进行扩充、标记化或筛选。

为获得最佳性能，请将 UCVolumeDataset 与 databricks.air.data.DataLoader 搭配使用，而不是使用默认的 PyTorch DataLoader。它针对 AI 运行时输入输出进行了调校，同时在 GPU 计算时同时获取和缓存文件。

将模型检查点保存到卷上

为了为模型创建检查点，以便从最新快照继续训练或在发生崩溃后恢复，你可以像使用本地文件系统一样使用 Unity Catalog 卷。

Databricks 建议使用分布式检查点（DCP），以提升单 GPU 和多 GPU 工作负载的性能。请参阅Databricks工程博客中关于 AI运行时的快速容错PyTorch训练内容。

import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict

import databricks.air.data

checkpoint_path = "/Volumes/my-catalog/my-schema/my-volume/checkpoints/step_1000"

# Save
model_sd, optim_sd = get_state_dict(model, optimizer)
state_dict = {"model": model_sd, "optim": optim_sd, "step": 1000}
dcp.async_save(
    state_dict,
    storage_writer=databricks.air.data.UCVolumeWriter(checkpoint_path))

# Load
model_sd, optim_sd = get_state_dict(model, optimizer)
state_dict = {"model": model_sd, "optim": optim_sd}
dcp.load(
    state_dict,
    storage_reader=databricks.air.data.UCVolumeReader(checkpoint_path))

set_state_dict(
    model,
    optimizer,
    model_state_dict=state_dict["model"],
    optim_state_dict=state_dict["optim"],
)

单体式的 torch.save 方法也可行。

对于单GPU模型检查点，

# The monolithic torch.save approach for single GPU chip

# Save
torch.save({"model": model.state_dict(), "opt": optimizer.state_dict()},
           "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt")


# Load
ckpt = torch.load(
    "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt",
    weights_only=True)
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["opt"])

对于通过Torchrun启动的分布式培训，

# The monolithic torch.save approach for multi-GPU distributed training.
# This snippet assumes your launcher has already called
# dist.init_process_group(...).

import os
import torch.distributed as dist

# Save only on rank 0.
if dist.get_rank() == 0:
    torch.save({"model": model.state_dict(), "opt": optimizer.state_dict()},
               "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt")

# Wait for rank 0 to finish writing before any rank reads.
dist.barrier()

# Load on ALL ranks (map to current rank's local GPU).
local_rank = int(os.environ["LOCAL_RANK"])
ckpt = torch.load(
    "/Volumes/<catalog-name>/<schema-name>/<volume-name>/sub-dir/ckpt.pt",
    map_location=f"cuda:{local_rank}",
    weights_only=True)
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["opt"])

加载表格数据

使用 Spark Connect 从 Delta 表加载表格机器学习数据。

对于单节点训练，可以使用 PySpark 方法toPandas()将 Apache Spark 数据帧转换为 pandas 数据帧，然后使用 PySpark 方法to_numpy()选择性地转换为 NumPy 格式。

注释

Spark Connect 将分析和名称解析推迟到执行时间，这可能会更改代码的行为。请参阅 “将 Spark 连接与 Spark 经典版进行比较”。

Spark Connect 支持大多数 PySpark API，包括 Spark SQL、Spark 上的 Pandas API、结构化流式处理和 MLlib（基于数据帧）。有关最新支持的 API，请参阅 PySpark API 参考文档。

有关其他限制，请参阅无服务器计算限制。

使用 Unity 目录卷加载大型 Delta 表

对于体积过大而无法通过 toPandas() 进行转换的大型增量表，请将数据导出到 Unity Catalog 卷中，并使用 PyTorch 或 Hugging Face 直接加载：

# Step 1: Export the Delta table to Parquet files in a UC volume
output_path = "/Volumes/catalog/schema/my_volume/training_data"
spark.table("catalog.schema.my_table").write.mode("overwrite").parquet(output_path)

# Step 2: Load the exported data directly using Hugging Face datasets
from datasets import load_dataset

dataset = load_dataset("parquet", data_files="/Volumes/catalog/schema/my_volume/training_data/*.parquet")

此方法可避免在训练期间产生 Spark 开销，并且适用于单 GPU 和分布式训练工作流。

反馈

此页面是否有帮助？

上次更新时间为 2026-10-05