Apache Spark
Gradient Boosted Trees
Performance Optimization
Machine Learning
Distributed Computing

Slow Performance with Apache Spark Gradient Boosted Tree training runs

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Spark Gradient Boosted Trees can be slow even on a healthy cluster because boosting is inherently sequential at the tree level. Each new tree depends on the residuals from the previous ones, so GBT never scales as perfectly as embarrassingly parallel workloads. That said, many slow training runs are slower than they need to be because the dataset is repeatedly recomputed, partitions are badly shaped, or the model parameters are much more expensive than the problem requires.

Know Where the Time Goes

A slow GBT training job usually spends time in one or more of these areas:

  • upstream feature pipeline recomputation
  • dataset shuffles and skew
  • expensive tree parameters such as deep trees or large maxBins
  • executor memory pressure and garbage collection
  • writing checkpoints or lineage recovery work

The first debugging step is to check the Spark UI and see whether the delay sits in the model stage itself or in the feature preparation that runs before each stage.

Cache the Training Data Before Fitting

One of the easiest mistakes is fitting the model on a DataFrame that is still backed by a long transformation chain. If Spark has to rebuild features repeatedly, the GBT stage inherits that cost every time.

A common fix is to persist the assembled training dataset right before fitting:

python
1from pyspark.storagelevel import StorageLevel
2from pyspark.ml.classification import GBTClassifier
3
4training = assembled.select("features", "label").persist(StorageLevel.MEMORY_AND_DISK)
5training.count()
6
7gbt = GBTClassifier(labelCol="label", featuresCol="features", maxIter=50)
8model = gbt.fit(training)

The count() materializes the cache so training does not pay for lazy recomputation later.

Tune the Expensive GBT Parameters First

GBT cost grows quickly with model complexity. The main parameters that drive training time are:

  • 'maxIter: number of boosting rounds'
  • 'maxDepth: tree depth'
  • 'maxBins: candidate split bins'
  • 'stepSize: learning rate'
  • 'subsamplingRate: fraction of rows used per tree'

A heavy configuration such as deep trees plus many iterations is often far more expensive than expected.

Example starter configuration:

python
1from pyspark.ml.classification import GBTClassifier
2
3gbt = GBTClassifier(
4    labelCol="label",
5    featuresCol="features",
6    maxIter=20,
7    maxDepth=5,
8    maxBins=32,
9    stepSize=0.1,
10    subsamplingRate=0.8,
11)

If this baseline trains reasonably, increase complexity only when the accuracy gain justifies it.

Watch Partitioning and Skew

Spark ML depends on balanced task execution. If one partition contains much more data than the others, one task becomes the long tail that determines the stage duration.

Useful checks:

python
print(training.rdd.getNumPartitions())
print(training.groupBy().count().collect())

And for repartitioning when needed:

python
training = training.repartition(200)

The right partition count depends on cluster size and data volume, but a wildly unbalanced or tiny partition layout can slow training substantially.

Reduce Feature Explosion

GBT can handle many features, but very wide vectors and high-cardinality categorical encodings increase memory use and split-search cost.

Common performance wins include:

  • removing irrelevant columns before vector assembly
  • reducing one-hot explosion where possible
  • capping cardinality of rare categories
  • checking whether a simpler model already performs well enough

A slow GBT run is often telling you as much about the feature pipeline as about the learner itself.

Make Sure the Cluster Shape Matches the Job

Tree training usually benefits from enough executor memory to hold cached data and enough cores to process partitions efficiently, but more is not always better. Overly large executors can increase garbage-collection pauses, while too many tiny executors create scheduling overhead.

A practical tuning pattern is:

  • keep executor sizes moderate
  • ensure the cached training set fits reasonably across executors
  • watch GC time in the Spark UI
  • avoid repeatedly spilling cached data to disk if possible

If the Spark UI shows heavy spilling or long GC pauses, the cluster memory configuration is likely part of the slowdown.

Compare Against Random Forests and Smaller Baselines

Sometimes a GBT run is slow because the task does not really need GBT. A smaller tree model or random forest baseline may reach acceptable quality faster.

That is worth checking because boosting is sequential by nature. If your use case values training speed and interpretability more than squeezing out the last bit of accuracy, a simpler model may be the correct engineering answer.

Common Pitfalls

The biggest mistake is fitting GBT on an unpersisted DataFrame that has to rebuild features repeatedly. That can make the model look slower than it really is.

Another mistake is turning maxIter, maxDepth, and maxBins up together. Those settings compound the cost quickly.

Teams also often ignore skew and partition shape because the data volume looks reasonable overall. In Spark, one bad partition can dominate stage time.

Finally, do not assume every slow GBT run needs more hardware. Many of them need less recomputation, smaller trees, or a better feature set.

Summary

  • Spark GBT training is slower than many ML jobs because boosting is sequential by design.
  • Cache the final training DataFrame before calling fit.
  • Tune maxIter, maxDepth, and maxBins conservatively before scaling up.
  • Watch partition skew, spill behavior, and executor memory pressure in the Spark UI.
  • If a simpler model is good enough, it may be the better production choice.

Related reading
Course
Beginner
27 lessons
10 hours
System Design Fundamentals

Build a strong foundation in designing scalable, reliable distributed systems.

View the course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.