Skip to content

[FEA] Multi-node Out of Core Streaming KMeans API#2066

Open
tarang-jain wants to merge 199 commits into
NVIDIA:mainfrom
tarang-jain:mnmg-streaming
Open

[FEA] Multi-node Out of Core Streaming KMeans API#2066
tarang-jain wants to merge 199 commits into
NVIDIA:mainfrom
tarang-jain:mnmg-streaming

Conversation

@tarang-jain

@tarang-jain tarang-jain commented May 7, 2026

Copy link
Copy Markdown
Contributor

Merge after #2015 and #2017

Allows a stream of input matrices per worker, that are further batched using the streaming_batch_size parameter. Reasoning: We should be able to supply dask partitions (on host) directly without having to concatenate them into one consolidated matrix.

As a part of this PR, we also unify the multi-GPU implementations into one (earlier the out of core implementation was separate).
Tests: We get rid of the separate out of core test file. The single MG testing unit is taking care of both out of core and on device matrices.

Comment thread cpp/src/cluster/detail/kmeans_mg_batched_init.cuh Outdated
Comment thread cpp/src/cluster/detail/kmeans_mg_batched_init.cuh
Comment thread cpp/src/cluster/detail/kmeans_mg_distributed_init.cuh Outdated
Comment thread cpp/src/cluster/detail/kmeans_mg_distributed_init.cuh Outdated
Comment thread cpp/src/cluster/detail/kmeans_mg_distributed_init.cuh Outdated
Comment thread cpp/CMakeLists.txt
@tarang-jain

Copy link
Copy Markdown
Contributor Author

/ok to test bd3e1de

@tarang-jain tarang-jain removed request for a team and msarahan June 23, 2026 13:55

@lowener lowener left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@tarang-jain

Copy link
Copy Markdown
Contributor Author

/merge

@tarang-jain

Copy link
Copy Markdown
Contributor Author

/ok to test 40c77ca

@tarang-jain

Copy link
Copy Markdown
Contributor Author

/ok to test 079c216

@tarang-jain

Copy link
Copy Markdown
Contributor Author

/ok to test 7a38f62

@tarang-jain

Copy link
Copy Markdown
Contributor Author

/ok to test 8ce1ade

@tarang-jain

Copy link
Copy Markdown
Contributor Author

/ok to test 53ca72b

@dantegd dantegd left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Recent changes look good to me. Should we open an issue to track whether to check the CTEST registration of other MG algos?

@tarang-jain tarang-jain requested review from a team as code owners June 26, 2026 16:40
@tarang-jain

Copy link
Copy Markdown
Contributor Author

/ok to test 7cb1397

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

breaking Introduces a breaking change feature request New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants