[FEA] Skip Device Concat for KmeansMG - #8084
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
tarang-jain
added
improvement
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
tarang-jain
changed the title
Enables KMeans to process datasets that do not fit in device memory by streaming batches through fit/predict. Scales inputs in batches, only allocates batch_size buffers for uniform sample weights, and wires the new streaming parameters through the C++ params, Cython bindings, and Python layer. Includes tests for the streaming path.
| # Coerce each partition to a C-contiguous array, preserving its memory | ||
| # residency (device stays device, host stays host) via ``mem_type=None`` | ||
| # so device-accessible partitions are never staged to host and | ||
| # host-resident partitions are never staged to device here. | ||
| # ``reset=True`` on the first partition records ``n_features_in_``; | ||
| # subsequent partitions are validated against it so mismatched shapes | ||
| # are rejected. | ||
| coerced_parts = [] | ||
| coerced_weights = [] if has_weights else None | ||
| for i, part in enumerate(parts): | ||
| sw_in = sample_weight_parts[i] if has_weights else None | ||
| part_c, sw_c = check_inputs( | ||
| self, | ||
| part, | ||
| sample_weight=sw_in, | ||
| dtype=("float32", "float64"), | ||
| order="C", | ||
| mem_type=None, | ||
| reset=(i == 0), | ||
| ) |
There was a problem hiding this comment.
device-accessible partitions are never staged to host and host-resident partitions are never staged to device here
check_inputs(..., mem_type=None)
Shouldn't the memory residency be uniform across partitions? It looks like the C++ code dispatches according to the residency of the first partition.
The Python code should be congruent with the C++ code. We should detect the residency of the first partition and then coerce the following ones accordingly.
Additionally, the check_inputs function has a default of ensure_min_samples=1 meaning that if Dask provides any empty partition it would fail despite it being handled correctly on the cuVS implementation side.
There was a problem hiding this comment.
I fixed it. We fetch the residency from the first partition and then load the subsequent ones using the same mem_type.
csadorf
left a comment
csadorf
left a comment
There was a problem hiding this comment.
No major concerns, but a few comments and questions.
|
/merge |
4 participants
NVIDIA/cuvs#2066 allows passing host partitions directly to cuVS KMeans (MNMG). Merge this PR after NVIDIA/cuvs#2066 and #8248
Closes #8201