XGBoost API by Method: fit, predict, inplace_predict, and Dask Entry Points

XGBoost’s Python API is easiest to understand as three related interfaces. The scikit-learn-style estimator API centers on XGBModel.fit() and XGBModel.predict(). The lower-level Booster exposes explicit prediction, persistence, and in-place prediction controls. For distributed work, xgboost.dask adds client-aware training and prediction entry points. These interfaces share XGBoost’s model format, but they do not expose identical parameters.

The Python API reference and scikit-learn estimator guide describe the estimator API primarily as a training interface. An estimator can return its underlying booster with get_booster(). Calling get_booster() before fitting raises an exception. When explicit booster behavior or cached prediction results are required, the corresponding lower-level tools are Booster.predict() and DMatrix.

fit: training, validation, and checkpoint semantics

The complete XGBModel.fit signature is:

fit(
    X,
    y,
    *,
    sample_weight=None,
    base_margin=None,
    eval_set=None,
    verbose=True,
    xgb_model=None,
    sample_weight_eval_set=None,
    base_margin_eval_set=None,
    feature_weights=None,
)

X is the input feature matrix and y is the target. Their accepted containers depend on the interface being called; the complete support rules appear later in this article.

The weighting parameters operate at different granularities:

  • sample_weight supplies per-instance weights for training. A row with a larger weight contributes more strongly to fitting than a row with a smaller weight.
  • sample_weight_eval_set supplies corresponding per-instance weights for validation data. Its weights must describe the instances in the associated evaluation matrix rather than the training matrix.
  • feature_weights supplies weights at the feature level, allowing the caller to express a weight for each feature rather than for each instance.

base_margin supplies an additive starting score before the objective’s prediction transformation. base_margin_eval_set provides the corresponding base margins for evaluation data. These parameters are separate from sample_weight: a base margin changes the starting prediction score, while sample weights change the contribution of instances during fitting.

eval_set is a sequence of (DMatrix, str) pairs. The string is the evaluation set’s name, and metrics are evaluated during training so that validation performance can be tracked. verbose controls whether training reports evaluation output; its documented default is True.

Starting again versus resuming training

Calling fit() again on the same estimator re-fits it from scratch. It does not silently continue from the previous model.

To resume training from a checkpoint, pass that previous model as xgb_model:

model.fit(
    X,
    y,
    xgb_model=previous_model,
    eval_set=[(validation_dmatrix, "validation")],
)

The resulting distinction is explicit:

  • Omitting xgb_model starts a new fit.
  • Supplying xgb_model resumes from the supplied checkpoint.

This prevents an accidental second call to fit() from being mistaken for incremental training.

The hist data optimization

When tree_method is set to hist, fitting internally uses QuantileDMatrix rather than constructing a regular DMatrix from supported matrix input. The documented purpose is memory conservation.

The device arrangement still matters. If the input is a NumPy array on the CPU while CUDA is used for training, for example, XGBoost first processes the data on the CPU and then transfers it to the GPU. Thus, QuantileDMatrix substitution does not make device placement irrelevant; an input/training-device mismatch can introduce a transfer and has performance implications.

predict: the estimator-level prediction path

The complete estimator prediction signature is:

predict(
    X,
    *,
    output_margin=False,
    validate_features=True,
    base_margin=None,
    iteration_range=None,
)

Transformed output versus the raw margin

With the documented default output_margin=False, XGBoost returns the objective’s normal transformed prediction. For an objective configured to produce classification probabilities, that output is the transformed probability.

Setting output_margin=True bypasses the objective’s prediction transformation and returns the raw, untransformed margin. This is useful when the downstream calculation needs the additive score rather than the model’s ordinary output:

probability = model.predict(X)

margin = model.predict(X, output_margin=True)

base_margin accepts the same kind of additive starting margin used during fitting.

Validation and iteration selection

validate_features controls feature validation and defaults to True. Keeping validation enabled helps detect an input whose features do not agree with the trained model. It can be disabled only when the caller has independently established that the feature arrangement is valid.

iteration_range selects a range of boosting iterations. If the estimator was trained with early stopping, model.predict() automatically uses best_iteration. Callers can supply an explicit range when they need to control which tree layers participate in prediction rather than relying on that automatic best-iteration selection.

The in-place path and its fallback

XGBModel.predict() uses inplace_predict() by default. If the input data and estimator are on different devices, it falls back to a DMatrix prediction path.

That routing gives the estimator API a convenient default while preserving operation when the local in-place path cannot use the input’s device directly. The lower-level details of the in-place method are described below.

Thread-safety boundary

The documented thread-safety guarantee for predict() is limited to the gbtree and dart booster types. It should not be generalized to every booster implementation.

Thread safety is also different from lock-free execution. A method may synchronize concurrent calls and still be thread safe; Booster.inplace_predict() has the stronger property only under the conditions described next.

inplace_predict: explicit local prediction without result caching

The complete signature is:

inplace_predict(
    data,
    *,
    iteration_range=(0, 0),
    predict_type="value",
    missing=nan,
    validate_features=True,
    base_margin=None,
    strict_shape=False,
)

This method was added in XGBoost 1.1.0.

Output controls

predict_type accepts:

  • "value", the default, for the model’s ordinary transformed prediction.
  • "margin" for the raw additive score before the objective transformation.

The local estimator API expresses the same choice through output_margin; the booster API calls it predict_type.

The default iteration_range=(0, 0) means all trees unless another range is specified. missing=nan marks the default missing-value sentinel. validate_features=True enables feature validation, while base_margin supplies an additive starting score.

strict_shape=False permits prediction output shape to follow the input-dependent shape behavior of the task. Setting strict_shape=True requests XGBoost’s consistent output shape rather than allowing shape differences caused by the arrangement of the prediction input. This matters particularly when downstream code expects the same result shape across different batches.

No cached prediction result

Unlike predict(), inplace_predict() does not cache the prediction result. The returned value is produced by the call, but inplace_predict() does not attach that result to a reusable prediction cache.

This makes the method appropriate when the caller explicitly wants prediction without retaining cached outputs, needs direct control over margin prediction, or already has an input container supported by the in-place path.

In-place does not mean “never copies.” If the input data’s device ordinal does not match the device configured for the booster, XGBoost copies the data to the booster’s device before prediction. The method avoids a result cache, but a device mismatch can still require a host-to-device or other cross-device transfer.

Lock-free use has a strict boundary

Calling only inplace_predict() from multiple threads is safe and lock free. That guarantee does not extend to concurrent use with other booster methods.

For example, one thread must not train the booster while another performs prediction through the same booster. The safe rule is therefore not merely “prediction is thread safe,” but “concurrent calls consisting only of in-place prediction are safe and lock free.”

Persistence: save_model and load_model

The persistence methods are:

save_model(fname)
load_model(fname)

save_model() writes the model in XGBoost’s internal format, which is shared across XGBoost interfaces. load_model() reads a model from a file or bytearray.

An estimator can expose its persisted booster through get_booster():

booster = model.get_booster()
booster.save_model("model.ubj")

loaded_booster = xgb.Booster()
loaded_booster.load_model("model.ubj")

Since XGBoost 2.1.0, the default serialization format is UBJSON. The default applies to model files, serialized buffers, and memory snapshots such as pickle-based mechanisms.

Persistence has an important scope. XGBoost saves the trees and model information needed to represent the learned predictor, including information such as the objective and input-column count. Some auxiliary attributes of the Python Booster object, including feature_names, are saved only in JSON or UBJSON; UBJSON is the default. Parameters that are not part of the saved model—such as metrics and max_depth—are not included.

These details come from the official model-saving tutorial.

Supported input containers by method

The official supported data structures reference distinguishes input to DMatrix, input to QuantileDMatrix, input through the scikit-learn interface, metadata support, in-place prediction, and multilabel support.

  • numpy.ndarray is directly supported in every listed category: DMatrix, QuantileDMatrix, the scikit-learn interface, metadata, in-place prediction, and multilabel prediction.
  • cudf.DataFrame, cupy.ndarray, torch.Tensor, and pyarrow.Table are also directly supported in every listed category.
  • pandas.DataFrame is supported in every category with the help of a NumPy array.
  • polars.DataFrame is supported in every category with the help of an Arrow table.
  • scipy.sparse.csr is supported as input to DMatrix, QuantileDMatrix, the scikit-learn interface, and in-place prediction. It is marked invalid for the metadata use case and unsupported for multilabel prediction.

For Dask, a separate rule applies: the Dask interface requires Dask collection inputs rather than selecting freely from this local matrix.

Dask entry points

Import the Dask interface as:

from xgboost import dask as dxgb

The official Dask tutorial requires calling the XGBoost Dask interface from the client side. Create a local cluster and client before constructing distributed matrices or invoking training:

import dask.array as da
import dask.distributed

if __name__ == "__main__":
    cluster = dask.distributed.LocalCluster()
    client = dask.distributed.Client(cluster)

The guard around cluster construction prevents the worker processes from independently recreating the cluster.

Distributed matrices

The two documented matrix choices are DaskDMatrix and DaskQuantileDMatrix. Both take the active Dask client as their leading argument:

dtrain = dxgb.DaskDMatrix(client, X, y)

dtrain = dxgb.DaskQuantileDMatrix(client, X, y)

Unlike the single-node interface, the data and labels must be Dask DataFrame or Dask Array instances.

Training with dxgb.train

The Dask training entry point mirrors the familiar booster-training structure, with the client supplied first:

output = dxgb.train(
    client,
    {
        "verbosity": 2,
        "tree_method": "hist",
        "objective": "reg:squarederror",
    },
    dtrain,
    num_boost_round=4,
    evals=[(dtrain, "train")],
)

This example uses the official quickstart’s num_boost_round=4. The returned object can be passed directly to prediction, or its booster member can be extracted.

Prediction with and without a distributed matrix

Prediction through dxgb.predict() accepts either the DaskDMatrix or the trained output directly:

prediction = dxgb.predict(client, output, dtrain)

# Equivalently:
prediction = dxgb.predict(client, output["booster"], dtrain)

A DaskDMatrix need not be constructed solely for prediction:

prediction = dxgb.predict(client, output, X)

Omitting the matrix can make the computation somewhat faster when metadata such as base_margin is not needed.

The Dask in-place entry point follows the same client-first convention:

prediction = dxgb.inplace_predict(client, output, X)

The tutorial notes that in-place prediction can help reduce both memory usage and prediction time. It does not promise a fixed speedup or memory reduction.

When the input is a DaskDMatrix or da.Array, prediction returns a Dask Array. Passing a Dask DataFrame directly produces a Dask Series instead. More generally, when a Dask collection is passed directly to prediction or to dxgb.inplace_predict(), the output type depends on the input collection.

Scikit-learn-compatible Dask wrappers

XGBoost also provides DaskXGBClassifier, DaskXGBRegressor, and DaskXGBRanker. These wrappers resemble their single-node scikit-learn counterparts but accept Dask collections as inputs.

They add a client attribute, which must be associated with the active distributed client before the wrapper performs cluster work. The wrappers provide a familiar estimator surface for Dask classification, regression, and ranking without replacing the lower-level dxgb.train, dxgb.predict, and dxgb.inplace_predict functions.

Choosing the method

Use XGBModel.fit() when a scikit-learn-style estimator is the intended interface, with explicit instance weights, validation matrices, feature weights, base margins, or checkpoint resumption. Its predict() method is the normal prediction API and delegates to in-place prediction when possible.

Use Booster.inplace_predict() when the booster is already available and the caller needs direct control over transformed-versus-margin output, iteration range, strict output shape, or prediction without a cached result. Check its device-copy behavior before assuming that in-place means zero-copy.

Use Booster.save_model() and load_model() for portable XGBoost persistence, relying on the documented UBJSON default while accounting for the distinction between saved model information and auxiliary Python attributes.

For distributed data, construct DaskDMatrix or DaskQuantileDMatrix, pass the active client to every Dask entry point, and choose dxgb.predict() or dxgb.inplace_predict() according to whether an explicit distributed matrix and cached prediction path are useful.