[FEA] Calculating distances with FP32 in NN Descent

[jinsolp]

I hope NN Descent can support calculating distances in local_join_kernel using fp32 instead of fp16 (__half).

Due to precision issues, it is difficult to grab the distances calculated by NN Descent and compare it against fp32-calculated distances for additional operations outside of NN Descent such as making predictions.

[jinsolp]

Summarizing details after investigating and comparing this to our brute force implementation.

TLDR;
The problem with our NN Descent implementation is that given FP32 data, we always convert it to FP16 to run wmma.
This causes problems when compared against the brute force implementation, which directly uses the FP32 data.
To solve this, we need to keep the FP32 data as FP32.

This writeup demonstrates this by showing

1. NN Descent vs Brute force show similar results when given fp16 data
2. NN Decent vs Brute force shows similar behavior when we modifying our Brute force implementation to do a fp32->fp16 conversion like NN Descent

FP16 data

Given fp16 data, our NN Descent calculations are done exactly the same as the brute force implementation

• Calculate norms in fp32 by converting fp16 to fp32 on the fly
• Calculate dot product in fp32 by converting fp16 to fp32 on the fly, accumulator is fp32.

from sklearn.datasets import make_moons
from cuvs.neighbors import brute_force, nn_descent
import numpy as np

# Data in fp16
X, y = make_moons(n_samples=1_000_000, noise=0.1, random_state=42)
X = X.astype(np.float16)
k = 16
metric="l2"

# NN Descent
build_params = nn_descent.IndexParams(metric=metric, return_distances=True, graph_degree=k)
index = nn_descent.build(build_params, X)
gpu_nnd_indices = index.graph
gpu_nnd_distances = index.distances

# Brute Force
cp_X = cp.asarray(X)
index = brute_force.build(cp_X, metric=metric)
gpu_bf_distances, gpu_bf_indices = brute_force.search(index, cp_X, k)

# Add self-loops: discard last column, shift right, and add row index as first column
def add_self_loops(indices):
    n_rows = indices.shape[0]
    # Create array with self-loops in first column
    self_loops = np.arange(n_rows).reshape(-1, 1)
    # Drop last column and concatenate with self-loops at the beginning
    return np.concatenate([self_loops, indices[:, :-1]], axis=1)

def knn_recall(pred: np.ndarray, gt: np.ndarray) -> float:
    # print(pred[0], gt[0])
    N, k = pred.shape
    correct = 0

    for i in range(N):
        pred_set = set(pred[i])
        gt_set = set(gt[i])
        correct += len(pred_set & gt_set)

    return correct / (N * k)

print(f"knn_recall: {knn_recall(add_self_loops(gpu_nnd_indices), gpu_bf_indices)}")

Running this script (data in fp16) results in a very high recall

knn_recall: 0.9928262727272728

Distances look similar too, with minor differences due to rounding errors arising from different computation order

# NN Descent nearest neighbors for 0:
indices: [603894 655574 476675 297473 186463 606604 872008  79939 668216 781306
 254310]
distances: [0.00097656 0.00097656 0.00218366 0.00239208 0.00258374 0.00292969
 0.00292969 0.00308816 0.00308816 0.00365396 0.00390625]

# Brute Force nearest neighbors for 0:
indices: [     0 655574 603894 476675 186463 297473 606604 872008  79939 668216
 781306]
distances: [0.         0.         0.00097656 0.00218366 0.00239208 0.00239208
 0.00276214 0.00276214 0.00292969 0.00308816 0.00365396]

Investigating some pairs in binary. First different value in bold:

pair (0, 655574)

NN Descent

• norm of vec 0: 4.4054 (0x408cf909)
• norm of vec 655574: 4.40576 (0x408cfbfb)
• dot product of vec 0 and vec 655574: 4.05579 (0x408cfa81)
• (norm vec 0) + (norm vec 655574) - 2 * (dot vec0 vec655574): 0.000001 (0x35800000)
• doing sqrt on this value (final l2 dist): 0: 0.000976562 (0x3a800000)

Brute Force

• norm of vec 0: 4.4054 (0x408cf909)
• norm of vec 655574: 4.40576 (0x408cfbfb)
• dot product of vec 0 and vec 655574: 4.40558 (0x408cfa82)
• (norm vec 0) + (norm vec 655574) - 2 * (dot vec0 vec655574) = 0 (0x00000000)
• doing sqrt on this value (final l2 dist): 0: 0 (0x00000000)

pair (0, 186463)

NN Descent

• norm of vec 0: 4.4054 (0x408cf909)
• norm of vec 186463: 4.4124 (0x408d3261)
• dot product of vec 0 and vec 186463: 4.408896 (0x408d15ae)
• (norm vec 0) + (norm vec 186463) - 2 * (dot vec0 vec186463): 0.000007 (0x36e00000)
• doing sqrt on this value (final l2 dist): 0: 0.00258374 (0x3b2953fd)

Brute Force

• norm of vec 0: 4.4054 (0x408cf909)
• norm of vec 186463: 4.4124 (0x408d3261)
• dot product of vec 0 and vec 186463: 4.408897 (0x408d15af)
• (norm vec 0) + (norm vec 186463) - 2 * (dot vec0 vec186463) = 0.000006 (0x36c00000)
• doing sqrt on this value (final l2 dist): 0: 0.00239208 (0x3b1cc471)

Tolerable differences propagate to final distance. This is okay.

FP32 data

Our NN Descent always converts the data to fp16 to use wmma .
This is different from our brute force implementation, which keeps the fp32 and uses it directly.
Using the same script as above with this line change

X = X.astype(np.float32)

The recall is really bad

knn_recall: 0.1783208181818182

Looking at the indices and distances, our NN Descent results in all-zero distances like this;

# NN Descent nearest neighbors for 0:
indices: [ 72843  85566 107950 133809 178352 190164 191978 209930 237327 241959
 264631]
distances: [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]

# Brute Force nearest neighbors for 0:
indices: [     0 655574 603894 297473 186463 476675 668216  79939 606604 781306
 872008]
distances: [0.         0.00097656 0.00138107 0.00218366 0.00239208 0.00239208
 0.00292969 0.00308816 0.00308816 0.00352105 0.00365396]

Reason for 0 distances in our NN Descent implementation:

pair (0, 655574)

NN Descent

• norm of vec 0 (calculated in fp32): 4.40435 (0x408cf072)
• norm of vec 655574 (calculated in fp32): 4.40251 (0x408ce153)
• dot product of vec 0 and vec 655574 (calculated after fp16 conversion): 4.405579 (0x408cfa81)
• (norm vec 0) + (norm vec 186463) - 2 * (dot vec0 vec186463): -0.004302 (0xbb8cf800) <= this is negative!
• clipped and squaring (final l2 dist): 0 (0x00000000)

Brute Force

• norm of vec 0 (calculated in fp32): 4.40435 (0x408cf072)
• norm of vec 655574 (calculated in fp32): 4.40251 (0x408ce153)
• dot product of vec 0 and vec 655574 (calculated in fp32): 4.403428 (0x408ce8e1)
• (norm vec 0) + (norm vec 655574) - 2 * (dot vec0 vec655574) = 0.000001 (0x35800000)
• doing sqrt on this value (final l2 dist): 0.00097656 (0x3a7fffd5)

This is because in NN Descent we do a fp32->fp16 conversion, and convert the fp16 values back to fp32 when adding to the accumulator to calculate the dot product.

Investigating Conversions

Original FP32 vec 0: [2.06613, 0.368037] = [0x40043b8, 0x3ebc6f56]
=> Converted to FP16 vec 0: [2.06641, 0.36792] = [0x4022, 0x35e3]

Original FP32 vec 65574: [2.06564, 0.36828923] = [0x40043372, 0x3ebc9068]
=> Converted to FP16 vec 65574: [2.06641, 0.368408] = [0x4022, 0x35e5]

After rounding, in this case we end up with larger elements in fp16, which makes the dot product bigger than it should be, resulting in negative distances

Mimic conversion behavior in brute force

To mimic a similar behavior in brute force, let's force a fp32 -> fp16 conversion in the dot product calculation step.

This is where the doc product is calculated.
https://github.com/jinsolp/cuvs/blob/9b371e49878bf1cffe2632

698a90aabf8ae0a1b3/cpp/src/distance/detail/distance_ops/l2_exp.cuh#L93-L100

If we add a case for when data is fp32 like this where we convert to fp16 and back to fp32 to do the calculation. This ends up doing the same computation as our NN Descent implementation;

  DI void core(AccT& acc, DataT& x, DataT& y) const
  {
    if constexpr ((std::is_same_v<AccT, float> && std::is_same_v<DataT, half>)) {
      acc += __half2float(x) * __half2float(y);
    } else if constexpr ((std::is_same_v<AccT, float> && std::is_same_v<DataT, float>)) {
      acc += __half2float(__float2half(x)) * __half2float(__float2half(y));
    } else {
      acc += x * y;
    }
  };

Running the same script after this change results in zero-distance brute force neighbors like our NN Descent implementation:

# Brute Force nearest neighbors for 0:
indices: [   0  642 2547 2830 3235 3363 3825 3916 3866 5496 6120]
distances: [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]