kmeans: cluster-parallel update via counting sort (O(n·d), thread-count-invariant) - #1
Conversation
|
Hey Rohan, thanks for this! The diagnosis is right and the I benchmarked it before commenting, on SIFT1M and MS MARCO Dragon. One note on the base: this branches off End-to-end k-means: HNSW centroid assignment (M=32, ef_c=200, ef_search=16, k=32768, n_iter=25):
We're measuring different things, so both sets of numbers are right. Your 176.8 → 41.1 ms is the update step in isolation, and it reproduces cleanly — at d=128 I get 3.6× at k=4096 and 6.5× at k=16384, so 4.3× at k=8192 sits right on that curve. On full Dragon the step alone is 4.1× at k=32768 and 5.0× at k=65536. The table above is the whole A clear win at low dimension, then; ~4% at Dragon's realistic sample size. Three things I'd like to resolve before merging. 1. What the accumulator cost comes to in practiceYou document the Peak RSS attributable to the step, full Dragon (8.8M × 768), one k per process:
That's more than a 20x increase, and it grows with core count, where main's 2. The page-fault traffic isn't freeFirst-touching those accumulators every call pushes most of the CPU time into the kernel: on full Dragon sys time goes 44% → 66% → 84% as k goes 16k → 32k → 64k, and it sits at ~90% at 1M scale. main stays under 2% at these It also caps the win rather than just costing memory, on SIFT the step-level speedup peaks at 7.5× (k=32768, 28% sys) and falls to 3.0× at k=65536 (90% sys). 3.
|
Replace the O(n*k) mean computation in update_and_split (which re-scanned all n points once per centroid) with a single O(n*d) scatter-add: each rayon worker owns a private (sums[k*d], counts[k]) accumulator over a disjoint point range, then partials are reduced in parallel over disjoint output ranges. The split, spherical renorm, and return path are unchanged, as is the function signature; same rayon primitives already in use, no new dependency. Benefits every caller (IVF build, PQ codebook training); the speedup grows with k. Measured on synthetic data (n=200k, d=128, k=8k, 16 cores): 176.8ms -> 41.1ms (4.3x), max centroid diff 1.8e-7 vs the brute-force reference (float reassociation only). Tests: update_and_split_scatter_matches_reference_means (equals hand-computed means) and update_and_split_fires_split_on_empty_cluster (split path intact).
…iant memory and results) Replace the point-parallel scatter-add with a counting sort over assignments (histogram -> prefix sum -> stable point-index grouping) followed by a cluster-parallel accumulation into one shared output. Still one O(n*d) pass, but scratch drops from n_threads*k*d floats of per-worker accumulators to one u32 per point, independent of thread count, and the fixed per-cluster summation order makes centroids bit-identical across thread counts (as the `seed` docs promise). Adds a bit-exactness test across 1/2/4/8-thread pools and a weighted means test; full suite passes (194 unit + 29 doc tests). Co-authored-by: Cursor <cursoragent@cursor.com>
aeff8a7 to
471ed14
Compare
robro612
changed the title
|
Thanks for the thorough review Silvio! You were right that all three issues trace back to parallelizing over points, so I took the direction you suggested: a counting sort groups point indices by cluster, then a cluster-parallel accumulation writes into a single shared output. That drops the scratch to one u32 per point (thread-count-independent), removes the per-iteration zero-page traffic, and fixes the summation order so centroids are bit-identical across thread counts — with a test enforcing it on non-integer data, since as you noted SIFT can't catch it. I benchmarked at your Dragon config's scale — synthetic uniform vectors with the same shape (8M × 768, k=32768), not the actual Dragon embeddings — on a 64-core node: the cluster-parallel update is 6.2× faster than the point-parallel version with ~0.2 GiB scratch vs ~6.1 GiB, which reproduces your 6302 MiB measurement almost exactly. Full numbers and methodology are in the updated PR description. The branch is updated in place: rebased onto 0.6.1 as requested, with the cluster-parallel change as a separate commit on top of the original scatter-add so the delta is easy to review. I'd take you up on the benchmark harness offer to confirm the end-to-end numbers on your setup. |
|
Hi Rohan, this is great, thanks for turning it around so fast. The cluster-parallel version fixes all three points, and the counting sort reads really cleanly. Approving. I re-ran the full benchmark suite on both collections. Everything below is total wall clock of the real build, n_iter=10, 64 threads, idle machine, one variant at a time. IVF build: k-means (HNSW assignment, M=32/ef_c=200/ef_search=16, k=32768) + final assignment of every vector:
PQ build:
So v2 beats main on all four builds, beats v1 on both IVF builds, and does it at main's memory footprint instead of v1's. The determinism test passes on real-valued data too, I confirmed bit-identical centroids at 1/4/16/64 threads on Dragon, where v1 gave four different answers. pr1_bench.tar.gz attached: both experiment binaries, the runner scripts, and the raw output behind every number above. kANNolo 0.9.0 already pins v0.6.1 so no patching is needed; just repoint the Merging this now. One heads-up: PR #2 is still based on the pre-rebase head of this branch, so it'll need a rebase now that this has landed, and since it changes |
Document the finite-range precondition on C (f16/bf16 do not saturate; unsigned fixed-point zeros negatives), port the per-redo init-RNG offset into train so n_redo is not a no-op on the flat path, and drop the ignored bench_kmeans_fixes leftover from TusKANNy#1.
2 participants
Summary
Updated after review (thanks @SilvioM97 — see discussion below). The O(n·k) update scan still goes away, but the strategy changes from parallelizing over points to parallelizing over clusters:
assignments(cluster-size histogram → prefix sum → stable permutation of point indices grouped by cluster), O(n), serial.par_chunks_mut(d)over one shared output, so each rayon task owns a disjoint set of clusters and gathers its members via the grouped index list. No per-thread accumulators, no reduce step.This addresses all three review points:
u32per point plus O(k) offsets, independent of thread count, instead ofn_threads × k × dfloats all live at once.n_threads × k × dfresh zero pages per iteration.seeddoc comment's "reproducible across runs and thread counts" now holds, and a new test enforces it.Branch is rebased onto current
main(0.6.1), with the cluster-parallel change as a separate commit on top of the original scatter-add so the delta is easy to review.Evidence
Both configs below use synthetic data (seeded uniform random values in [-1, 1]).
Scale 1: n=200k, d=128, k=8000, 32 cores
Minor page faults for the whole test process dropped 230k → 129k.
Scale 2: Dragon scale — n=8M, d=768, k=32768, random assignments, 64-core node
Synthetic vectors with the same shape as the MS MARCO Dragon config from the review, not the actual Dragon embeddings. Per-phase peak RSS measured via
VmHWMdeltas, both implementations in the same process:t·k·d·4Max centroid diff between the two: 1.6e-7.
n_threads × k × d × 4formula and reproduces the ~6.3 GiB at 64 cores measured in the review; cluster-parallel stays flat at ~0.2 GiB.One measurement note: peak-RSS comparisons of the point-parallel version need random (or real) assignments. With round-robin
i % kassignments a worker's contiguous point range touches at mostn/threadsdistinct clusters, so most accumulator pages stay untouched kernel zero pages and never count toward RSS.Tests
update_deterministic_across_thread_counts— centroids bit-identical (f32::to_bits) across 1/2/4/8-thread pools on non-integer random data (integer-valued data like SIFT is exactly representable in f32 and hides reassociation, per the review note).update_weighted_means_match_reference— non-uniform weights vs hand-computed means.