Databricks AI Research Leverages AI Runtime for Scalable, Efficient Model Development

The Challenge of Diverse AI Research Needs

The Databricks research team faces the complex task of developing state-of-the-art AI models for features such as ai_parse_document and ai_extract. This involves extensive experimentation across various architectures, datasets, training techniques, and hyperparameters, often resulting in tens of thousands of training and evaluation runs. The team utilizes methods like supervised fine-tuning (SFT) and reinforcement learning (RL) for models ranging from billions to nearly a trillion parameters, requiring access to powerful GPUs like NVIDIA H100 or B300. This diverse workload presents three key challenges:

• Resiliency for long-running training jobs: Failures in nodes or networking can halt progress on multi-day or multi-week training sessions, necessitating costly restarts from earlier checkpoints.
• Fast experimentation and prototyping: The need to rapidly test thousands of different approaches requires quick infrastructure provisioning and environment setup.
• GPU fleet management: Balancing multiple research projects with varying compute demands, priorities, and deadlines necessitates flexible resource allocation.

AI Runtime as the Solution

Databricks has been an early adopter of its own AI Runtime, using it to streamline model development and training. AIR addresses the identified challenges through several key features:

• Built-in GPU resiliency: AIR conducts preflight checks on GPU nodes and monitors their health during execution, detecting and isolating unhealthy nodes to prevent job failures and lost progress. This is crucial for projects like KARL, which require thousands of GPU hours.
• Fast iteration with AIR CLI and custom containers: The AIR Command Line Interface (CLI) allows for programmatic workload launches, specifying GPU requirements and enabling AIR to handle provisioning and scheduling, including high-bandwidth networking. Researchers can also use custom Docker containers for local development and seamlessly transition them to AIR for training, ensuring environment compatibility and reducing feedback loops.
• MLflow and logging integration: AIR automatically manages MLflow experiments for training jobs, simplifying experiment tracking. Researchers can monitor job status, stream logs, and retrieve metrics, providing essential data for AI agents to evaluate progress and plan next steps.
• Seamless, governed data access: Data stored in Unity Catalog is accessible via Unity Catalog FUSE mounting, allowing AIR to treat it as a local filesystem across different Databricks workspaces. This ensures data is accessible, tracked, shareable, and governed without altering training code.
• GPU capacity management: AIR's workload prioritization and preemption capabilities enable efficient allocation of GPU resources across projects with competing needs and deadlines, allowing for opportunistic use of spare capacity for exploratory experiments.

Scaling with Agent-Led Training

The rise of agent-driven research, where AI agents automate experiment proposal, execution, and evaluation, has dramatically increased experimentation volume. Databricks has seen an 8x increase in training jobs across its AI Functions research workloads due to this shift. AIR is well-suited to support this trend, offering fast job launches, programmatic interfaces, reliable execution, and accessible results. The autoresearch workflow, involving proposing, launching, learning, and updating, is facilitated by AIR's ability to handle concurrent workloads and integrate with MLflow metrics, allowing researchers to focus on innovation while agents manage the iterative process.

Databricks AI Runtime provides the robust infrastructure necessary to support both larger models and a faster pace of AI development, empowering researchers to spend less time on infrastructure management and more time on creating advanced AI capabilities.