CloudGPU

Posted on Nov 11

Build and Verify Chunked Cosine Search with Python and PyTorch on a Rented GPU

#machinelearning #pytorch #python #tutorial

A vector search function can return plausible document IDs while still being wrong: a chunk offset is missing, scores are sorted along the wrong axis, or a CPU fallback hides a broken GPU environment. Before connecting search to an application, give it a fixture with answers you can check.

We will use Python and PyTorch to find five cosine neighbors for each query, stream the corpus through a rented GPU in blocks, and validate the answers against a separate CPU calculation. The complete example creates its own numerical data, so there is no model download, external dataset, or service credential to reproduce.

I work on CloudGPU.

1. Define a small, checkable retrieval task

The fixture contains 8,192 vectors, each with 128 components. We normalize every vector to unit length. For normalized vectors, cosine similarity is their dot product; larger values mean a closer direction.

The 32 queries are copies of known corpus rows. Their nearest neighbor should therefore be the copied row itself. The remaining neighbors still require an actual search. This combines a planted-answer test with a broader numerical comparison.

These are synthetic vectors, not language embeddings. They test indexing and arithmetic, not whether a search system understands a sentence. The same search function can accept your own normalized embeddings later.

The GPU part uses a block matrix multiplication and merges local top-k candidates. It does not build a document-by-document matrix. The CPU validator deliberately uses a complete query-by-corpus matrix for this small fixture, making it an independent reference for the block-merging logic.

2. Prepare a Python environment on the GPU machine

Use a Linux GPU environment with Python 3.11 and a driver compatible with the CUDA 12.8 PyTorch wheel. The CloudGPU public catalog lists RTX 4090 and RTX 5090 cards; the code checks the actual device rather than assuming the selected environment works. An existing compatible GPU environment is also suitable.

Run the following commands in a terminal on that machine:

python3.11 -m venv .venv

Save this as requirements.txt:

--index-url https://download.pytorch.org/whl/cu128
torch==2.9.1

Then install and record the resolved environment:

.venv/bin/python -m pip install -r requirements.txt
.venv/bin/python -m pip freeze > environment.txt

This selects the CUDA 12.8 wheel and PyTorch version listed on the official previous versions page. Only torch is needed; the example does not use image or audio packages. environment.txt records transitive dependencies as well as the direct requirement.

3. Save the complete search and validation script

Save the following as search.py:

import argparse
import json
import os
import platform
from pathlib import Path

# Set before importing torch or creating a CUDA context.
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

import torch
import torch.nn.functional as F

SEED = 20261111
N_ROWS = 8192
N_QUERIES = 32
DIM = 128
K = 5


def fixture():
    generator = torch.Generator(device="cpu").manual_seed(SEED)
    raw = torch.randn(N_ROWS, DIM, generator=generator)
    corpus = F.normalize(raw, p=2, dim=1)
    expected_ids = torch.arange(N_QUERIES) * 137
    queries = corpus[expected_ids].clone()
    return corpus, queries, expected_ids


@torch.inference_mode()
def search(corpus, queries, device, chunk_rows):
    q = queries.to(device)
    scores = torch.empty((len(q), 0), device=device)
    ids = torch.empty((len(q), 0), dtype=torch.long, device=device)
    for offset in range(0, len(corpus), chunk_rows):
        block = corpus[offset:offset + chunk_rows].to(device)
        similarities = q @ block.T
        local_scores, local_ids = similarities.topk(
            min(K, len(block)), dim=1, sorted=True
        )
        local_ids = local_ids + offset
        candidates = torch.cat((scores, local_scores), dim=1)
        candidate_ids = torch.cat((ids, local_ids), dim=1)
        scores, positions = candidates.topk(
            min(K, candidates.shape[1]), dim=1, sorted=True
        )
        ids = candidate_ids.gather(1, positions)
        del block, similarities
    return scores.cpu(), ids.cpu()


@torch.inference_mode()
def validate(corpus, queries, expected_ids, scores, ids):
    assert scores.shape == ids.shape == (N_QUERIES, K)
    assert torch.isfinite(scores).all().item()
    assert ((ids >= 0) & (ids < len(corpus))).all().item()
    assert (scores[:, :-1] >= scores[:, 1:]).all().item()
    assert all(row.unique().numel() == K for row in ids)
    assert torch.equal(ids[:, 0], expected_ids)

    # A full matrix is acceptable only for this small CPU oracle.
    reference = queries.double() @ corpus.double().T
    expected_scores = reference.topk(K, dim=1, sorted=True).values
    selected_scores = reference.gather(1, ids)
    actual = scores.double()
    torch.testing.assert_close(
        actual, expected_scores, rtol=1e-4, atol=1e-6
    )
    torch.testing.assert_close(
        actual, selected_scores, rtol=1e-4, atol=1e-6
    )
    return (actual - expected_scores).abs().max().item()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--device", choices=("cuda", "cpu"), default="cuda")
    parser.add_argument("--chunk-rows", type=int, default=1024)
    parser.add_argument("--output", type=Path, default=Path("search-report.json"))
    args = parser.parse_args()
    if args.chunk_rows < 1:
        parser.error("--chunk-rows must be positive")
    if args.device == "cuda" and not torch.cuda.is_available():
        raise SystemExit("CUDA is unavailable; check the driver and torch build")

    torch.manual_seed(SEED)
    torch.use_deterministic_algorithms(True)
    torch.set_float32_matmul_precision("highest")
    device = torch.device(args.device)
    corpus, queries, expected_ids = fixture()
    scores, ids = search(corpus, queries, device, args.chunk_rows)
    error = validate(corpus, queries, expected_ids, scores, ids)
    report = {
        "assertions_passed": True,
        "seed": SEED,
        "python": platform.python_version(),
        "torch": torch.__version__,
        "cuda_runtime": torch.version.cuda,
        "device": str(device),
        "device_name": torch.cuda.get_device_name(0) if device.type == "cuda" else "cpu",
        "corpus_rows": N_ROWS,
        "query_rows": N_QUERIES,
        "dimensions": DIM,
        "k": K,
        "chunk_rows": args.chunk_rows,
        "max_score_error": error,
        "expected_first_ids": expected_ids.tolist(),
        "ids": ids.tolist(),
        "scores": scores.tolist(),
    }
    args.output.write_text(
        json.dumps(report, indent=2, allow_nan=False) + "\n", encoding="utf-8"
    )
    print(f"All correctness assertions passed; wrote {args.output}")


if __name__ == "__main__":
    main()

Three details are easy to miss:

Row offsets: topk returns indices within a block. Adding offset turns them into corpus IDs. Omitting this can produce believable scores paired with incorrect documents.
The final block: min(K, len(block)) handles blocks smaller than the requested number of neighbors. The merge also accepts fewer than five candidates during early iterations.
No silent CPU fallback: the default device is CUDA. torch.cuda.is_available must succeed before the search begins. CPU mode is an explicit diagnostic choice, not a substitute reported as GPU execution.

The normalize documentation explains the row normalization. Here it runs on the CPU before transferring the same input values to the GPU, keeping data generation separate from device-specific matrix multiplication.

Why can we discard everything except a block's top five? A discarded item already has at least five items in its own block scoring no lower than it. It cannot be required to construct a global top five by score. Repeatedly taking the top five of the accumulated candidates preserves that invariant. Tied items need separate attention, discussed below.

4. Run it and inspect the evidence

Check syntax, then run the CUDA path twice:

.venv/bin/python -m py_compile search.py
.venv/bin/python search.py --device cuda --output report-a.json
.venv/bin/python search.py --device cuda --output report-b.json
.venv/bin/python -m json.tool report-a.json

Each invocation writes its report only after every assertion succeeds. The checks establish that:

every query receives five distinct, valid corpus IDs;
scores are finite and sorted in descending order;
the first match is the known source row for every query;
returned scores agree with the float64 CPU top-k scores within explicit tolerances;
each returned ID actually corresponds to its claimed score in the CPU reference.

That last check catches a subtle failure: returning the correct scores alongside unrelated IDs. Comparing scores alone would miss it. The checks use torch.testing.assert_close with both relative and absolute tolerances, and compare float64 tensors on the same device.

The report includes Python, PyTorch, CUDA runtime, device name, seed, chunk configuration, IDs, and scores. max_score_error is calculated during your execution; it is not a published measurement. This tutorial presents correctness checks, not measured GPU performance or a benchmark transcript.

For a replay check, save this as check_replay.py and run it with .venv/bin/python check_replay.py:

import json
import math
from pathlib import Path

first = json.loads(Path("report-a.json").read_text())
second = json.loads(Path("report-b.json").read_text())
assert first["assertions_passed"] and second["assertions_passed"]
for field in ("seed", "corpus_rows", "query_rows", "dimensions", "k"):
    assert first[field] == second[field]
for report in (first, second):
    assert len(report["ids"]) == report["query_rows"]
    assert [row[0] for row in report["ids"]] == report["expected_first_ids"]
for row_a, row_b in zip(first["scores"], second["scores"], strict=True):
    for a, b in zip(row_a, row_b, strict=True):
        assert math.isclose(a, b, rel_tol=1e-4, abs_tol=1e-6)
print("Replay checks passed")

Do not run the scripts with Python's optimization flag: these tutorial checks use assert statements.

You can also exercise the tail-block path:

.venv/bin/python search.py --device cuda --chunk-rows 1000 --output report-tail.json

With this configuration, the corpus does not divide evenly into blocks. A diagnostic run using --chunk-rows 3 additionally exercises blocks smaller than K; it is an edge-case test, not a suggested operating configuration.

5. Understand what reproducibility means here

A fixed CPU generator seed makes the fixture repeatable within a matching environment. The script also enables deterministic algorithms and sets the documented matrix-multiplication workspace configuration before importing PyTorch. Unsupported deterministic operations fail rather than silently switching policy.

The float32 matrix multiplication precision setting is highest, keeping the arithmetic policy explicit. None of this promises identical floating-point bytes across different hardware or PyTorch releases. PyTorch's reproducibility notes explicitly describe that boundary. Preserve the environment record, script, and reports together.

There is another boundary: torch.topk can return different indices for tied scores. The validator therefore compares top-k score values and checks the selected IDs against those values instead of requiring every secondary ID to equal the CPU's choice. The planted first-match assertion applies to this synthetic fixture, not arbitrary duplicate vectors.

If production data requires a fixed ordering for exact ties, define a secondary document-ID ordering and implement it consistently in both block selection and merging. Do not perturb the scores with arbitrary offsets and assume the ranking is unchanged.

6. Move from the fixture to application data

For your own retrieval task, replace fixture() with corpus and query vectors from the same encoder and preprocessing pipeline. Reject non-finite inputs and zero-length vectors before normalization, preserve a separate mapping from row numbers to document IDs, and normalize along the component dimension. Cosine scores from unrelated vector spaces are not meaningful.

Keep the planted fixture as a regression test. The supplied validate() is fixture-specific: for application data, remove the planted-row assertion and fixed shape constants, and run a CPU reference on a representative query sample. Do not allocate the full reference matrix for an unrestricted corpus.

Streaming corpus blocks bounds the temporary score matrix, but the example still keeps the corpus in host memory and all queries on the selected device. For a larger query set, add an outer query-batching loop. If you get an allocation error, reduce --chunk-rows first; do not infer that the arithmetic or ID mapping is correct merely because the program completes.

If CUDA is unavailable, inspect the installed PyTorch build and the environment's GPU access before proceeding. --device cpu can isolate a logic problem, but a successful CPU run does not verify the GPU path.

DEV Community