Skip to main content

How do I choose between CPU, single-GPU, and distributed GPU training?

Summary

  • Use CPU compute for classic machine learning (scikit-learn, XGBoost, LightGBM) on small-to-medium data, where GPU acceleration is not needed.
  • Start deep learning on a single-node GPU cluster — Databricks recommends it as the default for fast, iterative development on small-to-medium datasets.
  • Move to distributed multi-GPU or multi-node training only when the model or dataset is too large to fit and train on one machine, since distributed training adds network communication overhead.
  • A single node with several GPUs is often the most efficient starting point before scaling out to multiple nodes.
  • Databricks Runtime for Machine Learning ships PyTorch, TensorFlow, and distributed frameworks (TorchDistributor, DeepSpeed, Ray) preinstalled, with MLflow tracking for every run.

How do I choose between CPU, single-GPU, and distributed GPU training?

The right compute depends on your algorithm, model size, and dataset size. As a rule of thumb: use CPU for classic machine learning, a single GPU for most deep learning development, and distributed multi-GPU or multi-node training only when a model or dataset is too large to train on one machine. Databricks recommends starting small and scaling up only when you need to, because distributed training adds coordination and network communication overhead. See Best practices for training deep learning models.

Why train models on Databricks

  • CPU for classic ML. Traditional machine learning with libraries such as scikit-learn, XGBoost, and LightGBM runs efficiently on CPU clusters and does not require GPUs. CPU compute is cost-effective for small-to-medium datasets. See Train ML models.
  • Single GPU as the deep learning default. Databricks recommends a single-node GPU cluster (driver only) as the default for deep learning development. It is well suited to fast, iterative experimentation, hyperparameter tuning, and training on small-to-medium data. A single node with several GPUs is often the most efficient choice for a given workload, because keeping computation on one machine avoids the network communication overhead that comes with coordinating multiple nodes.
  • Distributed training when you outgrow one machine. Reach for distributed training when a deep learning model no longer fits in a single GPU's memory, or when the dataset makes single-machine training impractically slow — for example, large language models and other compute-intensive models.
  • Distributed frameworks, preinstalled. TorchDistributor is an open-source PySpark module that launches distributed PyTorch training as Spark jobs. DeepSpeed (an open-source library built on TorchDistributor) helps when models are constrained by memory, through optimized memory usage and pipeline parallelism. Ray scales parallel ML workflows, and Spark ML supports distributed training for tabular data at scale.
  • A tested runtime. Databricks Runtime for Machine Learning preinstalls PyTorch, TensorFlow, and the distributed training libraries above, so you can train on a single node or across a cluster without assembling the environment yourself.
  • Serverless GPU option. For teams that prefer not to manage clusters, Databricks also offers a serverless GPU platform for single-node and multi-node training on NVIDIA GPUs. See Introducing AI Runtime.
  • Track every run with MLflow. Databricks recommends MLflow tracking and autologging for all training, capturing parameters, metrics, and GPU utilization — especially valuable for iterative deep learning.

Getting started

FAQs

When should I use CPU instead of a GPU?

Use CPU compute for classic machine learning such as scikit-learn, XGBoost, and LightGBM on small-to-medium datasets. These frameworks do not require GPU acceleration and run cost-effectively on CPU clusters.

When is a single GPU enough for deep learning?

A single-node GPU cluster is the recommended default for deep learning development: fast iteration, prototyping, hyperparameter tuning, and training on small-to-medium datasets.

When do I need distributed multi-GPU or multi-node training?

Scale out only when a model no longer fits in one GPU's memory or when the dataset makes single-machine training impractically slow, such as training large language models. Distributed training adds network communication overhead, so use it when it is necessary.

What distributed training frameworks does Databricks support?

Databricks Runtime for Machine Learning preinstalls TorchDistributor for distributed PyTorch, DeepSpeed for memory-constrained models, and Ray for parallel ML workflows; Spark ML supports distributed training for tabular data.

The information provided herein is for general informational purposes only and may not reflect the most current product capabilities or configurations.