Databricks AI Research Team Reports 8x More Training Jobs with AI Runtime

Many experiments, very different compute needs

According to Databricks, its research team develops the models behind Databricks AI Functions, including ai_parse_document and ai_extract. The company says reaching its document parsing quality required research across different architectures, datasets, training techniques and hyperparameters, spanning tens of thousands of training and evaluation runs.

The team uses 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, it runs on NVIDIA H100 or B300 GPUs, both available through AIR. Its largest runs can need hundreds of GPUs working together over days or weeks, with high-bandwidth networking.

Databricks names three challenges:

• Resiliency: a single node or RDMA network failure can leave hundreds of GPUs idle and force a restart from an earlier checkpoint, losing hours of progress.
• Fast experimentation: testing thousands of approaches requires launching experiments without repeatedly provisioning infrastructure or rebuilding environments.
• GPU fleet management: several projects compete with different requirements, priorities and deadlines.

How AI Runtime addresses them

The research team says it has been customer zero of AI Runtime since the product's inception last year. Databricks describes several features it relies on.

GPU resiliency. AIR runs preflight checks on assigned GPU nodes to confirm they are healthy and can communicate before training starts. It also monitors nodes during execution for problems such as NCCL communication failures, thermal or power throttling and GPU driver issues. If an unhealthy node causes a job to fail, AIR cordons it off so later jobs are not scheduled there. Databricks says this matters for projects like KARL, where rollout generation and post-training of models with more than 100 billion parameters can require thousands of GPU hours.

Fast iteration. With the AIR CLI, researchers launch workloads programmatically and specify the GPUs they need, while AIR handles provisioning and scheduling, including high-bandwidth networking for distributed training across clouds. The team uses custom Docker containers for local development and the same images for training, and shares them through Databricks Artifact Registry. According to the post, AIR's work on launch latency and container optimization helps start training jobs in seconds.

Tracking and data. AIR automatically creates and manages MLflow experiments for training jobs. Research data sits in Unity Catalog, and Unity Catalog FUSE mounting lets the team access it as if it were a local filesystem, including across different Databricks workspaces.

Capacity management. Workload prioritization and preemption let the team allocate resources across projects and favor time-sensitive experiments. Databricks gives the example of spare capacity supporting extra exploratory experiments over a weekend without disrupting higher-priority work.

Agent-led training raises the load

The post describes an increasing use of agents in a workflow popularized by Andrej Karpathy, called autoresearch. Instead of researchers launching and evaluating each experiment by hand, agents propose changes, run experiments, evaluate results and decide what to try next. Databricks says the number of training jobs across its AI Functions research workloads has grown by 8x as a result.

The company argues that this changes infrastructure requirements. A slow startup or confusing configuration error may be manageable when a researcher launches a handful of experiments, but becomes a bottleneck when agents launch hundreds. Agent-driven research needs fast job launches, programmatic interfaces, clear and actionable errors, reliable execution and accessible results.

The autoresearch cycle has four steps: agents propose improvements, launch multiple trials in parallel on AIR, evaluate training metrics, and update the plan with new experiments. Through the Databricks CLI and databricks-air-skills, agents can launch concurrent workloads, monitor job status, retrieve MLflow metrics and iterate on results.

What it means for teams training models

The post is a company account of its own internal use, so the 8x figure and the speed claims come from Databricks and are not independently verified. It does show how the company positions AI Runtime: as GPU infrastructure for teams that train or fine-tune their own models, with reliability, governed data access and programmatic interfaces that suit agent-driven experimentation.

Databricks says AIR lets its researchers spend less time managing infrastructure and more time developing better models. It points readers who train or fine-tune models to the AI Runtime documentation to get started.