Databricks beta streams Unity Catalog data directly to GPUs

Direct path from Unity Catalog to GPUs

UCDataset lets teams use a Unity Catalog table as a data source for distributed PyTorch training. Users specify the table, the columns required by the model and a row filter, then UCDataset streams training ready batches to the GPUs in the job.

According to Databricks, the loader reads the table's files directly from cloud storage. It splits matching rows evenly across GPUs and data loader workers, so teams do not need to write sampler or sharding code. UCDataset also produces whole batches, allowing it to be wrapped in a standard PyTorch DataLoader with batch_size=None.

The package is available on PyPI and can be installed with %pip install databricks-uc-dataset. The beta reads Delta tables registered in Unity Catalog, including tables with deletion vectors.

Databricks reports 10x faster benchmark

Databricks says UCDataset addresses a common training bottleneck: moving data to GPUs quickly enough to keep them busy. Teams often export Unity Catalog tables to Parquet before a run, but that creates a second copy, can become stale when new data arrives and falls outside Unity Catalog lineage.

In a simulated production workload, Databricks streamed a 300 GB feature table across 8 NVIDIA H100 GPUs to train a recommendation model. The company reports an average 10x improvement in training time compared with loading data directly from Parquet.

UCDataset can also apply filters and deletion vectors as data is read, meaning rows that do not match a filter do not reach the training code. Databricks says newly appended rows are picked up automatically on the next run.

Governance remains attached to the live table

Training jobs read data as the identity that runs them, and Unity Catalog permissions apply to the table, according to Databricks. This avoids the need to secure a separately exported copy of the data. In the beta, however, UCDataset cannot read tables that have row filters or column masks.

Databricks provides a cookbook example for a click through rate model using a table called main.recsys.ctr_features. The example trains on 16 GPUs across two nodes with a standard PyTorch distributed data parallel loop. To use more GPUs, users change num_accelerators in the workload configuration while leaving the training code unchanged.

For unstructured data in Unity Catalog volumes, Databricks directs users to UCVolumeDataset, which streams files with local caching. The company says it is also expanding UCDataset to support multimodal training through the FILE type, including images, audio and documents. For users with structured Delta tables, the beta offers a way to connect existing governance and live data directly to PyTorch training without an export step.