XGBoost in production: GPU training, out-of-core data, model persistence, and fast inference

Productionizing gradient-boosted trees is not one decision. It is a chain of decisions about compute, memory, serialization, and request handling:

• Does the training matrix fit in host or device memory?

• Is one GPU faster than the available CPU for this exact workload?

• Can the artifact survive a library upgrade?

• Can the service batch enough rows to amortize inference overhead?

XGBoost 3.3.0 is the current stable release. Its production path is clear: use device="cuda" with tree_method="hist", choose a quantized matrix before reaching for external memory, save a model file rather than a Python object snapshot, and serve contiguous batches through inplace_predict.

Train on the GPU with the current API

Older examples often use GPU-specific tree-method names. In XGBoost 3.x, processor selection belongs in device; histogram tree construction remains tree_method="hist".

from time import perf_counter

import numpy as np
import xgboost as xgb
from sklearn.datasets import make_regression
from sklearn.model_selection import train_test_split

X, y = make_regression(
    n_samples=250_000,
    n_features=64,
    n_informative=40,
    noise=0.5,
    random_state=42,
)
X = X.astype(np.float32)
y = y.astype(np.float32)

X_train, X_valid, y_train, y_valid = train_test_split(
    X,
    y,
    test_size=0.15,
    random_state=42,
)

model = xgb.XGBRegressor(
    n_estimators=800,
    max_depth=8,
    learning_rate=0.05,
    subsample=0.9,
    colsample_bytree=0.9,
    objective="reg:squarederror",
    eval_metric="rmse",
    tree_method="hist",
    device="cuda",
    random_state=42,
)

started = perf_counter()
model.fit(
    X_train,
    y_train,
    eval_set=[(X_valid, y_valid)],
    verbose=False,
)
print(f"Training time: {perf_counter() - started:.2f} seconds")

model.save_model("model.json")

Current GPU builds require CUDA 12.0 and compute capability 5.0 or newer. A model trained on a GPU can later run on a CPU-only inference host.

What speedup should you expect?

Do not capacity-plan around a universal multiplier. Use this practical expectation ladder:

• Small tables: expect little benefit, and sometimes a regression, because initialization and host-to-device transfer dominate.

• Large dense tables: a meaningful multi-fold speedup is a reasonable benchmark goal when the GPU stays busy.

• Sparse or narrow tables: measure carefully; memory layout and transfer cost can erase the advantage.

• Out-of-core GPU training: interconnect bandwidth can matter more than raw GPU compute.

The only defensible number is the one from an end-to-end benchmark that includes loading, conversion, training, evaluation, and artifact writing. XGBoost stores GPU training data in compressed, quantized ELLPACK form, while working memory still grows with rows, bins, features, and active tree nodes.

For a fair comparison, hold the train/validation split, max_bin, boosting rounds, subsampling, and evaluation cadence constant. Warm both paths, pin CPU thread counts, and report peak host memory and peak device memory alongside elapsed time. A GPU result that is faster but repeatedly approaches out-of-memory failure is not production-ready; neither is a CPU result measured with all cores while the GPU run includes slow data conversion inside the timed region.

Choose the matrix before scaling the machine

There are three important data paths:

• DMatrix: general-purpose training and prediction representation.

• QuantileDMatrix: quantized storage designed for histogram training with lower memory use.

• ExtMemQuantileDMatrix: iterator-fed, externally cached quantized pages for data that cannot fit in the main processor memory.

Try QuantileDMatrix first. Move to external memory only when the quantized representation still exceeds capacity. External memory is experimental in 3.x, supports histogram training, and trades memory pressure for repeated I/O.

The following CPU example writes synthetic server-telemetry shards and streams them through a custom iterator.

from pathlib import Path
from typing import Callable

import numpy as np
import xgboost as xgb


class NpyShardIterator(xgb.DataIter):
    def __init__(
        self,
        shards: list[tuple[Path, Path]],
        cache_prefix: Path | None,
    ) -> None:
        self.shards = shards
        self.position = 0
        super().__init__(
            cache_prefix=None if cache_prefix is None else str(cache_prefix),
            release_data=True,
        )

    def reset(self) -> None:
        self.position = 0

    def next(self, input_data: Callable) -> bool:
        if self.position == len(self.shards):
            return False

        features_path, labels_path = self.shards[self.position]
        features = np.load(features_path, mmap_mode="r")
        labels = np.load(labels_path, mmap_mode="r")
        input_data(data=features, label=labels)
        self.position += 1
        return True


root = Path("telemetry_shards")
root.mkdir(exist_ok=True)

rng = np.random.default_rng(42)
shards: list[tuple[Path, Path]] = []

for part in range(8):
    features = rng.normal(size=(25_000, 48)).astype(np.float32)
    labels = (
        100.0
        + 8.0 * features[:, 0]
        - 4.0 * features[:, 1]
        + rng.normal(scale=2.0, size=features.shape[0])
    ).astype(np.float32)

    features_path = root / f"features-{part:02d}.npy"
    labels_path = root / f"labels-{part:02d}.npy"
    np.save(features_path, features)
    np.save(labels_path, labels)
    shards.append((features_path, labels_path))

iterator = NpyShardIterator(shards, root / "xgb-cache")
dtrain = xgb.ExtMemQuantileDMatrix(iterator, max_bin=256)

booster = xgb.train(
    params={
        "objective": "reg:squarederror",
        "tree_method": "hist",
        "device": "cpu",
        "max_depth": 8,
        "max_bin": 256,
        "grow_policy": "depthwise",
    },
    dtrain=dtrain,
    num_boost_round=400,
)
booster.save_model("telemetry-model.json")

For an in-memory quantized matrix, create an iterator with cache_prefix=None and pass it to xgb.QuantileDMatrix. With external memory, avoid tiny batches, keep CPU cache files on fast local storage, and prefer grow_policy="depthwise"; loss-guided growth can require many more passes over cached pages. GPU external memory stages cache in host memory by default, and PCIe or chip-to-chip bandwidth can become the limiting resource.

Scale out with Dask or Spark

External memory increases the capacity of one machine. Distributed training divides work across workers. XGBoost 3.3 supports distributed GPU training through Dask and Spark/PySpark. In distributed mode, use device="cuda"; the framework assigns device ordinals.

A minimal Dask outline looks like this:

import dask.array as da
import numpy as np
from dask_cuda import LocalCUDACluster
from distributed import Client
from xgboost import dask as dxgb

X_train = da.random.standard_normal(
    size=(1_000_000, 64),
    chunks=(125_000, 64),
).astype(np.float32)
y_train = (
    3.0 * X_train[:, 0]
    - 1.5 * X_train[:, 1]
    + 0.2 * X_train[:, 2]
).astype(np.float32)

with LocalCUDACluster() as cluster, Client(cluster) as client:
    dtrain = dxgb.DaskQuantileDMatrix(client, X_train, y_train)
    result = dxgb.train(
        client,
        params={
            "objective": "reg:squarederror",
            "tree_method": "hist",
            "device": "cuda",
            "max_depth": 8,
            "learning_rate": 0.05,
        },
        dtrain=dtrain,
        num_boost_round=600,
    )
    result["booster"].save_model("dask-model.json")

For PySpark, feature columns can be numeric columns directly, and each XGBoost worker corresponds to one Spark task:

from pyspark.sql import SparkSession, functions as F
from xgboost.spark import SparkXGBRegressor

spark = SparkSession.builder.getOrCreate()

training_dataframe = (
    spark.range(1_000_000, numPartitions=8)
    .select(
        F.randn(42).alias("f0"),
        F.randn(43).alias("f1"),
        F.randn(44).alias("f2"),
    )
    .withColumn(
        "label",
        3.0 * F.col("f0")
        - 1.5 * F.col("f1")
        + 0.2 * F.col("f2")
        + 0.1 * F.randn(45),
    )
)

estimator = SparkXGBRegressor(
    features_col=["f0", "f1", "f2"],
    label_col="label",
    prediction_col="prediction",
    num_workers=4,
    n_estimators=600,
    max_depth=8,
    learning_rate=0.05,
    tree_method="hist",
    device="cuda",
)

spark_model = estimator.fit(training_dataframe)
predictions = spark_model.transform(training_dataframe.limit(1000))
predictions.select("label", "prediction").show(5)

Keep partitions large enough to amortize scheduling and communication. In Spark, configure spark.task.cpus and GPU resources at the cluster level instead of setting XGBoost thread counts manually.

Persist the model, not the Python process

Use XGBoost model I/O for deployment and store pipeline metadata separately.

import json
from pathlib import Path

import xgboost as xgb

model.save_model("model.json")

metadata = {
    "xgboost_version": xgb.__version__,
    "feature_order": [f"feature_{index}" for index in range(X_train.shape[1])],
    "input_dtype": "float32",
    "training_device": "cuda",
}
Path("model.metadata.json").write_text(
    json.dumps(metadata, indent=2),
    encoding="utf-8",
)

reloaded = xgb.XGBRegressor()
reloaded.load_model("model.json")

JSON and UBJSON model files preserve the trees, objective, and essential model parameters in a stable representation. Pickle, joblib, and cloudpickle capture a Python memory snapshot whose internal attributes can change between XGBoost or Python versions. XGBoost guarantees backward compatibility for model files, not memory snapshots.

Cherry on the cake: the pickle that broke after an upgrade

A real XGBoost issue documented a classifier trained with 1.2.1 and pickled, then loaded under 1.4.2. The restored estimator lacked the newer use_label_encoder attribute and failed during normal use. The durable recovery pattern is to recreate the original environment, load the snapshot there, and immediately export with save_model.

There is also a security reason to avoid untrusted pickle files: Python explicitly warns that malicious pickle data can execute arbitrary code during unpickling. Treat a model artifact as executable supply-chain input unless it is a validated XGBoost model file from a trusted build pipeline.

Serve contiguous batches with inplace_predict

Booster.inplace_predict accepts ordinary NumPy arrays and bypasses explicit DMatrix construction. Prediction is thread-safe for tree boosters as long as the same booster is not being mutated or trained concurrently.

import numpy as np
import xgboost as xgb

FEATURE_COUNT = 64

booster = xgb.Booster()
booster.load_model("model.json")


def predict_batch(rows: list[list[float]]) -> list[float]:
    batch = np.asarray(rows, dtype=np.float32, order="C")

    if batch.ndim != 2:
        raise ValueError("Input must be a two-dimensional batch.")
    if batch.shape[1] != FEATURE_COUNT:
        raise ValueError(
            f"Expected {FEATURE_COUNT} features, received {batch.shape[1]}."
        )

    predictions = booster.inplace_predict(batch)
    return predictions.tolist()

A service should combine requests for a short, bounded window, enforce a maximum batch size, and monitor p95 and p99 under concurrency. Preserve training-time feature order and dtype exactly.

Measured CPU latency

These measurements were produced in a Linux container with XGBoost 3.1.3, five CPU threads, 300 depth-eight trees, 64 float32 features, and warmed calls. The DMatrix route includes matrix construction; network, JSON parsing, and framework overhead are excluded.

Batch size	inplace_predict p50	inplace_predict p95	DMatrix route p50	p50 speedup	p50 per row
1	0.188 ms	0.295 ms	0.295 ms	1.56×	188.5 µs
32	0.416 ms	0.654 ms	0.482 ms	1.16×	13.0 µs
256	0.582 ms	0.998 ms	0.952 ms	1.63×	2.27 µs
4,096	6.216 ms	10.910 ms	6.916 ms	1.11×	1.52 µs

The main gain is batching. In this run, moving from one row to 256 rows reduced median compute cost per row by roughly 83 times; replacing the DMatrix path added a smaller, workload-dependent improvement.

Add monotonic constraints for sane behavior

Accuracy alone does not prevent implausible responses. For an e-commerce conversion model, you may require predictions to increase with review score and decrease with delivery time, all else equal.

import numpy as np
import pandas as pd
import xgboost as xgb

rng = np.random.default_rng(42)
row_count = 75_000

features = pd.DataFrame(
    {
        "discount_pct": rng.uniform(0.0, 40.0, row_count),
        "delivery_days": rng.integers(1, 10, row_count),
        "review_score": rng.uniform(1.0, 5.0, row_count),
        "price": rng.uniform(10.0, 500.0, row_count),
    }
)

score = (
    -2.0
    + 0.04 * features["discount_pct"]
    - 0.30 * features["delivery_days"]
    + 0.50 * features["review_score"]
    - 0.002 * features["price"]
)
probability = 1.0 / (1.0 + np.exp(-score))
labels = rng.binomial(1, probability).astype(np.int32)

constrained_model = xgb.XGBClassifier(
    n_estimators=500,
    max_depth=7,
    learning_rate=0.05,
    objective="binary:logistic",
    eval_metric="logloss",
    tree_method="hist",
    device="cuda",
    max_bin=512,
    monotone_constraints={
        "discount_pct": 1,
        "delivery_days": -1,
        "review_score": 1,
    },
    random_state=42,
)
constrained_model.fit(features, labels)

A value of 1 enforces an increasing relationship, -1 a decreasing relationship, and omitted features remain unconstrained. Histogram training can produce shallow trees when constraints eliminate available split candidates; increasing max_bin offers more candidate thresholds. Test the required direction with automated prediction sweeps before every release.

Production checklist

Before deployment, verify that:

• CPU and GPU training are benchmarked on the real feature matrix.

• QuantileDMatrix is tested before external-memory I/O is introduced.

• External-memory batches are large and cached on appropriate storage.

• Distributed partitions are balanced and large enough to amortize overhead.

• The artifact is saved with save_model("model.json"), not pickle.

• Feature order, dtypes, preprocessing version, and XGBoost version are recorded.

• A clean inference image can load the artifact without the training pipeline.

• Inference uses contiguous batches and measures p95/p99 under concurrency.

• Monotonic requirements have explicit regression tests.

Take one representative production dataset, benchmark CPU versus GPU end to end, export the chosen model as model.json, and reproduce the batching table against your own service-level objective. That exercise will expose most memory, persistence, and latency mistakes before they become incidents.