From b76c76bf1bea8364b5089cb3bcab5f2db3441492 Mon Sep 17 00:00:00 2001 From: Victor Lafargue Date: Thu, 30 Jul 2026 23:31:15 +0200 Subject: [PATCH] Fix Dask logistic regression with non-contiguous class labels (#8422) Should address #8398. (could not reproduce the hangs yet though) The multi-GPU logistic regression path passed original labels directly to the QN solver, which expects class indices from 0 to n_classes - 1. Labels such as [10, 50, 20] therefore produced an invalid optimization state. This may be the cause behind observed distributed training to hang. This PR: - Encodes labels using the globally sorted classes before training. - Preserves the original values in `classes_` and predictions. - Updates `test_n_classes_small` to verify predictions, not only class discovery. - Avoid empty sparse partitions in Dask logistic regression tests. Authors: - Victor Lafargue (https://github.com/viclafargue) Approvers: - Simon Adorf (https://github.com/csadorf) URL: https://github.com/rapidsai/cuml/pull/8422 --- .../cuml/linear_model/logistic_regression_mg.pyx | 12 +++++++++++- .../cuml/tests/dask/test_dask_logistic_regression.py | 12 +++++++++--- 2 files changed, 20 insertions(+), 4 deletions(-) diff --git a/python/cuml/cuml/linear_model/logistic_regression_mg.pyx b/python/cuml/cuml/linear_model/logistic_regression_mg.pyx index 8de5e02e16..da75fd4a9d 100644 --- a/python/cuml/cuml/linear_model/logistic_regression_mg.pyx +++ b/python/cuml/cuml/linear_model/logistic_regression_mg.pyx @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # import cupy as cp @@ -211,6 +211,16 @@ class LogisticRegressionMG(LogisticRegression): classes.sort() cdef int n_classes = len(classes) + # QN's classification losses expect labels to be dense class indices + # in [0, n_classes). Keep the original values in ``classes`` for + # predictions, but encode the solver-facing labels using the global + # class ordering shared by all ranks. + y = cp.searchsorted(cp.asarray(classes, dtype=X.dtype), y).astype( + X.dtype, copy=False + ) + opg.free_data_t(y_ptr, X.dtype) + y_ptr = opg.build_data_t([y]) + # Validate and initialize parameters l1_strength, l2_strength = self._get_l1_l2_strength() cdef qn_params params diff --git a/python/cuml/tests/dask/test_dask_logistic_regression.py b/python/cuml/tests/dask/test_dask_logistic_regression.py index 0a8c1c80b9..ad2bfda55a 100644 --- a/python/cuml/tests/dask/test_dask_logistic_regression.py +++ b/python/cuml/tests/dask/test_dask_logistic_regression.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -56,6 +56,9 @@ def _prep_training_data_sparse(c, X_train, y_train, partitions_per_worker): def cal_chunks(dataset, n_partitions): n_samples = dataset.shape[0] + # Avoid zero-row partitions when a small test dataset has fewer rows + # than the number of partitions requested by the available workers. + n_partitions = min(n_partitions, max(n_samples, 1)) n_samples_per_part = int(n_samples / n_partitions) chunk_sizes = [n_samples_per_part] * n_partitions samples_last_row = n_samples - ( @@ -315,13 +318,16 @@ def test_noreg(fit_intercept, client): def test_n_classes_small(client): - def assert_small(X, y, n_classes): + def assert_small(X, y, n_classes, check_predictions=False): X_df, y_df = _prep_training_data(client, X, y, partitions_per_worker=1) from cuml.dask.linear_model import LogisticRegression as cumlLBFGS_dask lr = cumlLBFGS_dask() lr.fit(X_df, y_df) assert len(lr.classes_) == n_classes + if check_predictions: + predictions = lr.predict(X_df, delayed=True).compute().to_numpy() + assert np.array_equal(predictions, y) return lr X = np.array([(1, 2), (1, 3)], np.float32) @@ -341,7 +347,7 @@ def assert_small(X, y, n_classes): X = np.array([(1, 2), (1, 3), (1, 2.5)], np.float32) y = np.array([10.0, 50.0, 20.0], np.float32) - lr = assert_small(X=X, y=y, n_classes=3) + lr = assert_small(X=X, y=y, n_classes=3, check_predictions=True) assert np.array_equal( lr.classes_, np.array([10.0, 20.0, 50.0], np.float32) )