Repository navigation
Compute pairwise_distance with matrix multiplications - #1365
Conversation
81711c4 to
4c7f998
Compare
jeromekelleher
left a comment
jeromekelleher
left a comment
There was a problem hiding this comment.
Looks like a straightforward win to me - nice job @Billyzhang1229 !
4c7f998 to
7615957
Compare
|
Thanks, there is something goes wrong about floating point precision (when numbers are small). I'm checking these. |
|
Looks like all tests are passing though? |
The existing tests don't cover this. The old code computed the euclidean distance from |
|
That sounds good. Perfectly fine to just say "possible loss of precision with float32" too though. |
7615957 to
41ed418
Compare
Thanks! I've mentioned it in the docstring. It also needs more memory, since the chunks are converted to float64 (details in the description), but I think that's a good trade-off for the speedup. |
jeromekelleher
left a comment
jeromekelleher
left a comment
There was a problem hiding this comment.
LGTM - one comment about CI
Write both metrics as matrix products of the values, with missing entries set to zero, and the mask of observed entries. For example, the squared Euclidean distance over the entries observed in both rows is sum(x_i^2) + sum(x_j^2) - 2 * sum(x_i * x_j), where each sum is a matrix product. Build them from Dask array operations instead of Numba gufuncs, CUDA kernels and hand-written blockwise and reduction steps, with each metric as a function in sgkit.distance.metrics, looked up by name in pairwise_distance, so that new metrics can be added there. This is several times faster on the CPU, and no longer passes any functions for Dask to tokenize (see issue 1267). The GPU path now uses CuPy instead of Numba's CUDA support: arrays with CuPy chunks are computed on the GPU, and device="gpu" moves the chunks to the GPU and the result back to NumPy. This fixes wrong results for int8 input on the GPU, where the partial sums were stored in an int8 array. The GPU workflow now installs CuPy for CUDA 12, with the CUDA libraries it uses from PyPI, instead of the CUDA toolkit from conda, which needs NVIDIA driver 525 or later on the runner, and checks that CuPy can use the GPU before running the tests, which are skipped otherwise. The distances are computed in float64, and returned in the floating point type of the input, or float64 for other types, so integer input now reports the float64 dtype of the result before it is computed. Subtracting the sums of squares and products loses small distances in float32: with float32 dosages, two rows 0.11 apart could come out 0 apart, and copies of a row could have a negative correlation distance. In float64 these sums are exact for small integers, such as allele counts. On GPUs with slow float64, such as the T4 that the GPU workflow runs on, float32 input now takes as long as float64 input. Infinite values are now missing, like negative and nan values, since an infinite value times zero is nan in the matrix products, which made every distance of its row nan. The correlation distance is clipped to [0, 2], and its denominator is sqrt(var_i * var_j), which unlike sqrt(var_i) * sqrt(var_j) keeps copies of a row of integers exactly zero apart. The diagonal is now exactly zero, or nan where the correlation is undefined. This needs Dask 2024.12 or later: the blockwise fusion in earlier releases computed wrong correlations from these expressions.
41ed418 to
ae5fff9
Compare
pairwise_distanceis the last function affected by #1267. It passes locally defined functions toda.blockwiseandda.reduction, and newer Dask tokenizes them by pickling, so the Numba gufuncs they reference are rebuilt and recompiled on every call (comment). #1275 fixed the other functions, but this one needed a different approach.This PR:
sum(x_i^2) + sum(x_j^2) - 2 * sum(x_i * x_j), where each sum is a matrix product.sgkit.distance.metrics, looked up by name inpairwise_distance, so new metrics can be added there.device="gpu"moves the chunks to the GPU and the result back to NumPy. This fixes wrong results for int8 input on the GPU, where the partial sums were stored in an int8 array.Other changes in behavior:
device="gpu"requires CuPy instead of Numba's CUDA support.With this, sgkit uses CuPy for
pairwise_distanceon the GPU, so the CuPy mention that #579 questioned now applies.Fixes #1267.
Testing:
GUFunc._rebuildwhile running the full test suite with Dask 2026.8.0: 98 before this change, all intest_distance.py, and none after.test_distance.pyon a T4 on Modal, withcupy-cuda12x[ctk]as the workflow installs it, and all 88 tests passed, including the 40 GPU tests. The benchmarks below ran the GPU path on an RTX 5090 with CuPy for CUDA 13, and on an A100 and an H100 withcupy-cuda12x[ctk]. It would still be good to dispatch the workflow to run the GPU tests on its T4.Precision
Both metrics are differences of large sums, such as
sum(x_i^2) + sum(x_j^2) - 2 * sum(x_i * x_j), so a small distance loses precision to rounding, and in float32 most of it. I checked this with six rows of float32 dosages (0 to 2, with 2% missing), where two rows are copies and two others differ in five entries by 0.05, with 100k, 1M and 10M variants in chunks of 100k, on an M4 Pro:On a T4, computing in float32 gave 0.15 for the pair 0.11 apart with 100k variants, and 0 with 1M and 10M. With int8 calls both were exact: integer input was already computed in float64, where the sums of small integers are exact.
So the distances are now computed in float64, as scikit-learn does for float32 input in
euclidean_distances, and returned in the floating point type of the input. The denominator of the correlation issqrt(var_i * var_j)rather thansqrt(var_i) * sqrt(var_j):sqrt(v * v)is exactlyv, so copies of a row of integers are exactly zero apart, while the product of square roots gave ±2.2e-16. A distance much smaller than the values it is computed from can still be wrong in float64, for example 0 instead of 1 between[1e8, 1e8]and[1e8, 1e8 + 1], which the docstring now mentions.float32 input now takes about as long as float64 input on the CPU, which is still 23× and 63× faster than before this PR at 2k × 20k on the M4 Pro, and on GPUs where float64 is much slower than float32. At 4k × 40k on Modal, median of 3, where the arrows go from float32 input computed in float32 to computed in float64:
I haven't rerun the RTX 5090 benchmarks below, but there float32 input should now take about as long as float64 input. At 8k × 80k, that's about 14.7 s for euclidean and 26.3 s for correlation, against 15.0 s and 17.2 s for float32 before this PR, so on GPUs like these float32 correlation is now slower than before. The A100 and the H100 compute float64 about as fast as float32, so there float32 input is still as fast as before: at 8k × 80k the A100 takes 1.50 s and 2.65 s, against 1.54 s and 2.73 s when float32 input was computed in float32.
Memory
The old code computed on the chunks as they were, for example on int8 calls, while the new code converts each chunk to float64 values with the missing entries set to zero, a float64 mask of the observed entries and the squares of the values, so it uses more memory. The peak memory of a call above the memory before it, with 1000 rows of int8 calls (0, 1 and 2, with 5% missing) in chunks of 1000 × 10k, on an M4 Pro with Dask's default 14 threads:
At 500k variants, the new code also takes 2.2 s instead of 41 s for euclidean, and 4.0 s instead of 163 s for correlation.
Its peak doesn't grow with the number of variants, since the chunks are converted as the matrix products use them: from 250k to 2M variants it stays between 3.2 and 3.4 GiB for euclidean, and between 2.6 and 2.8 GiB for correlation. It grows with the number of threads and the size of the chunks instead, by about 0.2 GiB for each thread with chunks of 1000 × 10k, so fewer threads or smaller chunks bring it down. At 500k variants:
Computing float32 input in float64 also doubles its memory. With float32 input at 100k variants, the peak is 2.9 GiB for euclidean and 2.6 GiB for correlation, against 1.4 GiB and 1.2 GiB when it was computed in float32, and 0.06 GiB and 0.48 GiB before this PR.
I haven't measured the memory on the GPU, where the CuPy chunks are converted in the same way.
Setup
pairwise_distance(x, metric=metric).compute(), fromru_maxrss, minus the resident set size just before it, in a new process for each run, after a warm-up call on a small array. The input was an in-memory NumPy array, generated row by row so that the call had no freed memory to reuse.Benchmarks
pairwise_distancebefore and after this PR, on an Apple M4 Pro and an NVIDIA A100 on Modal. The times before don't include the recompilation from #1267, so the speedups come from the new computation alone:The speedup grows with the input. On the CPU at 4k × 40k, euclidean goes from 72.5 s to 3.05 s (24×) and correlation from 498 s to 6.29 s (79×). Main's float64 correlation is unusually slow on the M4 Pro at 2k × 20k (123 s, against 47.9 s in float32, and 134 s in an earlier run), so the speedup there is 164×. On the A100 at 8k × 80k, float32 is 26× and 15× faster, and float64 24× and 16×. For the smallest input, whose 16 chunks are 125 × 1250, overhead dominates on the GPU, where correlation on the A100 is slower (179 ms before, 245 ms after in float32). For correlation, the new graph has 1113 tasks instead of 241. On the CPU, the smallest input is still 3.7× and 13× faster.
GPUs that compute float64 much more slowly than float32 gain much less. On an RTX 5090, float64 is only 2.2× and 2.4× faster at 8k × 80k. It computes float64 at 1/64 of its float32 rate, and the new code spends its time in matrix products: an 8192 × 8192 product with CuPy runs at 1.7 TFLOPS in float64 and 76 TFLOPS in float32. The A100 and the H100 compute float64 matrix products as fast as float32 ones (18.6 and 58.5 TFLOPS), and at 8k × 80k float64 is 24× and 16× faster on the A100, and 35× and 30× faster on the H100. float32 input is now computed in float64 too, so on the RTX 5090 it would be about as slow as float64. I haven't rerun the RTX 5090, so it only appears in float64 here:
The CPU times are with default settings, where Dask's threaded scheduler runs a task on every core and Accelerate, the BLAS that NumPy uses on macOS, also uses several cores for each matrix product. Limiting Accelerate to one thread with
VECLIB_MAXIMUM_THREADS=1makes the new code at most 8% faster. OpenBLAS is different: on an i9-12900K under WSL2,OPENBLAS_NUM_THREADS=1, which Dask's distributed workers set by default, made the new code about twice as fast in float64 at 4k × 40k (9.29 s → 4.99 s for euclidean and 17.2 s → 9.58 s for correlation). Dask's workers don't setVECLIB_MAXIMUM_THREADS. The old code doesn't use BLAS.Setup and all timings
pairwise_distance(x, metric=metric, device=device).compute(), including building the graph, with Dask's threaded scheduler. Each time is the median of 2 to 23 runs, after a warm-up call on a small array that compiles the Numba functions and the CUDA and CuPy kernels. On the GPU the input starts in NumPy chunks, so the times include copying the chunks to the GPU and the result back. The first call on a larger input is slower than these medians while CuPy's memory pool grows, for example 1.23 s instead of 0.63 s for euclidean in float32 at 8k × 80k on the H100. The A100, H100 and M4 Pro times leave out the first call at each size when it takes less than 5 s.cupy-cuda12x[ctk]14.2.0,numba-cuda[cu12]0.30.4), since the driver is older than the local CUDA 13.4 wheels. Dask used 16 threads, as on the local machine. The A100 is the SXM4 80GB model and the H100 the SXM5 80GB model.CPU, float64:
CPU, float64, with
VECLIB_MAXIMUM_THREADS=1for the new code (speedup over before in parentheses):Other input types on the CPU, at 2k × 20k:
Matrix product throughput, 8192 × 8192 with CuPy (no TF32):
GPU times on the three GPUs, before → after (speedup). The RTX 5090 wasn't rerun after the change to compute float32 input in float64, so it only has float64 times:
float32, euclidean:
float32, correlation:
float64, euclidean:
float64, correlation: