From 50ec050142b872fa72476c37ae6cff99ff78e90b Mon Sep 17 00:00:00 2001 From: Dante Gama Dessavre Date: Sun, 7 Jul 2024 23:41:03 -0500 Subject: [PATCH] Avoid extra memory copy when using cp.concatenate in cuml.dask kmeans (#5937) Partial solution for #5936 Issue was that concatenating when having a single array per worker was causing a memory copy (not sure if always, but often enough). This PR avoids the concatenation when a worker has a single partition of data. This is coming from a behavior from CuPy, where some testing reveals that sometimes it creates an extra allocation when concatenating lists that are comprised of a single array: ```python >>> import cupy as cp >>> a = cp.random.rand(2000000, 250).astype(cp.float32) # Memory occupied: 5936MB >>> b = [a] >>> c = cp.concatenate(b) # Memory occupied: 5936 MB <- no memory copy ``` ```python >>> import cupy as cp >>> a = cp.random.rand(1000000, 250) # Memory occupied: 2120 MB >>> b = [a] >>> c = cp.concatenate(b) # Memory occupied: 4028 MB <- memory copy was performed! ``` I'm not sure what are the exact rules that CuPy follows here, we could check, but in general avoiding the concatenate when we have a single partition is an easy fix that will not depend on the behavior outside of cuML's code. cc @tfeher @cjnolet Authors: - Dante Gama Dessavre (https://github.com/dantegd) Approvers: - Artem M. Chirkin (https://github.com/achirkin) - Tamas Bela Feher (https://github.com/tfeher) - Divye Gala (https://github.com/divyegala) URL: https://github.com/rapidsai/cuml/pull/5937 --- python/cuml/cluster/kmeans_mg.pyx | 4 ++-- python/cuml/dask/common/input_utils.py | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/python/cuml/cluster/kmeans_mg.pyx b/python/cuml/cluster/kmeans_mg.pyx index 8064cea374..51d415036b 100644 --- a/python/cuml/cluster/kmeans_mg.pyx +++ b/python/cuml/cluster/kmeans_mg.pyx @@ -1,5 +1,5 @@ # -# Copyright (c) 2019-2023, NVIDIA CORPORATION. +# Copyright (c) 2019-2024, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -214,7 +214,7 @@ class KMeansMG(KMeans): self.handle.sync() - self.labels_, _, _, _ = input_to_cuml_array(self.predict(X, + self.labels_, _, _, _ = input_to_cuml_array(self.predict(X_m, sample_weight=sample_weight), order='C', convert_to_dtype=self.dtype) diff --git a/python/cuml/dask/common/input_utils.py b/python/cuml/dask/common/input_utils.py index 5c4f7d0913..01c1a4e55f 100644 --- a/python/cuml/dask/common/input_utils.py +++ b/python/cuml/dask/common/input_utils.py @@ -203,6 +203,8 @@ def concatenate(objs, axis=0): else: return cudf.concat(objs) elif isinstance(objs[0], cp.ndarray): + if len(objs) == 1: + return objs[0] return cp.concatenate(objs, axis=axis) elif isinstance(objs[0], np.ndarray):