AI Runtime

Preview

This feature is in Public Preview.

AI Runtime is a serverless GPU compute offering at Databricks intended for deep learning workloads. You can use AI Runtime to train and fine-tune custom models using your favorite frameworks and get state-of-the-art efficiency, performance, and quality.

Get started

Run your first workload in minutes using our quickstart guides.

• • 🚀 CLI quickstart
  • Coming from a Slurm or Kubernetes cluster? Train on AI Runtime from your terminal in minutes.
• • 📓 Notebook quickstart
  • Attach your notebook to a GPU in seconds and develop interactively..
• • ⚡ Ray
  • Coming from a Ray cluster? Run Ray on AI Runtime with dashboard support.
• • 🧭 Overview
  • Learn the key things about AI Runtime in two minutes: available GPUs, how it works, and limitations.

Features

AI Runtime is a full-stack GPU platform, with many ways to connect to GPUs, built-in observability and debugging tools, and pre-built environments.

Connect to Serverless GPUs: Reach serverless GPUs from wherever you work.

• • Notebooks
  • Attach a notebook to serverless GPU compute and develop interactively, with no cluster setup.
• • IDE over SSH
  • Connect from your IDE or terminal over an SSH tunnel to work on a GPU node directly.
• • AI Runtime CLI
  • Submit and manage distributed GPU training jobs from your laptop using YAML job configs.
• • Ray
  • Run Ray Core, Ray Data, Ray Train, and Ray Tune on serverless GPU compute.

Managing Dependencies: Control the Python and system libraries your workload runs with.

• • Pre-built environments
  • Start from a minimal Standard environment or a Databricks AI environment pre-loaded with ML frameworks.
• • Docker support
  • Bring your own container image to run workloads with a custom dependency stack.

Loading data into GPUs: Feed training data to your GPUs efficiently.

• • Efficient data loaders
  • Stream data from Unity Catalog volumes with fault-tolerant loaders built for high GPU utilization.

Debugging and Observability: Track experiments, inspect output, and diagnose failures.

• • MLflow for deep learning
  • Track experiments, metrics, and runs with MLflow, and save model checkpoints to Unity Catalog volumes.
• • Logs viewer
  • View training output and monitor GPU resource usage while your code runs.
• • Debugging with agents
  • Use Genie Code to generate training code, resolve environment issues, and debug GPU failures.

Scheduling and real-time serving: Move from interactive development to scheduled jobs and endpoints.

• • Productionize your jobs
  • Deploy training code with Declarative Automation Bundles, schedule runs, and build multi-task GPU and CPU workflows.
• • Serve your models
  • Deploy your trained models behind a scalable endpoint with Model Serving.

Examples

Clone end-to-end examples and run them on AI Runtime in minutes, from the CLI or in notebooks.

• • Ray Examples
  • Take your existing Ray project and run it on AI Runtime in minutes.
• • Training a tabular recommendation model
  • Train deep-learning recommender models, such as two-tower architectures.
• • Computer vision
  • Object detection and image classification workloads on GPUs.
• • Classic machine learning
  • GPU-accelerated XGBoost, time series forecasting, and other classic ML tasks.
• • Fine-tuning an OSS LLM
  • Fine-tune open-source large language models with LoRA, QLoRA, or full fine-tuning.

Learn more

Find helpful content: the latest product updates and how AI Runtime works under the hood.

• • Product updates
  • See the latest AI Runtime product updates and announcements.
• • How AI Runtime works under the hood
  • Read how Databricks keeps GPUs reliable across AI Runtime.