How our AI Research team uses AI Runtime
Databricks’ AI Research team relies on AI Runtime to accelerate model development, scale training workloads, and support agent-led research, reducing infrastructure overhead and enabling faster experimentation.
Databricks’ AI Research team develops models for products like Genie, Smart Routing, and AI Functions, including state-of-the-art document parsing tools such as ai_parse_document and ai_extract. To achieve high-quality results, the team conducts extensive research across architectures, datasets, and training techniques, spanning tens of thousands of runs. These experiments require significant computational resources, including NVIDIA H100 or B300 GPUs, which are provisioned and managed through Databricks AI Runtime (AIR).
AIR serves as the research team’s primary infrastructure, handling GPU provisioning, scheduling, and monitoring for distributed training workloads. It performs preflight checks to ensure GPU node health, detects failures during execution, and automatically isolates problematic nodes to prevent repeated issues. For large-scale projects like KARL, which involve complex model rollouts and training runs exceeding 100 billion parameters, AIR’s reliability reduces downtime and allows researchers to focus on model improvements rather than infrastructure troubleshooting.
The team uses the AIR CLI to launch workloads programmatically, specifying GPU requirements and leveraging high-bandwidth networking for distributed training. Custom Docker containers developed locally are reused in AIR, ensuring consistency across environments. AIR’s optimizations for launch latency and container handling enable training jobs to start in seconds, shortening feedback loops and accelerating experimentation cycles.
Researchers track experiments using MLflow, with AIR automatically managing experiment creation and metrics collection. Data access is streamlined through Unity Catalog FUSE mounting, allowing seamless integration with local filesystems across workspaces. AIR’s workload prioritization and preemption features allocate resources dynamically, supporting both small debugging tasks and large-scale training runs without manual coordination, while agent-driven research has increased training jobs by 8x.