Add optional n_jobs to fit base learners in parallel

[loganprosser]

Title: Add optional n_jobs to fit base learners in parallel

What this does

NGBoost fits one base learner per distribution parameter on every boosting
round. For distributions with several parameters, like MultivariateNormal or
classification with many classes, that becomes a lot of independent model fits
running one after another, and it turns out to be most of the training time.

This adds an optional n_jobs argument. When you set it, those independent
fits run at the same time on a pool of threads instead of sequentially. The
default keeps everything serial, so nothing changes unless you opt in.

Usage

from ngboost import NGBRegressor
from ngboost.distns import MultivariateNormal

model = NGBRegressor(Dist=MultivariateNormal(10), n_jobs=-1)  # use every core
model.fit(X, Y)

Why threads instead of processes

The fits inside one boosting round do not depend on each other. Each one gets a
different gradient column but the same X, so running them together is safe.
scikit learn releases the GIL while it builds a tree, so a thread pool gives
real parallelism here without copying X out to worker processes on every round.
I tried processes first and the copying cost cancelled out the gain, so threads
it is.

The backend is set to threading explicitly rather than as a soft preference, so
an outer parallel context (for example a scikit learn GridSearchCV that is
itself parallel) cannot silently turn these inner fits into separate processes
and bring the copying cost back. If you do nest NGBoost inside another parallel
job, keep the product of the two thread counts near your core count so you do
not oversubscribe.

Because the threads share one X, sparse input is canonicalized once with a
single sort_indices before the fits, so concurrent tree fitting cannot race on
an in place index sort. This mirrors the guard scikit learn's own forests use.

Numbers

MultivariateNormal regression on an Apple M3 Max (14 core), 2000 rows, 50
estimators, comparing serial against all cores. With a fixed random_state on
the base learner, the parallel fit reproduces the serial fit exactly.

output dim 10 (65 trees):   16.6s serial    8.4s parallel    1.98x
output dim 20 (230 trees):  75.9s serial   45.1s parallel    1.68x

On a larger dataset (8000 rows, same output dimension 20) the speedup climbs to
1.93x, because larger fits hide the thread overhead better. Very high dimensions
(around 30 parameters per axis and above) get less out of it, since at that
point a serial linear algebra step, not the tree fitting, is the bottleneck.

A note on GPUs

I looked at pushing the work onto a GPU (Apple Metal and MPS) before landing on
this. It does not pay off. The runtime is dominated by scikit learn building
trees on the CPU, which a GPU cannot help with, and the linear algebra a GPU
could take on is a small share of the total and actually ran slower on MPS than
NumPy did on the CPU. The independent base learner fits are where the real time
goes, so that is what this parallelizes.

Prior discussion

Issue #156 raised this idea back in 2020. A maintainer noted that parallelism
here would only help multiparameter distributions and was not sure it was worth
doing, and the issue was closed without any implementation. This branch is the
first actual build of it, and the numbers match that prediction rather than
contradict it: the speedups above are real for MultivariateNormal and many class
classification, while a plain Normal regression (two learners per round) gets
essentially nothing. That is why the flag defaults to off and is opt in. The
claim: for the multiparameter case the win is large
and comes for free, since the results are identical to serial fitting. It also
answers issue #357, an open request for parallel NGBoost that has had no reply.

Changes

NGBoost takes an n_jobs argument, surfaces it in get_params, and passes it
into fit_base. fit_base runs the fits through joblib.Parallel with the
threading backend when n_jobs is set, and keeps the exact original serial loop
otherwise. NGBRegressor, NGBClassifier, and NGBSurvival forward the
argument to the base class. joblib is now listed in pyproject, though it
already comes in with scikit learn. Tests cover that n_jobs survives
get_params and clone, and that a parallel fit matches a serial fit exactly
on regression, multiclass classification, sparse float32 CSC input, and data
with missing values, given a fixed random_state.

Correctness

The full test suite passes, including the pickling and clone tests. The added
tests confirm that a parallel fit reproduces a serial fit exactly on regression,
classification, sparse float32 CSC input, and data with missing values, as long
as the base learner has a fixed random_state. The default DecisionTreeRegressor uses random_state=None,
which draws from NumPy's global random state; under threads that can make
results vary slightly, and vary from run to run, on data with missing values or
tied splits, so set a fixed random_state when you need exact reproducibility.

---

This change was authored with assistance from an AI coding agent... (っᵔ◡ᵔ)っ ᓚᘏᗢ

[alejandroschuler]

Thanks for this! I think it needs a few changes:

The X.sort_indices() call is gated on issparse, which is too coarse a check.
issparse is also true for LIL, COO, DOK, and DIA sparse formats, and none of those even have a sort_indices method. Only the compressed formats (CSR, CSC, BSR) store an indices array whose entries within a row or column may or may not be in
ascending order, so they're the only ones with anything to sort. The result is that a LIL input fits fine at n_jobs=1 and raises AttributeError at n_jobs=4, so the flag silently changes which inputs the estimator accepts. Check for the compressed formats specifically and let sklearn's own validation handle the rest, since it converts those and hands each thread a private copy anyway.

On reproducibility, I'd rather fix it than document it. The default base learner has random_state=None, so it draws from numpy's global state, and under threads the interleaving of those draws is nondeterministic. I recommend what sklearn's forests do: draw one integer seed per base learner from self.random_state in the calling thread before dispatch, and set it on each clone that doesn't already have a fixed random_state. It should apply on the serial path as well, not just the parallel branch, so the two agree. Output for the default base learner will shift relative to 0.5.x, so it needs a changelog note, but that's fine imo.

Last, the n_jobs docstring is duplicated verbatim in four places, there should be a way to have that just in one place and referenced. It should also note that for small data there is actually a slowdown. And, minor, but, I woulnd't say "exactly as before" anywhere in there b/c a reader should not have to understand the project history to read a docstring.