使用 UCDataset 加速 Unity Catalog 表上的 GPU 训练
内容提要
Databricks发布Unity Catalog Dataset（UCDataset）Beta版，这是一个原生PyTorch数据加载器，可将Unity Catalog表直接流式传输到训练任务的各GPU，无需导出或复制数据。在模拟生产测试中，读取300GB特征表训练推荐模型，训练时间比读取Parquet副本缩短10倍。它自动将数据均匀分片到各GPU，支持列选择和行过滤，并遵循Unity Catalog权限。目前仅支持Delta表，不支持行过滤或列掩码的表。
延伸解读
GPU闲置的根源与UCDataset的针对性
文章指出，GPU通常是训练中最昂贵的部分，但多数团队无法让其保持忙碌，瓶颈往往不在模型，而在于数据从存储到GPU的速度。传统做法需要导出或复制数据，并手动编写分片逻辑，容易导致GPU空闲或数据重复/遗漏。UCDataset直接流式读取Unity Catalog表，自动将数据均匀分片到各GPU，省去导出和自定义加载器代码，从而减少数据管道维护负担。
10倍加速的测试条件与适用边界
在模拟生产负载中，UCDataset流式读取300GB特征表，使用8块NVIDIA H100 GPU训练推荐模型，训练时间比直接读取导出的Parquet副本平均缩短10倍。需注意该结果基于特定模拟场景，实际收益可能因数据规模、集群配置和模型类型而异。此外，Beta版仅支持Delta表，且无法读取带有行过滤或列掩码的表，选型时需确认表属性。
权限与数据新鲜度的内置处理
UCDataset直接读取Unity Catalog中的实时表，训练任务以运行身份读取数据，Unity Catalog权限自动生效，无需为导出副本单独设置安全策略。读取时会应用过滤条件和表的删除向量，每次新运行能获取上次运行后追加的行，避免维护同步副本。但需留意，Beta版不支持行过滤或列掩码的表，若表启用了这些治理功能，则无法使用UCDataset。
Q&A
UCDataset 是什么？它主要解决什么问题？
UCDataset（Unity Catalog Dataset）是 Databricks 发布的 Beta 版原生 PyTorch 数据加载器，可将 Unity Catalog 表直接流式传输到训练任务的每个 GPU，无需导出或复制数据。它解决了分布式训练中数据加载瓶颈导致 GPU 空闲的问题，避免团队手动编写数据分片代码。
UCDataset 相比读取 Parquet 副本能提升多少训练速度？
在模拟生产测试中，使用 8 块 NVIDIA H100 GPU 流式读取 300GB 特征表训练推荐模型，UCDataset 相比直接读取 Parquet 副本平均缩短训练时间 10 倍。
UCDataset 如何将数据分配到多个 GPU 上？需要自己写分片代码吗？
UCDataset 自动将匹配的行均匀分片到每个 GPU 和数据加载器 worker，每个 GPU 只读取自己的份额。用户无需编写 sampler 或分片代码，只需用标准 PyTorch DataLoader 包装（batch_size=None）即可。
UCDataset 支持哪些数据源和操作？有哪些限制？
UCDataset 目前仅支持 Unity Catalog 中注册的 Delta 表（包括带删除向量的表），支持列选择和行过滤（过滤在读取文件时应用）。限制：Beta 版不支持有行过滤或列掩码的表。
使用 UCDataset 时，Unity Catalog 的权限控制如何生效？
在 AI Runtime 上，训练任务以运行它的身份读取数据，Unity Catalog 表上的权限自动应用。由于直接读取受治理的表，没有导出的副本需要单独保护。
如何开始使用 UCDataset？需要安装什么包？
可以通过 %pip install databricks-uc-dataset 安装 PyPI 包。在 train.py 中指定表、所需列和行过滤条件创建加载器，用 PyTorch DataLoader 包装，然后通过 workload.yaml 配置作业并用 AI Runtime CLI 启动。完整示例见 AI Runtime cookbook。
