Slow prediction with dask

[RAMitchell]

Prediction with Xgboost dask interface is much slower than expected, taking almost 5x training time.

from dask_cuda import LocalCUDACluster
from dask.distributed import Client
from dask import array as da
import xgboost as xgb
from xgboost.dask import DaskDMatrix
import time


def main(client):
    m = 10000000
    n = 100
    X = da.random.random(size=(m, n), chunks=100)
    y = da.random.random(size=(m,), chunks=100)
    dtrain = DaskDMatrix(client, X, y)

    start = time.time()
    output = xgb.dask.train(client,
                            {
                                'tree_method': 'gpu_hist'},
                            dtrain,
                            num_boost_round=500, evals=[(dtrain, 'train')])
    print("Train time: {}".format(time.time() - start))
    bst = output['booster']
    start = time.time()
    prediction = xgb.dask.predict(client, bst, dtrain)
    prediction = prediction.compute()
    print("Predict time: {}".format(time.time() - start))
    return prediction


if __name__ == '__main__':
    with LocalCUDACluster() as cluster:
        with Client(cluster) as client:
            main(client)

Train Time: 121s
Prediction time: 502s

[trivialfis]

Let me take a look.

[trivialfis]

The prediction is run for each partition/block so that there's no concatenation for input data (hence lower memory usage). When you set the chunk size to 100, you ended up running prediction 100000 times given you have 10000000 rows.

[cdeotte]

I see two issues here. First the chunk size is small. If you change to chunks=1_000_000, then train takes 27.2 seconds and prediction takes 151.2 seconds.

The reason prediction is 5.5x slower than train is because you need to add 'predictor' : 'gpu_predictor'. If you add this parameter then with chunks=1_000_000 train takes 27.2 seconds and predict takes 11.7 seconds.

[trivialfis]

I think the predictor is automatically running on GPU in this example.

[cdeotte]

No it is not. I just ran this on a DGX and posted my results above. By default XGB always uses CPU for prediction. Even if you use 'tree_method': 'gpu_hist', prediction is still on CPU. You must explicitly set the 'predictor' : 'gpu_predictor'.

[trivialfis]

Got it, took another look into the code. CPU predictor might be chosen because data comes from host. Thanks for correcting my mistake.

[cdeotte]

It is good you made your comment trivialfis. Many people including myself last week did not know this. As such I suggest making GPU prediction the default if a user activates the tree method gpu hist parameter.

(And yes, maybe the predictor depends on where the data comes from, I'm not sure).

[trivialfis]

It's a trade-off. We made some heuristics to avoid copying data into device from host. For training the data is not copied even you are using GPU. Currently DGX is not very accessible to wide public, so we can't make this setting.

[trivialfis]

The prediction speed has been improved quite significantly in recent PRs.