Memory consumption of fitting EBMs

[DerWeh]

I'm running out of memory trying to fit a ExplainableBoostingClassifier.

The order of magnitude of the features is 10, of the samples is $10^5$, and of the classes is 100. In the recent release, the feature ExplainableBoostingClassifier.estimate_mem added. The docstring is somewhat vague:

Estimate memory usage of the model.

Is this the memory necessary to store the model (unlikely, as X and y are arguments), the memory necessary to fit the model (most likely, as y is an argument), or the memory necessary to make predictions (unlikely, predictions aren't memory intensive).

However, ExplainableBoostingClassifier.extimate_mem indicates tiny memory usage while in reality I run out of memory (more than 100GiB needed). The functions indicate memory usage independent on the number of classes, while for 'small' numbers of 1 to 100 classes, I observe the memory consumption to be roughly an affine function in the number of classes. For larger number of classes it seems to 'saturate', I guess the system starts caching to disk.

So far, I used resource.getrusage(resource.RUSAGE_SELF).ru_maxrss to measure the memory consumption, so I am not really experienced in measuring memory. For once, I am not sure if this includes the memory of the parallel processes used by default for each bags.

But for both variants

ebm = ExplainableBoostingClassifier(interactions=0)
with joblib.parallel_config(backend="loki"):
    ebm.fit(X, y)

ebm = ExplainableBoostingClassifier(interactions=0)
with joblib.parallel_config(backend="threading"):
    ebm.fit(X, y)

I observed a rather linear increase of memory consumption with classes.
To me this seems not surprising, as for each class a shape functions is fit.

The question is threefold: Is estimate_mem correct? What is the required memory to fit a classifier? What can I do to fit the model with limited memory (largest machine I can get has around 512 GiB)?

Probably it's best to use threading instead of the default of multiprocessing, or are you making use of smart shared memory or memory mapping data from disk?

---

Example script for minimal tests:

import resource

import joblib
import numpy as np
from interpret.glassbox import ExplainableBoostingClassifier

backend = "loky"
n_samples = 10_000
n_features = 10
n_classes = 100
rng = np.random.default_rng(42)
X = rng.random([n_samples, n_features], dtype=np.float32)
y = rng.integers(n_classes, size=[n_samples], dtype=np.int32)
ebm = ExplainableBoostingClassifier(interactions=0, random_state=42, max_rounds=3, n_jobs=16)  # small max rounds for fast testing
print(f"Pre: {resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * 1024**-2:.3f} GiB")
with joblib.parallel_config(backend=backend):
    ebm.fit(X, y)
print(f"Post: {resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * 1024**-2:.3f} GiB")

[paulbkoch]

Thanks @DerWeh — You're right: the number of classes wasn’t being considered when estimating the memory required for fitting. I’ve just pushed a change that includes it in the estimate. I also updated the docstring to clarify that the memory estimate pertains specifically to the fit function.

Generally, for a dataset with a large number of samples, the dataset itself will dominate the memory requirements. We make multiple copies of the dataset in a compressed form (1 copy per outer bag). With smaller datasets, the size of the model can dominate. We make essentially 3 copies of the model during fitting. If you are observing that the model scales with the number of classes, then in your case it's likely dominated by the size of the resulting model. Does your larger dataset with 10^5 samples include any categoricals? Categoricals can make fairly large interaction tensors, since they are not limited by the max_interaction_bins parameter, unlike continuous features.

We do use shared memory. I would expect threading to use a little less memory, but not much less.

[DerWeh]

Thank you very much for your fast response and update, @paulbkoch.
The new version of the estimate_mem seems to reflect the scaling with classes properly, now.
The absolute value of memory, on the other hand, seems still very off (see below).

Categoricals can make fairly large interaction tensors, since they are not limited by the max_interaction_bins parameter, unlike continuous features.

Thanks for the warning. So far, I don't use any interactions (which aren't natively supported for multi-class yet). So we probably can rule out max_interactions_bins. But it's a good reminder for the future.

I can confirm your statement about multiprocessing vs multithreading, adding resource.getrusage(resource.RUSAGE_SELF).ru_maxrss and resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss,
both, multiprocessing and multithreading, consume roughly the same amount of memory.

However, there is still a huge difference between estimate_mem and the actual consumed memory.
I did a few test runs on a large dataset, for different number of bags, with one job per bag.
The estimate is consistently one to two orders of magnitude smaller than the actual consumed memory.

---

Click to sho memory results expand

(Most important parts: `delta` = `post` - `Base memory`, `Estimated` produced with new `estimate_mem`)

backend=loky, job=1, bag=1

Size of data 0.182 GiB
Base memory usage 0.409 GiB
Number of classes: 150
Estimated memory usage without data: 0.33288298547267914 GiB
Resource (pre): 1.152 GiB
Resource (post): 20.633 GiB
Resource (delta): 20.223 GiB
Resource (children): 1.152 GiB

backend=loky, jobs=2, bags=2

Size of data 0.182 GiB
Base memory usage 0.409 GiB
Number of classes: 150
Estimated memory usage without data: 0.45567476749420166 GiB
Resource (pre): 1.151 GiB
Resource (post): 20.645 GiB
Resource (delta): 20.236 GiB
Resource (children): 7.116 GiB

backend=thread, jobs=2, bags=2

Size of data 0.182 GiB
Base memory usage 0.407 GiB
Number of classes: 150
Estimated memory usage without data: 0.45567476749420166 GiB
Resource (pre): 1.150 GiB
Resource (post): 20.826 GiB
Resource (delta): 20.418 GiB
Resource (children): 1.150 GiB

backend=thread, jobs=4, bags=4

Size of data 0.182 GiB
Base memory usage 0.409 GiB
Number of classes: 150
Estimated memory usage without data: 0.7012583315372467 GiB
Resource (pre): 1.151 GiB
Resource (post): 27.942 GiB
Resource (delta): 27.532 GiB
Resource (children): 1.151 GiB

backend=loky, jobs=4, bags=4

Size of data 0.182 GiB
Base memory usage 0.405 GiB
Number of classes: 150
Estimated memory usage without data: 0.7012583315372467 GiB
Resource (pre): 1.150 GiB
Resource (post): 20.688 GiB
Resource (delta): 20.282 GiB
Resource (children): 7.116 GiB

---

Furthermore, I included MLflow to track memory usage over time.
(Code attached below.)
We observe to phases of memory consumption the first corresponds with high CPU usage indicating that this happens during multi-processing,
the second phase seems to happen during sequential calculations.

For small number of bags=jobs, the second phase dominates memory consumption:

System memory usage in MB in 1s steps for n_jobs=n_bags=2.

CPU utilization in 1s steps for n_jobs=n_bags=2.

For larger number of bags=jobs, the it is the first phase:

System memory usage in MB in 1s steps for n_jobs=n_bags=4.

CPU utilization in 1s steps for n_jobs=n_bags=4.

I tried to investigate the memory usage for smaller datasets, and the main culprit I could find is native._unsafe.CreateBooster.
It seems to roughly equal the difference between estimated and consumed memory (n_jobs=4, outer_bags=4).
For a smaller dataset with estimated 0.15 GiB, and 1.44 GiB used, it amounted to roughly 1 GiB.
For a slightly larger dataset with estimated 0.20 GiB, and 6.36 GiB used, it amounted to roughly 6 GiB.

You can use the code below to reproduce results or experiment yourself (requires mlflow, psutil for memory plots in MLflow):

"""Check the real memory consumption of EBMs."""
import gc
import os
import resource
import time

import joblib
import mlflow
import numpy as np
from interpret.glassbox import ExplainableBoostingClassifier
from sklearn import clone

os.environ["MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING"] = "true"
os.environ["MLFLOW_SYSTEM_METRICS_SAMPLING_INTERVAL"] = "1"

rng = np.random.default_rng(42)

clf = ExplainableBoostingClassifier(
    interactions=0,
    random_state=42,
    max_rounds=5,  # fast fitting
    smoothing_rounds=3,
    n_jobs=4,
    outer_bags=4
)

backend = "loky"
# backend = "threading"
n_features = 10
n_samples = 100_000
# n_samples = 300_000
# n_samples = 500_000

X = rng.random((n_samples, n_features), dtype=np.float32)
y = rng.integers(20, size=n_samples, dtype=np.int32)
print(f"Size of data {(X.nbytes + y.nbytes) * 1024**-3:.3f} GiB")
base_memory = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
print(f"Base memory usage {base_memory * 1024**-2:.3f} GiB")
del y
gc.collect()


def check_memory(n_classes):
    """Fit and check peak memory."""
    print(f"Number of classes: {n_classes}")
    y = rng.integers(n_classes, size=n_samples, dtype=np.int32)
    n_bytes = clf.estimate_mem(X, y, data_multiplier=0.0)
    print(f"Estimated memory usage without data: {n_bytes * 1024**-3} GiB")

    pre = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    print(f"Resource (pre):   {pre * 1024**-2:.3f} GiB")
    with joblib.parallel_config(backend=backend), mlflow.start_run(run_name=f"{n_classes=}"):
        mlflow.log_params(
            {
                "n_features": n_features,
                "n_samples": n_samples,
                "n_classes": n_classes,
                "joblib backend": backend,
                "n_jobs": clf.n_jobs,
            }
        )
        time.sleep(2)  # get at least one sample before fit
        clf_ = clone(clf)
        clf_.fit(X, y)
        time.sleep(2)  # get at least one sample after fit
        post = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
        print(f"Resource (post):  {post * 1024**-2:.3f} GiB")
        delta = post - base_memory
        print(f"Resource (delta): {delta * 1024**-2:.3f} GiB")
        del clf_
        gc.collect()
        mlflow.log_metrics(
            {
                "peak-memory_kiB": post,
                "consumed-memory_kiB": delta,
                "estimated-memory_kib": n_bytes * 1024**-1,
            }
        )
        time.sleep(2)  # get at least one sample after fit
    child = resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss
    print(f"Resource (children): {child * 1024**-2:.3f} GiB")


check_memory(150)

[paulbkoch]

Hi @DerWeh -- I pushed an update that adds some more sources of memory consumption to the estimate. I think it should work better in the sample you provided. Your note about the extra memory being allocated in CreateBooster was very helpful.

[paulbkoch]

Hi @DerWeh -- Just wanted to follow up and see if you checked whether the latest memory estimates were closer to what you observed. If so, I'll push out a release with the latest estimates.

[DerWeh]

Short answer @paulbkoch: I got inconsistent behavior, which is most likely my fault. I should find time on Friday to properly testing things again. I'll give you feedback then.

[DerWeh]

I tried checking my setup and re-running the code, and I am not quite sure.

First: it is quite dependent on the number of jobs and bags. The assumption is probably that all bags are run in parallel. For 4 = n_jobs = outer_bags, this seems to yield good results. The memory estimate is roughly 20% larger than the reported consumption, probably some memory is freed before other memory is allocated. The estimate being slightly large is no problem in practice. Typically, I need to check if my machine has enough memory to run a trying or reserve a certain amount of memory in a job scheduler. So a conservative estimate is what we want anyway.

For a small number n_jobs = outer_bags, the estimate is way too low. My estimate is that the "post-processing" after parallel boosting determines the peak memory in this case (see the plots above). I guess the CreateBooster memory was freed at this point and another "bottleneck" emerges.

---

For larger number of n_jobs = outer_bags I had to borrow a different machine, as this exceeds the limits of my hardware. In this regime I get inconsistent results. For 10 = n_jobs = outerbags the estimate is roughly 3 times the memory consumption reported by resource.getrusage, but MLflow (and htop) show a much higher memory usage more than twice as much as the estimate.

Another very surprising fact, is switching from the parallel (loky) backend to the threading backend, the memory consumption increases. My current assumption is that I fail to correctly measure the memory consumption in the multiprocessing backend. If this is true, the memory estimate is probably still too small...

---

Sorry @paulbkoch, but it seems my knowledge of memory is insufficient here. Probably I should also redo the memory profiling using the threading backend.

(And I haven't even started on my main problem: avoiding running out of memory when trying to fit a large dataset.)

[paulbkoch]

Well, it sounds like a big improvement at least.

We do currently assume that all jobs end up running at the same time in the estimate. If n_jobs=-1 is given, that's probably the right thing to do, but I guess we could improve it if the user specifies a number for n_jobs since then we get a guarantee that they won't run simultaneously (I think?).

I believe it should be rare that post-processing would be the dominant memory requirement. I could see this in the example above which has a larger than usual numbers of classes though since the final model ends up so big.

For multi-process, I think you might want to use resource.RUSAGE_BOTH?

[paulbkoch]

In terms of reducing memory, I believe we've already squeezed almost as much out of it as we can. You can reduce the memory footprint with n_jobs=1 because it'll allocate then deallocate each outer bag in sequence, but it's much slower.

In the future, if the parallel improvements I mentioned in the other thread get implemented, then we could see some memory improvements by doing the outer bags sequentially.

[paulbkoch]

The latest memory estimates are included in the 0v0.7.3 release. Should we close this issue now?

[DerWeh]

Sorry for the inactivity @paulbkoch, currently I'm tied up with other projects.
I wanted to run the checks again, last's time I still got some inconsistencies. If nothing else, we can improve the documentation a little.

But I have no problem in case you want to close the issue, I can always open a new issue/PR or reopen this one.

[paulbkoch]

No worries @DerWeh -- Since you got inconsistencies last time I'll keep it open for now.