Skip to main content

Serverless GPU compute

Beta

This feature is in Beta.

This article describes serverless GPU compute on Databricks and provides recommended use cases, guidance for how to set up GPU compute resources, and feature limitations.

What is serverless GPU compute?​

Serverless GPU compute is part of the Serverless compute offering. Serverless GPU compute is specialized for custom single and multi-node deep learning workloads. You can use serverless GPU compute to train and fine-tune custom models using your favorite frameworks and get state-of-the-art efficiency, performance, and quality.

Serverless GPU compute includes:

  • An integrated experience across Notebooks, Unity Catalog, and MLflow: You can develop your code interactively using Notebooks.
  • A10s and H100s GPU accelerators: Use A10 GPUs for cost-effective, small to medium machine learning and deep learning tasks, such as classic ML models or fine-tuning smaller language models. Choose H100 GPUs for large-scale AI workloads, including training or fine-tuning massive models or running advanced deep learning tasks.
  • Hyperparameter sweeps: You can perform hyperparameter sweep fine-tuning on A10 GPUs.
  • Multi-GPU and multi-node support: You can run distributed training workloads multiple GPUs (A10s and H100s) and multiple nodes (A10s only) using the Serverless GPU Python API. See Distributed training.

The pre-installed packages on serverless GPU compute are not a replacement for Databricks Runtime ML. While there are common packages, not all Databricks Runtime ML dependencies and libraries are reflected in the serverless GPU compute environment.

Databricks recommends serverless GPU compute for any model training use case that requires training customizations and GPUs.

For example:

  • LLM Fine-tuning
  • Computer vision
  • Recommender systems
  • Reinforcement learning
  • Deep-learning-based time series forecasting

Requirements​

  • A workspace in either us-west-2 or us-east-1.

What's installed​

Serverless GPU compute for notebooks uses environment versions, which provide a stable client API to ensure application compatibility. This allows Databricks to upgrade the server independently, delivering performance improvements, security enhancements, and bug fixes without requiring any code changes to workloads.

For more details, see the release notes:

note

Base environments are not supported for serverless GPU compute. In order to set up serverless GPU compute on your environment, specify the dependencies directly in the Environments side panel or pip install them.

Set up serverless GPU compute​

To connect your notebook to serverless GPU compute and configure the environment:

  1. From the notebook's compute selector, select Serverless GPU.
  2. Click the Environment icon. to open the Environment side panel.
  3. Select A10 or H100 from the Accelerator field.
  4. Select the Environment version.
  5. Click Apply and then Confirm that you want to apply the serverless GPU compute to your notebook environment.
note

Connection to your compute auto-terminates after 60 minutes of inactivity.

Add libraries to the environment​

You can install additional libraries to the serverless GPU compute environment. See Add dependencies to the notebook.

note

Adding dependencies using the Environments panel as seen in Add dependencies to the notebook is not supported for serverless GPU compute scheduled jobs.

Create and schedule a job​

The following steps show how to create and schedule jobs for your serverless GPU compute workloads. See Create and manage scheduled notebook jobs for more details.

After you open the notebook you want to use:

  1. Select the Schedule button on the top right.
  2. Select Add schedule.
  3. Populate the New schedule form with the Job name, Schedule _and _Compute.
  4. Select Create.

You can also create and schedule jobs from the Jobs and pipelines UI. See Create a new job for step-by-step guidance.

Distributed training​

note

Multi-GPU distributed training is supported on both H100s and A10s. Multi-node distributed training is only supported on A10 GPUs.

You can launch distributed training across multiple GPUs on a single node or multiple nodes using the Serverless GPU Python API. It has a simple, unified interface for running both single-GPU and multi-GPU jobs, automatically handling GPU provisioning, environment setup, and workload distribution behind the scenes. Whether you’re working with local GPUs or scaling out across nodes with remote GPUs, serverless_gpu makes distributed execution straightforward, requiring only minimal code changes.

For example, the code below uses serverless_gpu to distribute the hello_world function to execute across 8 remote A10 GPUs.

Python
# Import the distributed decorator
from serverless_gpu import distributed

# Decorate the function with @distributed and specify the number of GPUs, the GPU type, and whether or not the GPUs are remote
@distributed(gpus=8, gpu_type='A10', remote=True)
def hello_world(s: str) -> None:
  print('hello_world ', s)

# Trigger the distributed execution of the hello_world function
hello_world.distributed(s='abc')

After running the above code, the results and logs appear in the Experiment section of your workspace.

Start by importing the starter notebook to get hands-on with the API, then explore the notebook examples to see how it’s used in real distributed training scenarios.

For full details, refer to the Serverless GPU Python API documentation.

Limitations​

  • Serverless GPU compute only supports A10 and H100 accelerators.

  • H100 accelerators only support single-node workflows and jobs. Sweeps and multi-node workflows on H100s are not yet supported.

  • PrivateLink is not supported. Storage or pip repos behind PrivateLink are not supported.

  • Serverless GPU compute is not supported for compliance security profile workspaces (like HIPAA or PCT). Processing regulated data is not supported at this time.

  • Serverless GPU compute is only supported on interactive environments.

  • Scheduled jobs on Serverless GPU compute:

    • Only supported for a single task.
    • Auto recovery behavior for incompatible package versions that are associated with your notebook is not supported.

Best practices checklist​

Before you run a notebook using serverless GPU compute, check the following:

  • Environment: Ensure your libraries and packages are compatible with your selected serverless environment version.

  • Checkpoint storage: Check if you are saving checkpoints to DBFS or leave it unspecified to let MLflow default to DBFS.

    • Avoid using /Workspace, which has a 500 MB per file size limit.
    • Verify checkpointing sooner. For example, after 50 steps, instead of 1 epoch.

  • MLFlow logging: Set the logger step parameter to a sufficiently large number of batches to avoid logging every batch (default) and exceeding the 1M metric step limit.

  • Multi-node launch: Add retries or a longer timeout to avoid barrier timeout issues.

The following code shows how to implement these best practices:

Python
# Settings for a quick run to verify logging and checkpointing

# If using transformers
from transformers import TrainingArguments

training_args = TrainingArguments(
    # checkpoint to /Vol if no symlinks created
    output_dir = "/Volumes/your_catalog/your_schema/your_vol/your_model",
    logging_strategy = "steps",
    logging_steps = 10, # avoid exceeding mlflow 1M metric step limit
    # save checkpoints earlier after 100 steps to verify checkpointing
    save_strategy = "steps",
    save_steps = 100,
    # terminate job earlier after 200 steps as a trial run to verify logging and checkpointing
    max_steps = 200,
    ...
)

Notebook examples​

Below are various notebook examples that demonstrate how to use Serverless GPU compute for different scenarios.

Task

Description

Large language models (LLMs)

Examples for fine-tuning large language models including parameter-efficient methods like Low-Rank Adaptation (LoRA) and supervised fine-tuning approaches.

Computer vision

Examples for computer vision tasks including object detection and image classification.

Deep learning based recommender systems

Examples for building recommendation systems using modern deep learning approaches like two-tower models.

Classic ML

Examples for traditional machine learning tasks including XGBoost model training and time series forecasting.

Multi-GPU and multi-node distributed training

Examples for scaling training across multiple GPUs and nodes using the Serverless GPU API, including hyperparameter sweeps and distributed fine-tuning.

Last updated on