AI Runtime and serverless GPU compute

What it is

AI Runtime is serverless GPU compute for training and fine-tuning. You ask for a number of accelerators of a given type, Databricks provisions them, runs your Python on them, and releases them when the work finishes. There is no cluster to size, no driver to keep alive and no GPU quota sitting idle between experiments.

Three accelerator shapes exist: 1xA10 (one GPU, 24 GB) for small and medium jobs, 1xH100 (one GPU, 80 GB) for medium-scale work, and 8xH100 (eight 80 GB GPUs on a single node), which is the only shape that supports distributed training. The workspace has to be in us-west-2, us-west-1, us-east-1, us-east-2, ca-central-1 or sa-east-1.

Two entry points lead to the same compute. From a notebook you decorate a function and call it; from a terminal you write a YAML file and submit it with a CLI. Both produce MLflow runs, so a training job’s record lands exactly where MLflow tracking on Databricks already puts everything else.

Why it exists

This is where fine-tuning on Databricks now lives, because the previous answer is gone. Foundation Model Fine-tuning has reached end of life and is no longer supported. The databricks_genai package and the Foundation Model Fine-tuning UI are both no longer available, and the documentation for them now points at AI Runtime. Anyone whose fine-tuning pipeline imported databricks_genai is rewriting it, not maintaining it.

The older path was also narrower than people needed. It fine-tuned a fixed catalogue of base models through an API that accepted a task type and a training table. Anything outside that shape, a custom loss, a vision encoder, a reinforcement-learning loop, meant provisioning a GPU cluster by hand, installing CUDA-dependent wheels, and paying for the cluster while you read the stack trace. AI Runtime replaces both halves with one primitive: arbitrary Python on GPUs you do not manage.

How it works

Two managed environments

The Standard environment is minimal and leaves dependency choices to you. The Databricks AI environment is preloaded with the usual deep-learning stack, including PyTorch and Transformers. Either way you add pip packages declaratively, and dependencies are installed with uv, so pip-style flags such as --no-build-isolation and --no-deps work as list entries, which is how a compiled CUDA extension such as flash-attn can be built against the installed torch. If a dependency needs system libraries, you point the workload at a custom Docker image instead. A web terminal is available on the compute, so nvidia-smi works when you want to see what the GPUs are actually doing. Ray is supported for distributed workloads that are not plain PyTorch.

The notebook path: @distributed

from databricks.air import distributed gives you a decorator. On GPU environment version 5 and below the same decorator comes from the older serverless_gpu package, with the old argument names gpus and gpu_type. You wrap a training function, then submit it by calling .distributed() on the decorated function with the function’s own arguments:

Everything the function needs has to be defined inside it, data loading included. The arguments are pickled to reach the workers, and a dataset larger than pickle allows is the usual first failure. Each call creates an MLflow run, or a nested child run when one is already active, and prints a link to it.

The CLI path: databricks air

databricks air is the same compute driven from a laptop or an IDE, with the job description checked into Git rather than living in a notebook cell. It is a command group of the Databricks CLI, so there is nothing separate to install: it replaces the Python-based air tool from the databricks-air package, and old scripts need their commands renamed, not just a databricks prefix.

databricks auth login --host https://<workspace-url>
databricks air run -h config    # the workload YAML fields your installed version accepts

Six commands cover the lifecycle: databricks air run submits a workload YAML (--dry-run validates it locally without submitting), databricks air get <run-id> shows one run’s status and configuration, databricks air list shows your active runs (--all-status adds finished ones), databricks air logs streams or downloads output (--node for one worker, --retry for one attempt, --download-to for a file), databricks air cancel stops a run, and databricks air convert-to-dabs turns a workload YAML into a Declarative Automation Bundle. Every command takes -p for a Databricks CLI profile.

The workload YAML requires experiment_name, a compute block and a command. compute.accelerator_type is one of GPU_1xA10, GPU_1xH100, GPU_8xH100 or, in Public Preview, GPU_8xB300, and num_accelerators has to agree with it: exactly 1 for GPU_1xH100, a multiple of 8 for GPU_8xH100 and GPU_8xB300, any positive integer for GPU_1xA10. The optional blocks are where the useful details sit: environment (a version and a dependencies list, or instead a unity_catalog_image for a custom Docker image), code_source (a snapshot of a local directory, optionally pinned to a Git branch or commit, with include_paths to keep the upload small), parameters for structured hyperparameters, env_variables, a top-level secrets map referenced as scope/key so no token is written into the file (see Secrets and credentials), plus max_retries, timeout_minutes and usage_policy_name for a budget policy.

Connection limits, which are not runtime limits

An interactive notebook connection times out after 15 minutes of inactivity, while notebook jobs and CLI jobs hold theirs for 24 hours. The work itself runs longer: up to seven days for a connected notebook session or a notebook job, and up to 14 days for a CLI job. Losing the notebook connection is therefore not the same thing as losing the run.

Example: the same fine-tune from a notebook and from a terminal

In a notebook, on all eight GPUs of one node:

from databricks.air import distributed
import os, torch, torch.distributed as dist

@distributed(num_accelerators=8, accelerator_type="GPU_8xH100", timeout=7200)
def run_train(num_epochs: int, batch_size: int) -> None:
    import mlflow
    from torch.nn.parallel import DistributedDataParallel as DDP
    from torch.utils.data import DataLoader, DistributedSampler

    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
    dist.init_process_group("nccl")
    device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")

    dataset = load_dataset_from_volume("/Volumes/main/ml/training/support_tickets")
    sampler = DistributedSampler(dataset)
    loader = DataLoader(dataset, sampler=sampler, batch_size=batch_size)
    model = DDP(build_model().to(device), device_ids=[device])
    optimiser = torch.optim.AdamW(model.parameters(), lr=2e-5)

    for epoch in range(num_epochs):
        sampler.set_epoch(epoch)
        for step, (xb, yb) in enumerate(loader):
            loss = model(xb.to(device), labels=yb.to(device)).loss
            mlflow.log_metric("loss", loss.item(), step=step)  # lands in the run this call created
            loss.backward()
            optimiser.step()
            optimiser.zero_grad()

    dist.destroy_process_group()

run_train.distributed(num_epochs=3, batch_size=8)

The same job as a checked-in workload, train.yaml:

experiment_name: support-ticket-finetune
mlflow_experiment_directory: /Workspace/Shared/training-experiments
mlflow_artifact_location: /Volumes/main/ml/mlflow-artifacts/finetune
compute:
  num_accelerators: 8
  accelerator_type: GPU_8xH100
environment:
  version: 'databricks_ai_v5'
  dependencies:
    - transformers>=4.30
secrets:
  HF_TOKEN: 'ml_team/huggingface_token'
code_source:
  type: snapshot
  snapshot:
    root_path: /Users/me/support-finetune
    git:
      branch: main
    include_paths:
      - src
command: torchrun --nproc_per_node=8 $CODE_SOURCE_PATH/src/train.py
max_retries: 1
timeout_minutes: 180
usage_policy_name: ml_team_policy

databricks air run -f train.yaml -p prod --watch
databricks air logs <run-id> --download-to ./logs/

Either way the resulting model gets registered in Models in Unity Catalog and served from a serving endpoint, which is unchanged from before.

Common mistakes

• Loading the dataset outside the decorated function. The arguments are pickled to reach the eight workers, so a DataFrame passed in from the notebook either fails on size or silently costs you a serialisation round trip. Load inside the function.
• Asking for distributed training on 1xH100. Multi-GPU data parallelism needs an eight-GPU shape, 8xH100 or 8xB300. num_accelerators also has to match the shape: a multiple of 8 for GPU_8xH100, exactly 1 for GPU_1xH100.
• Treating a lost notebook connection as a lost run. The interactive connection drops after 15 minutes of inactivity, while the job keeps going for up to seven days. Use databricks air list or the MLflow run link rather than restarting.
• Porting a databricks_genai pipeline by changing imports. Foundation Model Fine-tuning is at end of life, not deprecated-but-working, and the package is gone. The training loop has to be written as ordinary PyTorch.
• Putting an API token in the workload YAML. secrets takes scope/key references precisely so the file can live in Git.
• Reading Public Preview as almost-GA. The notebook decorator is a step behind the runtime at Beta, and the whole feature needs a workspace admin to enable it. Prototype on it, do not promise on it.