Unpickling `RandomForestClassifier` with `cuml.accel`

[betatim]

Another discussion topic. Unpickling a random forest with cuml.accel enabled takes quite a bit more time than with it turned off.

I instructed cursor to create a script that trains forests with trees of different "shapes". Trying to get trees that are deep or shallow. Ideally having trees that are deep and narrow vs trees that are shallow and wide. I'm not sure how achievable that is without lots of doctoring of the training dataset. However, I think it doesn't matter too much, as the different tree shapes are just a guess of mine for what could influence pickling time. The script is in the "details" at the end. I've read through it to check it is not totally bananas and it seems alright. I've not (yet) experimented with changing the hyper-parameters to see what happens (eg do we really need the min_samples_split setting or is it all about max_depth?)

The average depth for the "deep" config is quite different in the two cases (~20 vs ~30). My guess would have been that less depth/less nodes means faster (un)pickling, so I don't think this difference matters wrt the slower (un)pickling we see for with cuml.accel.

scikit-learn with cuml.accel on:

====================================================================================================
BENCHMARK RESULTS
====================================================================================================
Config        Trees  AvgDepth  TotalNodes  AvgLeaves     Pickle(ms)   Unpickle(ms)   Size(MB)
----------------------------------------------------------------------------------------------------
shallow          50       5.0        1550       16.0    0.68±0.25     2.89±0.73        0.17
shallow         100       5.0        3100       16.0    0.88±0.17     3.92±0.48        0.33
shallow         200       5.0        6200       16.0    1.53±0.33     7.14±0.87        0.66
shallow         300       5.0        9300       16.0    2.31±0.48    21.96±22.78       0.99
deep             50      19.8     1056182    10562.3   48.23±22.05  505.60±40.33     104.77
deep            100      19.8     2113350    10567.2  135.02±7.35   889.58±6.10      209.63
deep            200      19.7     4228960    10572.9  215.88±88.97 2042.09±60.36     419.49
deep            300      19.7     6342006    10570.5  325.77±122.84 3044.53±39.58     629.09
medium           50      10.9       84526      845.8    2.34±0.25    20.78±0.56        8.40
medium          100      10.8      168594      843.5   11.47±13.38   41.65±0.80       16.75
medium          200      10.7      338528      846.8    9.39±1.20    86.75±0.72       33.63
medium          300      10.8      507654      846.6   14.92±2.07   126.44±1.47       50.43
====================================================================================================

Just scikit-learn:

====================================================================================================
BENCHMARK RESULTS
====================================================================================================
Config        Trees  AvgDepth  TotalNodes  AvgLeaves     Pickle(ms)   Unpickle(ms)   Size(MB)
----------------------------------------------------------------------------------------------------
shallow          50       4.0        1550       16.0    0.52±0.10     0.42±0.03        0.17
shallow         100       4.0        3100       16.0    0.99±0.17     1.00±0.37        0.34
shallow         200       4.0        6200       16.0    2.13±0.26     1.62±0.12        0.68
shallow         300       4.0        9300       16.0    3.38±0.54     2.76±0.73        1.02
deep             50      29.6      982974     9830.2   82.56±0.86    65.01±3.83       97.51
deep            100      29.6     1964450     9822.8  165.29±0.85   137.92±2.09      194.87
deep            200      29.5     3931726     9829.8  329.58±8.15   290.88±2.51      390.03
deep            300      29.6     5898986     9832.1  491.37±10.54  433.18±0.64      585.18
medium           50      10.0       79518      795.7    3.32±2.26     1.57±0.74        7.90
medium          100      10.0      157966      790.3    5.65±2.10     2.94±1.04       15.70
medium          200      10.0      316650      792.1   23.33±4.01    13.85±10.37      31.47
medium          300      10.0      475614      793.2   39.48±2.13    35.01±7.52       47.27
====================================================================================================

Cursor generated benchmark script

I've not made any edits to this, except remove a very small `n_estimators` entry and add the 300 one. I've read through it and it looks super good enough.

#!/usr/bin/env python3
"""Benchmark pickle/unpickle times for sklearn RandomForestClassifier

This script benchmarks forests with different tree shapes:
- Shallow/wide trees: low max_depth, high min_samples_split/leaf
- Deep/narrow trees: high max_depth, low min_samples_split/leaf
- Medium trees: balanced parameters

Each shape is tested with varying numbers of trees (n_estimators).
"""

import pickle
import time
import numpy as np
from sklearn.datasets import make_classification
from sklearn.ensemble import RandomForestClassifier


# Forest configurations controlling tree shape
FOREST_CONFIGS = {
    "shallow": {
        "max_depth": 4,
        "min_samples_split": 50,
        "min_samples_leaf": 20,
    },
    "deep": {
        "max_depth": 30,
        "min_samples_split": 2,
        "min_samples_leaf": 1,
    },
    "medium": {
        "max_depth": 10,
        "min_samples_split": 10,
        "min_samples_leaf": 5,
    },
}

# Number of trees to test
N_ESTIMATORS_LIST = [50, 100, 200, 300]

# Benchmark parameters
N_RUNS = 5
RANDOM_STATE = 42


def generate_dataset(n_samples=50_000, n_features=100, n_informative=50, n_classes=5):
    """Generate a classification dataset suitable for tree benchmarking."""
    print(f"Generating dataset: {n_samples} samples, {n_features} features, "
          f"{n_informative} informative, {n_classes} classes...")
    
    X, y = make_classification(
        n_samples=n_samples,
        n_features=n_features,
        n_informative=n_informative,
        n_redundant=n_features - n_informative - n_classes,
        n_classes=n_classes,
        random_state=RANDOM_STATE,
    )
    return X.astype(np.float32), y.astype(np.int32)


def get_tree_stats(model):
    """Extract statistics from a fitted random forest.
    
    Returns:
        tuple: (average_depth, total_nodes, avg_leaves)
    """
    depths = []
    total_nodes = 0
    total_leaves = 0
    
    for tree in model.estimators_:
        depths.append(tree.get_depth())
        total_nodes += tree.tree_.node_count
        total_leaves += tree.tree_.n_leaves
    
    return np.mean(depths), total_nodes, total_leaves / len(model.estimators_)


def benchmark_pickle(model, n_runs=N_RUNS):
    """Measure pickle and unpickle times using in-memory serialization.
    
    Returns:
        tuple: (mean_pickle_time, std_pickle_time, mean_unpickle_time, 
                std_unpickle_time, serialized_size_bytes)
    """
    pickle_times = []
    unpickle_times = []
    data = None
    
    for _ in range(n_runs):
        # Pickle to bytes
        start = time.perf_counter()
        data = pickle.dumps(model)
        pickle_times.append(time.perf_counter() - start)
        
        # Unpickle from bytes
        start = time.perf_counter()
        _ = pickle.loads(data)
        unpickle_times.append(time.perf_counter() - start)
    
    return (
        np.mean(pickle_times),
        np.std(pickle_times),
        np.mean(unpickle_times),
        np.std(unpickle_times),
        len(data),
    )


def train_and_benchmark(X, y, config_name, config_params, n_estimators):
    """Train a RandomForest and benchmark its pickle/unpickle times.
    
    Returns:
        dict: Results including timing and tree statistics
    """
    # Create and train model
    model = RandomForestClassifier(
        n_estimators=n_estimators,
        n_jobs=-1,
        random_state=RANDOM_STATE,
        **config_params,
    )
    
    train_start = time.perf_counter()
    model.fit(X, y)
    train_time = time.perf_counter() - train_start
    
    # Get tree statistics
    avg_depth, total_nodes, avg_leaves = get_tree_stats(model)
    
    # Benchmark pickle/unpickle
    pickle_mean, pickle_std, unpickle_mean, unpickle_std, size_bytes = benchmark_pickle(model)
    
    return {
        "config": config_name,
        "n_estimators": n_estimators,
        "train_time_s": train_time,
        "avg_depth": avg_depth,
        "total_nodes": total_nodes,
        "avg_leaves": avg_leaves,
        "pickle_mean_ms": pickle_mean * 1000,
        "pickle_std_ms": pickle_std * 1000,
        "unpickle_mean_ms": unpickle_mean * 1000,
        "unpickle_std_ms": unpickle_std * 1000,
        "size_mb": size_bytes / (1024 * 1024),
    }


def print_results_table(results):
    """Print results in a formatted table."""
    print("\n" + "=" * 100)
    print("BENCHMARK RESULTS")
    print("=" * 100)
    
    header = (
        f"{'Config':<12} {'Trees':>6} {'AvgDepth':>9} {'TotalNodes':>11} "
        f"{'AvgLeaves':>10} {'Pickle(ms)':>14} {'Unpickle(ms)':>14} {'Size(MB)':>10}"
    )
    print(header)
    print("-" * 100)
    
    for r in results:
        row = (
            f"{r['config']:<12} {r['n_estimators']:>6} {r['avg_depth']:>9.1f} "
            f"{r['total_nodes']:>11} {r['avg_leaves']:>10.1f} "
            f"{r['pickle_mean_ms']:>7.2f}±{r['pickle_std_ms']:<5.2f} "
            f"{r['unpickle_mean_ms']:>7.2f}±{r['unpickle_std_ms']:<5.2f} "
            f"{r['size_mb']:>10.2f}"
        )
        print(row)
    
    print("=" * 100)


def main():
    print("=" * 60)
    print("RandomForest Pickle/Unpickle Benchmark")
    print("=" * 60)
    print(f"Runs per config: {N_RUNS}")
    print()
    
    # Generate dataset
    X, y = generate_dataset()
    print(f"Dataset shape: X={X.shape}, y={y.shape}")
    print(f"Dataset memory: {X.nbytes / (1024**2):.2f} MB")
    print()
    
    # Run benchmarks
    results = []
    total_configs = len(FOREST_CONFIGS) * len(N_ESTIMATORS_LIST)
    current = 0
    
    for config_name, config_params in FOREST_CONFIGS.items():
        print(f"\n--- Testing '{config_name}' tree configuration ---")
        print(f"    max_depth={config_params['max_depth']}, "
              f"min_samples_split={config_params['min_samples_split']}, "
              f"min_samples_leaf={config_params['min_samples_leaf']}")
        
        for n_estimators in N_ESTIMATORS_LIST:
            current += 1
            print(f"\n[{current}/{total_configs}] {config_name} with {n_estimators} trees...")
            
            result = train_and_benchmark(X, y, config_name, config_params, n_estimators)
            results.append(result)
            
            print(f"    Train: {result['train_time_s']:.2f}s | "
                  f"Pickle: {result['pickle_mean_ms']:.2f}ms | "
                  f"Unpickle: {result['unpickle_mean_ms']:.2f}ms | "
                  f"Size: {result['size_mb']:.2f}MB")
    
    # Print final results table
    print_results_table(results)
    
    # Summary comparison
    print("\nSUMMARY: Shallow vs Deep comparison (at 100 trees)")
    shallow_100 = next(r for r in results if r["config"] == "shallow" and r["n_estimators"] == 100)
    deep_100 = next(r for r in results if r["config"] == "deep" and r["n_estimators"] == 100)
    
    print(f"  Shallow: {shallow_100['total_nodes']:,} nodes, "
          f"{shallow_100['pickle_mean_ms']:.2f}ms pickle, "
          f"{shallow_100['size_mb']:.2f}MB")
    print(f"  Deep:    {deep_100['total_nodes']:,} nodes, "
          f"{deep_100['pickle_mean_ms']:.2f}ms pickle, "
          f"{deep_100['size_mb']:.2f}MB")
    print(f"  Ratio:   {deep_100['total_nodes']/shallow_100['total_nodes']:.1f}x nodes, "
          f"{deep_100['pickle_mean_ms']/shallow_100['pickle_mean_ms']:.1f}x pickle time, "
          f"{deep_100['size_mb']/shallow_100['size_mb']:.1f}x size")


if __name__ == "__main__":
    main()

[ncclementi]

Am I reading this numbers right or cuml is worst at unpickling? Probably could be better at pickling too. But the pickling numbers are in some cases faster, now the unpickling is bad overall.

[jcrist]

With the way things are designed currently, serialization and deserialization should always take longer than sklearn (aside - the cases above where pickling is faster are because the cuml-fit models have fewer nodes).

This is because under cuml.accel, we always serialize the sklearn version of the model (with a thin wrapper) so that it can be unpickled in an environment without cuml.accel or cuml at all. On load we always first unpickle to sklearn, then optionally convert to cuml. This lets a user train and save a model with cuml.accel and load it later with only sklearn.

In pseudocode:

def pickle_like_cuml_accel_does(model):
    # coerce cuml model to sklearn
    sk_model = model.as_sklearn()
    # wrap in a thin magical wrapper to coerce to sklearn/cuml on load
    # depending on whether `cuml.accel` is enabled or installed.
    wrapped = magic_unpickling_wrapper(sk_model)
    # pickle the wrapper
    return pickle.dumps(wrapped)

def unpickle_like_cuml_accel_does(buffer):
    # load the sklearn model
    sk_model = pickle.load(buffer)
    if cuml_accel_not_enabled:
        # return sklearn directly
        return sk_model
    # Coerce sklearn model to cuml model
    return cuml_cls.from_sklearn(sk_model)

The sklearn paths are always taken on both ends, cuml can only add work here. For most models the conversion is pretty cheap - usually just coercing arrays to cupy. For cuml.ensemble things are less cheap - RandomForestClassifier.from_sklearn takes some time. The expensive steps on the load side are:

# Coerce the sklearn model to treelite (takes ~55% of the time on an example problem)
tl_model = treelite.sklearn.import_model(sk_model)
# Coerce the treelite model to bytes (takes ~45% of the time on an example problem)
tl_bytes = tl_model.serialize_bytes()

There's an additional cost on first inference later, but that's done lazily so wouldn't show up during pickle.loads.

A few options as I see them:

1. Look into whether these operations can be optimized in treelite itself. My guess is this is possible, that seems pretty slow for what it's doing.
2. Special case cuml.ensemble to include the treelite bytes as part of the serialized payload. Since treelite isn't a sklearn dependency, we'll also still need to pickle the sklearn version of the model as well, effectively doubling the binary size but saving the conversion costs. I don't like this option.
3. Live with it. Does this show up in a real world workflow? Having a motivating use case for optimizing this would be useful.

Closing with a simple, human written, pointed benchmark of just the operation that's slow:

bench.py

import time
from sklearn.datasets import make_classification
import cuml

X, y = make_classification(
    n_samples=50_000,
    n_features=100,
    n_informative=50,
    n_classes=5,
    random_state=42,
)
X = X.astype("float32")
y = y.astype("int32")

model = cuml.RandomForestClassifier(
    max_depth=30,
    min_samples_split=2,
    min_samples_leaf=1,
    n_estimators=300,
).fit(X, y)
sk_model = model.as_sklearn()

N_RUNS = 3
start = time.perf_counter()
for _ in range(N_RUNS):
    cuml.RandomForestClassifier.from_sklearn(sk_model)
stop = time.perf_counter()
print(f"{(stop - start) * 1000 / N_RUNS:.2f} ms")

[dantegd]

Thanks for the analysis @jcrist that tracks with a modified benchmark of @betatim benchmark I was running, if we compare the treelite times vs pickle and see where things are taking time we have:

==========================================================================================
COMPARISON: Standard Pickle vs Treelite Serialization
==========================================================================================
Config      Trees   Pickle(ms)   Treelite(ms)    Ratio   Pkl Size    TL Size
------------------------------------------------------------------------------------------
shallow        50         0.47           1.46    3.13x      0.17M      0.17M
shallow       100         1.94           5.86    3.02x      0.34M      0.35M
shallow       200         1.95           5.20    2.67x      0.68M      0.70M
shallow       300         5.74           8.70    1.52x      1.02M      1.04M
deep           50        59.65         342.47    5.74x     97.51M    104.07M
deep          100        75.17         668.97    8.90x    194.87M    207.97M
deep          200       161.20        1229.39    7.63x    390.03M    416.24M
deep          300       228.10        1915.08    8.40x    585.18M    624.51M
medium         50         2.12          32.67   15.40x      7.90M      8.43M
medium        100         4.88          40.46    8.29x     15.70M     16.74M
medium        200         9.29          76.96    8.29x     31.47M     33.56M
medium        300        13.33         113.80    8.54x     47.27M     50.41M
==========================================================================================

========================================================================================================================
TREELITE BENCHMARK RESULTS (sklearn -> treelite -> bytes)
========================================================================================================================
Config      Trees  TotalNodes     Import(ms)   Serialize(ms)    Total(ms)  Import%  Serial%  Size(MB)
------------------------------------------------------------------------------------------------------------------------
shallow        50        1550    1.42±0.27      0.04±0.01          1.46    97.3%     2.7%      0.17
shallow       100        3100    5.71±0.62      0.15±0.03          5.86    97.4%     2.6%      0.35
shallow       200        6200    5.00±0.26      0.20±0.13          5.20    96.2%     3.8%      0.70
shallow       300        9300    8.26±0.50      0.44±0.24          8.70    95.0%     5.0%      1.04
deep           50      982974  242.87±18.34    99.60±10.95       342.47    70.9%    29.1%    104.07
deep          100     1964450  485.04±24.97   183.93±35.83       668.97    72.5%    27.5%    207.97
deep          200     3931726  891.29±39.29   338.10±44.71      1229.39    72.5%    27.5%    416.24
deep          300     5898986 1335.64±76.48   579.44±75.36      1915.08    69.7%    30.3%    624.51
medium         50       79518   21.33±6.96     11.34±19.70        32.67    65.3%    34.7%      8.43
medium        100      157966   36.77±2.05      3.69±0.02         40.46    90.9%     9.1%     16.74
medium        200      316650   70.02±0.71      6.94±0.08         76.96    91.0%     9.0%     33.56
medium        300      475614  104.36±0.71      9.44±0.09        113.80    91.7%     8.3%     50.41
========================================================================================================================

From this, doing a quick initial analysis of the treelite code there are a few areas where the time is probably being spent

• For importing sklearn into treelite the main one being this for loop for sklearn models: https://github.com/dmlc/treelite/blob/a9ce6c352be315e7c170392e3f8c4c4f13ecfc7a/src/model_loader/sklearn.cc#L297. In other words the C++ function LoadSKLearnModel() builds treelite trees node-by-node with this rough pattern:

// sklearn.cc lines ~297-340
for (int tree_id = 0; tree_id < n_trees; ++tree_id) {
    builder->StartTree();
    for (int node_id = 0; node_id < n_nodes; ++node_id) {
        builder->StartNode(node_id);
        // ... set node properties ...
        builder->EndNode();
    }
    builder->EndTree();
}

So for example, for a forest with 6M total nodes (the largest in Tim's script), this results in ~30M function calls through the ModelBuilder API.

I think for solving this, we should be able to copy sklearn arrays directly into treelite's internal storage via memcpy instead of iterating nodes, which would save a lot of function calls.

• For treelite to sklearn, probably less of an overhead but still, the export_model() function processes trees one at a time:

# exporter.py lines 209-222
for tree_id in range(n_trees):  # 7200 trees in our case
    estimators.append(_export_tree(model, tree_id=tree_id, ...))

I think here we probably can create a single C++ call that returns all tree data as concatenated arrays, then slice in Python, but need to think it through a bit more.

There's probably some additional loop optimizations we can do around as well.

[betatim]

Sorry for the missing context. We found this as part of investigating using cuml.accel with autogluon. What happens there is that cuml.accel accelerates the fit part but measuring the "end to end" time things are slower. One reason for this is that the end-to-end time is only a few seconds (less than 10s) so accelerating the fitting only buys you a small amount of time and autogluon does a bunch of pickle operations. I've not thoroughly investigated when and where they pickle, from a first look some of it could be changed/not done. However, In the spirit of meeting our users where they are I'm thinking that we shouldn't say "yeah stop doing that!".

I'm not a fan of option (2) either. It would be nice if the pickle was something like 10MB - doubling that is "in the noise". But for larger models I think it is too much.

Let's see how much we can move treelite optimisation forward. It sounds like there is potential.

---

As an aside, are we tied to using treelite? Or would it be possible to investigate using something like onnx (insert name of other serialisation library)? I'm wondering if we can buy into a different serialisation ecosystem with a large community of users and developers - with the benefit that others help push on the performance of the library/ecosystem. Maybe something for a new thread/issue.

[betatim]

xref dmlc/treelite#651