Serverless GPU & AI Runtime: Running GPU Workloads Without Cluster Management on Databricks

GPU provisioning has been the biggest friction point in ML/AI on Databricks. You want to fine-tune a model? First, configure a GPU cluster. Pick an instance type. Choose a runtime version. Wait for it to start. Hope your spot instances don’t get reclaimed mid-training. Get the bill — and realize you paid for 6 hours of idle time because you forgot to terminate the cluster after your run finished.

Databricks AI Runtime (formerly Serverless GPU Compute) removes all of that. You select an accelerator, click Apply, and you’re running on GPUs. No cluster configuration. No instance selection. No idle costs. One bill that includes both compute and infrastructure.

In this article, we’ll go hands-on: run a GPU workload on a single A10, scale to 8xH100 with distributed training, fine-tune Llama 3.2 with LoRA, and compare the real costs against traditional GPU clusters. With code you can run today.

Traditional GPU Clusters vs. Serverless GPU

#

Here’s what running a GPU workload on Databricks looked like before AI Runtime:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

[You want to fine-tune a model]
      |
  Create a GPU cluster (pick instance type, driver, workers)
      |
  Wait 5-10 minutes for cluster startup
      |
  Install dependencies (pip install, conda, etc.)
      |
  Run your training job
      |
  Remember to terminate the cluster (or pay for idle time)
      |
[Done — maybe $200 in idle costs later]

With Serverless GPU:

1
2
3
4
5
6
7

[You want to fine-tune a model]
      |
  Select "Serverless GPU" → pick A10 or H100 → Apply
      |
  Run your training job
      |
[Done — you paid only for what you used]

The difference isn’t just convenience. It’s structural:

Getting Started: Your First Serverless GPU Notebook

#

Step 1: Connect to Serverless GPU

#

1. Open a Databricks notebook
2. Click the compute selector (top-right)
3. Select Serverless GPU
4. In the Environment tab, choose your accelerator:
   • A10 — 1 GPU, good for small-to-medium ML/DL tasks
   • 8xH100 — 8 GPUs, for large-scale training and fine-tuning
5. Click Apply

That’s it. No cluster creation page. No instance type matrix. No driver selection.

Step 2: Verify GPU Access

#

1

%sh nvidia-smi

You should see your GPU(s) listed with driver version, CUDA version, and memory info. If you selected A10, you’ll see one GPU. If you selected 8xH100, you’ll see eight.

Step 3: Run a Quick GPU Workload

#

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

import torch

# Verify CUDA is available
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"GPU count: {torch.cuda.device_count()}")
print(f"GPU name: {torch.cuda.get_device_name(0)}")

# Quick tensor operation on GPU
x = torch.randn(1000, 1000, device='cuda')
y = torch.randn(1000, 1000, device='cuda')
z = torch.mm(x, y)
print(f"Matrix multiply result shape: {z.shape}, device: {z.device}")

If you’re using the AI environment (recommended), PyTorch comes pre-installed. If you’re on the base environment (v5), you’ll need to install it yourself — more on that in the Environment section below.

Distributed Training on 8xH100

#

Single-GPU workloads are straightforward — just write normal PyTorch code. The real power of Serverless GPU comes when you need to scale across multiple GPUs for fine-tuning or training larger models.

Databricks provides the @distributed decorator from the serverless_gpu package. It handles rank assignment, process spawning, and communication setup — you write the training logic.

Hello World: Distributed

#

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

from serverless_gpu import distributed
from serverless_gpu import runtime as rt

@distributed(gpus=8, gpu_type='h100')
def hello_distributed(name: str) -> list[int]:
    rank = rt.get_global_rank()
    local_rank = rt.get_local_rank()
    
    if local_rank == 0:
        print(f"Hello from rank 0, {name}!")
    
    return rank

# Call with .distributed() — NOT directly
result = hello_distributed.distributed('Serverless GPU')
assert result == [0, 1, 2, 3, 4, 5, 6, 7]
print(f"All 8 GPUs reported in: {result}")

Key Rules for @distributed

#

1. gpus must be 8 for H100 (that’s the node configuration)
2. gpu_type must match the accelerator your notebook is connected to
3. A10 does NOT support distributed training — it’s single-GPU only
4. Move data loading inside the decorated function to avoid pickle serialization issues
5. Multi-node (multiple 8xH100 nodes) is currently in Private Preview

Real-World: Fine-Tuning Llama 3.2 1B with LoRA

#

Let’s do something practical — fine-tune Meta’s Llama 3.2 1B model using LoRA on 8xH100 GPUs with DeepSpeed ZeRO-3.

Install Dependencies

#

1
2
3
4
5
6

%pip install --upgrade transformers==4.56.1
%pip install peft==0.17.1
%pip install trl==0.18.1
%pip install deepspeed>=0.15.4
%pip install mlflow>=3.6.0
%restart_python

The Training Script

#

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93

import json
import os
import tempfile

from serverless_gpu import distributed

UC_CATALOG = "main"
UC_SCHEMA = "default"
CHECKPOINT_DIR = f"/Volumes/{UC_CATALOG}/{UC_SCHEMA}/checkpoints/llama3_2-1b"

@distributed(gpus=8, gpu_type='H100')
def fine_tune_llama():
    import torch
    import mlflow
    from datasets import load_dataset
    from transformers import (
        AutoTokenizer,
        AutoModelForCausalLM,
        TrainingArguments,
    )
    from trl import SFTTrainer
    from peft import LoraConfig

    # Load tokenizer and model
    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct")
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        "meta-llama/Llama-3.2-1B-Instruct",
        torch_dtype="auto",
    )

    # LoRA configuration
    peft_config = LoraConfig(
        r=16,
        lora_alpha=32,
        target_modules=[
            "q_proj", "v_proj", "k_proj", "o_proj",
            "gate_proj", "up_proj", "down_proj",
        ],
        lora_dropout=0.1,
        bias="none",
        task_type="CAUSAL_LM",
    )

    # Dataset
    dataset = load_dataset("trl-lib/Capybara")

    # DeepSpeed ZeRO-3 config
    ds_config = {
        "bf16": {"enabled": True},
        "zero_optimization": {
            "stage": 3,
            "overlap_comm": True,
            "contiguous_gradients": True,
            "reduce_bucket_size": "auto",
            "stage3_prefetch_bucket_size": "auto",
            "stage3_param_persistence_threshold": "auto",
        },
        "gradient_accumulation_steps": "auto",
        "gradient_clipping": "auto",
    }
    ds_path = os.path.join(tempfile.mkdtemp(), "ds_config.json")
    with open(ds_path, "w") as f:
        json.dump(ds_config, f)

    # Training arguments
    training_args = TrainingArguments(
        output_dir=CHECKPOINT_DIR,
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        learning_rate=2e-4,
        max_steps=60,
        bf16=True,
        gradient_checkpointing=True,
        report_to="mlflow",
        deepspeed=ds_path,
    )

    # Train
    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset["train"],
        processing_class=tokenizer,
        peft_config=peft_config,
    )
    trainer.train()
    trainer.save_model()


results = fine_tune_llama.distributed()

This runs across all 8 H100 GPUs with DeepSpeed ZeRO-3 for memory-efficient distributed training. MLflow automatically logs metrics, and checkpoints go to a Unity Catalog Volume.

Environment v5: What’s Under the Hood

#

Serverless GPU runs on Environment v5, which ships with a modern stack:

Two Environments: Base vs. AI

#

You choose between two pre-built environments:

Base environment (~180 packages): Minimal. Good for custom setups where you control every dependency. Breaking change from v4: PyTorch is NOT included in the base environment. You must install it yourself or switch to the AI environment.

AI environment (~220 packages): Everything in base, plus 40+ ML/AI libraries:

Costs: Real Numbers

#

Serverless GPU uses the Model Training SKU at $0.65/DBU (AWS, US East, Premium tier). This single rate includes cloud VM infrastructure — no separate EC2 charges.

Fine-Tuning Cost Estimates

#

The Hidden Cost of Traditional Clusters

#

Consider a typical workflow: spin up a p3.8xlarge cluster (4x V100), spend 2 hours debugging data loading, run a 3-hour training job, forget to terminate the cluster overnight.

With traditional clusters, you pay for all 13+ hours (2 debugging + 3 training + 8 idle). With Serverless GPU, you pay for 5 hours (2 + 3). The idle time literally costs zero because there’s no cluster to forget about.

For teams running GPU workloads intermittently — a few training runs per week, periodic fine-tuning, experimentation — Serverless GPU can cut GPU spend by 50-70% just by eliminating idle time.

Bonus: ai_parse_document — GPU-Powered Document Intelligence

#

While not technically a Serverless GPU workload (it runs on Foundation Model APIs), ai_parse_document is worth mentioning because it’s the kind of GPU-accelerated capability that “just works” without any infrastructure:

1
2
3
4
5
6
7
8

-- Parse PDFs into structured elements
SELECT 
    path,
    ai_parse_document(content, MAP('version', '2.0')) AS parsed
FROM READ_FILES(
    '/Volumes/finance/invoices/',
    format => 'binaryFile'
);

Combine it with ai_extract for structured extraction:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

WITH parsed_docs AS (
    SELECT
        path,
        ai_parse_document(content) AS parsed_content
    FROM READ_FILES('/Volumes/finance/invoices/', format => 'binaryFile')
)
SELECT
    path,
    ai_extract(parsed_content, '["invoice_id", "vendor_name", "total_amount"]') AS invoice_data
FROM parsed_docs;

This handles PDFs, images (JPG, PNG), Word docs, and PowerPoint files — up to 500 pages per document. No GPU cluster needed, no model deployment, no inference endpoint. Just a SQL function.

Limitations You Should Know

#

Before going all-in on Serverless GPU, understand the constraints:

Hard limits:

• 7-day maximum runtime — implement checkpointing for longer training runs
• 10 million MLflow metric steps — don’t log every batch on large runs
• Only A10 and 8xH100 — no other GPU types
• Distributed training requires H100 — A10 is single-GPU only
• Multi-node (multiple 8xH100 nodes) is Private Preview — you’re limited to 8 GPUs max today

Environment restrictions:

• Python notebooks only — no Scala, no R
• No Spark UI — use query profile instead
• No environment variables — use notebook widgets for job parameters
• Notebook-scoped libraries are NOT cached across sessions (cold installs every time)

Data access:

• Unity Catalog is mandatory — no direct DBFS access
• Move data loading inside @distributed functions to avoid serialization issues
• Cache datasets to /tmp (local NVMe SSD) for multi-epoch training — UC Volumes are networked storage

Compliance:

• No support for compliance security profiles (HIPAA, PCI DSS)
• Cross-region GPU provisioning can happen during high demand, which may incur egress costs

Checkpointing Best Practice

#

Since you have a 7-day limit, always checkpoint to UC Volumes:

1
2
3
4
5
6
7
8
9

import torch

# Save to Unity Catalog Volume — survives session termination
checkpoint_path = "/Volumes/main/default/checkpoints/model_epoch_5.pt"
torch.save(model.state_dict(), checkpoint_path)

# For distributed training, use torch.distributed.checkpoint
import torch.distributed.checkpoint as dcp
dcp.save(model.state_dict(), checkpoint_dir=checkpoint_path)

When to Use Serverless GPU vs. Traditional Clusters

#

Use Serverless GPU when:

• You need GPUs intermittently (experimentation, periodic fine-tuning)
• You want zero idle costs
• You’re fine with A10 or H100 accelerators
• Your training runs fit within 7 days
• You don’t need compliance security profiles

Stick with traditional GPU clusters when:

• You need specific GPU types (V100, A100, etc.)
• Your workloads run continuously (always-on serving)
• You require HIPAA/PCI compliance
• You need more than 8 GPUs (until multi-node GA)
• You need Spark RDD APIs or non-Python languages

What’s Coming Next

#

Databricks has signaled several expansions for AI Runtime:

• Multi-node distributed training moving from Private Preview toward GA
• NVIDIA RTX PRO 4500 Blackwell GPU support (announced at GTC 2026)
• Broader region availability — currently limited to select AWS regions (us-west-2, us-west-1, us-east-1, us-east-2, ca-central-1, sa-east-1)

The trajectory is clear: Databricks wants GPU compute to be as invisible as serverless SQL warehouses already are. You describe the workload, the platform handles the infrastructure.

Wrapping Up

#

Serverless GPU on Databricks isn’t just a convenience feature — it fundamentally changes the economics of GPU workloads. By eliminating cluster management and idle costs, it makes GPU compute accessible to teams that previously couldn’t justify dedicated GPU clusters.

The key takeaway: if you’re running intermittent GPU workloads on Databricks today — fine-tuning runs, model experimentation, batch inference — switch to Serverless GPU and stop paying for idle VMs. The migration is trivial (your PyTorch code doesn’t change), and the cost savings are immediate.

For distributed training on 8xH100, the @distributed decorator is remarkably simple. The hard part isn’t the infrastructure anymore — it’s the ML engineering. Which is exactly how it should be.