Databricks AI Runtime reaches GA for training models on NVIDIA GPUs

Serverless GPUs for training and fine-tuning

According to Databricks, AIR lets developers run their existing code and containers at scale on serverless NVIDIA GPUs, while Databricks manages the underlying clusters and networking. The service supports multi-node distributed training, so larger workloads can be spread across GPU nodes. NVIDIA Quantum InfiniBand provides high-bandwidth, low-latency networking with remote direct memory access (RDMA), and the NVIDIA Collective Communications Library (NCCL) handles communication between GPUs.

Developers can prepare Lakehouse data and experiment in GPU-powered notebooks or through the AIR CLI. Standard environments include PyTorch and CUDA, and teams can use distributed training frameworks or bring their own libraries and custom training loops. For repeatable runs, Lakeflow Jobs can orchestrate training alongside data preparation and evaluation, and Declarative Automation Bundles support CI/CD for those workflows.

MLflow tracks experiments and GPU utilization. Unity Catalog governs access to training data, and native system tables show CPU and GPU usage.

NVIDIA software, containers and Nemotron models

AIR now supports Docker and private registries, so developers can bring custom containerized environments. GPU-optimized containers from the NVIDIA NGC Catalog can be used with the Databricks Artifact Registry. Environment caching is meant to cut setup time when teams rerun experiments, and NVIDIA profiling tools can be installed to look into GPU performance bottlenecks.

Databricks says that with the release of Nemotron 3.5 Lightning and Super VL, customers can load, experiment with, adapt and serve these models on the platform from day one. AIR supports customization with NVIDIA NeMo and NeMo Automodel. Databricks and NVIDIA also published Nemotron cookbooks to help teams get started.

Customization is not limited to language models. AIR includes a validated fine-tuning workflow for NVIDIA Kumo Anomaly, an open research model for time-series anomaly detection, with MLflow tracking and Unity Catalog checkpoints. Databricks says pairing it with the Genie Ontology offers a path to agents that connect unusual patterns in time-series data with business meaning.

Genie Code and the support agent example

Genie Code can help developers write and adapt training scripts in a notebook. Developers describe the changes they need, ask for help configuring a fine-tuning run and work through errors, then review the suggested code before running it. Databricks says Genie Code can also debug model code, identify GPU performance bottlenecks such as memory issues and communication overhead, and apply NVIDIA-specific guidance on CUDA, cuDNN, NCCL and profiling tools.

As an example, Databricks describes an enterprise building a technical support agent. The team could start with a Nemotron model, fine-tune it on labeled support tickets with NeMo on AIR, track results in MLflow, and then deploy it through Databricks Model Serving for real-time inference. The resulting agent would interpret incoming support requests and generate responses informed by the company's products and troubleshooting workflows.

What it means for enterprise teams

The release targets a familiar problem: getting GPU access is only part of the work, and teams also have to configure training jobs, govern data and manage infrastructure. By keeping training, evaluation and serving inside the Databricks platform, the companies aim to reduce that overhead and leave more time for evaluating models against business needs.

Hanlin Tang, CTO for Neural Networks at Databricks, said AIR "will accelerate the rise of customized models in enterprise." Pat Lee, VP of Strategic Partnerships at NVIDIA, said the companies are bringing open models like Nemotron and Kumo to governed enterprise data. Pricing and regional availability are not given in the announcement.