Databricks 发布 UCDataset：将 Unity Catalog 表直接流式送入 GPU 训练

GPU 通常是任何训练任务中最昂贵的部分，而大多数团队都无法让它们保持忙碌。瓶颈通常不在模型本身，而在于如何足够快地把数据从它所在的位置送到 GPU——这些芯片每空等一秒，你都在同时浪费时间和金钱。

今天，我们在 AI Runtime 上发布 Beta 版的 Unity Catalog Dataset (UCDataset)。UCDataset 是一个原生 PyTorch 数据加载器，可将 Unity Catalog 表直接流式传输到训练任务中的每个 GPU，无需导出、无需复制，也无需自定义加载器代码。

在模拟的生产工作负载上，与读取同一数据的导出 Parquet 副本相比，UCDataset 将训练时间缩短了 10 倍。

问题所在：你的 GPU 在等待你的数据

你模型的大多数训练数据已经存放在 Unity Catalog 表中。要将其送入分布式训练任务，意味着要从三种变通方案中选择一种，而每一种都会耗费你的时间和金钱：

每次运行前缓慢而繁琐的导出：最常见的模式是在每次运行前将表导出为卷中的 Parquet（或其他分片格式），然后让数据集指向这些文件。你需要为存储第二份副本付费，新数据一落地副本就会过时，而且它不在 Unity Catalog 血缘追踪范围内。
将整张表拉入内存无法扩展：用 toPandas() 将表读入 pandas 对小表有效，对大表则会失败。它还会把所有数据都放在一个节点上，这对在 32 个 GPU 上训练的任务毫无帮助。
通过仓库流式传输很慢：仓库并非原生为向 GPU 供数而构建，未经优化时，当许多 GPU 同时读取就会被限流。

然后是加载器本身。每个 GPU 和每个 worker 都需要自己均衡的数据分片。团队要花时间手动构建这一点，而一个细微的 bug 就可能导致模型在重复或跳过的行上训练，却无人察觉。

结果：昂贵的 GPU 闲置，训练耗时超出应有水平，ML 工程师维护数据管道，而不是改进模型。

隆重推出 UCDataset

UCDataset 使 Unity Catalog 表成为分布式 PyTorch 训练的一等数据源。你只需描述要读取的内容（一张表、所需的列以及行过滤器），UCDataset 就会将可直接用于训练的批次流式传输到每个 GPU。优势包括：

为速度而构建：UCDataset 直接从云存储读取表的文件，因此你的数据与 GPU 之间没有任何中间环节。它将匹配的行均匀拆分到每个 GPU 和数据加载器 worker。
在湖上训练，而非副本上：将其指向实时表。过滤器和删除向量在读取数据时生效，下一次运行会自动获取新追加的行。
端到端治理：训练任务以运行它们的身份读取数据，Unity Catalog 权限的适用方式与 Databricks 上其他任何地方完全一致。

基准测试：在推荐模型工作负载上实现更快的训练

我们在模拟的生产工作负载上对 UCDataset 进行了基准测试，将一张 300 GB 的特征表流式传输到 8 块 NVIDIA H100 GPU 上以训练推荐模型，结果与直接从 Parquet 加载数据相比，训练时间平均缩短了 10 倍。

如何使用：

假设一个团队将用户和物品特征以及点击标签存储在 main.recsys.ctr_features 表中。每天都有新的曝光数据写入该表。他们希望直接在两个节点的 16 块 GPU 上基于该表训练一个点击率模型。

1. 描述要读取的内容并创建加载器。在 train.py 中，选择该表，保留带标签的训练行，并仅选择模型所需的列。过滤会在读取表文件时应用，因此不匹配的行永远不会进入你的训练代码。UCDataset 已经按整批产出数据，因此将其包装在标准的 PyTorch DataLoader 中，并设置 batch_size=None。

2. 在两个节点上训练并启动。使用标准的 PyTorch DDP 循环进行训练。UCDataset 会将匹配的行分配给 16 块 GPU 中的每一块，因此无需编写采样器或分片代码。在 workload.yaml 中描述该作业，它会在每个节点上通过 torchrun 启动 train.py，并使用 AI Runtime CLI 运行：

要在更多 GPU 上训练，请更改 workload.yaml 中的 num_accelerators。训练代码保持不变。完整的 train.py 和 workload.yaml，以及可供试用的示例表，都在 AI Runtime cookbook 中：Unity Catalog 表上的点击率模型。

开始使用

包现已发布在 PyPI 上 - 因此你可以通过简单的 %pip install databricks-uc-dataset 进行安装

UCDataset 在 AI Runtime 上处于 Beta 阶段。要试用它：

阅读文档：使用 UCDataset 将 Unity Catalog 表加载到 PyTorch 中
从可运行的示例开始：cookbook 配方
联系你的客户团队，我们很乐意听取你的反馈。

对于存储在 Unity Catalog 卷中的非结构化数据，请使用 UCVolumeDataset 以本地缓存方式流式读取文件。UCDataset 处理结构化表，同时我们也在扩展它，以通过 FILE 类型支持多模态训练——让你可以直接从表中训练图像、音频和文档。Beta 版本可读取 Unity Catalog 中的 Delta 表。

常见问题

什么是 UCDataset？

UCDataset（Unity Catalog Dataset）是一个原生数据加载器，可将 Unity Catalog 中的 Unity Catalog 表直接流式传输到训练作业中的 GPU。它已在 Databricks AI Runtime 上以 Beta 版提供。你选择一张表以及所需的列，UCDataset 就会提供可直接用于训练的批次，无需导出步骤或自定义加载器代码。

UCDataset 如何加速 GPU 训练？

UCDataset 直接从云存储读取表的数据文件，中间没有导出的副本，从而让 GPU 保持忙碌。它将匹配的行均匀分配给每块 GPU 和每个数据加载器工作进程，因此每块 GPU 只读取自己的那部分数据。

我需要将 Unity Catalog 表导出为 Parquet 才能在其上训练吗？

不需要。UCDataset 读取 Unity Catalog 中的实时表，因此没有导出步骤，也没有需要存储或保持同步的第二份副本。它在读取时应用你的过滤条件和表的删除向量，并且每次新运行都会获取自上次运行以来追加的行。

UCDataset 是否强制执行 Unity Catalog 权限？

是的。在 AI Runtime 上，训练作业以运行它们的身份读取数据，并且表上的 Unity Catalog 权限适用。由于作业直接读取受治理的表，因此没有需要单独保护的导出副本。在 Beta 版中，UCDataset 无法读取带有行过滤器或列掩码的表。

UCDataset 可以读取哪些表？

UCDataset 可读取在 Unity Catalog 中注册的 Delta 表，包括带有删除向量的表。