Accelerate GPU training on Unity Catalog tables with UCDataset

Databricks released the Unity Catalog Dataset (UCDataset) in Beta on AI Runtime, a native PyTorch data loader that streams Unity Catalog tables directly to every GPU in a training job without exports or custom loader code. In a benchmark streaming a 300 GB feature table across 8 NVIDIA H100 GPUs to train a recommendation model, Databricks measured an average 10x improvement in training time versus loading data directly from an exported Parquet copy. The package is available on PyPI as databricks-uc-dataset, and Databricks said it is expanding UCDataset to multimodal training via the FILE type for images, audio and documents.

by John Spencer https://www.databricks.com/blog/author/john-spencer , Milan Stefanovic https://www.databricks.com/blog/author/milan-stefanovic , Bruce Fontaine https://www.databricks.com/blog/author/bruce-fontaine and Tejas Sundaresan https://www.databricks.com/blog/author/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 loader that streams Unity Catalog tables directly to every GPU in your training job, with no exports, no copies, and no custom loader 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 loader 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 loading 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 loader. 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 DataLoader 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 https://github.com/databricks-solutions/ai-runtime-cookbook/tree/main/recommendation-and-ranking/click-through-rate-uc-dataset . 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 https://docs.databricks.com/aws/en/machine-learning/ai-runtime/develop-workloads/dataloading load-unstructured-data-with-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 https://docs.databricks.com/aws/en/unstructured/file - 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 loader 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 loader 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 loader 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 https://www.databricks.com/blog/tag/announcements , Unity Catalog https://www.databricks.com/blog/tag/unity-catalog , Data Science and ML https://www.databricks.com/blog/tag/data-science-machine-learning and AI Research https://www.databricks.com/blog/tag/ai-research Subscribe to our blog and get the latest posts delivered to your inbox.