- 将训练时间加速 10倍 并通过 UCDataset 降低计算成本，这是一个目前处于 Beta 阶段的原生数据加载器，可直接将 Unity Catalog 表中的数据流式传输到你的 GPU
- UCDataset 可直接接入标准 PyTorch 训练循环，并自动处理多 GPU 数据分片，无需自定义代码
- 通过直接从 Unity Catalog 读取训练数据，你可以确保治理策略被自动应用
GPU 通常是任何训练运行中最昂贵的部分，而大多数团队无法让它们保持忙碌。瓶颈通常不是模型，而是如何足够快地将数据从其所在位置传输到 GPU——这些芯片每等待一秒，你都在同时浪费时间和金钱。
今天我们发布 Unity Catalog Dataset (UCDataset) 的 Beta 版，登陆 AI Runtime。UCDataset 是一个原生 PyTorch 数据加载器，可将 Unity Catalog 表直接流式传输到训练任务中的每块 GPU，无需导出、无需复制，也无需自定义加载器代码。
在模拟生产工作负载中，与读取同一数据的已导出 Parquet 副本相比，UCDataset 将训练时间缩短了 10倍 。
问题：你的 GPU 正在等待你的数据
你模型的大多数训练数据已经存放在 Unity Catalog 表中。要将其送入分布式训练任务，意味着要从三种变通方案中选择一种，而每一种都会耗费你的时间和金钱：
- 每次运行前缓慢而繁琐的导出： 最常见的模式是在每次运行前将表导出为卷中的 Parquet（或其他分片格式），然后让数据集指向这些文件。你需要为存储第二份副本付费，新数据一到达它就会过时，而且它落在 Unity Catalog 血缘之外。
- 将整张表拉入内存无法扩展： 使用 toPandas() 将表读入 pandas 对小表有效，对大表则会失败。它还会将所有内容放在一个节点上，这对在 32 块 GPU 上训练的任务没有帮助。
- 通过仓库流式传输很慢： 仓库并非原生为向 GPU 供数而构建，如果没有优化，当许多 GPU 同时读取时会被限流。
然后是加载器本身。每块 GPU 和每个工作进程都需要自己均衡的数据切片。团队花时间手工构建这一点，而一个细微的 bug 可能意味着模型在重复或跳过的行上训练，却无人察觉。
结果： 昂贵的 GPU 闲置，训练耗时超过应有水平，ML 工程师维护数据管道而不是改进模型。
隆重推出 UCDataset
UCDataset 使 Unity Catalog 表成为分布式 PyTorch 训练的一等数据源。你描述要读取的内容（一张表、你需要的列以及行过滤器），UCDataset 就会将可直接用于训练的批次流式传输到每块 GPU。优势包括：
- 为速度而构建： UCDataset 直接从云存储读取表的文件，因此你的数据与 GPU 之间没有任何中间层。它将匹配的行均匀拆分到每块 GPU 和每个数据加载器工作进程。
- 在湖上训练，而不是在副本上： 将其指向实时表。过滤器和删除向量在读取数据时应用，下一次运行会自动获取新追加的行。
- 端到端治理： 训练任务以运行它们的身份读取数据，Unity Catalog 权限的应用方式与 Databricks 上其他任何地方完全一致。
基准测试：在推荐模型工作负载上实现更快的训练
我们在模拟生产工作负载上对 UCDataset 进行了基准测试，在 8 块 NVIDIA H100 GPU 上流式传输一张 300 GB 的特征表来训练推荐模型，结果发现与直接从 Parquet 加载数据相比，训练时间平均 提升 10 倍 。
如何使用：
假设一个团队将用户和物品特征以及点击标签存储在 main.recsys.ctr_features 表中。每天都有新的曝光落到该表上。他们希望跨两个节点、在 16 块 GPU 上直接从该表训练一个点击率模型。
1. 描述要读取的内容并创建加载器。 在 train.py 中，选择表、保留带标签的训练行，并只选择模型需要的列。过滤器在读取表的文件时应用，因此不匹配的行永远不会到达你的训练代码。UCDataset 已经产出完整批次，因此用 batch_size=None 将其包装在标准 PyTorch DataLoader 中。
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 表，包括带有删除向量的表。
相关标签：公告, Unity Catalog, 数据科学与机器学习 和 AI 研究
在收件箱中获取最新文章
订阅我们的博客，最新文章将直接发送到你的收件箱。