Databricks Machine Learning Associate study materials
Databricks Machine Learning Associate study resources

Hyperparameter Sweep That Got Slower With More Parallelism

Databricks Machine Learning Associate · Put your knowledge to work in context.

All case studies

Hyperparameter Sweep That Got Slower With More Parallelism

A team is tuning an XGBoost churn model on a 400M-row Delta table (~60 GB after feature expansion). The cluster has 8 workers, each 16 cores/64 GB, and a 32-core/128 GB driver. They use Hyperopt with SparkTrials(parallelism=8) and 300 evaluations; each trial takes ~18 minutes, so the sweep runs about three days. To speed it up they raise parallelism to 32. The sweep now takes four days. Driver logs show long GC pauses, repeated CollectLimit operators, and shuffle spill on workers; several trials fail with org.apache.spark.SparkException: Job aborted due to stage failure. MLflow logging is enabled for every trial.

1. Explain why raising `SparkTrials` parallelism from 8 to 32 made the sweep slower, and propose a concrete configuration change.

Read the worked answer

The cluster has 8 x 16 = 128 worker cores. At parallelism 8, each trial can use ~16 cores, which matches one worker and keeps each XGBoost/Spark job's shuffle local. At parallelism 32, four trials share each worker's 16 cores, so every trial's Spark job gets a small slice, creates more partitions, and spills. The driver also becomes the bottleneck: Hyperopt bookkeeping, MLflow run creation, metric logging, and CollectLimit actions all serialize there, and 32 concurrent runs amplify GC.

A better configuration is SparkTrials(parallelism=8) and fewer evaluations, e.g. 120 with early stopping or a timeout, plus a per-trial worker cap (num_workers) so one trial does not consume the whole cluster. If the driver remains the bottleneck, reduce logged metrics per trial and avoid collecting large prediction tables to the driver. More parallelism only helps when each trial is small relative to the cluster and driver-side work is negligible.

2. The team proposes abandoning Spark and training single-node XGBoost on the driver with all 60 GB. Evaluate that proposal and give a decision rule.

Read the worked answer

A 60 GB feature matrix will not fit comfortably in a 128 GB driver once Spark, Python, XGBoost's internal histogram structures, and conversion overhead are counted. Converting the Delta table to a pandas DataFrame also requires a driver-side collect, which is exactly the failure mode already visible. Single-node training becomes attractive only after the data is reduced.

A practical decision rule: if feature selection and sampling bring the training matrix below roughly 30-40% of driver memory and the job does not need distributed joins, use single-node XGBoost on a dedicated single-node cluster; it is simpler and often faster. Otherwise keep distributed training and control per-trial resources. The team should first run feature importance on a sample, drop low-value one-hot columns, and test a 5-10 GB sample on one node before deciding.