Overview

AI Runtime is serverless GPU compute on Azure Databricks for deep learning workloads. This page covers how it works, example use cases, hardware, key features, and limitations.

How AI Runtime works

AI Runtime lets you train and fine-tune models on Azure Databricks without managing GPU infrastructure. You bring your code, and Azure Databricks handles GPU capacity, procurement, and management.

For example, the following command submits a workload from your terminal:

databricks air run --file workload.yaml --watch

When you run databricks air run, AI Runtime packages the code you submit and orchestrates it on serverless GPU compute managed by Azure Databricks. To try it, see AI Runtime quickstart.

Example use cases

You can use AI Runtime to train or fine-tune many types of AI models. For example:

• LLM fine-tuning (LoRA, QLoRA, full fine-tuning). See Large language models (LLMs).
• Computer vision (object detection, image classification). See Computer vision.
• Deep-learning-based recommender systems. See Deep learning based recommender systems.
• GPU-accelerated classic ML and deep-learning-based time series forecasting. See Classic machine learning.
• Large-scale batch inference. See Batch inference.
• Multi-GPU distributed training. See Multi-GPU distributed training.
• Reinforcement learning. See Large language models (LLMs).

For all notebook tutorials and Databricks CLI recipes, see AI Runtime examples.

Hardware

All AI Runtime accelerators provision a single node. The number of GPUs on that node depends on the accelerator type:

Accelerator	GPUs per node	GPU memory	Best for	Distributed training
1xA10	1	24 GB	Small to medium ML and deep learning tasks such as classic ML models or fine-tuning smaller language models	Not supported (single GPU)
1xH100	1	80 GB	Workloads that need more GPU memory than 1xA10 provides but do not require multi-GPU distributed training, such as fine-tuning medium-size language models	Not supported (single GPU)
8xH100	8	80 GB per GPU	Large-scale AI workloads including training or fine-tuning massive models or running advanced deep learning tasks	Supported using the @distributed decorator, which uses all eight GPUs by default

Tip

Only 8xH100 supports multi-GPU distributed training. For single-GPU workloads, choose 1xA10 or 1xH100 based on the GPU memory your model needs.

Key features

Category	Features
Connect to GPUs	Databricks CLI: Submit and manage GPU workloads from your terminal with YAML configs. Notebooks: Attach a notebook to a GPU and develop interactively. Distributed training in notebooks: Run PyTorch DDP, FSDP, or DeepSpeed across all GPUs on a node with the @distributed decorator. Jobs: Schedule a GPU workload as a recurring job. IDE over SSH: Connect from your IDE or terminal to a GPU node. Ray: Run Ray Core, Ray Data, Ray Train, and Ray Tune.
Develop workloads	Pre-built environments: Choose the minimal Standard environment or the Databricks AI environment with PyTorch and Transformers preinstalled. Custom Docker images: Bring your own image, and push it to the artifact registry. Load and save data: Stream data from Unity Catalog volumes and tables. Coding agents: Use Genie Code and agent skills to write, submit, and debug GPU workloads.
Monitor and deploy	Track runs, logs, and metrics: Track experiments with MLflow, view logs, and monitor GPU usage. Schedule GPU workloads: Deploy with Declarative Automation Bundles and build multi-task GPU and CPU jobs. Convert a workload to a bundle: Turn a workload YAML into a Declarative Automation Bundles job. Configure bundle tasks: Package code, select GPU compute, and pass configuration to tasks.

Requirements

Before you start, confirm your workspace meets these requirements:

• Check AI Runtime availability by Geo for supported Geos and cross-Geo processing requirements. GPU availability varies by workspace.
• If your workspace has the compliance security profile enabled, a workspace admin must enable the AI Runtime and AI Runtime multi-node H100 previews. See Manage Databricks previews. For details about default enablement for these workspaces, see the advance notice.

Limitations

Keep the following limitations in mind when you plan an AI Runtime workload:

• AI Runtime only supports A10 and H100 accelerators.

• AI Runtime doesn't support the compliance security profile for the FedRAMP High or DoD IL5 standards.
• Adding dependencies using the Environments panel is not supported for AI Runtime scheduled jobs. Install dependencies programmatically using %pip install in your notebook instead.
• For scheduled jobs on AI Runtime, auto recovery behavior for incompatible package versions that are associated with your notebook is not supported.
• AI Runtime connection attempts terminate after 15 minutes for interactive notebooks and after 24 hours for notebook jobs or CLI jobs. Connected notebook sessions and notebook jobs can run for up to two days, and CLI jobs can run for up to 14 days. For long-running model training jobs, implement checkpointing and restart the job when the maximum runtime is reached. See Checkpoint with Distributed Checkpoint (DCP).
• Running a scheduled job that uses a workspace base environment, including the Databricks AI environment, is in Beta. A workspace admin must turn on the Serverless workspace base environment support in Jobs preview. See Manage Azure Databricks previews.
• AI Runtime provides on-demand access to GPU resources. While this leads to easy, flexible access to GPUs, there may be periods where capacity is constrained or unavailable in your region.
• AI Runtime leverages cross-region GPUs in certain cases during moments of high demand. There may be egress costs associated with such usage, and cross-region network connectivity might be limited at certain times.

Guides

• Guides for AI Runtime: Track usage and costs, and troubleshoot common errors.
• Migrate classic GPU workloads to serverless: Move a deep learning workload from a classic cluster to AI Runtime.
• Migrate from Slurm: Map Slurm batch scripts, launchers, and shared storage onto the Databricks CLI.
• Improve training performance and resiliency on AI Runtime: Speed up training and make long-running jobs resilient to failures.