How our AI Research team uses AI Runtime

How our AI Research team uses AI Runtime | Databricks Blog Skip to main content

Summary

Databricks AI Research uses AI Runtime to train and fine-tune models at scale, from small experiments to workloads spanning hundreds of GPUs.

Built-in GPU reliability, fast job launches, and integrations with MLflow and Unity Catalog help researchers experiment faster.

Agent-driven research has increased training jobs 8x, with AI Runtime supporting faster, automated experimentation.

At Databricks, we continuously develop and evaluate new AI research techniques to make our products more intelligent, efficient, and useful for customers. Our research spans products like Genie, Smart Routing, and AI Functions. Recently, we've published how we made  search 3x faster , developed  new techniques for offline reinforcement learning , and pushed the  frontier of complex document extraction . To build great AI models and manage hundreds of parallel research workstreams, we need infrastructure that can keep pace. That's why we use Databricks AI Runtime (AIR). In this post, we'll share how AIR helps us experiment quickly, scale to large training workloads, and adapt to agent-led training, along with what we've learned about managing GPU resources across research projects. The challenge: many experiments, very different compute needs Our research team develops models that power  Databricks AI Functions , including ai_parse_document and  ai_extract . For example, Databricks AI Functions  deliver state-of-the-art document parsing quality , backed by models trained to understand enterprise documents. Achieving that quality requires extensive research across different architectures, datasets, training techniques, and hyperparameters, spanning tens of thousands of training and evaluation runs. We use supervised fine-tuning (SFT), offline reinforcement learning (RL), and online RL to train models ranging from a few billion parameters to nearly a trillion. Depending on the workload, we use NVIDIA H100 or B300 GPUs, both available through AIR. Our largest training runs can require hundreds of GPUs working together over days or weeks, with high-bandwidth networking to keep training efficient. Supporting these workloads creates three challenges: Resiliency for long-running training jobs. Training can require thousands of GPU hours. A single node or RDMA network failure can leave hundreds of GPUs idle and force us to restart from an earlier checkpoint, losing hours of progress. For reinforcement learning  projects like KARL , reliable infrastructure is essential. Fast experimentation and prototyping. Improving AI Functions means testing thousands of approaches, from different models and training techniques to data recipes. We need to launch experiments quickly without repeatedly provisioning infrastructure or rebuilding environments. GPU fleet management. We balance multiple research projects with different compute requirements, priorities, and deadlines, while remaining flexible enough to respond to new research techniques or open-source models.

Adopting AI Runtime to solve these challenges Our research team has been customer zero of AI Runtime since the product's inception last year. We use AIR to develop, train, and test models that power Databricks AI products, relying on its GPU infrastructure and integration with our existing tools. Built-in GPU resiliency AIR performs  preflight checks on assigned GPU nodes to ensure they're healthy and can communicate effectively before training starts. It also monitors node health during execution, detecting issues such as NCCL communication failures, thermal or power throttling, and GPU driver problems. If a node becomes unhealthy and causes a job to fail, AIR detects the issue and cordons off the affected node so subsequent jobs aren't scheduled on the same problematic infrastructure. This is especially important for projects like KARL, where complex rollout generation and post-training of models with more than 100 billion parameters can require thousands of GPU hours. Rather than repeatedly troubleshooting infrastructure failures, we can focus on improving our models. Fast iteration with the AIR CLI and custom containers Research requires us to change model configurations, training approaches, and code frequently. With the AIR CLI, we can launch workloads programmatically, specify the GPUs we need, and let AIR handle provisioning and scheduling, including high-bandwidth networking for distributed training across clouds. We also rely on custom Docker containers for local development. AIR lets us use those same images for training, ensuring compatibility without recreating our environments. We can upload our codebase, execute custom code in those containers, and use Databricks Artifact Registry to register and share images across experiments. AIR's investments in launch latency and container optimization help us start training jobs in seconds, shortening feedback loops as we test new ideas. MLflow and logging for experiment tracking We use MLflow to track our research experiments. AIR automatically creates and manages MLflow experiments for our training jobs, eliminating separate configuration. We can monitor job status, stream and inspect training logs, and retrieve metrics to understand how experiments are progressing. This also gives our AI agents access to the information they need to evaluate results and determine what to try next. Seamless, governed access to research data All of our research data is stored in Unity Catalog. Through Unity Catalog FUSE mounting, AIR lets us access that data as though it were part of a local filesystem, even when working across different Databricks workspaces. This keeps our data accessible, tracked, shareable, and governed without changing how our training code interacts with files. Combined with MLflow and our Docker-based environments, it gives us a consistent workflow from local development to distributed training. Managing GPU capacity across research projects Our workloads have very different resource needs. A debugging experiment may require only a few GPUs, while a large training run might need hundreds for days or weeks. Different projects also have competing timelines and priorities. AIR's workload prioritization and preemption capabilities help us allocate resources across projects, prioritize time-sensitive experiments, and use excess capacity when it's available. For...