Databricks releases UCDataset beta for direct GPU training

Databricks released UCDataset in beta on AI Runtime. The native PyTorch data loader streams rows from Unity Catalog tables directly into distributed GPU training jobs.

UCDataset reads a table’s underlying files in cloud storage. That removes the need to export a second Parquet copy or feed training through a Spark session or SQL warehouse. The loader applies selected columns, supported query filters and Delta deletion vectors during the read. A later run can also pick up rows appended to the live table.

Databricks’ documentation identifies UCDataset as a PyTorch IterableDataset. It divides matching rows across every distributed rank and DataLoader worker without a separate distributed sampler; by default, shard sizes differ by no more than one row. Multi-GPU jobs still need a launcher such as torchrun to set RANK and WORLD_SIZE for each process.

The identity running the workload determines access. Reading a table requires USE CATALOG on its catalog, USE SCHEMA on its schema and SELECT on the table. Unity Catalog permissions therefore govern the direct read rather than a separately secured exported copy.

Databricks reported an average 10× improvement in training time on a simulated recommendation-model workload using a 300 GB feature table and eight NVIDIA H100 GPUs, compared with loading the same data from Parquet. That figure is an issuer benchmark; Databricks did not publish enough methodology or raw results in the announcement for independent reproduction, and the result has not been independently verified.

The company’s November release notes and product documentation record the beta on November 6, three days before the November 9 launch post. The available official sources do not explain the timing difference.

The beta is limited to Delta Lake tables registered in Unity Catalog. It does not support views, materialized views, streaming tables, shallow clones, tables containing a VARIANT column, or tables protected by row filters or column masks. Databricks also warns that nearly equal shards can still produce unequal batch counts under DistributedDataParallel. UCDataset does not checkpoint its read position, so a restarted job begins again at the start of the table. Databricks has not provided a general-availability timetable.