by John Spencer, Milan Stefanovic, Bruce Fontaine and Tejas Sundaresan
GPUs are usually the most expensive part of any training run, and most teams can't keep them busy. The bottleneck usually isn't the model. It's getting data from where it lives to the GPU fast enough - and every second those chips sit waiting, you're burning both time and money.
Today we're releasing the Unity Catalog Dataset (UCDataset) in Beta on AI Runtime. UCDataset is a native PyTorch data that streams Unity Catalog tables directly to every GPU in your training job, with no exports, no copies, and no custom code.
On simulated production workloads, UCDataset cut training time by 10x compared with reading an exported Parquet copy of the same data.
Most training data for your models already lives in Unity Catalog tables. Getting it into a distributed training job means picking one of three workarounds, and each one costs you time and money:
Then there's the itself. Every GPU and every worker needs its own balanced slice of the data. Teams spend time building this by hand, and a subtle bug can mean the model trains on duplicated or skipped rows without anyone noticing.
The result: expensive GPUs sit idle, training takes longer than it should, and ML engineers maintain data plumbing instead of improving models.
UCDataset makes Unity Catalog tables a first-class data source for distributed PyTorch training. You describe what to read (a table, the columns you need, and a row filter) and UCDataset streams training-ready batches to every GPU. Benefits include:
We benchmarked UCDataset on a simulated production workload, streaming a 300 GB feature table across 8 NVIDIA H100 GPUs to train a recommendation model and saw on average a 10x improvement in training time compared to data directly from Parquet.
How to use it:
Say a team keeps user and item features with a click label stored in a main.recsys.ctr_features table. New impressions land on the table every day. They want to train a click-through-rate model on 16 GPUs across two nodes, straight from the table.
1. Describe what to read and create the . In train.py, pick the table, keep the labeled training rows, and select only the columns the model needs. The filter is applied as the table's files are read, so rows that don't match never reach your training code. UCDataset already yields whole batches, so wrap it in a standard PyTorch Data with batch_size=None.
2. Train and launch on two nodes. Train with a standard PyTorch DDP loop. UCDataset gives each of the 16 GPUs its own share of the matching rows, so there's no sampler or sharding code to write. Describe the job in workload.yaml, which starts train.py with torchrun on every node, and run it with the AI Runtime CLI:
To train on more GPUs, change num_accelerators in workload.yaml. The training code stays the same. The complete train.py and workload.yaml, with a sample table to try them on, are in the AI Runtime cookbook: Click-through-rate model on a Unity Catalog table.
Package is on PyPI now - so you can install with simple %pip install databricks-uc-dataset UCDataset is available in Beta on AI Runtime. To try it:
For unstructured data stored in Unity Catalog volumes, use UCVolumeDataset to stream files with local caching. While UCDataset handles structured tables, we are also expanding it to support multimodal training via the FILE type - allowing you to train on images, audio, and documents directly from a table. The Beta reads Delta tables in Unity Catalog. UCDataset (Unity Catalog Dataset) is a native data that streams Unity Catalog tables in Unity Catalog directly to the GPUs in a training job. It's available in Beta on Databricks AI Runtime. You choose a table, the columns you need, and UCDataset delivers training-ready batches with no export step or custom code.
UCDataset keeps GPUs busy by reading a table's data files directly from cloud storage, without an exported copy in between. It splits the matching rows evenly across every GPU and data worker, so each GPU reads only its own share.
No. UCDataset reads the live table in Unity Catalog, so there's no export step and no second copy to store or keep in sync. It applies your filter and the table's deletion vectors as it reads, and each new run picks up rows appended since the last one.
Yes. On AI Runtime, training jobs read data as the identity that runs them, and Unity Catalog permissions on the table apply. Because the job reads the governed table directly, there's no exported copy to secure separately. In the Beta, UCDataset can't read tables that have row filters or column masks.
UCDataset reads Delta tables registered in Unity Catalog, including tables with deletion vectors.
Announcements, Unity Catalog, Data Science and ML and AI Research
Subscribe to our blog and get the latest posts delivered to your inbox.