Scaling ML Training on Databricks: Single Node, Spark, Ray and pandas Function APIs

1.Vertical first: why Databricks starts you on one machine

You can scale a training workload in two ways. Vertical scaling means moving it to a bigger single machine. Horizontal scaling means spreading it across several machines. Databricks recommends a default: when possible, train neural networks on a single machine. Distributed code is harder to write and debug, and the workers have to exchange gradients, parameters or data over the network. That traffic is overhead a single machine never pays.

The reason to scale horizontally is size, not impatience. Databricks says to consider distributed training and inference when your model or your data is too large to fit in memory on one machine. If a job already fits comfortably on one node, buying a larger instance or simply waiting is usually the better trade than adding nodes. When a job really has outgrown one machine, Databricks Runtime ML ships three options: TorchDistributor, the DeepSpeed distributor and Ray. TorchDistributor is an open-source PySpark module that launches PyTorch training jobs as Spark jobs. It sets up the communication channels between workers and runs torch.distributed.run across the worker nodes.

Sources1

2.Splitting the data or splitting the model

After you decide to scale out, the next question is what gets split. In data parallelism, every worker holds a full copy of the model and processes a different slice of the data. Distributed Data Parallel (DDP) is the standard example. The full model is replicated on each GPU and the data batches are split across the GPUs. This only works if the model fits on one device. Adding workers then increases throughput.

Once the model itself outgrows a single device's memory, you need a strategy that spreads the model across devices. Databricks points you to FSDP or DeepSpeed for this case. The DeepSpeed distributor is built on top of TorchDistributor. Databricks recommends it for models that need more compute but are limited by memory, because it offers optimized memory usage, reduced communication overhead and advanced pipeline parallelism.

Sources21

3.Spark or Ray for distributing the work

Databricks runs Spark and Ray in the same execution environment, so you rarely have to pick just one. What you do need to know is which kinds of work each engine is built for. Spark is optimized for large-scale data processing such as joins, filters, aggregations, ETL and feature engineering, and its MLlib and Spark ML libraries are built for large-scale ML algorithms. Ray is built for task parallelism: many independent, compute-heavy tasks running at once. Databricks lists the workloads where Spark is less optimized and Ray is the better fit. They include reinforcement learning, hierarchical time series forecasting, simulation modeling, hyperparameter search, deep learning training and HPC.

For a single task that uses both engines, the main pattern Databricks recommends is Spark for data handling and Ray for computation. Spark reads the data. databricks.ray.data.from_spark passes it to Ray Data. The resulting model is saved to MLflow, or the output data to Unity Catalog tables. Two advanced patterns are also documented: running Ray inside Spark functions such as UDFs, and running Spark and Ray operations concurrently. You can also put the ETL step in its own Databricks Workflow subtask, so each step gets a cluster type that suits it and the two engines don't compete for resources.

Sources3

4.pandas function APIs: per-group training and batch inference

For classic ML on large tables, pyspark.ml.connect provides distributed training and inference for Spark ML models. In Databricks Runtime 17.0 and above, Spark ML is enabled by default on Standard compute. Many teams, though, need to train many small single-node models, such as one per store or one per device. pandas function APIs handle that case. They apply a native Python function that takes and returns pandas objects to a PySpark DataFrame. Like pandas UDFs, they use Apache Arrow to move the data, and Python type hints are optional.

groupBy().applyInPandas() follows the split-apply-combine pattern. Spark splits the rows by key, calls your function once per group with a pandas.DataFrame that contains all the rows and columns for that group, and combines the returned DataFrames. You must supply the function and an output schema. The documentation's example subtracts a group mean rather than training a model, and it shows no per-group training example. The contract is the same either way, though: each function call sees one complete group, which is what per-group model fitting needs.

def subtract_mean(pdf):
    # pdf is a pandas.DataFrame
    v = pdf.v
    return pdf.assign(v=v - v.mean())

df.groupby("id"). ? (subtract_mean, schema="id long, v double").show()

That design has a cost. All of a group's data is loaded into memory before the function runs, and maxRecordsPerBatch does not apply to groups. If the keys are skewed, one very large group can cause an out-of-memory error, so it is up to you to make sure each group fits. For inference you usually don't need grouping. DataFrame.mapInPandas() passes your function an iterator of pandas DataFrames, and the function yields DataFrames back. The output can be any length. The third API type, cogrouped map, pairs the groups of two DataFrames that share a key.

df = spark.createDataFrame([(1, 21), (2, 30)], ("id", "age"))

def filter_func(iterator):
    for pdf in iterator:
        yield pdf[pdf.id == 1]

df.mapInPandas(filter_func, schema=df.schema).show()

Sources41

5.Distributed hyperparameter tuning with Optuna and MLflow

Hyperparameter search is a task-parallel job: each trial is independent, so many can run at once. Optuna is an open-source Python library for hyperparameter tuning that can be scaled horizontally across multiple compute resources. MLflow 3.0 integrates with it through two classes. MlflowStorage lets Optuna use the MLflow Tracking Server as its storage backend, which gives all the workers one shared place to record trials. MlflowSparkStudy launches parallel Optuna studies using PySpark executors.

The workflow has four steps. First, define an objective function. Inside it you define the search space with trial.suggest_* methods such as suggest_float, suggest_int and suggest_categorical, and you return the value to optimize, such as the validation error. Second, create shared storage with MlflowStorage(experiment_id=...). Third, create an MlflowSparkStudy with a study name and that storage. Fourth, call optimize with the objective, the number of trials and n_jobs. By default the study uses optuna.samplers.TPESampler and optuna.pruners.MedianPruner. MLflow 3.0 is pre-installed in Databricks Runtime 17.0 ML and above. On older runtimes you install the latest MLflow and Optuna with %pip.

This is also the current answer where you might once have used Hyperopt with SparkTrials. Hyperopt is not included in Databricks Runtime ML after 16.4 LTS ML. Databricks recommends Optuna, or Ray Tune for a similar experience to the deprecated Hyperopt distributed tuning.

from mlflow.pyspark.optuna.study import MlflowSparkStudy

mlflow_study = MlflowSparkStudy(
    study_name="spark-mlflow-tuning",
    storage=mlflow_storage,
)

mlflow_study.optimize(objective, n_trials=8, n_jobs=4)

Sources56

Exam traps

Each one states something that sounds right. Open it to see what is actually true.

1. 1.Adding worker nodes always speeds up a training job.Why is that wrong?

   Distributed training adds complexity and communication overhead. Databricks recommends a single machine unless the model or data no longer fits in memory.

   Covered in Vertical first: why Databricks starts you on one machine

2. 2.Setting maxRecordsPerBatch keeps a skewed group in applyInPandas from running out of memory.Why is that wrong?

   A whole group is loaded into memory before the function runs, and maxRecordsPerBatch is not applied to groups. Keeping each group small enough to fit is your job.

   Covered in pandas function APIs: per-group training and batch inference

3. 3.Distributed tuning on current Databricks Runtime ML should use Hyperopt with SparkTrials.Why is that wrong?

   Hyperopt is no longer included after 16.4 LTS ML. Databricks recommends Optuna, or Ray Tune for a similar experience to the deprecated Hyperopt distributed tuning.

   Covered in Distributed hyperparameter tuning with Optuna and MLflow