TL;DR
Databricks AI Runtime, launched March 2026, is serverless NVIDIA GPU infrastructure built into the Databricks platform so ML teams can attach on-demand A10 and H100 GPUs in a few clicks and pay only for usage, instead of spending days provisioning clusters. This guide covers how it works, what it supports, and where it fits against alternatives.
GPU infrastructure has been the bottleneck for enterprise ML teams for years. Provisioning clusters takes days. Configuring distributed training environments takes weeks. Debugging data loading bottlenecks consumes engineering time that should be going into modeling. The result is that deep learning teams spend more time fighting infrastructure than building the models their organizations need.
Databricks launched AI Runtime in March 2026 as its answer to this problem. Built on serverless NVIDIA GPUs and integrated across the full Databricks platform, AI Runtime lets ML teams attach on-demand A10 and H100 GPUs in two to three clicks, pay only for what they use, and work within the same governance and orchestration layer their data engineering teams already run. Hundreds of customers, including Rivian, FactSet, and YipitData, have used it to train and deploy deep learning models into production.
In this article, we cover what Databricks AI Runtime is, how it works, what it supports, how to get started, where it fits against alternatives, and what teams should know before deploying it at scale.
Key Takeaways Databricks AI Runtime is a serverless GPU training platform offering on-demand NVIDIA A10 and H100 GPUs with no cluster configuration, driver management, or idle time costs AI Runtime is purpose-built for deep learning: LLM training, fine-tuning, computer vision, recommendation systems, and reinforcement learning workflows It integrates natively with MLflow for experiment tracking, Unity Catalog for data governance , Lakeflow for job orchestration, and Genie Code for ML-specific agent assistance The platform is the same infrastructure Databricks used to train its own foundation models including DBRX and KARL Multi-node distributed training became generally available at Data + AI Summit 2026, extending AI Runtime from single-node GPU work to large-scale parallel training The AI Runtime CLI (air) is now in Beta, enabling programmatic access and CI/CD integration alongside the notebook interface
Build Smarter AI Workflows with Databricks AI Runtime. Kanerika Enables Enterprises to Optimize AI Development from Training to Production.
Book a Meeting
What is Databricks AI Runtime? Databricks AI Runtime, also referred to as AIR, is a managed deep learning compute platform embedded in the Databricks workspace. It provides on-demand access to serverless NVIDIA GPU infrastructure for training, fine-tuning, and running inference on complex AI models, without requiring teams to provision clusters, install drivers, or manage GPU scheduling themselves.
AI Runtime sits above the Databricks Serverless GPU Compute layer. When a notebook connects to serverless GPU compute, the AI Runtime environment provides the pre-configured software stack including PyTorch, Transformers, and CUDA, optimized for distributed training and fast data loading from the Lakehouse.
Before AI Runtime, deep learning teams on Databricks had two options, both with real costs:
Classic GPU clusters: Required cluster configuration, driver selection, autoscaling policy management, and paying for idle time between training runsExternal GPU platforms: Moving work outside Databricks broke data lineage, removed Unity Catalog governance, and created fragmentation between where data lived and where models trained
Databricks AI Runtime closes both gaps. Teams get GPU access in seconds, pay only for the compute they consume, and stay within the same workspace where their data, pipelines, and experiment tracking already live. The governance model, access controls, and lineage tracking from Unity Catalog apply to AI Runtime workloads the same way they apply to everything else on the platform.
The End-to-End AI Workflow with Databricks AI Runtime AI Runtime is accessed through Databricks Notebooks by selecting Serverless GPU from the compute selector. From there, teams choose their GPU type, either A10 or H100, and select an environment: a minimal base environment for teams that want full control over their dependency stack, or the full-featured AI environment pre-loaded with PyTorch, Transformers, and other popular ML frameworks.
Serverless NVIDIA GPUs: A10 and H100 The two GPU options in AI Runtime serve different use cases:
NVIDIA A10 GPUs (24GB VRAM): Suited for smaller model fine-tuning, computer vision training, recommendation system development, and experimentation where H100 memory capacity is not required. A10s are the standard choice for iterative development and debugging workflowsNVIDIA H100 GPUs (80GB HBM3 VRAM): Suited for large model training, multi-billion parameter fine-tuning, and workloads where high throughput and large GPU memory are both required. H100s are available in 8-GPU configurations on a single node , enabling substantial single-node training capacity before moving to multi-node distributed setups
Both GPU types connect through the same serverless interface. Teams do not manage driver versions, CUDA compatibility, or instance provisioning. Databricks handles that layer entirely.
Multi-Node Distributed Training At Data + AI Summit 2026, Databricks announced that AI Runtime now supports high-performance multi-node training . This extends AI Runtime’s capabilities from single-node GPU work, which was the initial preview scope, to large-scale parallel training across multiple nodes. Multi-node support is what makes AI Runtime viable for teams training foundation models, running large-scale LLM fine-tuning, or working with model architectures that exceed single-node GPU memory capacity.
The distributed training API follows a straightforward pattern:
python
from serverless_gpu import distributed
from serverless_gpu import runtime as rt
@distributed(
gpus=8,
gpu_type='h100',
)
def train_model(config: dict) -> dict:
if rt.get_local_rank() == 0:
print('Training started on rank 0')
# Training logic here
return {"rank": rt.get_global_rank()}
result = train_model.distributed(config)The @distributed decorator handles node coordination, rank assignment, and result aggregation. Teams write training logic once and the runtime scales it across the configured number of GPUs and nodes without additional infrastructure code.
The AI Runtime Environments AI Runtime offers two managed Python environments:
Minimal base environment: A clean Python environment with CUDA and basic GPU support pre-installed. Teams add their own dependencies, giving full control over the software stack without inheriting framework versions they do not needDatabricks AI environment (AI v5): A full-featured environment pre-loaded with PyTorch, Transformers, Accelerate, DeepSpeed, Ray, and other common deep learning libraries. The correct approach for teams that want to start training immediately without environment setup time
Both environments handle package caching, custom module imports, and version compatibility within the serverless GPU context. The AI v5 environment is what Databricks recommends for most production training use cases.
Core Capabilities of Databricks AI Runtime 1. Deep Learning Model Training AI Runtime supports the full range of deep learning training use cases:
LLM pre-training and fine-tuning: Training or fine-tuning large language models using PyTorch FSDP, DeepSpeed ZeRO, or QLoRA for memory-efficient fine-tuning on A10 instances. The platform supports quantization-based approaches like Unsloth for fitting large models within single-GPU memory constraintsComputer vision: Training convolutional neural networks , vision transformers, and object detection models with standard computer vision frameworksRecommendation systems: Training embedding-based and deep learning recommendation models at scale, a common enterprise use case across retail, media, and financial servicesReinforcement learning: RL workloads including the reinforcement learning from human feedback patterns used in LLM alignment. Databricks AI Research used AI Runtime for RL training on the KARL modelForecasting models: Deep learning-based time series forecasting for demand forecasting, financial prediction, and operations planning
2. MLflow Integration for Experiment Tracking Every AI Runtime training run integrates natively with MLflow for experiment tracking. Runs, artifacts, hyperparameters, metrics, and model lineage are all logged automatically. Teams can compare runs across experiments, register successful models to Unity Catalog, and deploy them through Databricks Model Serving without leaving the workspace.
At Data + AI Summit 2026, Databricks expanded Genie Code’s ML capabilities to include reading MLflow experimentation data, allowing the ML agent to answer questions grounded in the team’s real past experiments rather than generic documentation. That integration lets a team ask Genie Code why a previous training run underperformed, and it can retrieve the actual MLflow data to diagnose the issue.
3. Lakeflow Orchestration for Production Jobs For production training pipelines, AI Runtime integrates with Lakeflow Jobs and Declarative Automation Bundles (DABs) for scheduling, dependency management, and CI/CD workflows. The maximum runtime for a single AI Runtime workload is seven days, which covers most training runs. For workloads that exceed this limit, Databricks recommends implementing checkpointing and restarting the job to continue from the last saved state.
The production pattern most teams use:
Develop and debug training code interactively in notebooks connected to serverless GPU Parameterize the training job and validate it on a small data subset Submit the production training job through Lakeflow with the full dataset and configured resources Track results through MLflow and promote successful models through Unity Catalog to Model Serving
4. Unity Catalog Governance All AI Runtime workloads operate within the same Unity Catalog governance model as the rest of the Databricks platform. Data access for training follows existing permissions. Models trained and registered through MLflow carry lineage back to the training data and experiment. Access controls apply to who can trigger training jobs, who can access model artifacts, and who can deploy models to serving endpoints.
This is the governance integration that most external GPU training platforms lack. When a model trained on AI Runtime is deployed, the full lineage from raw data through feature engineering to the trained artifact is traceable through Unity Catalog. For regulated industries where model provenance is a compliance requirement, that lineage is not optional infrastructure.
5. Genie Code Integration for ML Engineering When a notebook connects to AI Runtime serverless GPU compute, Genie Code switches to an ML-optimized mode. It gains awareness of the GPU context and can help with:
Diagnosing model convergence issues and suggesting architectural adjustments Resolving CUDA errors, NCCL communication failures, and distributed training configuration problems Optimizing data loading pipelines to reduce GPU idle time from I/O bottlenecks Generating distributed training code patterns for specific model architectures Reviewing MLflow experiment data to explain why a particular run underperformed
The ML-specific intelligence is what distinguishes Genie Code’s behavior in AI Runtime from its general workspace behavior. It has access to NVIDIA-specific knowledge around CUDA, NCCL, and common deep learning framework behaviors that general-purpose coding agents lack.
Common Enterprise Use Cases for Databricks AI Runtime Databricks AI Research used AI Runtime to train DBRX and KARL, Databricks’ own foundation models. That makes AI Runtime the same infrastructure currently handling frontier model development, not a simplified consumer tier built separately.
Externally, the customer use cases across the hundreds of organizations in preview span several patterns:
Rivian and Volkswagen Group Technologies used AI Runtime to train advanced audio and multimodal models for automotive applications without infrastructure overhead, according to Arjuna Siva, VP of Infotainment and ConnectivityFactSet uses AI Runtime for financial data model training, where the governance integration with Unity Catalog supports compliance with financial data handling requirementsYipitData uses AI Runtime for training models on alternative data sources, where training speed and experiment iteration velocity are the primary performance drivers
The common thread across these use cases is that GPU compute was previously managed separately from the data platform, and the overhead of maintaining that separation was consuming engineering capacity. AI Runtime collapses the separation.
AI Runtime vs Classic GPU Clusters on Databricks Teams already on Databricks who are evaluating AI Runtime against their existing classic GPU cluster workflows will find the tradeoffs are real on both sides.
Dimension AI Runtime (Serverless GPU) Classic GPU Clusters Setup time 2 to 3 clicks, seconds to attach Cluster configuration, minutes to hours Cost model Pay per use, no idle charges Pay for cluster uptime including idle Driver management Fully managed Team-managed Distributed training Native multi-node support Custom configuration required Max job duration 7 days Cluster lifetime CLI access Beta (air CLI) Full CLI and API support Genie Code integration Optimized for ML tasks Standard workspace Genie Code Custom instances A10 and H100 only Full cloud GPU instance catalog
The primary reason to stay on classic GPU clusters is access to GPU instance types outside the A10 and H100 catalog. Teams with specific workloads requiring V100s, A100s, or specialized instances available on their cloud provider but not yet in the AI Runtime catalog may need classic clusters for those specific cases. For most standard deep learning workloads, AI Runtime is the lower-overhead path.
AI Runtime vs External GPU Training Platforms The more relevant comparison for teams considering AI Runtime is against external GPU training platforms: AWS SageMaker Training, Google Cloud Vertex AI, and standalone GPU cloud providers.
The core advantage of AI Runtime over external platforms is integration depth for teams already on Databricks. Training on SageMaker or Vertex AI requires moving data out of the Lakehouse, configuring separate governance, managing two experiment tracking systems, and stitching together orchestration across platform boundaries. AI Runtime eliminates all of those handoffs.
The advantage of external platforms over AI Runtime is GPU catalog breadth and compute scale for very large training runs. For teams training at the scale of hundreds of GPU nodes over weeks, dedicated training infrastructure on AWS or Google Cloud may provide more flexible capacity options than the current AI Runtime preview supports.
Dimension AI Runtime AWS SageMaker Training Google Vertex AI Training Data access Native Lakehouse, Unity Catalog S3 native, requires data movement from Databricks GCS native, requires data movement Governance Unity Catalog lineage and access controls IAM-based, separate from data governance IAM-based, separate from data governance Experiment tracking MLflow native SageMaker Experiments Vertex AI Experiments Orchestration Lakeflow native SageMaker Pipelines Vertex AI Pipelines Agent assistance Genie Code ML integration CodeWhisperer Gemini Code Assist Pricing model DBU consumption Per-instance, per-hour Per-instance, per-hour Setup overhead Low, within existing workspace Medium, separate service setup Medium, separate service setup
Getting Started With AI Runtime Step 1: Enable Serverless GPU Compute Serverless GPU Compute is enabled through workspace admin settings. Once enabled, it becomes available to all workspace users with the appropriate permissions. Unity Catalog permissions govern which data sources training jobs can access, consistent with the broader workspace governance model.
Step 2: Connect a Notebook to Serverless GPU From any notebook, select Serverless GPU from the compute selector. Choose the GPU type: A10 for smaller workloads and experimentation, H100 for larger models and high-throughput training. Select the AI v5 environment if using standard frameworks, or the minimal base environment if managing your own dependency stack.
python
# Verify GPU connection
import subprocess
result = subprocess.run(['nvidia-smi'], capture_output=True, text=True)
print(result.stdout)Step 3: Mount Data From the Lakehouse AI Runtime reads training data directly from Unity Catalog tables and volumes. For large datasets, Databricks recommends storing data in Delta Lake or Parquet format and using the platform’s optimized data loading utilities for efficient GPU feeding:
python
from databricks.sdk.runtime import *
import torch
from torch.utils.data import DataLoader
# Load data from Unity Catalog
df = spark.table("my_catalog.training_data.features")
pandas_df = df.toPandas()
# Convert to PyTorch dataset
dataset = YourDataset(pandas_df)
loader = DataLoader(dataset, batch_size=32, num_workers=4)
Step 4: Configure MLflow Tracking MLflow tracking is available automatically in the AI Runtime environment. Start a run before training and log metrics, parameters, and artifacts throughout:
python
import mlflow
mlflow.set_experiment("/experiments/my_model_training")
with mlflow.start_run() as run:
mlflow.log_params({"learning_rate": 1e-4, "epochs": 10, "model": "bert-base"})
for epoch in range(num_epochs):
train_loss = train_epoch(model, loader, optimizer)
mlflow.log_metric("train_loss", train_loss, step=epoch)
mlflow.pytorch.log_model(model, "model")Step 5: Submit Production Jobs Through Lakeflow Once the training logic is validated interactively, submit it as a Lakeflow job for scheduled or event-triggered execution:
python
# Using the Databricks SDK to submit an AI Runtime job
from databricks.sdk import WorkspaceClient
w = WorkspaceClient()
job = w.jobs.create(
name="production_training_run",
tasks=[{
"task_key": "train",
"notebook_task": {
"notebook_path": "/Workspace/training/train_model",
},
"new_cluster": {
"spark_version": "ai-runtime:latest",
"num_workers": 0,
"custom_tags": {"gpu_type": "h100"}
}
}]
)Fine-Tuning LLMs on AI Runtime Fine-tuning is the most common GPU use case in 2026. Most teams are adapting existing open-source models, Llama, Mistral, Qwen, to their domain data rather than training from scratch. A10 GPUs with QLoRA quantization handle models up to 20B parameters on a single node. H100s extend that to 70B and beyond.
Standard single-node fine-tuning on an A10 using Unsloth for memory efficiency:
python
from unsloth import FastLanguageModel
import mlflow
model, tokenizer = FastLanguageModel.from_pretrained(
model_name="unsloth/Llama-3.1-8B-Instruct",
max_seq_length=2048,
load_in_4bit=True,
)
model = FastLanguageModel.get_peft_model(
model,
r=16,
target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
lora_alpha=16,
)
mlflow.set_experiment("/experiments/llama_finetune")
with mlflow.start_run():
mlflow.log_params({"model": "Llama-3.1-8B", "quantization": "4bit", "lora_rank": 16})
trainer = SFTTrainer(model=model, tokenizer=tokenizer, train_dataset=dataset)
trainer.train()
mlflow.pytorch.log_model(model, "finetuned_model")For 70B models, use the 8xH100 configuration with the @distributed decorator to shard across GPUs. Three things to do in any fine-tuning run on AI Runtime: implement checkpointing every 500 to 1000 steps given the 7-day job limit, store checkpoints to a Unity Catalog volume so they persist across runs, and register final adapter weights to Unity Catalog via MLflow so lineage from training data to deployed model is captured.
Governance and Security in AI Runtime 1. Unity Catalog Permissions AI Runtime workloads respect the same Unity Catalog permissions as all other Databricks operations. A training job running under a given user’s credentials can only access data that user is authorized to read. Model artifacts registered to Unity Catalog inherit the governance model of the catalog they are registered to, including access controls, lineage tracking, and sensitivity labels.
2. Model Lineage Every model trained through AI Runtime and registered via MLflow carries full lineage back to its training data. The lineage record captures which Unity Catalog tables contributed to the training dataset, which experiment produced the registered model, and which version of the training code was used. For organizations operating under regulations that require model provenance documentation, this lineage is captured automatically rather than requiring manual documentation.
3. Data Residency AI Runtime is a designated service with region-specific processing rules. Teams with data residency requirements should verify that their workspace region supports AI Runtime before enabling it, and should confirm that GPU workloads process within the required geographic boundaries.
Databricks AI Runtime Pricing AI Runtime operates on a consumption-based model. Teams pay for the GPU DBUs consumed during training runs and do not pay for idle time between runs. The specific DBU rates for A10 and H100 compute are published on the Databricks pricing page and vary by cloud provider and region.
Databricks is also bringing serverless NVIDIA GPUs to its Free Edition, broadening access for developers, students, and startups who want to experiment with AI Runtime without an enterprise contract. For production deployments, Unity AI Gateway provides spend controls including hard caps, rate limits, and workload routing rules that allow organizations to set and enforce GPU spend budgets before training workloads run.
How Kanerika Works With Databricks AI Runtime Kanerika is a Databricks Consulting Partner with hands-on experience deploying ML workflows and AI agents into governed, production-grade Databricks environments. Our data engineering and AI/ML services cover the full stack from Unity Catalog architecture through model training, serving, and monitoring.
For organizations evaluating AI Runtime, Kanerika helps teams with three areas where most deployments get stuck:
Data foundation setup: Structuring Unity Catalog tables and Delta Lake pipelines to enable efficient data loading into AI Runtime training jobs, which is where most teams find performance bottlenecks before model architecture becomes the variableMLflow and governance configuration: Setting up experiment tracking, model registry, and Unity Catalog lineage so training runs produce auditable, reproducible artifacts that meet enterprise governance requirements from day oneProduction pipeline design: Moving from interactive notebook development to Lakeflow-orchestrated production training pipelines with monitoring, alerting, and cost controls configured before the first production run
A manufacturing client needed real-time inventory intelligence across distributed operations. Manual data analysis was consuming team capacity, insights were delayed, and inventory reconciliation cycles were adding overhead to every planning cycle. The team needed a solution that could surface patterns and anomalies from large operational datasets without requiring SQL expertise from the business users who needed the insights.
Challenge Inventory and operational data was spread across multiple systems with no unified layer for real-time querying. Operations teams were waiting on reports that needed to inform same-day decisions, and ad hoc analysis requests to the data team were creating a backlog that slowed operational response times.
Solution Kanerika deployed Karl , its AI data insights agent built natively on Databricks , trained on the client’s inventory and operational data using the Databricks ML platform. Karl surfaces inventory patterns, flags anomalies, and delivers business insights directly to operations teams in plain language. The deployment connected to the client’s existing data environment through Unity Catalog without requiring restructuring, and the model training workflow used Databricks’ managed ML infrastructure to keep governance consistent with the rest of the data estate.
Results 30% faster inventory reconciliation across distributed operations 50% faster reporting cycles for operations and planning teams 28% improvement in customer satisfaction driven by faster fulfillment decisions
Wrapping Up Databricks AI Runtime removes the infrastructure layer that has historically sat between ML teams and productive deep learning work. Serverless NVIDIA GPUs, native integration with MLflow and Unity Catalog, multi-node distributed training, and Genie Code ML assistance combine into a platform where teams spend their time on modeling rather than infrastructure management.
The tradeoffs are real: AI Runtime is workspace-bound, the GPU catalog is currently limited to A10 and H100, and the CLI is still in Beta. For teams whose deep learning work is primarily conducted within the Databricks platform, those constraints are worth accepting for the integration depth that AI Runtime provides.
For teams evaluating how to build or modernize their ML infrastructure on Databricks, Kanerika’s AI Maturity Assessment maps current data infrastructure against the requirements for AI Runtime deployment and produces a sequenced implementation plan. Talk to our team to discuss how AI Runtime fits your specific ML workflows and data environment.
Ready to Scale Enterprise AI Faster? Partner with Kanerika to Modernize AI Infrastructure and Accelerate Enterprise Innovation.
Book a Meeting
FAQs 1. What is Databricks AI Runtime? Databricks AI Runtime is an optimized runtime environment designed to simplify AI and machine learning development on the Databricks Data Intelligence Platform. It comes with pre-installed libraries, frameworks, and performance optimizations that help data scientists and engineers build, train, and deploy AI models more efficiently. By reducing manual setup and configuration, teams can focus on developing AI solutions instead of managing infrastructure.
2. How is Databricks AI Runtime different from the standard Databricks Runtime? Unlike the standard Databricks Runtime, Databricks AI Runtime includes built-in support for popular AI and machine learning frameworks such as PyTorch, TensorFlow, Hugging Face Transformers, and XGBoost. It is specifically optimized for AI workloads, providing better performance for model training, fine-tuning, and inference while reducing the need to install and maintain additional libraries.
3. What are the key benefits of Databricks AI Runtime? Databricks AI Runtime accelerates AI development by offering optimized compute performance, pre-configured machine learning libraries, GPU support, and seamless integration with MLflow. These capabilities help reduce development time, improve model performance, simplify collaboration, and enable faster deployment of AI applications in enterprise environments.
4. Which AI and machine learning frameworks are supported by Databricks AI Runtime? Databricks AI Runtime supports a wide range of widely used AI and machine learning frameworks, including TensorFlow, PyTorch, Scikit-learn, Hugging Face Transformers, XGBoost, LightGBM, and MLflow. This broad ecosystem allows developers to build, train, evaluate, and deploy models using the tools they are already familiar with while benefiting from Databricks’ performance optimizations.
5. Can Databricks AI Runtime be used for generative AI and large language models? Yes. Databricks AI Runtime is designed to support modern generative AI workloads, including fine-tuning, evaluating, and serving large language models (LLMs). It provides optimized GPU environments, integrates with open-source AI frameworks, and works alongside Databricks tools for vector search, model serving, and AI governance , making it suitable for enterprise generative AI applications.
6. How does Databricks AI Runtime improve machine learning performance? Databricks AI Runtime improves performance through optimized compute resources, pre-configured libraries, distributed processing, and efficient GPU utilization. These enhancements reduce training times, accelerate experimentation, and improve resource efficiency, enabling organizations to develop and deploy machine learning models faster while managing infrastructure more effectively.
7. Who should use Databricks AI Runtime? Databricks AI Runtime is designed for data scientists, machine learning engineers, data engineers, and AI development teams working on enterprise-scale analytics and AI initiatives. It is particularly valuable for organizations building predictive models, recommendation systems, generative AI applications, and large-scale machine learning pipelines that require high performance and simplified workflows.
8. Is Databricks AI Runtime suitable for enterprise AI projects? Yes. Databricks AI Runtime is built to support enterprise AI initiatives by combining scalability, security, governance, and high-performance computing in a single platform. It integrates with Unity Catalog, MLflow, and other Databricks services, allowing organizations to manage AI development, model lifecycle, and production deployment while maintaining governance and compliance standards.
