Skip to main content
Last updated on

Migrate classic GPU workloads to serverless

If you are moving an existing deep learning workload from a classic Databricks cluster (with Databricks Runtime ML) to serverless (with AI Runtime), follow these steps:

  1. Replace cluster-dependent code. Remove any references to Spark-based distributed training (for example, TorchDistributor) and replace them with the @distributed decorator from serverless_gpu.
  2. Update data loading. Replace direct DBFS paths with Unity Catalog volumes paths (/Volumes/...). Replace local Spark DataFrame operations with Spark Connect. For streaming file-based data from volumes, use UCVolumeDataset from serverless_gpu.data. See Load data on AI Runtime.
  3. Set up your environment. Some Databricks Runtime ML pre-installed libraries aren't available on AI Runtime. Set up a serverless GPU environment and install the packages your workload needs with %pip install. See Set up your environment.
  4. Update checkpoint paths. Move checkpoints from DBFS or local storage to Unity Catalog volumes (/Volumes/<catalog>/<schema>/<volume>/...). For distributed checkpointing, use UCVolumeWriter and UCVolumeReader from serverless_gpu.data, which stage I/O through local NVMe. See Model checkpointing.
  5. Review your MLflow configuration. On AI Runtime, the .distributed() API creates MLflow runs automatically. To customize the experiment, set it with an absolute path. To resume interrupted runs, set a run name. See MLflow integration.
  6. Test interactively first. Validate your workload in an interactive notebook before scheduling it as a job.