How our AI Research team uses AI Runtime

來源原標題

How our AI Research team uses AI Runtime

機構

Databricks

原始標籤

Engineering · Research

原始連結

https://www.databricks.com/blog/how-our-ai-research-team-uses-ai-runtime

摘要原文

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. 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: 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. 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 example, spare capacity can support additional exploratory experiments over a weekend without disrupting higher-priority work. Our goal is to keep GPU resources available for the experiments that matter, without requiring researchers to manually coordinate infrastructure. Popularized by Andrej Karpathy, autoresearch is changing how we develop AI models. Instead of researchers manually launching and evaluating every experiment, we're increasingly using agents to propose changes, run experiments, evaluate results, and determine what to try next. This shift has dramatically increased our experimentation volume. Across our AI Functions research workloads, the number of training jobs has grown by 8x . As agents take on more of the research process, our infrastructure needs to keep up. A slow startup or confusing configuration error might be manageable when a researcher launches a handful of experiments. When agents launch hundreds, those same issues quickly become bottlenecks. Agent-driven research requires fast job launches, programmatic interfaces, clear and actionable errors, reliable execution, and accessible training results. Our autoresearch workflow follows a continuous cycle: AIR is particularly well suited to this workflow because agents can interact with it programmatically and benefit from its fast job launches, built-in reliability, and ability to scale. Through the Databricks CLI and databricks-air-skills , our agents can launch concurrent workloads, monitor job status, retrieve MLflow metrics, and iterate based on the results. This lets us run more experiments while our researchers focus on developing new ideas and improving model quality. As we move further toward agent-driven research, AIR gives us the infrastructure to support both larger models and a faster pace of experimentation. Our experience using AI Runtime has shown us how much reliable, flexible GPU infrastructure matters for research. From large-scale training and governed data access to agent-driven experimentation, AIR lets us spend less time managing infrastructure and more time developing better models. If you're training or fine-tuning your own models, check out the AI Runtime documentation to get started. AI Runtime has also powered our research across areas like faster search , reinforcement learning , and document understanding . Explore more of our work at Databricks AI Research , Subscribe to our blog and get the latest posts delivered to your inbox.