merge changes Jun092023

.github/workflows/update_rapids.yml

@@ -1,11 +1,13 @@
 name: update-rapids
 
 on:
+  workflow_dispatch:
   schedule:
-    - cron: "20 20 * * *"  # Run once daily
+    - cron: "0 20 * * *"  # Run once daily
 
 permissions:
-  contents: read  # to fetch code (actions/checkout)
+  pull-requests: write
+  contents: write
 
 defaults:
   run:
@@ -37,4 +39,6 @@ jobs:
           tests/buildkite
         branch: create-pull-request/update-rapids
         base: master
+        title: "[CI] Update RAPIDS to latest stable"
+        commit-message: "[CI] Update RAPIDS to latest stable"
       

demo/guide-python/learning_to_rank.py

@@ -0,0 +1,212 @@
+"""
+Getting started with learning to rank
+=====================================
+
+  .. versionadded:: 2.0.0
+
+This is a demonstration of using XGBoost for learning to rank tasks using the
+MSLR_10k_letor dataset. For more infomation about the dataset, please visit its
+`description page <https://www.microsoft.com/en-us/research/project/mslr/>`_.
+
+This is a two-part demo, the first one contains a basic example of using XGBoost to
+train on relevance degree, and the second part simulates click data and enable the
+position debiasing training.
+
+For an overview of learning to rank in XGBoost, please see
+:doc:`Learning to Rank </tutorials/learning_to_rank>`.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import pickle as pkl
+
+import numpy as np
+import pandas as pd
+from sklearn.datasets import load_svmlight_file
+
+import xgboost as xgb
+from xgboost.testing.data import RelDataCV, simulate_clicks, sort_ltr_samples
+
+
+def load_mlsr_10k(data_path: str, cache_path: str) -> RelDataCV:
+    """Load the MSLR10k dataset from data_path and cache a pickle object in cache_path.
+
+    Returns
+    -------
+
+    A list of tuples [(X, y, qid), ...].
+
+    """
+    root_path = os.path.expanduser(args.data)
+    cacheroot_path = os.path.expanduser(args.cache)
+    cache_path = os.path.join(cacheroot_path, "MSLR_10K_LETOR.pkl")
+
+    # Use only the Fold1 for demo:
+    # Train,      Valid, Test
+    # {S1,S2,S3}, S4,    S5
+    fold = 1
+
+    if not os.path.exists(cache_path):
+        fold_path = os.path.join(root_path, f"Fold{fold}")
+        train_path = os.path.join(fold_path, "train.txt")
+        valid_path = os.path.join(fold_path, "vali.txt")
+        test_path = os.path.join(fold_path, "test.txt")
+        X_train, y_train, qid_train = load_svmlight_file(
+            train_path, query_id=True, dtype=np.float32
+        )
+        y_train = y_train.astype(np.int32)
+        qid_train = qid_train.astype(np.int32)
+
+        X_valid, y_valid, qid_valid = load_svmlight_file(
+            valid_path, query_id=True, dtype=np.float32
+        )
+        y_valid = y_valid.astype(np.int32)
+        qid_valid = qid_valid.astype(np.int32)
+
+        X_test, y_test, qid_test = load_svmlight_file(
+            test_path, query_id=True, dtype=np.float32
+        )
+        y_test = y_test.astype(np.int32)
+        qid_test = qid_test.astype(np.int32)
+
+        data = RelDataCV(
+            train=(X_train, y_train, qid_train),
+            test=(X_test, y_test, qid_test),
+            max_rel=4,
+        )
+
+        with open(cache_path, "wb") as fd:
+            pkl.dump(data, fd)
+
+    with open(cache_path, "rb") as fd:
+        data = pkl.load(fd)
+
+    return data
+
+
+def ranking_demo(args: argparse.Namespace) -> None:
+    """Demonstration for learning to rank with relevance degree."""
+    data = load_mlsr_10k(args.data, args.cache)
+
+    # Sort data according to query index
+    X_train, y_train, qid_train = data.train
+    sorted_idx = np.argsort(qid_train)
+    X_train = X_train[sorted_idx]
+    y_train = y_train[sorted_idx]
+    qid_train = qid_train[sorted_idx]
+
+    X_test, y_test, qid_test = data.test
+    sorted_idx = np.argsort(qid_test)
+    X_test = X_test[sorted_idx]
+    y_test = y_test[sorted_idx]
+    qid_test = qid_test[sorted_idx]
+
+    ranker = xgb.XGBRanker(
+        tree_method="gpu_hist",
+        lambdarank_pair_method="topk",
+        lambdarank_num_pair_per_sample=13,
+        eval_metric=["ndcg@1", "ndcg@8"],
+    )
+    ranker.fit(
+        X_train,
+        y_train,
+        qid=qid_train,
+        eval_set=[(X_test, y_test)],
+        eval_qid=[qid_test],
+        verbose=True,
+    )
+
+
+def click_data_demo(args: argparse.Namespace) -> None:
+    """Demonstration for learning to rank with click data."""
+    data = load_mlsr_10k(args.data, args.cache)
+    train, test = simulate_clicks(data)
+    assert test is not None
+
+    assert train.X.shape[0] == train.click.size
+    assert test.X.shape[0] == test.click.size
+    assert test.score.dtype == np.float32
+    assert test.click.dtype == np.int32
+
+    X_train, clicks_train, y_train, qid_train = sort_ltr_samples(
+        train.X,
+        train.y,
+        train.qid,
+        train.click,
+        train.pos,
+    )
+    X_test, clicks_test, y_test, qid_test = sort_ltr_samples(
+        test.X,
+        test.y,
+        test.qid,
+        test.click,
+        test.pos,
+    )
+
+    class ShowPosition(xgb.callback.TrainingCallback):
+        def after_iteration(
+            self,
+            model: xgb.Booster,
+            epoch: int,
+            evals_log: xgb.callback.TrainingCallback.EvalsLog,
+        ) -> bool:
+            config = json.loads(model.save_config())
+            ti_plus = np.array(config["learner"]["objective"]["ti+"])
+            tj_minus = np.array(config["learner"]["objective"]["tj-"])
+            df = pd.DataFrame({"ti+": ti_plus, "tj-": tj_minus})
+            print(df)
+            return False
+
+    ranker = xgb.XGBRanker(
+        n_estimators=512,
+        tree_method="gpu_hist",
+        learning_rate=0.01,
+        reg_lambda=1.5,
+        subsample=0.8,
+        sampling_method="gradient_based",
+        # LTR specific parameters
+        objective="rank:ndcg",
+        # - Enable bias estimation
+        lambdarank_unbiased=True,
+        # - normalization (1 / (norm + 1))
+        lambdarank_bias_norm=1,
+        # - Focus on the top 12 documents
+        lambdarank_num_pair_per_sample=12,
+        lambdarank_pair_method="topk",
+        ndcg_exp_gain=True,
+        eval_metric=["ndcg@1", "ndcg@3", "ndcg@5", "ndcg@10"],
+        callbacks=[ShowPosition()],
+    )
+    ranker.fit(
+        X_train,
+        clicks_train,
+        qid=qid_train,
+        eval_set=[(X_test, y_test), (X_test, clicks_test)],
+        eval_qid=[qid_test, qid_test],
+        verbose=True,
+    )
+    ranker.predict(X_test)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Demonstration of learning to rank using XGBoost."
+    )
+    parser.add_argument(
+        "--data",
+        type=str,
+        help="Root directory of the MSLR-WEB10K data.",
+        required=True,
+    )
+    parser.add_argument(
+        "--cache",
+        type=str,
+        help="Directory for caching processed data.",
+        required=True,
+    )
+    args = parser.parse_args()
+
+    ranking_demo(args)
+    click_data_demo(args)

doc/contrib/coding_guide.rst

@@ -16,8 +16,10 @@ C++ Coding Guideline
   * Each line of text may contain up to 100 characters.
   * The use of C++ exceptions is allowed.
 
-- Use C++11 features such as smart pointers, braced initializers, lambda functions, and ``std::thread``.
+- Use C++17 features such as smart pointers, braced initializers, lambda functions, and ``std::thread``.
 - Use Doxygen to document all the interface code.
+- We have some comments around symbols imported by headers, some of those are hinted by `include-what-you-use <https://include-what-you-use.org>`_. It's not required.
+- We use clang-tidy and clang-format. You can check their configuration in the root directory of the XGBoost source tree.
 - We have a series of automatic checks to ensure that all of our codebase complies with the Google style. Before submitting your pull request, you are encouraged to run the style checks on your machine. See :ref:`running_checks_locally`.
 
 ***********************

doc/tutorials/index.rst

@@ -21,6 +21,7 @@ See `Awesome XGBoost <https://github.com/dmlc/xgboost/tree/master/demo>`_ for mo
   monotonic
   rf
   feature_interaction_constraint
+  learning_to_rank
   aft_survival_analysis
   c_api_tutorial
   input_format

doc/tutorials/learning_to_rank.rst

@@ -0,0 +1,201 @@
+################
+Learning to Rank
+################
+
+**Contents**
+
+.. contents::
+  :local:
+  :backlinks: none
+
+********
+Overview
+********
+Often in the context of information retrieval, learning-to-rank aims to train a model that arranges a set of query results into an ordered list `[1] <#references>`__. For surprivised learning-to-rank, the predictors are sample documents encoded as feature matrix, and the labels are relevance degree for each sample. Relevance degree can be multi-level (graded) or binary (relevant or not). The training samples are often grouped by their query index with each query group containing multiple query results.
+
+XGBoost implements learning to rank through a set of objective functions and performance metrics. The default objective is ``rank:ndcg`` based on the ``LambdaMART`` `[2] <#references>`__ algorithm, which in turn is an adaptation of the ``LambdaRank`` `[3] <#references>`__ framework to gradient boosting trees. For a history and a summary of the algorithm, see `[5] <#references>`__. The implementation in XGBoost features deterministic GPU computation, distributed training, position debiasing and two different pair construction strategies.
+
+************************************
+Training with the Pariwise Objective
+************************************
+``LambdaMART`` is a pairwise ranking model, meaning that it compares the relevance degree for every pair of samples in a query group and calculate a proxy gradient for each pair. The default objective ``rank:ndcg`` is using the surrogate gradient derived from the ``ndcg`` metric. To train a XGBoost model, we need an additional sorted array called ``qid`` for specifying the query group of input samples. An example input would look like this:
+
++-------+-----------+---------------+
+|   QID |   Label   |   Features    |
++=======+===========+===============+
+|   1   |   0       |   :math:`x_1` |
++-------+-----------+---------------+
+|   1   |   1       |   :math:`x_2` |
++-------+-----------+---------------+
+|   1   |   0       |   :math:`x_3` |
++-------+-----------+---------------+
+|   2   |   0       |   :math:`x_4` |
++-------+-----------+---------------+
+|   2   |   1       |   :math:`x_5` |
++-------+-----------+---------------+
+|   2   |   1       |   :math:`x_6` |
++-------+-----------+---------------+
+|   2   |   1       |   :math:`x_7` |
++-------+-----------+---------------+
+
+Notice that the samples are sorted based on their query index in a non-decreasing order. In the above example, the first three samples belong to the first query and the next four samples belong to the second. For the sake of simplicity, we will use a synthetic binary learning-to-rank dataset in the following code snippets, with binary labels representing whether the result is relevant or not, and randomly assign the query group index to each sample. For an example that uses a real world dataset, please see :ref:`sphx_glr_python_examples_learning_to_rank.py`.
+
+.. code-block:: python
+
+  from sklearn.datasets import make_classification
+  import numpy as np
+
+  import xgboost as xgb
+
+  # Make a synthetic ranking dataset for demonstration
+  X, y = make_classification(random_state=rng)
+  rng = np.random.default_rng(1994)
+  n_query_groups = 3
+  qid = rng.integers(0, 3, size=X.shape[0])
+
+  # Sort the inputs based on query index
+  sorted_idx = np.argsort(qid)
+  X = X[sorted_idx, :]
+  y = y[sorted_idx]
+
+The simpliest way to train a ranking model is by using the scikit-learn estimator interface. Continuing the previous snippet, we can train a simple ranking model without tuning:
+
+.. code-block:: python
+
+  ranker = xgb.XGBRanker(tree_method="hist", lambdarank_num_pair_per_sample=8, objective="rank:ndcg", lambdarank_pair_method="topk")
+  ranker.fit(X, y, qid=qid)
+
+Please note that, as of writing, there's no learning-to-rank interface in scikit-learn. As a result, the :py:class:`xgboost.XGBRanker` class does not fully conform the scikit-learn estimator guideline and can not be directly used with some of its utility functions. For instances, the ``auc_score`` and ``ndcg_score`` in scikit-learn don't consider query group information nor the pairwise loss. Most of the metrics are implemented as part of XGBoost, but to use scikit-learn utilities like :py:func:`sklearn.model_selection.cross_validation`, we need to make some adjustments in order to pass the ``qid`` as an additional parameter for :py:meth:`xgboost.XGBRanker.score`. Given a data frame ``X`` (either pandas or cuDF), add the column ``qid`` as follows:
+
+.. code-block:: python
+
+  df = pd.DataFrame(X, columns=[str(i) for i in range(X.shape[1]))
+  df["qid"] = qid
+  ranker.fit(df, y)  # No need to pass qid as a separate argument
+
+  from sklearn.model_selection import StratifiedGroupKFold, cross_val_score
+  # Works with cv in scikit-learn, along with HPO utilities like GridSearchCV
+  kfold = StratifiedGroupKFold(shuffle=False)
+  cross_val_score(ranker, df, y, cv=kfold, groups=df.qid)
+
+The above snippets build a model using ``LambdaMART`` with the ``NDCG@8`` metric. The outputs of a ranker are relevance scores:
+
+.. code-block:: python
+
+  scores = ranker.predict(X)
+  sorted_idx = np.argsort(scores)[::-1]
+  # Sort the relevance scores from most relevant to least relevant
+  scores = scores[sorted_idx]
+
+
+*************
+Position Bias
+*************
+
+.. versionadded:: 2.0.0
+
+.. note::
+
+   The feature is considered experimental. This is a heated research area, and your input is much appreciated!
+
+Obtaining real relevance degrees for query results is an expensive and strenuous, requiring human labelers to label all results one by one. When such labeling task is infeasible, we might want to train the learning-to-rank model on user click data instead, as it is relatively easy to collect. Another advantage of using click data directly is that it can reflect the most up-to-date user preferences `[1] <#references>`__. However, user clicks are often biased,  as users tend to choose results that are displayed in higher positions. User clicks are also noisy, where users might accidentally click on irrelevant documents. To ameliorate these issues, XGBoost implements the ``Unbiased LambdaMART`` `[4] <#references>`__ algorithm to debias the position-dependent click data. The feature can be enabled by the ``lambdarank_unbiased`` parameter; see :ref:`ltr-param` for related options and :ref:`sphx_glr_python_examples_learning_to_rank.py` for a worked example with simulated user clicks.
+
+****
+Loss
+****
+
+XGBoost implements different ``LambdaMART`` objectives based on different metrics. We list them here as a reference. Other than those used as objective function, XGBoost also implements metrics like ``pre`` (for precision) for evaluation. See :doc:`parameters </parameter>` for available options and the following sections for how to choose these objectives based of the amount of effective pairs.
+
+* NDCG
+
+`Normalized Discounted Cumulative Gain` ``NDCG`` can be used with both binary relevance and multi-level relevance. If you are not sure about your data, this metric can be used as the default. The name for the objective is ``rank:ndcg``.
+
+
+* MAP
+
+`Mean average precision` ``MAP`` is a binary measure. It can be used when the relevance label is 0 or 1. The name for the objective is ``rank:map``.
+
+
+* Pairwise
+
+The `LambdaMART` algorithm scales the logistic loss with learning to rank metrics like ``NDCG`` in the hope of including ranking information into the loss function. The ``rank:pairwise`` loss is the original version of the pairwise loss, also known as the `RankNet loss` `[7] <#references>`__ or the `pairwise logistic loss`. Unlike the ``rank:map`` and the ``rank:ndcg``, no scaling is applied (:math:`|\Delta Z_{ij}| = 1`).
+
+Whether scaling with a LTR metric is actually more effective is still up for debate; `[8] <#references>`__ provides a theoretical foundation for general lambda loss functions and some insights into the framework.
+
+******************
+Constructing Pairs
+******************
+
+There are two implemented strategies for constructing document pairs for :math:`\lambda`-gradient calculation. The first one is the ``mean`` method, another one is the ``topk`` method. The preferred strategy can be specified by the ``lambdarank_pair_method`` parameter.
+
+For the ``mean`` strategy, XGBoost samples ``lambdarank_num_pair_per_sample`` pairs for each document in a query list. For example, given a list of 3 documents and ``lambdarank_num_pair_per_sample`` is set to 2, XGBoost will randomly sample 6 pairs, assuming the labels for these documents are different. On the other hand, if the pair method is set to ``topk``, XGBoost constructs about :math:`k \times |query|` number of pairs with :math:`|query|` pairs for each sample at the top :math:`k = lambdarank\_num\_pair` position. The number of pairs counted here is an approximation since we skip pairs that have the same label.
+
+*********************
+Obtaining Good Result
+*********************
+
+Learning to rank is a sophisticated task and an active research area. It's not trivial to train a model that generalizes well. There are multiple loss functions available in XGBoost along with a set of hyperparameters. This section contains some hints for how to choose hyperparameters as a starting point. One can further optimize the model by tuning these hyperparameters.
+
+The first question would be how to choose an objective that matches the task at hand. If your input data has multi-level relevance degrees, then either ``rank:ndcg`` or ``rank:pairwise`` should be used. However, when the input has binary labels, we have multiple options based on the target metric. `[6] <#references>`__ provides some guidelines on this topic and users are encouraged to see the analysis done in their work. The choice should be based on the number of `effective pairs`, which refers to the number of pairs that can generate non-zero gradient and contribute to training. `LambdaMART` with ``MRR`` has the least amount of effective pairs as the :math:`\lambda`-gradient is only non-zero when the pair contains a non-relevant document ranked higher than the top relevant document. As a result, it's not implemented in XGBoost. Since ``NDCG`` is a multi-level metric, it usually generate more effective pairs than ``MAP``.
+
+However, when there are sufficiently many effective pairs, it's shown in `[6] <#references>`__ that matching the target metric with the objective is of significance. When the target metric is ``MAP`` and you are using a large dataset that can provide a sufficient amount of effective pairs, ``rank:map`` can in theory yield higher ``MAP`` value than ``rank:ndcg``.
+
+The consideration of effective pairs also applies to the choice of pair method (``lambdarank_pair_method``) and the number of pairs for each sample (``lambdarank_num_pair_per_sample``). For example, the mean-``NDCG`` considers more pairs than ``NDCG@10``, so the former generates more effective pairs and provides more granularity than the latter. Also, using the ``mean`` strategy can help the model generalize with random sampling. However, one might want to focus the training on the top :math:`k` documents instead of using all pairs, to better fit their real-world application.
+
+When using the mean strategy for generating pairs, where the target metric (like ``NDCG``) is computed over the whole query list, users can specify how many pairs should be generated per each document, by setting the ``lambdarank_num_pair_per_sample``. XGBoost will randomly sample ``lambdarank_num_pair_per_sample`` pairs for each element in the query group (:math:`|pairs| = |query| \times num\_pairsample`). Often, setting it to 1 can produce reasonable results. In cases where performance is inadequate due to insufficient number of effective pairs being generated, set ``lambdarank_num_pair_per_sample`` to a higher value. As more document pairs are generated, more effective pairs will be generated as well.
+
+On the other hand, if you are prioritizing the top :math:`k` documents, the ``lambdarank_num_pair_per_sample`` should be set slightly higher than :math:`k` (with a few more documents) to obtain a good training result.
+
+**Summary** If you have large amount of training data:
+
+* Use the target-matching objective.
+* Choose the ``topk`` strategy for generating document pairs (if it's appropriate for your application).
+
+On the other hand, if you have comparatively small amount of training data:
+
+* Select ``NDCG`` or the RankNet loss (``rank:pairwise``).
+* Choose the ``mean`` strategy for generating document pairs, to obtain more effective pairs.
+
+For any method chosen, you can modify ``lambdarank_num_pair_per_sample`` to control the amount of pairs generated.
+
+********************
+Distributed Training
+********************
+XGBoost implements distributed learning-to-rank with integration of multiple frameworks including Dask, Spark, and PySpark. The interface is similar to the single-node counterpart. Please refer to document of the respective XGBoost interface for details. Scattering a query group onto multiple workers is theoretically sound but can affect the model accuracy. For most of the use cases, the small discrepancy is not an issue, as the amount of training data is usually large when distributed training is used. As a result, users don't need to partition the data based on query groups. As long as each data partition is correctly sorted by query IDs, XGBoost can aggregate sample gradients accordingly.
+
+*******************
+Reproducible Result
+*******************
+
+Like any other tasks, XGBoost should generate reproducible results given the same hardware and software environments (and data partitions, if distributed interface is used). Even when the underlying environment has changed, the result should still be consistent. However, when the ``lambdarank_pair_method`` is set to ``mean``, XGBoost uses random sampling, and results may differ depending on the platform used. The random number generator used on Windows (Microsoft Visual C++) is different from the ones used on other platforms like Linux (GCC, Clang) [#f0]_, so the output varies significantly between these platforms.
+
+.. [#f0] `minstd_rand` implementation is different on MSVC. The implementations from GCC and Thrust produce the same output.
+
+**********
+References
+**********
+
+[1] Tie-Yan Liu. 2009. "`Learning to Rank for Information Retrieval`_". Found. Trends Inf. Retr. 3, 3 (March 2009), 225–331.
+
+[2] Christopher J. C. Burges, Robert Ragno, and Quoc Viet Le. 2006. "`Learning to rank with nonsmooth cost functions`_". In Proceedings of the 19th International Conference on Neural Information Processing Systems (NIPS'06). MIT Press, Cambridge, MA, USA, 193–200.
+
+[3] Wu, Q., Burges, C.J.C., Svore, K.M. et al. "`Adapting boosting for information retrieval measures`_". Inf Retrieval 13, 254–270 (2010).
+
+[4] Ziniu Hu, Yang Wang, Qu Peng, Hang Li. "`Unbiased LambdaMART: An Unbiased Pairwise Learning-to-Rank Algorithm`_". Proceedings of the 2019 World Wide Web Conference.
+
+[5] Burges, Chris J.C. "`From RankNet to LambdaRank to LambdaMART: An Overview`_". MSR-TR-2010-82
+
+[6] Pinar Donmez, Krysta M. Svore, and Christopher J.C. Burges. 2009. "`On the local optimality of LambdaRank`_". In Proceedings of the 32nd international ACM SIGIR conference on Research and development in information retrieval (SIGIR '09). Association for Computing Machinery, New York, NY, USA, 460–467.
+
+[7] Chris Burges, Tal Shaked, Erin Renshaw, Ari Lazier, Matt Deeds, Nicole Hamilton, and Greg Hullender. 2005. "`Learning to rank using gradient descent`_". In Proceedings of the 22nd international conference on Machine learning (ICML '05). Association for Computing Machinery, New York, NY, USA, 89–96.
+
+[8] Xuanhui Wang and Cheng Li and Nadav Golbandi and Mike Bendersky and Marc Najork. 2018. "`The LambdaLoss Framework for Ranking Metric Optimization`_". Proceedings of The 27th ACM International Conference on Information and Knowledge Management (CIKM '18).
+
+.. _`Learning to Rank for Information Retrieval`: https://doi.org/10.1561/1500000016
+.. _`Learning to rank with nonsmooth cost functions`: https://dl.acm.org/doi/10.5555/2976456.2976481
+.. _`Adapting boosting for information retrieval measures`: https://doi.org/10.1007/s10791-009-9112-1
+.. _`Unbiased LambdaMART: An Unbiased Pairwise Learning-to-Rank Algorithm`: https://dl.acm.org/doi/10.1145/3308558.3313447
+.. _`From RankNet to LambdaRank to LambdaMART: An Overview`: https://www.microsoft.com/en-us/research/publication/from-ranknet-to-lambdarank-to-lambdamart-an-overview/
+.. _`On the local optimality of LambdaRank`: https://doi.org/10.1145/1571941.1572021
+.. _`Learning to rank using gradient descent`:  https://doi.org/10.1145/1102351.1102363
+.. _`The LambdaLoss Framework for Ranking Metric Optimization`: https://dl.acm.org/doi/10.1145/3269206.3271784

jvm-packages/pom.xml

@@ -431,7 +431,7 @@
             <plugin>
                 <groupId>org.apache.maven.plugins</groupId>
                 <artifactId>maven-surefire-plugin</artifactId>
-                <version>3.1.0</version>
+                <version>3.1.2</version>
                 <configuration>
                     <skipTests>false</skipTests>
                     <useSystemClassLoader>false</useSystemClassLoader>
@@ -454,7 +454,7 @@
         <plugins>
             <plugin>
                 <artifactId>maven-project-info-reports-plugin</artifactId>
-                <version>3.4.4</version>
+                <version>3.4.5</version>
             </plugin>
             <plugin>
                 <groupId>net.alchim31.maven</groupId>

jvm-packages/xgboost4j-spark-gpu/src/test/scala/ml/dmlc/xgboost4j/scala/rapids/spark/GpuXGBoostRegressorSuite.scala

@@ -220,7 +220,7 @@ class GpuXGBoostRegressorSuite extends GpuTestSuite {
 
   test("Ranking: train with Group") {
     withGpuSparkSession(enableCsvConf()) { spark =>
-      val xgbParam = Map("eta" -> 0.1f, "max_depth" -> 2, "objective" -> "rank:pairwise",
+      val xgbParam = Map("eta" -> 0.1f, "max_depth" -> 2, "objective" -> "rank:ndcg",
         "num_round" -> 10, "num_workers" -> 1, "tree_method" -> "gpu_hist",
         "features_cols" -> featureNames, "label_col" -> labelName)
       val Array(trainingDf, testDf) = spark.read.option("header", "true").schema(schema)

jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/params/LearningTaskParams.scala

@@ -25,7 +25,7 @@ private[spark] trait LearningTaskParams extends Params {
   /**
    * Specify the learning task and the corresponding learning objective.
    * options: reg:squarederror, reg:squaredlogerror, reg:logistic, binary:logistic, binary:logitraw,
-   * count:poisson, multi:softmax, multi:softprob, rank:pairwise, reg:gamma.
+   * count:poisson, multi:softmax, multi:softprob, rank:ndcg, reg:gamma.
    * default: reg:squarederror
    */
   final val objective = new Param[String](this, "objective",

jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostGeneralSuite.scala

@@ -201,7 +201,7 @@ class XGBoostGeneralSuite extends AnyFunSuite with TmpFolderPerSuite with PerTes
       sc,
       buildTrainingRDD,
       List("eta" -> "1", "max_depth" -> "6",
-        "objective" -> "rank:pairwise", "num_round" -> 5, "num_workers" -> numWorkers,
+        "objective" -> "rank:ndcg", "num_round" -> 5, "num_workers" -> numWorkers,
         "custom_eval" -> null, "custom_obj" -> null, "use_external_memory" -> false,
         "missing" -> Float.NaN).toMap)
 
@@ -268,7 +268,7 @@ class XGBoostGeneralSuite extends AnyFunSuite with TmpFolderPerSuite with PerTes
     val training = buildDataFrameWithGroup(Ranking.train, 5)
     val Array(train, eval1, eval2) = training.randomSplit(Array(0.6, 0.2, 0.2), 0)
     val paramMap1 = Map("eta" -> "1", "max_depth" -> "6",
-      "objective" -> "rank:pairwise",
+      "objective" -> "rank:ndcg",
       "num_round" -> 5, "num_workers" -> numWorkers, "group_col" -> "group")
     val xgb1 = new XGBoostRegressor(paramMap1).setEvalSets(Map("eval1" -> eval1, "eval2" -> eval2))
     val model1 = xgb1.fit(train)
@@ -281,7 +281,7 @@ class XGBoostGeneralSuite extends AnyFunSuite with TmpFolderPerSuite with PerTes
     assert(model1.summary.trainObjectiveHistory !== model1.summary.validationObjectiveHistory(1))
 
     val paramMap2 = Map("eta" -> "1", "max_depth" -> "6",
-      "objective" -> "rank:pairwise",
+      "objective" -> "rank:ndcg",
       "num_round" -> 5, "num_workers" -> numWorkers, "group_col" -> "group",
       "eval_sets" -> Map("eval1" -> eval1, "eval2" -> eval2))
     val xgb2 = new XGBoostRegressor(paramMap2)

jvm-packages/xgboost4j-spark/src/test/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostRegressorSuite.scala

@@ -121,7 +121,7 @@ class XGBoostRegressorSuite extends AnyFunSuite with PerTest with TmpFolderPerSu
 
   test("ranking: use group data") {
     val paramMap = Map("eta" -> "1", "max_depth" -> "6", "silent" -> "1",
-      "objective" -> "rank:pairwise", "num_workers" -> numWorkers, "num_round" -> 5,
+      "objective" -> "rank:ndcg", "num_workers" -> numWorkers, "num_round" -> 5,
       "group_col" -> "group", "tree_method" -> treeMethod)
 
     val trainingDF = buildDataFrameWithGroup(Ranking.train)

python-package/xgboost/_typing.py

@@ -31,7 +31,7 @@ ArrayLike = Any
 PathLike = Union[str, os.PathLike]
 CupyT = ArrayLike  # maybe need a stub for cupy arrays
 NumpyOrCupy = Any
-NumpyDType = Union[str, Type[np.number]]
+NumpyDType = Union[str, Type[np.number]]  # pylint: disable=invalid-name
 PandasDType = Any  # real type is pandas.core.dtypes.base.ExtensionDtype
 
 FloatCompatible = Union[float, np.float32, np.float64]

python-package/xgboost/sklearn.py

@@ -1796,7 +1796,11 @@ def _get_qid(
 
 
 @xgboost_model_doc(
-    """Implementation of the Scikit-Learn API for XGBoost Ranking.""",
+    """Implementation of the Scikit-Learn API for XGBoost Ranking.
+
+See :doc:`Learning to Rank </tutorials/learning_to_rank>` for an introducion.
+
+    """,
     ["estimators", "model"],
     end_note="""
         .. note::
@@ -1845,7 +1849,7 @@ def _get_qid(
 class XGBRanker(XGBModel, XGBRankerMixIn):
     # pylint: disable=missing-docstring,too-many-arguments,invalid-name
     @_deprecate_positional_args
-    def __init__(self, *, objective: str = "rank:pairwise", **kwargs: Any):
+    def __init__(self, *, objective: str = "rank:ndcg", **kwargs: Any):
         super().__init__(objective=objective, **kwargs)
         if callable(self.objective):
             raise ValueError("custom objective function not supported by XGBRanker")
@@ -2029,7 +2033,7 @@ class XGBRanker(XGBModel, XGBRankerMixIn):
             self._Booster = train(
                 params,
                 train_dmatrix,
-                self.get_num_boosting_rounds(),
+                num_boost_round=self.get_num_boosting_rounds(),
                 early_stopping_rounds=early_stopping_rounds,
                 evals=evals,
                 evals_result=evals_result,

python-package/xgboost/testing/dask.py

@@ -1,6 +1,8 @@
 """Tests for dask shared by different test modules."""
 import numpy as np
+import pandas as pd
 from dask import array as da
+from dask import dataframe as dd
 from distributed import Client
 
 import xgboost as xgb
@@ -52,3 +54,22 @@ def check_init_estimation(tree_method: str, client: Client) -> None:
     """Test init estimation."""
     check_init_estimation_reg(tree_method, client)
     check_init_estimation_clf(tree_method, client)
+
+
+def check_uneven_nan(client: Client, tree_method: str, n_workers: int) -> None:
+    """Issue #9271, not every worker has missing value."""
+    assert n_workers >= 2
+
+    with client.as_current():
+        clf = xgb.dask.DaskXGBClassifier(tree_method=tree_method)
+        X = pd.DataFrame({"a": range(10000), "b": range(10000, 0, -1)})
+        y = pd.Series([*[0] * 5000, *[1] * 5000])
+
+        X["a"][:3000:1000] = np.NaN
+
+        client.wait_for_workers(n_workers=n_workers)
+
+        clf.fit(
+            dd.from_pandas(X, npartitions=n_workers),
+            dd.from_pandas(y, npartitions=n_workers),
+        )

python-package/xgboost/testing/data.py

@@ -1,11 +1,14 @@
+# pylint: disable=invalid-name
 """Utilities for data generation."""
 import os
 import zipfile
-from typing import Any, Generator, List, Tuple, Union
+from dataclasses import dataclass
+from typing import Any, Generator, List, NamedTuple, Optional, Tuple, Union
 from urllib import request
 
 import numpy as np
 import pytest
+from numpy import typing as npt
 from numpy.random import Generator as RNG
 from scipy import sparse
 
@@ -340,3 +343,263 @@ def get_mq2008(
         y_valid,
         qid_valid,
     )
+
+
+RelData = Tuple[sparse.csr_matrix, npt.NDArray[np.int32], npt.NDArray[np.int32]]
+
+
+@dataclass
+class ClickFold:
+    """A structure containing information about generated user-click data."""
+
+    X: sparse.csr_matrix
+    y: npt.NDArray[np.int32]
+    qid: npt.NDArray[np.int32]
+    score: npt.NDArray[np.float32]
+    click: npt.NDArray[np.int32]
+    pos: npt.NDArray[np.int64]
+
+
+class RelDataCV(NamedTuple):
+    """Simple data struct for holding a train-test split of a learning to rank dataset."""
+
+    train: RelData
+    test: RelData
+    max_rel: int
+
+    def is_binary(self) -> bool:
+        """Whether the label consists of binary relevance degree."""
+        return self.max_rel == 1
+
+
+class PBM:  # pylint: disable=too-few-public-methods
+    """Simulate click data with position bias model. There are other models available in
+    `ULTRA <https://github.com/ULTR-Community/ULTRA.git>`_ like the cascading model.
+
+    References
+    ----------
+    Unbiased LambdaMART: An Unbiased Pairwise Learning-to-Rank Algorithm
+
+    """
+
+    def __init__(self, eta: float) -> None:
+        # click probability for each relevance degree. (from 0 to 4)
+        self.click_prob = np.array([0.1, 0.16, 0.28, 0.52, 1.0])
+        exam_prob = np.array(
+            [0.68, 0.61, 0.48, 0.34, 0.28, 0.20, 0.11, 0.10, 0.08, 0.06]
+        )
+        # Observation probability, encoding positional bias for each position
+        self.exam_prob = np.power(exam_prob, eta)
+
+    def sample_clicks_for_query(
+        self, labels: npt.NDArray[np.int32], position: npt.NDArray[np.int64]
+    ) -> npt.NDArray[np.int32]:
+        """Sample clicks for one query based on input relevance degree and position.
+
+        Parameters
+        ----------
+
+        labels :
+            relevance_degree
+
+        """
+        labels = np.array(labels, copy=True)
+
+        click_prob = np.zeros(labels.shape)
+        # minimum
+        labels[labels < 0] = 0
+        # maximum
+        labels[labels >= len(self.click_prob)] = -1
+        click_prob = self.click_prob[labels]
+
+        exam_prob = np.zeros(labels.shape)
+        assert position.size == labels.size
+        ranks = np.array(position, copy=True)
+        # maximum
+        ranks[ranks >= self.exam_prob.size] = -1
+        exam_prob = self.exam_prob[ranks]
+
+        rng = np.random.default_rng(1994)
+        prob = rng.random(size=labels.shape[0], dtype=np.float32)
+
+        clicks: npt.NDArray[np.int32] = np.zeros(labels.shape, dtype=np.int32)
+        clicks[prob < exam_prob * click_prob] = 1
+        return clicks
+
+
+def rlencode(x: npt.NDArray[np.int32]) -> Tuple[npt.NDArray, npt.NDArray, npt.NDArray]:
+    """Run length encoding using numpy, modified from:
+    https://gist.github.com/nvictus/66627b580c13068589957d6ab0919e66
+
+    """
+    x = np.asarray(x)
+    n = x.size
+    starts = np.r_[0, np.flatnonzero(~np.isclose(x[1:], x[:-1], equal_nan=True)) + 1]
+    lengths = np.diff(np.r_[starts, n])
+    values = x[starts]
+    indptr = np.append(starts, np.array([x.size]))
+
+    return indptr, lengths, values
+
+
+def init_rank_score(
+    X: sparse.csr_matrix,
+    y: npt.NDArray[np.int32],
+    qid: npt.NDArray[np.int32],
+    sample_rate: float = 0.1,
+) -> npt.NDArray[np.float32]:
+    """We use XGBoost to generate the initial score instead of SVMRank for
+    simplicity. Sample rate is set to 0.1 by default so that we can test with small
+    datasets.
+
+    """
+    # random sample
+    rng = np.random.default_rng(1994)
+    n_samples = int(X.shape[0] * sample_rate)
+    index = np.arange(0, X.shape[0], dtype=np.uint64)
+    rng.shuffle(index)
+    index = index[:n_samples]
+
+    X_train = X[index]
+    y_train = y[index]
+    qid_train = qid[index]
+
+    # Sort training data based on query id, required by XGBoost.
+    sorted_idx = np.argsort(qid_train)
+    X_train = X_train[sorted_idx]
+    y_train = y_train[sorted_idx]
+    qid_train = qid_train[sorted_idx]
+
+    ltr = xgboost.XGBRanker(objective="rank:ndcg", tree_method="hist")
+    ltr.fit(X_train, y_train, qid=qid_train)
+
+    # Use the original order of the data.
+    scores = ltr.predict(X)
+    return scores
+
+
+def simulate_one_fold(
+    fold: Tuple[sparse.csr_matrix, npt.NDArray[np.int32], npt.NDArray[np.int32]],
+    scores_fold: npt.NDArray[np.float32],
+) -> ClickFold:
+    """Simulate clicks for one fold."""
+    X_fold, y_fold, qid_fold = fold
+    assert qid_fold.dtype == np.int32
+
+    qids = np.unique(qid_fold)
+
+    position = np.empty((y_fold.size,), dtype=np.int64)
+    clicks = np.empty((y_fold.size,), dtype=np.int32)
+    pbm = PBM(eta=1.0)
+
+    # Avoid grouping by qid as we want to preserve the original data partition by
+    # the dataset authors.
+    for q in qids:
+        qid_mask = q == qid_fold
+        qid_mask = qid_mask.reshape(qid_mask.shape[0])
+        query_scores = scores_fold[qid_mask]
+        # Initial rank list, scores sorted to decreasing order
+        query_position = np.argsort(query_scores)[::-1]
+        position[qid_mask] = query_position
+        # get labels
+        relevance_degrees = y_fold[qid_mask]
+        query_clicks = pbm.sample_clicks_for_query(relevance_degrees, query_position)
+        clicks[qid_mask] = query_clicks
+
+    assert X_fold.shape[0] == qid_fold.shape[0], (X_fold.shape, qid_fold.shape)
+    assert X_fold.shape[0] == clicks.shape[0], (X_fold.shape, clicks.shape)
+
+    return ClickFold(X_fold, y_fold, qid_fold, scores_fold, clicks, position)
+
+
+# pylint: disable=too-many-locals
+def simulate_clicks(cv_data: RelDataCV) -> Tuple[ClickFold, Optional[ClickFold]]:
+    """Simulate click data using position biased model (PBM)."""
+    X, y, qid = list(zip(cv_data.train, cv_data.test))
+
+    # ptr to train-test split
+    indptr = np.array([0] + [v.shape[0] for v in X])
+    indptr = np.cumsum(indptr)
+
+    assert len(indptr) == 2 + 1  # train, test
+    X_full = sparse.vstack(X)
+    y_full = np.concatenate(y)
+    qid_full = np.concatenate(qid)
+
+    # Obtain initial relevance score for click simulation
+    scores_full = init_rank_score(X_full, y_full, qid_full)
+    # partition it back to (train, test) tuple
+    scores = [scores_full[indptr[i - 1] : indptr[i]] for i in range(1, indptr.size)]
+
+    X_lst, y_lst, q_lst, s_lst, c_lst, p_lst = [], [], [], [], [], []
+    for i in range(indptr.size - 1):
+        fold = simulate_one_fold((X[i], y[i], qid[i]), scores[i])
+        X_lst.append(fold.X)
+        y_lst.append(fold.y)
+        q_lst.append(fold.qid)
+        s_lst.append(fold.score)
+        c_lst.append(fold.click)
+        p_lst.append(fold.pos)
+
+    scores_check_1 = [s_lst[i] for i in range(indptr.size - 1)]
+    for i in range(2):
+        assert (scores_check_1[i] == scores[i]).all()
+
+    if len(X_lst) == 1:
+        train = ClickFold(X_lst[0], y_lst[0], q_lst[0], s_lst[0], c_lst[0], p_lst[0])
+        test = None
+    else:
+        train, test = (
+            ClickFold(X_lst[i], y_lst[i], q_lst[i], s_lst[i], c_lst[i], p_lst[i])
+            for i in range(len(X_lst))
+        )
+    return train, test
+
+
+def sort_ltr_samples(
+    X: sparse.csr_matrix,
+    y: npt.NDArray[np.int32],
+    qid: npt.NDArray[np.int32],
+    clicks: npt.NDArray[np.int32],
+    pos: npt.NDArray[np.int64],
+) -> Tuple[
+    sparse.csr_matrix,
+    npt.NDArray[np.int32],
+    npt.NDArray[np.int32],
+    npt.NDArray[np.int32],
+]:
+    """Sort data based on query index and position."""
+    sorted_idx = np.argsort(qid)
+    X = X[sorted_idx]
+    clicks = clicks[sorted_idx]
+    qid = qid[sorted_idx]
+    pos = pos[sorted_idx]
+
+    indptr, _, _ = rlencode(qid)
+
+    for i in range(1, indptr.size):
+        beg = indptr[i - 1]
+        end = indptr[i]
+
+        assert beg < end, (beg, end)
+        assert np.unique(qid[beg:end]).size == 1, (beg, end)
+
+        query_pos = pos[beg:end]
+        assert query_pos.min() == 0, query_pos.min()
+        assert query_pos.max() >= query_pos.size - 1, (
+            query_pos.max(),
+            query_pos.size,
+            i,
+            np.unique(qid[beg:end]),
+        )
+        sorted_idx = np.argsort(query_pos)
+
+        X[beg:end] = X[beg:end][sorted_idx]
+        clicks[beg:end] = clicks[beg:end][sorted_idx]
+        y[beg:end] = y[beg:end][sorted_idx]
+        # not necessary
+        qid[beg:end] = qid[beg:end][sorted_idx]
+
+    data = X, clicks, y, qid
+
+    return data

python-package/xgboost/testing/params.py

@@ -67,3 +67,17 @@ cat_parameter_strategy = strategies.fixed_dictionaries(
         "max_cat_threshold": strategies.integers(1, 128),
     }
 )
+
+lambdarank_parameter_strategy = strategies.fixed_dictionaries(
+    {
+        "lambdarank_unbiased": strategies.sampled_from([True, False]),
+        "lambdarank_pair_method": strategies.sampled_from(["topk", "mean"]),
+        "lambdarank_num_pair_per_sample": strategies.integers(1, 8),
+        "lambdarank_bias_norm": strategies.floats(0.5, 2.0),
+        "objective": strategies.sampled_from(
+            ["rank:ndcg", "rank:map", "rank:pairwise"]
+        ),
+    }
+).filter(
+    lambda x: not (x["lambdarank_unbiased"] and x["lambdarank_pair_method"] == "mean")
+)

src/metric/elementwise_metric.cu

@@ -499,9 +499,13 @@ class QuantileError : public MetricNoCache {
 
   const char* Name() const override { return "quantile"; }
   void LoadConfig(Json const& in) override {
-    auto const& name = get<String const>(in["name"]);
+    auto const& obj = get<Object const>(in);
-    CHECK_EQ(name, "quantile");
+    auto it = obj.find("quantile_loss_param");
-    FromJson(in["quantile_loss_param"], &param_);
+    if (it != obj.cend()) {
+      FromJson(it->second, &param_);
+      auto const& name = get<String const>(in["name"]);
+      CHECK_EQ(name, "quantile");
+    }
   }
   void SaveConfig(Json* p_out) const override {
     auto& out = *p_out;

src/predictor/cpu_predictor.cc

@@ -152,7 +152,7 @@ void PredictByAllTrees(gbm::GBTreeModel const &model, std::uint32_t const tree_b
       } else {
         for (std::size_t i = 0; i < block_size; ++i) {
           out_predt(predict_offset + i, gid) +=
-              scalar::PredValueByOneTree<true>(thread_temp[offset + i], tree, cats);
+              scalar::PredValueByOneTree<false>(thread_temp[offset + i], tree, cats);
         }
       }
     }
@@ -430,8 +430,7 @@ class ColumnSplitHelper {
         << "column-split prediction is only supported for distributed training";
 
     for (auto const &batch : p_fmat->GetBatches<SparsePage>()) {
-      CHECK_EQ(out_preds->size(),
+      CHECK_EQ(out_preds->size(), p_fmat->Info().num_row_ * (tree_end_ - tree_begin_));
-               p_fmat->Info().num_row_ * model_.learner_model_param->num_output_group);
       PredictBatchKernel<SparsePageView, kBlockOfRowsSize, true>(SparsePageView{&batch}, out_preds);
     }
   }
@@ -543,8 +542,12 @@ class ColumnSplitHelper {
     for (size_t tree_id = tree_begin_; tree_id < tree_end_; ++tree_id) {
       auto const gid = model_.tree_info[tree_id];
       for (size_t i = 0; i < block_size; ++i) {
-        preds[(predict_offset + i) * num_group + gid] +=
+        auto const result = PredictOneTree<predict_leaf>(tree_id, batch_offset + i);
-            PredictOneTree<predict_leaf>(tree_id, batch_offset + i);
+        if constexpr (predict_leaf) {
+          preds[(predict_offset + i) * (tree_end_ - tree_begin_) + tree_id] = result;
+        } else {
+          preds[(predict_offset + i) * num_group + gid] += result;
+        }
       }
     }
   }
@@ -645,6 +648,9 @@ class CPUPredictor : public Predictor {
   void PredictDMatrix(DMatrix *p_fmat, std::vector<bst_float> *out_preds,
                       gbm::GBTreeModel const &model, int32_t tree_begin, int32_t tree_end) const {
     if (p_fmat->Info().IsColumnSplit()) {
+      CHECK(!model.learner_model_param->IsVectorLeaf())
+          << "Predict DMatrix with column split" << MTNotImplemented();
+
       ColumnSplitHelper helper(this->ctx_->Threads(), model, tree_begin, tree_end);
       helper.PredictDMatrix(p_fmat, out_preds);
       return;
@@ -743,6 +749,8 @@ class CPUPredictor : public Predictor {
                       unsigned tree_end) const override {
     auto proxy = dynamic_cast<data::DMatrixProxy *>(p_m.get());
     CHECK(proxy)<< "Inplace predict accepts only DMatrixProxy as input.";
+    CHECK(!p_m->Info().IsColumnSplit())
+        << "Inplace predict support for column-wise data split is not yet implemented.";
     auto x = proxy->Adapter();
     if (x.type() == typeid(std::shared_ptr<data::DenseAdapter>)) {
       this->DispatchedInplacePredict<data::DenseAdapter, kBlockOfRowsSize>(
@@ -773,6 +781,9 @@ class CPUPredictor : public Predictor {
     out_preds->resize(model.learner_model_param->num_output_group);
 
     if (is_column_split) {
+      CHECK(!model.learner_model_param->IsVectorLeaf())
+          << "Predict instance with column split" << MTNotImplemented();
+
       ColumnSplitHelper helper(this->ctx_->Threads(), model, 0, ntree_limit);
       helper.PredictInstance(inst, out_preds);
       return;
@@ -802,6 +813,9 @@ class CPUPredictor : public Predictor {
     preds.resize(info.num_row_ * ntree_limit);
 
     if (p_fmat->Info().IsColumnSplit()) {
+      CHECK(!model.learner_model_param->IsVectorLeaf())
+          << "Predict leaf with column split" << MTNotImplemented();
+
       ColumnSplitHelper helper(n_threads, model, 0, ntree_limit);
       helper.PredictLeaf(p_fmat, &preds);
       return;

src/tree/updater_gpu_hist.cu

@@ -302,7 +302,7 @@ struct GPUHistMakerDevice {
         matrix.feature_segments,
         matrix.gidx_fvalue_map,
         matrix.min_fvalue,
-        matrix.is_dense
+        matrix.is_dense && !collective::IsDistributed()
     };
     auto split = this->evaluator_.EvaluateSingleSplit(inputs, shared_inputs);
     return split;
@@ -316,11 +316,11 @@ struct GPUHistMakerDevice {
     std::vector<bst_node_t> nidx(2 * candidates.size());
     auto h_node_inputs = pinned2.GetSpan<EvaluateSplitInputs>(2 * candidates.size());
     auto matrix = page->GetDeviceAccessor(ctx_->gpu_id);
-    EvaluateSplitSharedInputs shared_inputs{
+    EvaluateSplitSharedInputs shared_inputs{GPUTrainingParam{param}, *quantiser, feature_types,
-        GPUTrainingParam{param}, *quantiser, feature_types,     matrix.feature_segments,
+                                            matrix.feature_segments, matrix.gidx_fvalue_map,
-        matrix.gidx_fvalue_map,  matrix.min_fvalue,
+                                            matrix.min_fvalue,
-        matrix.is_dense
+                                            // is_dense represents the local data
-    };
+                                            matrix.is_dense && !collective::IsDistributed()};
     dh::TemporaryArray<GPUExpandEntry> entries(2 * candidates.size());
     // Store the feature set ptrs so they dont go out of scope before the kernel is called
     std::vector<std::shared_ptr<HostDeviceVector<bst_feature_t>>> feature_sets;

src/tree/updater_quantile_hist.cc

@@ -419,6 +419,7 @@ class HistBuilder {
 
   CPUExpandEntry InitRoot(DMatrix *p_fmat, linalg::MatrixView<GradientPair const> gpair,
                           RegTree *p_tree) {
+    monitor_->Start(__func__);
     CPUExpandEntry node(RegTree::kRoot, p_tree->GetDepth(0));
 
     std::size_t page_id = 0;
@@ -434,7 +435,7 @@ class HistBuilder {
 
     {
       GradientPairPrecise grad_stat;
-      if (p_fmat->IsDense()) {
+      if (p_fmat->IsDense() && !collective::IsDistributed()) {
         /**
          * Specialized code for dense data: For dense data (with no missing value), the sum
          * of gradient histogram is equal to snode[nid]
@@ -475,12 +476,14 @@ class HistBuilder {
       node = entries.front();
     }
 
+    monitor_->Stop(__func__);
     return node;
   }
 
   void BuildHistogram(DMatrix *p_fmat, RegTree *p_tree,
                       std::vector<CPUExpandEntry> const &valid_candidates,
                       linalg::MatrixView<GradientPair const> gpair) {
+    monitor_->Start(__func__);
     std::vector<CPUExpandEntry> nodes_to_build(valid_candidates.size());
     std::vector<CPUExpandEntry> nodes_to_sub(valid_candidates.size());
 
@@ -508,6 +511,7 @@ class HistBuilder {
                                     nodes_to_sub, gpair.Values());
       ++page_id;
     }
+    monitor_->Stop(__func__);
   }
 
   void UpdatePosition(DMatrix *p_fmat, RegTree const *p_tree,
@@ -525,6 +529,7 @@ class HistBuilder {
                      std::vector<bst_node_t> *p_out_position) {
     monitor_->Start(__func__);
     if (!task_->UpdateTreeLeaf()) {
+    monitor_->Stop(__func__);
       return;
     }
     for (auto const &part : partitioner_) {

tests/buildkite/conftest.sh

@@ -24,7 +24,7 @@ set -x
 
 CUDA_VERSION=11.8.0
 NCCL_VERSION=2.16.5-1
-RAPIDS_VERSION=23.04
+RAPIDS_VERSION=23.06
 SPARK_VERSION=3.4.0
 JDK_VERSION=8
 

tests/ci_build/lint_python.py

@@ -36,6 +36,7 @@ class LintersPaths:
         "demo/guide-python/individual_trees.py",
         "demo/guide-python/quantile_regression.py",
         "demo/guide-python/multioutput_regression.py",
+        "demo/guide-python/learning_to_rank.py",
         # CI
         "tests/ci_build/lint_python.py",
         "tests/ci_build/test_r_package.py",
@@ -76,6 +77,7 @@ class LintersPaths:
         "demo/guide-python/individual_trees.py",
         "demo/guide-python/quantile_regression.py",
         "demo/guide-python/multioutput_regression.py",
+        "demo/guide-python/learning_to_rank.py",
         # CI
         "tests/ci_build/lint_python.py",
         "tests/ci_build/test_r_package.py",

tests/cpp/common/test_algorithm.cu

@@ -88,7 +88,7 @@ TEST(Algorithm, GpuArgSort) {
 TEST(Algorithm, SegmentedSequence) {
   dh::device_vector<std::size_t> idx(16);
   dh::device_vector<std::size_t> ptr(3);
-  Context ctx = CreateEmptyGenericParam(0);
+  Context ctx = MakeCUDACtx(0);
   ptr[0] = 0;
   ptr[1] = 4;
   ptr[2] = idx.size();

tests/cpp/common/test_column_matrix.cc

@@ -14,7 +14,7 @@ TEST(DenseColumn, Test) {
   int32_t max_num_bins[] = {static_cast<int32_t>(std::numeric_limits<uint8_t>::max()) + 1,
                             static_cast<int32_t>(std::numeric_limits<uint16_t>::max()) + 1,
                             static_cast<int32_t>(std::numeric_limits<uint16_t>::max()) + 2};
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
   BinTypeSize last{kUint8BinsTypeSize};
   for (int32_t max_num_bin : max_num_bins) {
     auto dmat = RandomDataGenerator(100, 10, 0.0).GenerateDMatrix();
@@ -63,7 +63,7 @@ TEST(SparseColumn, Test) {
   int32_t max_num_bins[] = {static_cast<int32_t>(std::numeric_limits<uint8_t>::max()) + 1,
                             static_cast<int32_t>(std::numeric_limits<uint16_t>::max()) + 1,
                             static_cast<int32_t>(std::numeric_limits<uint16_t>::max()) + 2};
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
   for (int32_t max_num_bin : max_num_bins) {
     auto dmat = RandomDataGenerator(100, 1, 0.85).GenerateDMatrix();
     GHistIndexMatrix gmat{&ctx, dmat.get(), max_num_bin, 0.5f, false};
@@ -92,7 +92,7 @@ TEST(DenseColumnWithMissing, Test) {
   int32_t max_num_bins[] = {static_cast<int32_t>(std::numeric_limits<uint8_t>::max()) + 1,
                             static_cast<int32_t>(std::numeric_limits<uint16_t>::max()) + 1,
                             static_cast<int32_t>(std::numeric_limits<uint16_t>::max()) + 2};
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
   for (int32_t max_num_bin : max_num_bins) {
     auto dmat = RandomDataGenerator(100, 1, 0.5).GenerateDMatrix();
     GHistIndexMatrix gmat(&ctx, dmat.get(), max_num_bin, 0.2, false);

tests/cpp/common/test_hist_util.cc

@@ -156,28 +156,28 @@ TEST(CutsBuilder, SearchGroupInd) {
 }
 
 TEST(HistUtil, DenseCutsCategorical) {
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
-   int categorical_sizes[] = {2, 6, 8, 12};
+  int categorical_sizes[] = {2, 6, 8, 12};
-   int num_bins = 256;
+  int num_bins = 256;
-   int sizes[] = {25, 100, 1000};
+  int sizes[] = {25, 100, 1000};
-   for (auto n : sizes) {
+  for (auto n : sizes) {
-     for (auto num_categories : categorical_sizes) {
+    for (auto num_categories : categorical_sizes) {
-       auto x = GenerateRandomCategoricalSingleColumn(n, num_categories);
+      auto x = GenerateRandomCategoricalSingleColumn(n, num_categories);
-       std::vector<float> x_sorted(x);
+      std::vector<float> x_sorted(x);
-       std::sort(x_sorted.begin(), x_sorted.end());
+      std::sort(x_sorted.begin(), x_sorted.end());
-       auto dmat = GetDMatrixFromData(x, n, 1);
+      auto dmat = GetDMatrixFromData(x, n, 1);
-       HistogramCuts cuts = SketchOnDMatrix(&ctx, dmat.get(), num_bins);
+      HistogramCuts cuts = SketchOnDMatrix(&ctx, dmat.get(), num_bins);
-       auto cuts_from_sketch = cuts.Values();
+      auto cuts_from_sketch = cuts.Values();
-       EXPECT_LT(cuts.MinValues()[0], x_sorted.front());
+      EXPECT_LT(cuts.MinValues()[0], x_sorted.front());
-       EXPECT_GT(cuts_from_sketch.front(), x_sorted.front());
+      EXPECT_GT(cuts_from_sketch.front(), x_sorted.front());
-       EXPECT_GE(cuts_from_sketch.back(), x_sorted.back());
+      EXPECT_GE(cuts_from_sketch.back(), x_sorted.back());
-       EXPECT_EQ(cuts_from_sketch.size(), static_cast<size_t>(num_categories));
+      EXPECT_EQ(cuts_from_sketch.size(), static_cast<size_t>(num_categories));
-     }
+    }
-   }
+  }
 }
 
 TEST(HistUtil, DenseCutsAccuracyTest) {
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
   int bin_sizes[] = {2, 16, 256, 512};
   int sizes[] = {100};
   int num_columns = 5;
@@ -195,7 +195,7 @@ TEST(HistUtil, DenseCutsAccuracyTestWeights) {
   int bin_sizes[] = {2, 16, 256, 512};
   int sizes[] = {100, 1000, 1500};
   int num_columns = 5;
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
   for (auto num_rows : sizes) {
     auto x = GenerateRandom(num_rows, num_columns);
     auto dmat = GetDMatrixFromData(x, num_rows, num_columns);
@@ -218,7 +218,7 @@ void TestQuantileWithHessian(bool use_sorted) {
   int bin_sizes[] = {2, 16, 256, 512};
   int sizes[] = {1000, 1500};
   int num_columns = 5;
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
   for (auto num_rows : sizes) {
     auto x = GenerateRandom(num_rows, num_columns);
     auto dmat = GetDMatrixFromData(x, num_rows, num_columns);
@@ -257,7 +257,7 @@ TEST(HistUtil, DenseCutsExternalMemory) {
   int bin_sizes[] = {2, 16, 256, 512};
   int sizes[] = {100, 1000, 1500};
   int num_columns = 5;
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
   for (auto num_rows : sizes) {
     auto x = GenerateRandom(num_rows, num_columns);
     dmlc::TemporaryDirectory tmpdir;
@@ -278,7 +278,7 @@ TEST(HistUtil, IndexBinBound) {
                                            kUint32BinsTypeSize};
   size_t constexpr kRows = 100;
   size_t constexpr kCols = 10;
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
   size_t bin_id = 0;
   for (auto max_bin : bin_sizes) {
     auto p_fmat = RandomDataGenerator(kRows, kCols, 0).GenerateDMatrix();
@@ -303,7 +303,7 @@ TEST(HistUtil, IndexBinData) {
                                      static_cast<uint64_t>(std::numeric_limits<uint16_t>::max()) + 2 };
   size_t constexpr kRows = 100;
   size_t constexpr kCols = 10;
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
 
   for (auto max_bin : kBinSizes) {
     auto p_fmat = RandomDataGenerator(kRows, kCols, 0).GenerateDMatrix();
@@ -331,7 +331,7 @@ void TestSketchFromWeights(bool with_group) {
   size_t constexpr kRows = 300, kCols = 20, kBins = 256;
   size_t constexpr kGroups = 10;
   auto m = RandomDataGenerator{kRows, kCols, 0}.Device(0).GenerateDMatrix();
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
   common::HistogramCuts cuts = SketchOnDMatrix(&ctx, m.get(), kBins);
 
   MetaInfo info;
@@ -397,7 +397,7 @@ TEST(HistUtil, SketchFromWeights) {
 }
 
 TEST(HistUtil, SketchCategoricalFeatures) {
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
   TestCategoricalSketch(1000, 256, 32, false, [&ctx](DMatrix* p_fmat, int32_t num_bins) {
     return SketchOnDMatrix(&ctx, p_fmat, num_bins);
   });

tests/cpp/common/test_hist_util.cu

@@ -324,7 +324,7 @@ TEST(HistUtil, AdapterDeviceSketch) {
   data::CupyAdapter adapter(str);
 
   auto device_cuts = MakeUnweightedCutsForTest(adapter, num_bins, missing);
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
   auto host_cuts = GetHostCuts(&ctx, &adapter, num_bins, missing);
 
   EXPECT_EQ(device_cuts.Values(), host_cuts.Values());

tests/cpp/common/test_quantile.cc

@@ -302,7 +302,7 @@ namespace {
 void TestSameOnAllWorkers() {
   auto const world = collective::GetWorldSize();
   constexpr size_t kRows = 1000, kCols = 100;
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
 
   RunWithSeedsAndBins(
       kRows, [=, &ctx](int32_t seed, size_t n_bins, MetaInfo const&) {

tests/cpp/data/test_gradient_index.cc

@@ -21,7 +21,7 @@
 #include "../../../src/data/adapter.h"          // for SparsePageAdapterBatch
 #include "../../../src/data/gradient_index.h"   // for GHistIndexMatrix
 #include "../../../src/tree/param.h"            // for TrainParam
-#include "../helpers.h"                         // for CreateEmptyGenericParam, GenerateRandomCa...
+#include "../helpers.h"                         // for GenerateRandomCategoricalSingleColumn...
 #include "xgboost/base.h"                       // for bst_bin_t
 #include "xgboost/context.h"                    // for Context
 #include "xgboost/host_device_vector.h"         // for HostDeviceVector
@@ -29,7 +29,7 @@
 namespace xgboost {
 namespace data {
 TEST(GradientIndex, ExternalMemory) {
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
   std::unique_ptr<DMatrix> dmat = CreateSparsePageDMatrix(10000);
   std::vector<size_t> base_rowids;
   std::vector<float> hessian(dmat->Info().num_row_, 1);
@@ -58,7 +58,7 @@ TEST(GradientIndex, FromCategoricalBasic) {
   size_t max_bins = 8;
   auto x = GenerateRandomCategoricalSingleColumn(kRows, kCats);
   auto m = GetDMatrixFromData(x, kRows, 1);
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
 
   auto &h_ft = m->Info().feature_types.HostVector();
   h_ft.resize(kCols, FeatureType::kCategorical);

tests/cpp/data/test_sparse_page_dmatrix.cc

@@ -67,7 +67,7 @@ void TestSparseDMatrixLoadFile(Context const* ctx) {
 }
 
 TEST(SparsePageDMatrix, LoadFile) {
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
   TestSparseDMatrixLoadFile<SparsePage>(&ctx);
   TestSparseDMatrixLoadFile<CSCPage>(&ctx);
   TestSparseDMatrixLoadFile<SortedCSCPage>(&ctx);
@@ -77,7 +77,7 @@ TEST(SparsePageDMatrix, LoadFile) {
 template <typename Page>
 void TestRetainPage() {
   auto m = CreateSparsePageDMatrix(10000);
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
   auto batches = m->GetBatches<Page>(&ctx);
   auto begin = batches.begin();
   auto end = batches.end();
@@ -145,7 +145,7 @@ TEST(SparsePageDMatrix, ColAccess) {
   const std::string tmp_file = tempdir.path + "/simple.libsvm";
   CreateSimpleTestData(tmp_file);
   xgboost::DMatrix *dmat = xgboost::DMatrix::Load(UriSVM(tmp_file, tmp_file));
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
 
   // Loop over the batches and assert the data is as expected
   size_t iter = 0;
@@ -224,7 +224,7 @@ TEST(SparsePageDMatrix, ColAccessBatches) {
   // Create multiple sparse pages
   std::unique_ptr<xgboost::DMatrix> dmat{xgboost::CreateSparsePageDMatrix(kEntries)};
   ASSERT_EQ(dmat->Ctx()->Threads(), AllThreadsForTest());
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
   for (auto const &page : dmat->GetBatches<xgboost::CSCPage>(&ctx)) {
     ASSERT_EQ(dmat->Info().num_col_, page.Size());
   }

tests/cpp/data/test_sparse_page_dmatrix.cu

@@ -112,7 +112,7 @@ TEST(SparsePageDMatrix, RetainEllpackPage) {
 }
 
 TEST(SparsePageDMatrix, EllpackPageContent) {
-  auto ctx = CreateEmptyGenericParam(0);
+  auto ctx = MakeCUDACtx(0);
   constexpr size_t kRows = 6;
   constexpr size_t kCols = 2;
   constexpr size_t kPageSize = 1;

tests/cpp/helpers.h

@@ -382,13 +382,6 @@ std::unique_ptr<GradientBooster> CreateTrainedGBM(std::string name, Args kwargs,
                                                   LearnerModelParam const* learner_model_param,
                                                   Context const* generic_param);
 
-inline Context CreateEmptyGenericParam(int gpu_id) {
-  xgboost::Context tparam;
-  std::vector<std::pair<std::string, std::string>> args{{"gpu_id", std::to_string(gpu_id)}};
-  tparam.Init(args);
-  return tparam;
-}
-
 inline std::unique_ptr<HostDeviceVector<GradientPair>> GenerateGradients(
     std::size_t rows, bst_target_t n_targets = 1) {
   auto p_gradients = std::make_unique<HostDeviceVector<GradientPair>>(rows * n_targets);
@@ -407,9 +400,14 @@ inline std::unique_ptr<HostDeviceVector<GradientPair>> GenerateGradients(
 }
 
 /**
- * \brief Make a context that uses CUDA.
+ * \brief Make a context that uses CUDA if device >= 0.
  */
-inline Context MakeCUDACtx(std::int32_t device) { return Context{}.MakeCUDA(device); }
+inline Context MakeCUDACtx(std::int32_t device) {
+  if (device == Context::kCpuId) {
+    return Context{};
+  }
+  return Context{}.MakeCUDA(device);
+}
 
 inline HostDeviceVector<GradientPair> GenerateRandomGradients(const size_t n_rows,
                                                               float lower= 0.0f, float upper = 1.0f) {

tests/cpp/linear/test_json_io.h

@@ -12,19 +12,19 @@
 
 namespace xgboost {
 inline void TestUpdaterJsonIO(std::string updater_str) {
-  auto runtime = xgboost::CreateEmptyGenericParam(GPUIDX);
+  Context ctx{MakeCUDACtx(GPUIDX)};
   Json config_0 {Object() };
 
   {
-    auto updater = std::unique_ptr<xgboost::LinearUpdater>(
+    auto updater =
-        xgboost::LinearUpdater::Create(updater_str, &runtime));
+        std::unique_ptr<xgboost::LinearUpdater>(xgboost::LinearUpdater::Create(updater_str, &ctx));
     updater->Configure({{"eta", std::to_string(3.14)}});
     updater->SaveConfig(&config_0);
   }
 
   {
-    auto updater = std::unique_ptr<xgboost::LinearUpdater>(
+    auto updater =
-        xgboost::LinearUpdater::Create(updater_str, &runtime));
+        std::unique_ptr<xgboost::LinearUpdater>(xgboost::LinearUpdater::Create(updater_str, &ctx));
     updater->LoadConfig(config_0);
     Json config_1 { Object() };
     updater->SaveConfig(&config_1);

tests/cpp/linear/test_linear.cc

@@ -17,7 +17,7 @@ TEST(Linear, Shotgun) {
 
   auto p_fmat = xgboost::RandomDataGenerator(kRows, kCols, 0).GenerateDMatrix();
 
-  auto ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   LearnerModelParam mparam{MakeMP(kCols, .5, 1)};
 
   {
@@ -49,7 +49,7 @@ TEST(Linear, coordinate) {
 
   auto p_fmat = xgboost::RandomDataGenerator(kRows, kCols, 0).GenerateDMatrix();
 
-  auto ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   LearnerModelParam mparam{MakeMP(kCols, .5, 1)};
 
   auto updater = std::unique_ptr<xgboost::LinearUpdater>(

tests/cpp/linear/test_linear.cu

@@ -13,7 +13,7 @@ TEST(Linear, GPUCoordinate) {
   size_t constexpr kCols = 10;
 
   auto mat = xgboost::RandomDataGenerator(kRows, kCols, 0).GenerateDMatrix();
-  auto ctx = CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(0);
 
   LearnerModelParam mparam{MakeMP(kCols, .5, 1)};
   auto updater = std::unique_ptr<xgboost::LinearUpdater>(

tests/cpp/metric/test_auc.h

@@ -11,7 +11,7 @@ namespace xgboost {
 namespace metric {
 
 inline void VerifyBinaryAUC(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   std::unique_ptr<Metric> uni_ptr{Metric::Create("auc", &ctx)};
   Metric* metric = uni_ptr.get();
   ASSERT_STREQ(metric->Name(), "auc");
@@ -54,7 +54,7 @@ inline void VerifyBinaryAUC(DataSplitMode data_split_mode = DataSplitMode::kRow)
 }
 
 inline void VerifyMultiClassAUC(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   std::unique_ptr<Metric> uni_ptr{Metric::Create("auc", &ctx)};
   auto metric = uni_ptr.get();
 
@@ -115,7 +115,7 @@ inline void VerifyMultiClassAUC(DataSplitMode data_split_mode = DataSplitMode::k
 }
 
 inline void VerifyRankingAUC(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   std::unique_ptr<Metric> metric{Metric::Create("auc", &ctx)};
 
   // single group
@@ -149,7 +149,7 @@ inline void VerifyRankingAUC(DataSplitMode data_split_mode = DataSplitMode::kRow
 }
 
 inline void VerifyPRAUC(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
 
   xgboost::Metric* metric = xgboost::Metric::Create("aucpr", &ctx);
   ASSERT_STREQ(metric->Name(), "aucpr");
@@ -186,7 +186,7 @@ inline void VerifyPRAUC(DataSplitMode data_split_mode = DataSplitMode::kRow) {
 }
 
 inline void VerifyMultiClassPRAUC(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
 
   std::unique_ptr<Metric> metric{Metric::Create("aucpr", &ctx)};
 
@@ -210,7 +210,7 @@ inline void VerifyMultiClassPRAUC(DataSplitMode data_split_mode = DataSplitMode:
 }
 
 inline void VerifyRankingPRAUC(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
 
   std::unique_ptr<Metric> metric{Metric::Create("aucpr", &ctx)};
 

tests/cpp/metric/test_elementwise_metric.h

@@ -15,7 +15,7 @@ namespace xgboost {
 namespace metric {
 
 inline void CheckDeterministicMetricElementWise(StringView name, int32_t device) {
-  auto ctx = CreateEmptyGenericParam(device);
+  auto ctx = MakeCUDACtx(device);
   std::unique_ptr<Metric> metric{Metric::Create(name.c_str(), &ctx)};
 
   HostDeviceVector<float> predts;
@@ -46,7 +46,7 @@ inline void CheckDeterministicMetricElementWise(StringView name, int32_t device)
 }
 
 inline void VerifyRMSE(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   xgboost::Metric * metric = xgboost::Metric::Create("rmse", &ctx);
   metric->Configure({});
   ASSERT_STREQ(metric->Name(), "rmse");
@@ -75,7 +75,7 @@ inline void VerifyRMSE(DataSplitMode data_split_mode = DataSplitMode::kRow) {
 }
 
 inline void VerifyRMSLE(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   xgboost::Metric * metric = xgboost::Metric::Create("rmsle", &ctx);
   metric->Configure({});
   ASSERT_STREQ(metric->Name(), "rmsle");
@@ -104,7 +104,7 @@ inline void VerifyRMSLE(DataSplitMode data_split_mode = DataSplitMode::kRow) {
 }
 
 inline void VerifyMAE(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   xgboost::Metric * metric = xgboost::Metric::Create("mae", &ctx);
   metric->Configure({});
   ASSERT_STREQ(metric->Name(), "mae");
@@ -133,7 +133,7 @@ inline void VerifyMAE(DataSplitMode data_split_mode = DataSplitMode::kRow) {
 }
 
 inline void VerifyMAPE(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   xgboost::Metric * metric = xgboost::Metric::Create("mape", &ctx);
   metric->Configure({});
   ASSERT_STREQ(metric->Name(), "mape");
@@ -162,7 +162,7 @@ inline void VerifyMAPE(DataSplitMode data_split_mode = DataSplitMode::kRow) {
 }
 
 inline void VerifyMPHE(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   std::unique_ptr<xgboost::Metric> metric{xgboost::Metric::Create("mphe", &ctx)};
   metric->Configure({});
   ASSERT_STREQ(metric->Name(), "mphe");
@@ -197,7 +197,7 @@ inline void VerifyMPHE(DataSplitMode data_split_mode = DataSplitMode::kRow) {
 }
 
 inline void VerifyLogLoss(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   xgboost::Metric * metric = xgboost::Metric::Create("logloss", &ctx);
   metric->Configure({});
   ASSERT_STREQ(metric->Name(), "logloss");
@@ -230,7 +230,7 @@ inline void VerifyLogLoss(DataSplitMode data_split_mode = DataSplitMode::kRow) {
 }
 
 inline void VerifyError(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   xgboost::Metric * metric = xgboost::Metric::Create("error", &ctx);
   metric->Configure({});
   ASSERT_STREQ(metric->Name(), "error");
@@ -292,7 +292,7 @@ inline void VerifyError(DataSplitMode data_split_mode = DataSplitMode::kRow) {
 }
 
 inline void VerifyPoissonNegLogLik(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   xgboost::Metric * metric = xgboost::Metric::Create("poisson-nloglik", &ctx);
   metric->Configure({});
   ASSERT_STREQ(metric->Name(), "poisson-nloglik");
@@ -332,7 +332,7 @@ inline void VerifyMultiRMSE(DataSplitMode data_split_mode = DataSplitMode::kRow)
 
   HostDeviceVector<float> predt(n_samples * n_targets, 0);
 
-  auto ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   std::unique_ptr<Metric> metric{Metric::Create("rmse", &ctx)};
   metric->Configure({});
 
@@ -347,7 +347,7 @@ inline void VerifyMultiRMSE(DataSplitMode data_split_mode = DataSplitMode::kRow)
 }
 
 inline void VerifyQuantile(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   std::unique_ptr<Metric> metric{Metric::Create("quantile", &ctx)};
 
   HostDeviceVector<float> predts{0.1f, 0.9f, 0.1f, 0.9f};

tests/cpp/metric/test_metric.cc

@@ -2,10 +2,10 @@
 #include <xgboost/metric.h>
 
 #include "../helpers.h"
-
+namespace xgboost {
 TEST(Metric, UnknownMetric) {
-  auto ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
-  xgboost::Metric * metric = nullptr;
+  xgboost::Metric* metric = nullptr;
   EXPECT_ANY_THROW(metric = xgboost::Metric::Create("unknown_name", &ctx));
   EXPECT_NO_THROW(metric = xgboost::Metric::Create("rmse", &ctx));
   if (metric) {
@@ -18,3 +18,4 @@ TEST(Metric, UnknownMetric) {
     delete metric;
   }
 }
+}  // namespace xgboost

tests/cpp/metric/test_multiclass_metric.h

@@ -8,7 +8,7 @@ namespace xgboost {
 namespace metric {
 
 inline void CheckDeterministicMetricMultiClass(StringView name, int32_t device) {
-  auto ctx = CreateEmptyGenericParam(device);
+  auto ctx = MakeCUDACtx(device);
   std::unique_ptr<Metric> metric{Metric::Create(name.c_str(), &ctx)};
 
   HostDeviceVector<float> predts;
@@ -45,7 +45,7 @@ inline void CheckDeterministicMetricMultiClass(StringView name, int32_t device)
 }
 
 inline void TestMultiClassError(int device, DataSplitMode data_split_mode) {
-  auto ctx = xgboost::CreateEmptyGenericParam(device);
+  auto ctx = MakeCUDACtx(device);
   ctx.gpu_id = device;
   xgboost::Metric * metric = xgboost::Metric::Create("merror", &ctx);
   metric->Configure({});
@@ -66,7 +66,7 @@ inline void VerifyMultiClassError(DataSplitMode data_split_mode = DataSplitMode:
 }
 
 inline void TestMultiClassLogLoss(int device, DataSplitMode data_split_mode) {
-  auto ctx = xgboost::CreateEmptyGenericParam(device);
+  auto ctx = MakeCUDACtx(device);
   ctx.gpu_id = device;
   xgboost::Metric * metric = xgboost::Metric::Create("mlogloss", &ctx);
   metric->Configure({});

tests/cpp/metric/test_rank_metric.cc

@@ -22,7 +22,7 @@ namespace metric {
 
 #if !defined(__CUDACC__) && !defined(__HIP_PLATFORM_AMD__)
 TEST(Metric, AMS) {
-  auto ctx = CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   EXPECT_ANY_THROW(Metric::Create("ams", &ctx));
   Metric* metric = Metric::Create("ams@0.5f", &ctx);
   ASSERT_STREQ(metric->Name(), "ams@0.5");

tests/cpp/metric/test_rank_metric.h

@@ -20,7 +20,7 @@
 namespace xgboost::metric {
 
 inline void VerifyPrecision(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   std::unique_ptr<xgboost::Metric> metric{Metric::Create("pre", &ctx)};
   ASSERT_STREQ(metric->Name(), "pre");
   EXPECT_NEAR(GetMetricEval(metric.get(), {0, 1}, {0, 1}, {}, {}, data_split_mode), 0.5, 1e-7);
@@ -44,7 +44,7 @@ inline void VerifyPrecision(DataSplitMode data_split_mode = DataSplitMode::kRow)
 }
 
 inline void VerifyNDCG(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   Metric * metric = xgboost::Metric::Create("ndcg", &ctx);
   ASSERT_STREQ(metric->Name(), "ndcg");
   EXPECT_ANY_THROW(GetMetricEval(metric, {0, 1}, {}, {}, {}, data_split_mode));
@@ -102,7 +102,7 @@ inline void VerifyNDCG(DataSplitMode data_split_mode = DataSplitMode::kRow) {
 }
 
 inline void VerifyMAP(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   Metric * metric = xgboost::Metric::Create("map", &ctx);
   ASSERT_STREQ(metric->Name(), "map");
   EXPECT_NEAR(GetMetricEval(metric, {0, 1}, {0, 1}, {}, {}, data_split_mode), 1, kRtEps);
@@ -150,7 +150,7 @@ inline void VerifyMAP(DataSplitMode data_split_mode = DataSplitMode::kRow) {
 }
 
 inline void VerifyNDCGExpGain(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  Context ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  Context ctx = MakeCUDACtx(GPUIDX);
 
   auto p_fmat = xgboost::RandomDataGenerator{0, 0, 0}.GenerateDMatrix();
   MetaInfo& info = p_fmat->Info();

tests/cpp/metric/test_survival_metric.cu

@@ -31,7 +31,7 @@ TEST_F(DeclareUnifiedDistributedTest(MetricTest), IntervalRegressionAccuracyColu
 
 // Test configuration of AFT metric
 TEST(AFTNegLogLikMetric, DeclareUnifiedTest(Configuration)) {
-  auto ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   std::unique_ptr<Metric> metric(Metric::Create("aft-nloglik", &ctx));
   metric->Configure({{"aft_loss_distribution", "normal"}, {"aft_loss_distribution_scale", "10"}});
 

tests/cpp/metric/test_survival_metric.h

@@ -13,7 +13,7 @@
 namespace xgboost {
 namespace common {
 inline void CheckDeterministicMetricElementWise(StringView name, int32_t device) {
-  auto ctx = CreateEmptyGenericParam(device);
+  auto ctx = MakeCUDACtx(device);
   std::unique_ptr<Metric> metric{Metric::Create(name.c_str(), &ctx)};
   metric->Configure(Args{});
 
@@ -48,7 +48,7 @@ inline void CheckDeterministicMetricElementWise(StringView name, int32_t device)
 }
 
 inline void VerifyAFTNegLogLik(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
 
   /**
    * Test aggregate output from the AFT metric over a small test data set.
@@ -79,7 +79,7 @@ inline void VerifyAFTNegLogLik(DataSplitMode data_split_mode = DataSplitMode::kR
 }
 
 inline void VerifyIntervalRegressionAccuracy(DataSplitMode data_split_mode = DataSplitMode::kRow) {
-  auto ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
 
   auto p_fmat = EmptyDMatrix();
   MetaInfo& info = p_fmat->Info();

tests/cpp/objective/test_aft_obj.cc

@@ -16,7 +16,7 @@ namespace xgboost {
 namespace common {
 
 TEST(Objective, DeclareUnifiedTest(AFTObjConfiguration)) {
-  auto ctx = CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   std::unique_ptr<ObjFunction> objective(ObjFunction::Create("survival:aft", &ctx));
   objective->Configure({ {"aft_loss_distribution", "logistic"},
                           {"aft_loss_distribution_scale", "5"} });
@@ -77,7 +77,7 @@ static inline void CheckGPairOverGridPoints(
 }
 
 TEST(Objective, DeclareUnifiedTest(AFTObjGPairUncensoredLabels)) {
-  auto ctx = CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   std::unique_ptr<ObjFunction> obj(ObjFunction::Create("survival:aft", &ctx));
 
   CheckGPairOverGridPoints(obj.get(), 100.0f, 100.0f, "normal",
@@ -101,7 +101,7 @@ TEST(Objective, DeclareUnifiedTest(AFTObjGPairUncensoredLabels)) {
 }
 
 TEST(Objective, DeclareUnifiedTest(AFTObjGPairLeftCensoredLabels)) {
-  auto ctx = CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   std::unique_ptr<ObjFunction> obj(ObjFunction::Create("survival:aft", &ctx));
 
   CheckGPairOverGridPoints(obj.get(), 0.0f, 20.0f, "normal",
@@ -122,7 +122,7 @@ TEST(Objective, DeclareUnifiedTest(AFTObjGPairLeftCensoredLabels)) {
 }
 
 TEST(Objective, DeclareUnifiedTest(AFTObjGPairRightCensoredLabels)) {
-  auto ctx = CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   std::unique_ptr<ObjFunction> obj(ObjFunction::Create("survival:aft", &ctx));
 
   CheckGPairOverGridPoints(obj.get(), 60.0f, std::numeric_limits<float>::infinity(), "normal",
@@ -146,7 +146,7 @@ TEST(Objective, DeclareUnifiedTest(AFTObjGPairRightCensoredLabels)) {
 }
 
 TEST(Objective, DeclareUnifiedTest(AFTObjGPairIntervalCensoredLabels)) {
-  auto ctx = CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   std::unique_ptr<ObjFunction> obj(ObjFunction::Create("survival:aft", &ctx));
 
   CheckGPairOverGridPoints(obj.get(), 16.0f, 200.0f, "normal",

tests/cpp/objective/test_hinge.cc

@@ -4,14 +4,12 @@
 #include <limits>
 
 #include "../helpers.h"
-
+namespace xgboost {
 TEST(Objective, DeclareUnifiedTest(HingeObj)) {
-  xgboost::Context ctx = xgboost::CreateEmptyGenericParam(GPUIDX);
+  Context ctx = MakeCUDACtx(GPUIDX);
-  std::unique_ptr<xgboost::ObjFunction> obj {
+  std::unique_ptr<ObjFunction> obj{ObjFunction::Create("binary:hinge", &ctx)};
-    xgboost::ObjFunction::Create("binary:hinge", &ctx)
-  };
 
-  xgboost::bst_float eps = std::numeric_limits<xgboost::bst_float>::min();
+  float eps = std::numeric_limits<xgboost::bst_float>::min();
   CheckObjFunction(obj,
                    {-1.0f, -0.5f, 0.5f, 1.0f, -1.0f, -0.5f,  0.5f, 1.0f},
                    { 0.0f,  0.0f, 0.0f, 0.0f,  1.0f,  1.0f,  1.0f, 1.0f},
@@ -27,3 +25,4 @@ TEST(Objective, DeclareUnifiedTest(HingeObj)) {
 
   ASSERT_NO_THROW(obj->DefaultEvalMetric());
 }
+}  // namespace xgboost

tests/cpp/objective/test_multiclass_obj.cc

@@ -9,7 +9,7 @@
 namespace xgboost {
 
 TEST(Objective, DeclareUnifiedTest(SoftmaxMultiClassObjGPair)) {
-  Context ctx = CreateEmptyGenericParam(GPUIDX);
+  Context ctx = MakeCUDACtx(GPUIDX);
   std::vector<std::pair<std::string, std::string>> args {{"num_class", "3"}};
   std::unique_ptr<ObjFunction> obj {
     ObjFunction::Create("multi:softmax", &ctx)
@@ -36,7 +36,7 @@ TEST(Objective, DeclareUnifiedTest(SoftmaxMultiClassObjGPair)) {
 }
 
 TEST(Objective, DeclareUnifiedTest(SoftmaxMultiClassBasic)) {
-  auto ctx = CreateEmptyGenericParam(GPUIDX);
+  auto ctx = MakeCUDACtx(GPUIDX);
   std::vector<std::pair<std::string, std::string>> args{
       std::pair<std::string, std::string>("num_class", "3")};
 
@@ -57,7 +57,7 @@ TEST(Objective, DeclareUnifiedTest(SoftmaxMultiClassBasic)) {
 }
 
 TEST(Objective, DeclareUnifiedTest(SoftprobMultiClassBasic)) {
-  Context ctx = CreateEmptyGenericParam(GPUIDX);
+  Context ctx = MakeCUDACtx(GPUIDX);
   std::vector<std::pair<std::string, std::string>> args {
     std::pair<std::string, std::string>("num_class", "3")};
 

tests/cpp/objective/test_quantile_obj.cc

@@ -10,11 +10,11 @@
 #include <memory>               // std::unique_ptr
 #include <vector>               // std::vector
 
-#include "../helpers.h"         // CheckConfigReload,CreateEmptyGenericParam,DeclareUnifiedTest
+#include "../helpers.h"         // CheckConfigReload,MakeCUDACtx,DeclareUnifiedTest
 
 namespace xgboost {
 TEST(Objective, DeclareUnifiedTest(Quantile)) {
-  Context ctx = CreateEmptyGenericParam(GPUIDX);
+  Context ctx = MakeCUDACtx(GPUIDX);
 
   {
     Args args{{"quantile_alpha", "[0.6, 0.8]"}};
@@ -37,7 +37,7 @@ TEST(Objective, DeclareUnifiedTest(Quantile)) {
 }
 
 TEST(Objective, DeclareUnifiedTest(QuantileIntercept)) {
-  Context ctx = CreateEmptyGenericParam(GPUIDX);
+  Context ctx = MakeCUDACtx(GPUIDX);
   Args args{{"quantile_alpha", "[0.6, 0.8]"}};
   std::unique_ptr<ObjFunction> obj{ObjFunction::Create("reg:quantileerror", &ctx)};
   obj->Configure(args);

tests/cpp/objective/test_regression_obj.cc

@@ -17,7 +17,7 @@
 namespace xgboost {
 
 TEST(Objective, DeclareUnifiedTest(LinearRegressionGPair)) {
-  Context ctx = CreateEmptyGenericParam(GPUIDX);
+  Context ctx = MakeCUDACtx(GPUIDX);
   std::vector<std::pair<std::string, std::string>> args;
 
   std::unique_ptr<ObjFunction> obj{ObjFunction::Create("reg:squarederror", &ctx)};
@@ -39,7 +39,7 @@ TEST(Objective, DeclareUnifiedTest(LinearRegressionGPair)) {
 }
 
 TEST(Objective, DeclareUnifiedTest(SquaredLog)) {
-  Context ctx = CreateEmptyGenericParam(GPUIDX);
+  Context ctx = MakeCUDACtx(GPUIDX);
   std::vector<std::pair<std::string, std::string>> args;
 
   std::unique_ptr<ObjFunction> obj{ObjFunction::Create("reg:squaredlogerror", &ctx)};
@@ -62,7 +62,7 @@ TEST(Objective, DeclareUnifiedTest(SquaredLog)) {
 }
 
 TEST(Objective, DeclareUnifiedTest(PseudoHuber)) {
-  Context ctx = CreateEmptyGenericParam(GPUIDX);
+  Context ctx = MakeCUDACtx(GPUIDX);
   Args args;
 
   std::unique_ptr<ObjFunction> obj{ObjFunction::Create("reg:pseudohubererror", &ctx)};
@@ -91,7 +91,7 @@ TEST(Objective, DeclareUnifiedTest(PseudoHuber)) {
 }
 
 TEST(Objective, DeclareUnifiedTest(LogisticRegressionGPair)) {
-  Context ctx = CreateEmptyGenericParam(GPUIDX);
+  Context ctx = MakeCUDACtx(GPUIDX);
   std::vector<std::pair<std::string, std::string>> args;
   std::unique_ptr<ObjFunction> obj{ObjFunction::Create("reg:logistic", &ctx)};
 
@@ -107,7 +107,7 @@ TEST(Objective, DeclareUnifiedTest(LogisticRegressionGPair)) {
 }
 
 TEST(Objective, DeclareUnifiedTest(LogisticRegressionBasic)) {
-  Context ctx = CreateEmptyGenericParam(GPUIDX);
+  Context ctx = MakeCUDACtx(GPUIDX);
   std::vector<std::pair<std::string, std::string>> args;
   std::unique_ptr<ObjFunction> obj{ObjFunction::Create("reg:logistic", &ctx)};
 
@@ -136,7 +136,7 @@ TEST(Objective, DeclareUnifiedTest(LogisticRegressionBasic)) {
 }
 
 TEST(Objective, DeclareUnifiedTest(LogisticRawGPair)) {
-  Context ctx = CreateEmptyGenericParam(GPUIDX);
+  Context ctx = MakeCUDACtx(GPUIDX);
   std::vector<std::pair<std::string, std::string>> args;
   std::unique_ptr<ObjFunction>  obj {
     ObjFunction::Create("binary:logitraw", &ctx)
@@ -152,7 +152,7 @@ TEST(Objective, DeclareUnifiedTest(LogisticRawGPair)) {
 }
 
 TEST(Objective, DeclareUnifiedTest(PoissonRegressionGPair)) {
-  Context ctx = CreateEmptyGenericParam(GPUIDX);
+  Context ctx = MakeCUDACtx(GPUIDX);
   std::vector<std::pair<std::string, std::string>> args;
   std::unique_ptr<ObjFunction> obj {
     ObjFunction::Create("count:poisson", &ctx)
@@ -176,7 +176,7 @@ TEST(Objective, DeclareUnifiedTest(PoissonRegressionGPair)) {
 }
 
 TEST(Objective, DeclareUnifiedTest(PoissonRegressionBasic)) {
-  Context ctx = CreateEmptyGenericParam(GPUIDX);
+  Context ctx = MakeCUDACtx(GPUIDX);
   std::vector<std::pair<std::string, std::string>> args;
   std::unique_ptr<ObjFunction> obj {
     ObjFunction::Create("count:poisson", &ctx)
@@ -205,7 +205,7 @@ TEST(Objective, DeclareUnifiedTest(PoissonRegressionBasic)) {
 }
 
 TEST(Objective, DeclareUnifiedTest(GammaRegressionGPair)) {
-  Context ctx = CreateEmptyGenericParam(GPUIDX);
+  Context ctx = MakeCUDACtx(GPUIDX);
   std::vector<std::pair<std::string, std::string>> args;
   std::unique_ptr<ObjFunction> obj {
     ObjFunction::Create("reg:gamma", &ctx)
@@ -227,7 +227,7 @@ TEST(Objective, DeclareUnifiedTest(GammaRegressionGPair)) {
 }
 
 TEST(Objective, DeclareUnifiedTest(GammaRegressionBasic)) {
-  Context ctx = CreateEmptyGenericParam(GPUIDX);
+  Context ctx = MakeCUDACtx(GPUIDX);
   std::vector<std::pair<std::string, std::string>> args;
   std::unique_ptr<ObjFunction> obj{ObjFunction::Create("reg:gamma", &ctx)};
 
@@ -256,7 +256,7 @@ TEST(Objective, DeclareUnifiedTest(GammaRegressionBasic)) {
 }
 
 TEST(Objective, DeclareUnifiedTest(TweedieRegressionGPair)) {
-  Context ctx = CreateEmptyGenericParam(GPUIDX);
+  Context ctx = MakeCUDACtx(GPUIDX);
   std::vector<std::pair<std::string, std::string>> args;
   std::unique_ptr<ObjFunction> obj{ObjFunction::Create("reg:tweedie", &ctx)};
 
@@ -280,7 +280,7 @@ TEST(Objective, DeclareUnifiedTest(TweedieRegressionGPair)) {
 
 #if defined(__CUDACC__) || defined(__HIP_PLATFORM_AMD__)
 TEST(Objective, CPU_vs_CUDA) {
-  Context ctx = CreateEmptyGenericParam(GPUIDX);
+  Context ctx = MakeCUDACtx(GPUIDX);
 
   ObjFunction* obj = ObjFunction::Create("reg:squarederror", &ctx);
   HostDeviceVector<GradientPair> cpu_out_preds;
@@ -331,7 +331,7 @@ TEST(Objective, CPU_vs_CUDA) {
 #endif
 
 TEST(Objective, DeclareUnifiedTest(TweedieRegressionBasic)) {
-  Context ctx = CreateEmptyGenericParam(GPUIDX);
+  Context ctx = MakeCUDACtx(GPUIDX);
   std::vector<std::pair<std::string, std::string>> args;
   std::unique_ptr<ObjFunction> obj{ObjFunction::Create("reg:tweedie", &ctx)};
 
@@ -360,7 +360,7 @@ TEST(Objective, DeclareUnifiedTest(TweedieRegressionBasic)) {
 // CoxRegression not implemented in GPU code, no need for testing.
 #if !defined(__CUDACC__) && !defined(__HIP_PLATFORM_AMD__)
 TEST(Objective, CoxRegressionGPair) {
-  Context ctx = CreateEmptyGenericParam(GPUIDX);
+  Context ctx = MakeCUDACtx(GPUIDX);
   std::vector<std::pair<std::string, std::string>> args;
   std::unique_ptr<ObjFunction> obj{ObjFunction::Create("survival:cox", &ctx)};
 
@@ -375,7 +375,7 @@ TEST(Objective, CoxRegressionGPair) {
 #endif
 
 TEST(Objective, DeclareUnifiedTest(AbsoluteError)) {
-  Context ctx = CreateEmptyGenericParam(GPUIDX);
+  Context ctx = MakeCUDACtx(GPUIDX);
   std::unique_ptr<ObjFunction> obj{ObjFunction::Create("reg:absoluteerror", &ctx)};
   obj->Configure({});
   CheckConfigReload(obj, "reg:absoluteerror");
@@ -419,7 +419,7 @@ TEST(Objective, DeclareUnifiedTest(AbsoluteError)) {
 }
 
 TEST(Objective, DeclareUnifiedTest(AbsoluteErrorLeaf)) {
-  Context ctx = CreateEmptyGenericParam(GPUIDX);
+  Context ctx = MakeCUDACtx(GPUIDX);
   bst_target_t constexpr kTargets = 3, kRows = 16;
   std::unique_ptr<ObjFunction> obj{ObjFunction::Create("reg:absoluteerror", &ctx)};
   obj->Configure({});

tests/cpp/plugin/test_example_objective.cc

@@ -4,12 +4,10 @@
 #include "../helpers.h"
 
 namespace xgboost {
-
 TEST(Plugin, ExampleObjective) {
-  xgboost::Context ctx = CreateEmptyGenericParam(GPUIDX);
+  xgboost::Context ctx = MakeCUDACtx(GPUIDX);
   auto* obj = xgboost::ObjFunction::Create("mylogistic", &ctx);
   ASSERT_EQ(obj->DefaultEvalMetric(), std::string{"logloss"});
   delete obj;
 }
-
 }  // namespace xgboost

tests/cpp/plugin/test_predictor_oneapi.cc

@@ -12,7 +12,7 @@
 
 namespace xgboost {
 TEST(Plugin, OneAPIPredictorBasic) {
-  auto lparam = CreateEmptyGenericParam(0);
+  auto lparam = MakeCUDACtx(0);
   std::unique_ptr<Predictor> oneapi_predictor =
       std::unique_ptr<Predictor>(Predictor::Create("oneapi_predictor", &lparam));
 
@@ -82,7 +82,7 @@ TEST(Plugin, OneAPIPredictorExternalMemory) {
   dmlc::TemporaryDirectory tmpdir;
   std::string filename = tmpdir.path + "/big.libsvm";
   std::unique_ptr<DMatrix> dmat = CreateSparsePageDMatrix(12, 64, filename);
-  auto lparam = CreateEmptyGenericParam(0);
+  auto lparam = MakeCUDACtx(0);
 
   std::unique_ptr<Predictor> oneapi_predictor =
       std::unique_ptr<Predictor>(Predictor::Create("oneapi_predictor", &lparam));

tests/cpp/plugin/test_regression_obj_oneapi.cc

@@ -9,7 +9,7 @@
 namespace xgboost {
 
 TEST(Plugin, LinearRegressionGPairOneAPI) {
-  Context tparam = CreateEmptyGenericParam(0);
+  Context tparam = MakeCUDACtx(0);
   std::vector<std::pair<std::string, std::string>> args;
 
   std::unique_ptr<ObjFunction> obj {
@@ -33,7 +33,7 @@ TEST(Plugin, LinearRegressionGPairOneAPI) {
 }
 
 TEST(Plugin, SquaredLogOneAPI) {
-  Context tparam = CreateEmptyGenericParam(0);
+  Context tparam = MakeCUDACtx(0);
   std::vector<std::pair<std::string, std::string>> args;
 
   std::unique_ptr<ObjFunction> obj { ObjFunction::Create("reg:squaredlogerror_oneapi", &tparam) };
@@ -56,7 +56,7 @@ TEST(Plugin, SquaredLogOneAPI) {
 }
 
 TEST(Plugin, LogisticRegressionGPairOneAPI) {
-  Context tparam = CreateEmptyGenericParam(0);
+  Context tparam = MakeCUDACtx(0);
   std::vector<std::pair<std::string, std::string>> args;
   std::unique_ptr<ObjFunction> obj { ObjFunction::Create("reg:logistic_oneapi", &tparam) };
 
@@ -72,7 +72,7 @@ TEST(Plugin, LogisticRegressionGPairOneAPI) {
 }
 
 TEST(Plugin, LogisticRegressionBasicOneAPI) {
-  Context lparam = CreateEmptyGenericParam(0);
+  Context lparam = MakeCUDACtx(0);
   std::vector<std::pair<std::string, std::string>> args;
   std::unique_ptr<ObjFunction> obj {
     ObjFunction::Create("reg:logistic_oneapi", &lparam)
@@ -103,7 +103,7 @@ TEST(Plugin, LogisticRegressionBasicOneAPI) {
 }
 
 TEST(Plugin, LogisticRawGPairOneAPI) {
-  Context lparam = CreateEmptyGenericParam(0);
+  Context lparam = MakeCUDACtx(0);
   std::vector<std::pair<std::string, std::string>> args;
   std::unique_ptr<ObjFunction>  obj {
     ObjFunction::Create("binary:logitraw_oneapi", &lparam)
@@ -120,7 +120,7 @@ TEST(Plugin, LogisticRawGPairOneAPI) {
 }
 
 TEST(Plugin, CPUvsOneAPI) {
-  Context ctx = CreateEmptyGenericParam(0);
+  Context ctx = MakeCUDACtx(0);
 
   ObjFunction * obj_cpu =
       ObjFunction::Create("reg:squarederror", &ctx);
@@ -140,8 +140,8 @@ TEST(Plugin, CPUvsOneAPI) {
   }
   auto& info = pdmat->Info();
 
-  info.labels_.Resize(kRows);
+  info.labels.Reshape(kRows, 1);
-  auto& h_labels = info.labels_.HostVector();
+  auto& h_labels = info.labels.Data()->HostVector();
   for (size_t i = 0; i < h_labels.size(); ++i) {
     h_labels[i] = 1 / static_cast<float>(i+1);
   }

tests/cpp/predictor/test_cpu_predictor.cc

@@ -20,16 +20,15 @@ namespace xgboost {
 
 namespace {
 void TestBasic(DMatrix* dmat) {
-  auto lparam = CreateEmptyGenericParam(GPUIDX);
+  Context ctx;
   std::unique_ptr<Predictor> cpu_predictor =
-      std::unique_ptr<Predictor>(Predictor::Create("cpu_predictor", &lparam));
+      std::unique_ptr<Predictor>(Predictor::Create("cpu_predictor", &ctx));
 
   size_t const kRows = dmat->Info().num_row_;
   size_t const kCols = dmat->Info().num_col_;
 
   LearnerModelParam mparam{MakeMP(kCols, .0, 1)};
 
-  Context ctx;
   ctx.UpdateAllowUnknown(Args{});
   gbm::GBTreeModel model = CreateTestModel(&mparam, &ctx);
 
@@ -117,7 +116,7 @@ void TestColumnSplit() {
 }
 }  // anonymous namespace
 
-TEST(CpuPredictor, ColumnSplitBasic) {
+TEST(CpuPredictor, BasicColumnSplit) {
   auto constexpr kWorldSize = 2;
   RunWithInMemoryCommunicator(kWorldSize, TestColumnSplit);
 }
@@ -126,6 +125,10 @@ TEST(CpuPredictor, IterationRange) {
   TestIterationRange("cpu_predictor");
 }
 
+TEST(CpuPredictor, IterationRangeColmnSplit) {
+  TestIterationRangeColumnSplit("cpu_predictor");
+}
+
 TEST(CpuPredictor, ExternalMemory) {
   size_t constexpr kPageSize = 64, kEntriesPerCol = 3;
   size_t constexpr kEntries = kPageSize * kEntriesPerCol * 2;
@@ -223,10 +226,18 @@ TEST(CPUPredictor, CategoricalPrediction) {
   TestCategoricalPrediction("cpu_predictor");
 }
 
+TEST(CPUPredictor, CategoricalPredictionColumnSplit) {
+  TestCategoricalPredictionColumnSplit("cpu_predictor");
+}
+
 TEST(CPUPredictor, CategoricalPredictLeaf) {
   TestCategoricalPredictLeaf(StringView{"cpu_predictor"});
 }
 
+TEST(CPUPredictor, CategoricalPredictLeafColumnSplit) {
+  TestCategoricalPredictLeafColumnSplit(StringView{"cpu_predictor"});
+}
+
 TEST(CpuPredictor, UpdatePredictionCache) {
   TestUpdatePredictionCache(false);
   TestUpdatePredictionCache(true);
@@ -236,11 +247,20 @@ TEST(CpuPredictor, LesserFeatures) {
   TestPredictionWithLesserFeatures("cpu_predictor");
 }
 
+TEST(CpuPredictor, LesserFeaturesColumnSplit) {
+  TestPredictionWithLesserFeaturesColumnSplit("cpu_predictor");
+}
+
 TEST(CpuPredictor, Sparse) {
   TestSparsePrediction(0.2, "cpu_predictor");
   TestSparsePrediction(0.8, "cpu_predictor");
 }
 
+TEST(CpuPredictor, SparseColumnSplit) {
+  TestSparsePredictionColumnSplit(0.2, "cpu_predictor");
+  TestSparsePredictionColumnSplit(0.8, "cpu_predictor");
+}
+
 TEST(CpuPredictor, Multi) {
   Context ctx;
   ctx.nthread = 1;

tests/cpp/predictor/test_gpu_predictor.cu

@@ -23,8 +23,8 @@ namespace xgboost {
 namespace predictor {
 
 TEST(GPUPredictor, Basic) {
-  auto cpu_lparam = CreateEmptyGenericParam(-1);
+  auto cpu_lparam = MakeCUDACtx(-1);
-  auto gpu_lparam = CreateEmptyGenericParam(0);
+  auto gpu_lparam = MakeCUDACtx(0);
 
   std::unique_ptr<Predictor> gpu_predictor =
       std::unique_ptr<Predictor>(Predictor::Create("gpu_predictor", &gpu_lparam));
@@ -88,7 +88,7 @@ TEST(GPUPredictor, EllpackTraining) {
 }
 
 TEST(GPUPredictor, ExternalMemoryTest) {
-  auto lparam = CreateEmptyGenericParam(0);
+  auto lparam = MakeCUDACtx(0);
   std::unique_ptr<Predictor> gpu_predictor =
       std::unique_ptr<Predictor>(Predictor::Create("gpu_predictor", &lparam));
   gpu_predictor->Configure({});
@@ -165,7 +165,7 @@ TEST(GPUPredictor, ShapStump) {
   trees.push_back(std::unique_ptr<RegTree>(new RegTree));
   model.CommitModelGroup(std::move(trees), 0);
 
-  auto gpu_lparam = CreateEmptyGenericParam(0);
+  auto gpu_lparam = MakeCUDACtx(0);
   std::unique_ptr<Predictor> gpu_predictor = std::unique_ptr<Predictor>(
       Predictor::Create("gpu_predictor", &gpu_lparam));
   gpu_predictor->Configure({});
@@ -193,8 +193,8 @@ TEST(GPUPredictor, Shap) {
   trees[0]->ExpandNode(0, 0, 0.5, true, 1.0, -1.0, 1.0, 0.0, 5.0, 2.0, 3.0);
   model.CommitModelGroup(std::move(trees), 0);
 
-  auto gpu_lparam = CreateEmptyGenericParam(0);
+  auto gpu_lparam = MakeCUDACtx(0);
-  auto cpu_lparam = CreateEmptyGenericParam(-1);
+  auto cpu_lparam = MakeCUDACtx(-1);
   std::unique_ptr<Predictor> gpu_predictor = std::unique_ptr<Predictor>(
       Predictor::Create("gpu_predictor", &gpu_lparam));
   std::unique_ptr<Predictor> cpu_predictor = std::unique_ptr<Predictor>(
@@ -228,7 +228,7 @@ TEST(GPUPredictor, CategoricalPredictLeaf) {
 TEST(GPUPredictor, PredictLeafBasic) {
   size_t constexpr kRows = 5, kCols = 5;
   auto dmat = RandomDataGenerator(kRows, kCols, 0).Device(0).GenerateDMatrix();
-  auto lparam = CreateEmptyGenericParam(GPUIDX);
+  auto lparam = MakeCUDACtx(GPUIDX);
   std::unique_ptr<Predictor> gpu_predictor =
       std::unique_ptr<Predictor>(Predictor::Create("gpu_predictor", &lparam));
   gpu_predictor->Configure({});

tests/cpp/predictor/test_predictor.cc

@@ -153,28 +153,32 @@ void TestInplacePrediction(std::shared_ptr<DMatrix> x, std::string predictor, bs
   learner->Configure();
 }
 
-void TestPredictionWithLesserFeatures(std::string predictor_name) {
+namespace {
-  size_t constexpr kRows = 256, kTrainCols = 256, kTestCols = 4, kIters = 4;
+std::unique_ptr<Learner> LearnerForTest(std::shared_ptr<DMatrix> dmat, size_t iters,
-  auto m_train = RandomDataGenerator(kRows, kTrainCols, 0.5).GenerateDMatrix(true);
+                                        size_t forest = 1) {
-  auto m_test = RandomDataGenerator(kRows, kTestCols, 0.5).GenerateDMatrix(false);
+  std::unique_ptr<Learner> learner{Learner::Create({dmat})};
-  std::unique_ptr<Learner> learner{Learner::Create({m_train})};
+  learner->SetParams(Args{{"num_parallel_tree", std::to_string(forest)}});
-
+  for (size_t i = 0; i < iters; ++i) {
-  for (size_t i = 0; i < kIters; ++i) {
+    learner->UpdateOneIter(i, dmat);
-    learner->UpdateOneIter(i, m_train);
   }
+  return learner;
+}
 
+void VerifyPredictionWithLesserFeatures(Learner *learner, std::string const &predictor_name,
+                                        size_t rows, std::shared_ptr<DMatrix> const &m_test,
+                                        std::shared_ptr<DMatrix> const &m_invalid) {
   HostDeviceVector<float> prediction;
   learner->SetParam("predictor", predictor_name);
   learner->Configure();
   Json config{Object()};
   learner->SaveConfig(&config);
-  ASSERT_EQ(get<String>(config["learner"]["gradient_booster"]["gbtree_train_param"]["predictor"]), predictor_name);
+  ASSERT_EQ(get<String>(config["learner"]["gradient_booster"]["gbtree_train_param"]["predictor"]),
+            predictor_name);
 
   learner->Predict(m_test, false, &prediction, 0, 0);
-  ASSERT_EQ(prediction.Size(), kRows);
+  ASSERT_EQ(prediction.Size(), rows);
 
-  auto m_invalid = RandomDataGenerator(kRows, kTrainCols + 1, 0.5).GenerateDMatrix(false);
+  ASSERT_THROW({ learner->Predict(m_invalid, false, &prediction, 0, 0); }, dmlc::Error);
-  ASSERT_THROW({learner->Predict(m_invalid, false, &prediction, 0, 0);}, dmlc::Error);
 
 #if defined(XGBOOST_USE_CUDA) || defined(XGBOOST_USE_HIP)
   HostDeviceVector<float> from_cpu;
@@ -185,13 +189,49 @@ void TestPredictionWithLesserFeatures(std::string predictor_name) {
   learner->SetParam("predictor", "gpu_predictor");
   learner->Predict(m_test, false, &from_cuda, 0, 0);
 
-  auto const& h_cpu = from_cpu.ConstHostVector();
+  auto const &h_cpu = from_cpu.ConstHostVector();
-  auto const& h_gpu = from_cuda.ConstHostVector();
+  auto const &h_gpu = from_cuda.ConstHostVector();
   for (size_t i = 0; i < h_cpu.size(); ++i) {
     ASSERT_NEAR(h_cpu[i], h_gpu[i], kRtEps);
   }
 #endif  // defined(XGBOOST_USE_CUDA) || defined(XGBOOST_USE_HIP)
 }
+}  // anonymous namespace
+
+void TestPredictionWithLesserFeatures(std::string predictor_name) {
+  size_t constexpr kRows = 256, kTrainCols = 256, kTestCols = 4, kIters = 4;
+  auto m_train = RandomDataGenerator(kRows, kTrainCols, 0.5).GenerateDMatrix(true);
+  auto learner = LearnerForTest(m_train, kIters);
+  auto m_test = RandomDataGenerator(kRows, kTestCols, 0.5).GenerateDMatrix(false);
+  auto m_invalid = RandomDataGenerator(kRows, kTrainCols + 1, 0.5).GenerateDMatrix(false);
+  VerifyPredictionWithLesserFeatures(learner.get(), predictor_name, kRows, m_test, m_invalid);
+}
+
+namespace {
+void VerifyPredictionWithLesserFeaturesColumnSplit(Learner *learner,
+                                                   std::string const &predictor_name, size_t rows,
+                                                   std::shared_ptr<DMatrix> m_test,
+                                                   std::shared_ptr<DMatrix> m_invalid) {
+  auto const world_size = collective::GetWorldSize();
+  auto const rank = collective::GetRank();
+  std::shared_ptr<DMatrix> sliced_test{m_test->SliceCol(world_size, rank)};
+  std::shared_ptr<DMatrix> sliced_invalid{m_invalid->SliceCol(world_size, rank)};
+
+  VerifyPredictionWithLesserFeatures(learner, predictor_name, rows, sliced_test, sliced_invalid);
+}
+}  // anonymous namespace
+
+void TestPredictionWithLesserFeaturesColumnSplit(std::string predictor_name) {
+  size_t constexpr kRows = 256, kTrainCols = 256, kTestCols = 4, kIters = 4;
+  auto m_train = RandomDataGenerator(kRows, kTrainCols, 0.5).GenerateDMatrix(true);
+  auto learner = LearnerForTest(m_train, kIters);
+  auto m_test = RandomDataGenerator(kRows, kTestCols, 0.5).GenerateDMatrix(false);
+  auto m_invalid = RandomDataGenerator(kRows, kTrainCols + 1, 0.5).GenerateDMatrix(false);
+
+  auto constexpr kWorldSize = 2;
+  RunWithInMemoryCommunicator(kWorldSize, VerifyPredictionWithLesserFeaturesColumnSplit,
+                              learner.get(), predictor_name, kRows, m_test, m_invalid);
+}
 
 void GBTreeModelForTest(gbm::GBTreeModel *model, uint32_t split_ind,
                         bst_cat_t split_cat, float left_weight,
@@ -212,7 +252,7 @@ void GBTreeModelForTest(gbm::GBTreeModel *model, uint32_t split_ind,
   model->CommitModelGroup(std::move(trees), 0);
 }
 
-void TestCategoricalPrediction(std::string name) {
+void TestCategoricalPrediction(std::string name, bool is_column_split) {
   size_t constexpr kCols = 10;
   PredictionCacheEntry out_predictions;
 
@@ -236,6 +276,9 @@ void TestCategoricalPrediction(std::string name) {
 
   std::vector<FeatureType> types(10, FeatureType::kCategorical);
   m->Info().feature_types.HostVector() = types;
+  if (is_column_split) {
+    m = std::shared_ptr<DMatrix>{m->SliceCol(collective::GetWorldSize(), collective::GetRank())};
+  }
 
   predictor->InitOutPredictions(m->Info(), &out_predictions.predictions, model);
   predictor->PredictBatch(m.get(), &out_predictions, model, 0);
@@ -246,13 +289,21 @@ void TestCategoricalPrediction(std::string name) {
 
   row[split_ind] = split_cat + 1;
   m = GetDMatrixFromData(row, 1, kCols);
+  if (is_column_split) {
+    m = std::shared_ptr<DMatrix>{m->SliceCol(collective::GetWorldSize(), collective::GetRank())};
+  }
   out_predictions.version = 0;
   predictor->InitOutPredictions(m->Info(), &out_predictions.predictions, model);
   predictor->PredictBatch(m.get(), &out_predictions, model, 0);
   ASSERT_EQ(out_predictions.predictions.HostVector()[0], left_weight + score);
 }
 
-void TestCategoricalPredictLeaf(StringView name) {
+void TestCategoricalPredictionColumnSplit(std::string name) {
+  auto constexpr kWorldSize = 2;
+  RunWithInMemoryCommunicator(kWorldSize, TestCategoricalPrediction, name, true);
+}
+
+void TestCategoricalPredictLeaf(StringView name, bool is_column_split) {
   size_t constexpr kCols = 10;
   PredictionCacheEntry out_predictions;
 
@@ -275,6 +326,9 @@ void TestCategoricalPredictLeaf(StringView name) {
   std::vector<float> row(kCols);
   row[split_ind] = split_cat;
   auto m = GetDMatrixFromData(row, 1, kCols);
+  if (is_column_split) {
+    m = std::shared_ptr<DMatrix>{m->SliceCol(collective::GetWorldSize(), collective::GetRank())};
+  }
 
   predictor->PredictLeaf(m.get(), &out_predictions.predictions, model);
   CHECK_EQ(out_predictions.predictions.Size(), 1);
@@ -283,25 +337,25 @@ void TestCategoricalPredictLeaf(StringView name) {
 
   row[split_ind] = split_cat + 1;
   m = GetDMatrixFromData(row, 1, kCols);
+  if (is_column_split) {
+    m = std::shared_ptr<DMatrix>{m->SliceCol(collective::GetWorldSize(), collective::GetRank())};
+  }
   out_predictions.version = 0;
   predictor->InitOutPredictions(m->Info(), &out_predictions.predictions, model);
   predictor->PredictLeaf(m.get(), &out_predictions.predictions, model);
   ASSERT_EQ(out_predictions.predictions.HostVector()[0], 1);
 }
 
+void TestCategoricalPredictLeafColumnSplit(StringView name) {
+  auto constexpr kWorldSize = 2;
+  RunWithInMemoryCommunicator(kWorldSize, TestCategoricalPredictLeaf, name, true);
+}
 
 void TestIterationRange(std::string name) {
-  size_t constexpr kRows = 1000, kCols = 20, kClasses = 4, kForest = 3;
+  size_t constexpr kRows = 1000, kCols = 20, kClasses = 4, kForest = 3, kIters = 10;
   auto dmat = RandomDataGenerator(kRows, kCols, 0).GenerateDMatrix(true, true, kClasses);
-  std::unique_ptr<Learner> learner{Learner::Create({dmat})};
+  auto learner = LearnerForTest(dmat, kIters, kForest);
-
+  learner->SetParams(Args{{"predictor", name}});
-  learner->SetParams(Args{{"num_parallel_tree", std::to_string(kForest)},
-                          {"predictor", name}});
-
-  size_t kIters = 10;
-  for (size_t i = 0; i < kIters; ++i) {
-    learner->UpdateOneIter(i, dmat);
-  }
 
   bool bound = false;
   std::unique_ptr<Learner> sliced {learner->Slice(0, 3, 1, &bound)};
@@ -363,15 +417,82 @@ void TestIterationRange(std::string name) {
   }
 }
 
-void TestSparsePrediction(float sparsity, std::string predictor) {
+namespace {
-  size_t constexpr kRows = 512, kCols = 128;
+void VerifyIterationRangeColumnSplit(DMatrix *dmat, Learner *learner, Learner *sliced,
-  auto Xy = RandomDataGenerator(kRows, kCols, sparsity).GenerateDMatrix(true);
+                                     std::vector<float> const &expected_margin_ranged,
-  std::unique_ptr<Learner> learner{Learner::Create({Xy})};
+                                     std::vector<float> const &expected_margin_sliced,
-  learner->Configure();
+                                     std::vector<float> const &expected_leaf_ranged,
-  for (size_t i = 0; i < 4; ++i) {
+                                     std::vector<float> const &expected_leaf_sliced) {
-    learner->UpdateOneIter(i, Xy);
+  auto const world_size = collective::GetWorldSize();
+  auto const rank = collective::GetRank();
+  std::shared_ptr<DMatrix> Xy{dmat->SliceCol(world_size, rank)};
+
+  HostDeviceVector<float> out_predt_sliced;
+  HostDeviceVector<float> out_predt_ranged;
+
+  // margin
+  {
+    sliced->Predict(Xy, true, &out_predt_sliced, 0, 0, false, false, false, false, false);
+    learner->Predict(Xy, true, &out_predt_ranged, 0, 3, false, false, false, false, false);
+    auto const &h_sliced = out_predt_sliced.HostVector();
+    auto const &h_range = out_predt_ranged.HostVector();
+    ASSERT_EQ(h_sliced.size(), expected_margin_sliced.size());
+    ASSERT_EQ(h_sliced, expected_margin_sliced);
+    ASSERT_EQ(h_range.size(), expected_margin_ranged.size());
+    ASSERT_EQ(h_range, expected_margin_ranged);
   }
 
+  // Leaf
+  {
+    sliced->Predict(Xy, false, &out_predt_sliced, 0, 0, false, true, false, false, false);
+    learner->Predict(Xy, false, &out_predt_ranged, 0, 3, false, true, false, false, false);
+    auto const &h_sliced = out_predt_sliced.HostVector();
+    auto const &h_range = out_predt_ranged.HostVector();
+    ASSERT_EQ(h_sliced.size(), expected_leaf_sliced.size());
+    ASSERT_EQ(h_sliced, expected_leaf_sliced);
+    ASSERT_EQ(h_range.size(), expected_leaf_ranged.size());
+    ASSERT_EQ(h_range, expected_leaf_ranged);
+  }
+}
+}  // anonymous namespace
+
+void TestIterationRangeColumnSplit(std::string name) {
+  size_t constexpr kRows = 1000, kCols = 20, kClasses = 4, kForest = 3, kIters = 10;
+  auto dmat = RandomDataGenerator(kRows, kCols, 0).GenerateDMatrix(true, true, kClasses);
+  auto learner = LearnerForTest(dmat, kIters, kForest);
+  learner->SetParams(Args{{"predictor", name}});
+
+  bool bound = false;
+  std::unique_ptr<Learner> sliced{learner->Slice(0, 3, 1, &bound)};
+  ASSERT_FALSE(bound);
+
+  // margin
+  HostDeviceVector<float> margin_predt_sliced;
+  HostDeviceVector<float> margin_predt_ranged;
+  sliced->Predict(dmat, true, &margin_predt_sliced, 0, 0, false, false, false, false, false);
+  learner->Predict(dmat, true, &margin_predt_ranged, 0, 3, false, false, false, false, false);
+  auto const &margin_sliced = margin_predt_sliced.HostVector();
+  auto const &margin_ranged = margin_predt_ranged.HostVector();
+
+  // Leaf
+  HostDeviceVector<float> leaf_predt_sliced;
+  HostDeviceVector<float> leaf_predt_ranged;
+  sliced->Predict(dmat, false, &leaf_predt_sliced, 0, 0, false, true, false, false, false);
+  learner->Predict(dmat, false, &leaf_predt_ranged, 0, 3, false, true, false, false, false);
+  auto const &leaf_sliced = leaf_predt_sliced.HostVector();
+  auto const &leaf_ranged = leaf_predt_ranged.HostVector();
+
+  auto constexpr kWorldSize = 2;
+  RunWithInMemoryCommunicator(kWorldSize, VerifyIterationRangeColumnSplit, dmat.get(),
+                              learner.get(), sliced.get(), margin_ranged, margin_sliced,
+                              leaf_ranged, leaf_sliced);
+}
+
+void TestSparsePrediction(float sparsity, std::string predictor) {
+  size_t constexpr kRows = 512, kCols = 128, kIters = 4;
+  auto Xy = RandomDataGenerator(kRows, kCols, sparsity).GenerateDMatrix(true);
+  auto learner = LearnerForTest(Xy, kIters);
+
   HostDeviceVector<float> sparse_predt;
 
   Json model{Object{}};
@@ -419,6 +540,43 @@ void TestSparsePrediction(float sparsity, std::string predictor) {
   }
 }
 
+namespace {
+void VerifySparsePredictionColumnSplit(DMatrix *dmat, Learner *learner,
+                                       std::vector<float> const &expected_predt) {
+  std::shared_ptr<DMatrix> sliced{
+      dmat->SliceCol(collective::GetWorldSize(), collective::GetRank())};
+  HostDeviceVector<float> sparse_predt;
+  learner->Predict(sliced, false, &sparse_predt, 0, 0);
+
+  auto const &predt = sparse_predt.HostVector();
+  ASSERT_EQ(predt.size(), expected_predt.size());
+  for (size_t i = 0; i < predt.size(); ++i) {
+    ASSERT_FLOAT_EQ(predt[i], expected_predt[i]);
+  }
+}
+}  // anonymous namespace
+
+void TestSparsePredictionColumnSplit(float sparsity, std::string predictor) {
+  size_t constexpr kRows = 512, kCols = 128, kIters = 4;
+  auto Xy = RandomDataGenerator(kRows, kCols, sparsity).GenerateDMatrix(true);
+  auto learner = LearnerForTest(Xy, kIters);
+
+  HostDeviceVector<float> sparse_predt;
+
+  Json model{Object{}};
+  learner->SaveModel(&model);
+
+  learner.reset(Learner::Create({Xy}));
+  learner->LoadModel(model);
+
+  learner->SetParam("predictor", predictor);
+  learner->Predict(Xy, false, &sparse_predt, 0, 0);
+
+  auto constexpr kWorldSize = 2;
+  RunWithInMemoryCommunicator(kWorldSize, VerifySparsePredictionColumnSplit, Xy.get(),
+                              learner.get(), sparse_predt.HostVector());
+}
+
 void TestVectorLeafPrediction(Context const *ctx) {
   std::unique_ptr<Predictor> cpu_predictor =
       std::unique_ptr<Predictor>(Predictor::Create("cpu_predictor", ctx));

tests/cpp/predictor/test_predictor.h

@@ -37,10 +37,10 @@ void TestPredictionFromGradientIndex(std::string name, size_t rows, size_t cols,
   constexpr size_t kClasses { 3 };
 
   LearnerModelParam mparam{MakeMP(cols, .5, kClasses)};
-  auto lparam = CreateEmptyGenericParam(0);
+  auto cuda_ctx = MakeCUDACtx(0);
 
   std::unique_ptr<Predictor> predictor =
-      std::unique_ptr<Predictor>(Predictor::Create(name, &lparam));
+      std::unique_ptr<Predictor>(Predictor::Create(name, &cuda_ctx));
   predictor->Configure({});
 
   Context ctx;
@@ -86,14 +86,24 @@ void TestInplacePrediction(std::shared_ptr<DMatrix> x, std::string predictor, bs
 
 void TestPredictionWithLesserFeatures(std::string preditor_name);
 
-void TestCategoricalPrediction(std::string name);
+void TestPredictionWithLesserFeaturesColumnSplit(std::string preditor_name);
 
-void TestCategoricalPredictLeaf(StringView name);
+void TestCategoricalPrediction(std::string name, bool is_column_split = false);
+
+void TestCategoricalPredictionColumnSplit(std::string name);
+
+void TestCategoricalPredictLeaf(StringView name, bool is_column_split = false);
+
+void TestCategoricalPredictLeafColumnSplit(StringView name);
 
 void TestIterationRange(std::string name);
 
+void TestIterationRangeColumnSplit(std::string name);
+
 void TestSparsePrediction(float sparsity, std::string predictor);
 
+void TestSparsePredictionColumnSplit(float sparsity, std::string predictor);
+
 void TestVectorLeafPrediction(Context const* ctx);
 }  // namespace xgboost
 

tests/cpp/tree/hist/test_histogram.cc

@@ -25,7 +25,7 @@ void InitRowPartitionForTest(common::RowSetCollection *row_set, size_t n_samples
 }  // anonymous namespace
 
 void TestAddHistRows(bool is_distributed) {
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
   std::vector<CPUExpandEntry> nodes_for_explicit_hist_build_;
   std::vector<CPUExpandEntry> nodes_for_subtraction_trick_;
   int starting_index = std::numeric_limits<int>::max();
@@ -74,7 +74,7 @@ TEST(CPUHistogram, AddRows) {
 void TestSyncHist(bool is_distributed) {
   size_t constexpr kNRows = 8, kNCols = 16;
   int32_t constexpr kMaxBins = 4;
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
 
   std::vector<CPUExpandEntry> nodes_for_explicit_hist_build_;
   std::vector<CPUExpandEntry> nodes_for_subtraction_trick_;
@@ -229,7 +229,7 @@ TEST(CPUHistogram, SyncHist) {
 void TestBuildHistogram(bool is_distributed, bool force_read_by_column, bool is_col_split) {
   size_t constexpr kNRows = 8, kNCols = 16;
   int32_t constexpr kMaxBins = 4;
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
   auto p_fmat =
       RandomDataGenerator(kNRows, kNCols, 0.8).Seed(3).GenerateDMatrix();
   if (is_col_split) {
@@ -330,7 +330,7 @@ void TestHistogramCategorical(size_t n_categories, bool force_read_by_column) {
   auto x = GenerateRandomCategoricalSingleColumn(kRows, n_categories);
   auto cat_m = GetDMatrixFromData(x, kRows, 1);
   cat_m->Info().feature_types.HostVector().push_back(FeatureType::kCategorical);
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
 
   BatchParam batch_param{0, static_cast<int32_t>(kBins)};
 
@@ -475,7 +475,7 @@ void TestHistogramExternalMemory(Context const *ctx, BatchParam batch_param, boo
 
 TEST(CPUHistogram, ExternalMemory) {
   int32_t constexpr kBins = 256;
-  auto ctx = CreateEmptyGenericParam(Context::kCpuId);
+  Context ctx;
 
   TestHistogramExternalMemory(&ctx, BatchParam{kBins, common::Span<float>{}, false}, true, false);
   TestHistogramExternalMemory(&ctx, BatchParam{kBins, common::Span<float>{}, false}, true, true);

tests/cpp/tree/test_gpu_hist.cu

@@ -102,7 +102,7 @@ void TestBuildHist(bool use_shared_memory_histograms) {
 
   auto page = BuildEllpackPage(kNRows, kNCols);
   BatchParam batch_param{};
-  Context ctx{CreateEmptyGenericParam(0)};
+  Context ctx{MakeCUDACtx(0)};
   GPUHistMakerDevice<GradientSumT> maker(&ctx, page.get(), {}, kNRows, param, kNCols, kNCols,
                                          batch_param);
   xgboost::SimpleLCG gen;
@@ -186,7 +186,7 @@ void TestHistogramIndexImpl() {
   int constexpr kNRows = 1000, kNCols = 10;
 
   // Build 2 matrices and build a histogram maker with that
-  Context ctx(CreateEmptyGenericParam(0));
+  Context ctx(MakeCUDACtx(0));
   ObjInfo task{ObjInfo::kRegression};
   tree::GPUHistMaker hist_maker{&ctx, &task}, hist_maker_ext{&ctx, &task};
   std::unique_ptr<DMatrix> hist_maker_dmat(
@@ -279,7 +279,7 @@ TEST(GpuHist, UniformSampling) {
   // Build a tree using the in-memory DMatrix.
   RegTree tree;
   HostDeviceVector<bst_float> preds(kRows, 0.0, 0);
-  Context ctx(CreateEmptyGenericParam(0));
+  Context ctx(MakeCUDACtx(0));
   UpdateTree(&ctx, &gpair, dmat.get(), 0, &tree, &preds, 1.0, "uniform", kRows);
   // Build another tree using sampling.
   RegTree tree_sampling;
@@ -309,7 +309,7 @@ TEST(GpuHist, GradientBasedSampling) {
   // Build a tree using the in-memory DMatrix.
   RegTree tree;
   HostDeviceVector<bst_float> preds(kRows, 0.0, 0);
-  Context ctx(CreateEmptyGenericParam(0));
+  Context ctx(MakeCUDACtx(0));
   UpdateTree(&ctx, &gpair, dmat.get(), 0, &tree, &preds, 1.0, "uniform", kRows);
 
   // Build another tree using sampling.
@@ -344,7 +344,7 @@ TEST(GpuHist, ExternalMemory) {
 
   // Build a tree using the in-memory DMatrix.
   RegTree tree;
-  Context ctx(CreateEmptyGenericParam(0));
+  Context ctx(MakeCUDACtx(0));
   HostDeviceVector<bst_float> preds(kRows, 0.0, 0);
   UpdateTree(&ctx, &gpair, dmat.get(), 0, &tree, &preds, 1.0, "uniform", kRows);
   // Build another tree using multiple ELLPACK pages.
@@ -382,7 +382,7 @@ TEST(GpuHist, ExternalMemoryWithSampling) {
   // Build a tree using the in-memory DMatrix.
   auto rng = common::GlobalRandom();
 
-  Context ctx(CreateEmptyGenericParam(0));
+  Context ctx(MakeCUDACtx(0));
   RegTree tree;
   HostDeviceVector<bst_float> preds(kRows, 0.0, 0);
   UpdateTree(&ctx, &gpair, dmat.get(), 0, &tree, &preds, kSubsample, kSamplingMethod, kRows);
@@ -403,7 +403,7 @@ TEST(GpuHist, ExternalMemoryWithSampling) {
 }
 
 TEST(GpuHist, ConfigIO) {
-  Context ctx(CreateEmptyGenericParam(0));
+  Context ctx(MakeCUDACtx(0));
   ObjInfo task{ObjInfo::kRegression};
   std::unique_ptr<TreeUpdater> updater{TreeUpdater::Create("grow_gpu_hist", &ctx, &task)};
   updater->Configure(Args{});
@@ -421,7 +421,7 @@ TEST(GpuHist, ConfigIO) {
 }
 
 TEST(GpuHist, MaxDepth) {
-  Context ctx(CreateEmptyGenericParam(0));
+  Context ctx(MakeCUDACtx(0));
   size_t constexpr kRows = 16;
   size_t constexpr kCols = 4;
   auto p_mat = RandomDataGenerator{kRows, kCols, 0}.GenerateDMatrix();

tests/cpp/tree/test_prune.cc

@@ -29,7 +29,7 @@ TEST(Updater, Prune) {
   std::shared_ptr<DMatrix> p_dmat {
     RandomDataGenerator{32, 10, 0}.GenerateDMatrix() };
 
-  auto ctx = CreateEmptyGenericParam(GPUIDX);
+  Context ctx;
 
   // prepare tree
   RegTree tree = RegTree{1u, kCols};

tests/cpp/tree/test_refresh.cc

@@ -29,7 +29,7 @@ TEST(Updater, Refresh) {
       {"reg_lambda", "1"}};
 
   RegTree tree = RegTree{1u, kCols};
-  auto ctx = CreateEmptyGenericParam(GPUIDX);
+  Context ctx;
   std::vector<RegTree*> trees{&tree};
 
   ObjInfo task{ObjInfo::kRegression};

tests/cpp/tree/test_tree_stat.cc

@@ -33,8 +33,7 @@ class UpdaterTreeStatTest : public ::testing::Test {
     ObjInfo task{ObjInfo::kRegression};
     param.Init(Args{});
 
-    Context ctx(updater == "grow_gpu_hist" ? CreateEmptyGenericParam(0)
+    Context ctx(updater == "grow_gpu_hist" ? MakeCUDACtx(0) : MakeCUDACtx(Context::kCpuId));
-                                           : CreateEmptyGenericParam(Context::kCpuId));
     auto up = std::unique_ptr<TreeUpdater>{TreeUpdater::Create(updater, &ctx, &task)};
     up->Configure(Args{});
     RegTree tree{1u, kCols};
@@ -79,8 +78,7 @@ class UpdaterEtaTest : public ::testing::Test {
   void RunTest(std::string updater) {
     ObjInfo task{ObjInfo::kClassification};
 
-    Context ctx(updater == "grow_gpu_hist" ? CreateEmptyGenericParam(0)
+    Context ctx(updater == "grow_gpu_hist" ? MakeCUDACtx(0) : MakeCUDACtx(Context::kCpuId));
-                                           : CreateEmptyGenericParam(Context::kCpuId));
 
     float eta = 0.4;
     auto up_0 = std::unique_ptr<TreeUpdater>{TreeUpdater::Create(updater, &ctx, &task)};
@@ -156,8 +154,7 @@ class TestMinSplitLoss : public ::testing::Test {
     param.UpdateAllowUnknown(args);
     ObjInfo task{ObjInfo::kRegression};
 
-    Context ctx(updater == "grow_gpu_hist" ? CreateEmptyGenericParam(0)
+    Context ctx{MakeCUDACtx(updater == "grow_gpu_hist" ? 0 : Context::kCpuId)};
-                                           : CreateEmptyGenericParam(Context::kCpuId));
     auto up = std::unique_ptr<TreeUpdater>{TreeUpdater::Create(updater, &ctx, &task)};
     up->Configure({});
 

tests/python/test_eval_metrics.py

@@ -299,7 +299,9 @@ class TestEvalMetrics:
     def run_pr_auc_ltr(self, tree_method):
         from sklearn.datasets import make_classification
         X, y = make_classification(128, 4, n_classes=2, random_state=1994)
-        ltr = xgb.XGBRanker(tree_method=tree_method, n_estimators=16)
+        ltr = xgb.XGBRanker(
+            tree_method=tree_method, n_estimators=16, objective="rank:pairwise"
+        )
         groups = np.array([32, 32, 64])
         ltr.fit(
             X,

tests/python/test_ranking.py

@@ -1,12 +1,57 @@
 import itertools
+import json
 import os
 import shutil
+from typing import Optional
 
 import numpy as np
+import pytest
+from hypothesis import given, note, settings
 from scipy.sparse import csr_matrix
 
 import xgboost
 from xgboost import testing as tm
+from xgboost.testing.data import RelDataCV, simulate_clicks, sort_ltr_samples
+from xgboost.testing.params import lambdarank_parameter_strategy
+
+
+def test_ndcg_custom_gain():
+    def ndcg_gain(y: np.ndarray) -> np.ndarray:
+        return np.exp2(y.astype(np.float64)) - 1.0
+
+    X, y, q, w = tm.make_ltr(n_samples=1024, n_features=4, n_query_groups=3, max_rel=3)
+    y_gain = ndcg_gain(y)
+
+    byxgb = xgboost.XGBRanker(tree_method="hist", ndcg_exp_gain=True, n_estimators=10)
+    byxgb.fit(
+        X,
+        y,
+        qid=q,
+        sample_weight=w,
+        eval_set=[(X, y)],
+        eval_qid=(q,),
+        sample_weight_eval_set=(w,),
+        verbose=True,
+    )
+    byxgb_json = json.loads(byxgb.get_booster().save_raw(raw_format="json"))
+
+    bynp = xgboost.XGBRanker(tree_method="hist", ndcg_exp_gain=False, n_estimators=10)
+    bynp.fit(
+        X,
+        y_gain,
+        qid=q,
+        sample_weight=w,
+        eval_set=[(X, y_gain)],
+        eval_qid=(q,),
+        sample_weight_eval_set=(w,),
+        verbose=True,
+    )
+    bynp_json = json.loads(bynp.get_booster().save_raw(raw_format="json"))
+
+    # Remove the difference in parameter for comparison
+    byxgb_json["learner"]["objective"]["lambdarank_param"]["ndcg_exp_gain"] = "0"
+    assert byxgb.evals_result() == bynp.evals_result()
+    assert byxgb_json == bynp_json
 
 
 def test_ranking_with_unweighted_data():
@@ -73,8 +118,77 @@ def test_ranking_with_weighted_data():
         assert all(p <= q for p, q in zip(is_sorted, is_sorted[1:]))
 
 
-class TestRanking:
+def test_error_msg() -> None:
+    X, y, qid, w = tm.make_ltr(10, 2, 2, 2)
+    ranker = xgboost.XGBRanker()
+    with pytest.raises(ValueError, match=r"equal to the number of query groups"):
+        ranker.fit(X, y, qid=qid, sample_weight=y)
 
+
+@given(lambdarank_parameter_strategy)
+@settings(deadline=None, print_blob=True)
+def test_lambdarank_parameters(params):
+    if params["objective"] == "rank:map":
+        rel = 1
+    else:
+        rel = 4
+    X, y, q, w = tm.make_ltr(4096, 3, 13, rel)
+    ranker = xgboost.XGBRanker(tree_method="hist", n_estimators=64, **params)
+    ranker.fit(X, y, qid=q, sample_weight=w, eval_set=[(X, y)], eval_qid=[q])
+    for k, v in ranker.evals_result()["validation_0"].items():
+        note(v)
+        assert v[-1] >= v[0]
+        assert ranker.n_features_in_ == 3
+
+
+@pytest.mark.skipif(**tm.no_pandas())
+@pytest.mark.skipif(**tm.no_sklearn())
+def test_unbiased() -> None:
+    import pandas as pd
+    from sklearn.model_selection import train_test_split
+
+    X, y, q, w = tm.make_ltr(8192, 2, n_query_groups=6, max_rel=4)
+    X, Xe, y, ye, q, qe = train_test_split(X, y, q, test_size=0.2, random_state=3)
+    X = csr_matrix(X)
+    Xe = csr_matrix(Xe)
+    data = RelDataCV((X, y, q), (Xe, ye, qe), max_rel=4)
+
+    train, _ = simulate_clicks(data)
+    x, c, y, q = sort_ltr_samples(
+        train.X, train.y, train.qid, train.click, train.pos
+    )
+    df: Optional[pd.DataFrame] = None
+
+    class Position(xgboost.callback.TrainingCallback):
+        def after_training(self, model) -> bool:
+            nonlocal df
+            config = json.loads(model.save_config())
+            ti_plus = np.array(config["learner"]["objective"]["ti+"])
+            tj_minus = np.array(config["learner"]["objective"]["tj-"])
+            df = pd.DataFrame({"ti+": ti_plus, "tj-": tj_minus})
+            return model
+
+    ltr = xgboost.XGBRanker(
+        n_estimators=8,
+        tree_method="hist",
+        lambdarank_unbiased=True,
+        lambdarank_num_pair_per_sample=12,
+        lambdarank_pair_method="topk",
+        objective="rank:ndcg",
+        callbacks=[Position()],
+        boost_from_average=0,
+    )
+    ltr.fit(x, c, qid=q, eval_set=[(x, c)], eval_qid=[q])
+
+    assert df is not None
+    # normalized
+    np.testing.assert_allclose(df["ti+"].iloc[0], 1.0)
+    np.testing.assert_allclose(df["tj-"].iloc[0], 1.0)
+    # less biased on low ranks.
+    assert df["ti+"].iloc[-1] < df["ti+"].iloc[0]
+
+
+class TestRanking:
     @classmethod
     def setup_class(cls):
         """

tests/python/test_with_sklearn.py

@@ -130,11 +130,11 @@ def test_ranking():
 
     params = {
         "tree_method": "exact",
+        "objective": "rank:pairwise",
         "learning_rate": 0.1,
         "gamma": 1.0,
         "min_child_weight": 0.1,
         "max_depth": 6,
-        "eval_metric": "ndcg",
         "n_estimators": 4,
     }
     model = xgb.sklearn.XGBRanker(**params)
@@ -163,7 +163,6 @@ def test_ranking():
         "gamma": 1.0,
         "min_child_weight": 0.1,
         "max_depth": 6,
-        "eval_metric": "ndcg",
     }
     xgb_model_orig = xgb.train(
         params_orig, train_data, num_boost_round=4, evals=[(valid_data, "validation")]

tests/test_distributed/test_gpu_with_dask/test_gpu_with_dask.py

@@ -44,7 +44,7 @@ try:
     from dask_cuda import LocalCUDACluster
 
     from xgboost import dask as dxgb
-    from xgboost.testing.dask import check_init_estimation
+    from xgboost.testing.dask import check_init_estimation, check_uneven_nan
 except ImportError:
     pass
 
@@ -224,6 +224,12 @@ class TestDistributedGPU:
     def test_init_estimation(self, local_cuda_client: Client) -> None:
         check_init_estimation("gpu_hist", local_cuda_client)
 
+    def test_uneven_nan(self) -> None:
+        n_workers = 2
+        with LocalCUDACluster(n_workers=n_workers) as cluster:
+            with Client(cluster) as client:
+                check_uneven_nan(client, "gpu_hist", n_workers)
+
     @pytest.mark.skipif(**tm.no_dask_cudf())
     def test_dask_dataframe(self, local_cuda_client: Client) -> None:
         run_with_dask_dataframe(dxgb.DaskDMatrix, local_cuda_client)

tests/test_distributed/test_with_dask/test_with_dask.py

@@ -4,7 +4,6 @@ import json
 import os
 import pickle
 import socket
-import subprocess
 import tempfile
 from concurrent.futures import ThreadPoolExecutor
 from functools import partial
@@ -41,7 +40,7 @@ from distributed import Client, LocalCluster
 from toolz import sliding_window  # dependency of dask
 
 from xgboost.dask import DaskDMatrix
-from xgboost.testing.dask import check_init_estimation
+from xgboost.testing.dask import check_init_estimation, check_uneven_nan
 
 dask.config.set({"distributed.scheduler.allowed-failures": False})
 
@@ -2014,6 +2013,14 @@ def test_init_estimation(client: Client) -> None:
     check_init_estimation("hist", client)
 
 
+@pytest.mark.parametrize("tree_method", ["hist", "approx"])
+def test_uneven_nan(tree_method) -> None:
+    n_workers = 2
+    with LocalCluster(n_workers=n_workers) as cluster:
+        with Client(cluster) as client:
+            check_uneven_nan(client, tree_method, n_workers)
+
+
 class TestDaskCallbacks:
     @pytest.mark.skipif(**tm.no_sklearn())
     def test_early_stopping(self, client: "Client") -> None: