Distributed Embedding Generation

"A GPU batch of ten thousand chunks, all becoming vectors at once. I never read any of them; I just turned meaning into geometry and moved on to the next ten thousand."

An Embedding Worker That Saw Only Its Shard

By the time a corpus reaches this stage it has already been gathered, filtered, deduplicated, and segmented into retrievable chunks. What remains is the step that converts text into the representation the retriever actually searches: a dense vector for every chunk. A billion source documents, segmented into a handful of chunks each, becomes several billion short forward passes through an embedding model. None of those forward passes depends on any other, which is the defining property of the job. There is no gradient to synchronize, no parameter to update, no collective to wait on; each worker can embed its slice of the corpus in complete isolation and the union of the results is the finished index. That independence is what makes the job both simple to reason about and ruthless to optimize, because every inefficiency on one GPU is multiplied by the thousands of GPUs running beside it.

1. From Chunks to Forward Passes Beginner

The unit of embedding is the chunk, not the document. A retriever returns passages, so the corpus is first segmented into spans short enough to be a coherent answer and to fit the encoder's context window, typically a few hundred tokens with a small overlap so that a fact straddling a boundary survives in at least one chunk. Segmentation strategy belongs to the chunking section that precedes this one; what matters here is the multiplier it produces. A corpus of $N_{\text{docs}}$ documents averaging $c$ chunks each yields

forward passes. For a billion documents at six chunks each, that is six billion forward passes, and the entire cost model of this section follows from that single number. The embedding model itself is a compact encoder, usually a sentence-transformer in the 100-to-500-million-parameter range that maps a chunk to a vector of dimension $d$ (commonly 384, 768, or 1024). It is small enough to fit in one accelerator's memory with room to spare, which is precisely why this job never needs the model-sharding machinery of Chapter 16: the model is replicated, not split. What is split is the data.

2. Sizing the Job: GPU-Hours, Wall-Clock, and Cost Intermediate

The most useful thing to compute before launching a billion-document embedding run is its bill. Let $r$ be the throughput of one accelerator in chunks per second at the chosen batch size. The total accelerator work is fixed by the corpus and the per-GPU rate, independent of how many GPUs you use:

Add $G$ workers and the wall-clock time divides, because the chunks split cleanly across them with no coordination overhead to claw any of it back:

where $p$ is the price per GPU-hour. The cost depends only on the total GPU-hours and the price, not on $G$: renting four thousand GPUs for half an hour costs the same as renting eight GPUs for ten days, give or take the overhead of spinning the larger fleet up. That decoupling, total cost set by the workload and wall-clock set by the fleet size, is the practical meaning of an embarrassingly parallel job, and it is what lets a team trade money against time freely. Code 36.5.1 turns these three identities into a sizing calculator and sweeps the worker count to show the linear scale-out explicitly.

import math

# --- Corpus and model parameters ---
N_docs = 1_000_000_000          # one billion documents
chunks_per_doc = 6              # average chunks after segmentation
N_chunks = N_docs * chunks_per_doc
tput_per_gpu = 800             # chunks/sec embedded per GPU at the chosen batch size
gpu_price = 1.10               # USD per GPU-hour (cloud accelerator)

# --- Core sizing identities ---
# total GPU-seconds is independent of how many GPUs run in parallel
gpu_seconds = N_chunks / tput_per_gpu
gpu_hours = gpu_seconds / 3600.0
job_cost = gpu_hours * gpu_price

print(f"corpus: {N_docs:,} docs x {chunks_per_doc} chunks = {N_chunks:,} chunks")
print(f"per-GPU throughput : {tput_per_gpu:,} chunks/sec")
print(f"total GPU-hours    : {gpu_hours:,.0f}")
print(f"job cost (@ ${gpu_price}/gpu-hr): ${job_cost:,.0f}")
print()

# --- Linear scale-out: sweep worker count, wall-clock = GPU-hours / num_gpus ---
print(f"{'GPUs':>8} {'wall-clock (h)':>16} {'speedup':>10} {'efficiency':>12}")
base = None
for num_gpus in [8, 64, 256, 1024, 4096]:
    wall_h = gpu_hours / num_gpus
    if base is None:
        base = wall_h * 8         # reference: wall-clock that 1 GPU would take
    speedup = base / wall_h
    eff = speedup / num_gpus
    print(f"{num_gpus:>8} {wall_h:>16,.1f} {speedup:>10,.0f}x {eff:>11.0%}")

print()
# --- Throughput vs batch size (illustrative saturation curve) ---
# small batches underuse the GPU; throughput rises then plateaus at the FLOPs ceiling
print(f"{'batch':>8} {'chunks/sec':>14}")
peak = 900.0
for bs in [8, 32, 128, 512, 2048]:
    # saturating model: half-saturation batch size = 96
    tput = peak * bs / (bs + 96)
    print(f"{bs:>8} {tput:>14,.0f}")

corpus: 1,000,000,000 docs x 6 chunks = 6,000,000,000 chunks
per-GPU throughput : 800 chunks/sec
total GPU-hours    : 2,083
job cost (@ $1.1/gpu-hr): $2,292

    GPUs   wall-clock (h)    speedup   efficiency
       8            260.4          8x        100%
      64             32.6         64x        100%
     256              8.1        256x        100%
    1024              2.0      1,024x        100%
    4096              0.5      4,096x        100%

   batch     chunks/sec
       8             69
      32            225
     128            514
     512            758
    2048            860

The efficiency column staying pinned at 100% is the signature of the embarrassingly parallel regime, and it is exactly what data-parallel training cannot achieve, because there the all-reduce tax erodes efficiency as workers multiply. The honest caveat is that this idealization ignores the input pipeline and the write path; a real run loses a few percent to reading chunks from storage and committing vectors to the index, which is why Sections 3 and 5 treat those two boundaries as the places where the linear curve actually bends.

3. Driving the GPU: Batch Size, Mixed Precision, and Sharding Intermediate

The per-GPU throughput $r$ that the cost model treats as a constant is in fact the most important knob in the job, and three settings determine it. The first is batch size. An encoder forward pass is a sequence of matrix multiplications, and a modern accelerator reaches its arithmetic peak only when those multiplications are wide enough to keep every core busy. A batch of eight chunks leaves most of the hardware idle; a batch of several hundred or a few thousand saturates it, which is the rising curve in the bottom of Output 36.5.1. The relationship is the same per-node arithmetic-intensity argument developed in Chapter 22: throughput climbs with batch size until the kernel becomes compute-bound, then plateaus. Offline embedding is the ideal place to exploit this, because, unlike online serving, there is no latency budget to respect; you may use the largest batch that fits in memory and pay nothing for the queuing delay that would be unacceptable in production.

The second knob is numerical precision. Embedding does not need full 32-bit floating point; the encoder runs comfortably in half precision (FP16 or BF16), which doubles arithmetic throughput and halves the memory the activations consume, so a larger batch fits and the saturation point arrives sooner. The output vectors are typically stored at lower precision still, since retrieval quality is robust to vector quantization, a point Section 36.6 and Chapter 25 develop when they compress the index. The third knob is corpus sharding: the chunked corpus is partitioned into $G$ disjoint shards, one per worker, sized so that no worker finishes long before the others and leaves a GPU idle at the tail of the job. Uneven shards reintroduce the straggler problem from Chapter 3, so shards are balanced by chunk count, and a work-stealing queue lets an idle worker pull from a slow one's remaining backlog. The aggregate effect of these three settings is captured by the throughput identity that feeds the cost model:

where $b$ is the batch size, $r_{\max}$ the saturated per-GPU rate, and $b_{1/2}$ the batch at which the GPU reaches half its peak. The curve in Output 36.5.1 is this function with $b_{1/2} = 96$; reading it tells you the smallest batch that still earns most of the hardware's throughput, which is the batch you actually run.

# Launched once per worker; rank and world_size come from the cluster launcher.
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("BAAI/bge-base-en-v1.5", device="cuda")
model.half()                                   # FP16: ~2x throughput, larger batch fits

for batch_ids, batch_texts in shard_batches(rank, world_size, batch_size=512):
    vecs = model.encode(batch_texts,           # one forward pass over the whole batch
                        batch_size=512,
                        normalize_embeddings=True,   # unit vectors for cosine retrieval
                        convert_to_numpy=True)
    vector_store.write(batch_ids, vecs)        # stream straight to the index shard

4. Incremental and Streaming Embedding Intermediate

The bulk job embeds the corpus once, but a corpus is never finished. New documents arrive continuously, old ones are revised, and a few are deleted, so the index must be kept current without re-embedding the billions of chunks that did not change. The principle is to embed only the delta. A change-data-capture feed identifies which documents are new or modified since the last run; those are re-chunked and embedded by a small standing pool of workers, and the resulting vectors are upserted into the store while the vectors of deleted documents are tombstoned. This is the offline embedding pipeline reframed as a stream-processing job, and it inherits its design directly from the online-AI and stream-processing material of Chapter 9: the same windowing, exactly-once delivery, and watermark discipline that govern any streaming pipeline govern this one.

The economics flip in the incremental regime. The bulk job is throughput-bound and runs at maximum batch size, because a billion chunks are waiting and latency does not matter. The streaming job is latency-bound: a freshly published document should become retrievable within seconds or minutes, so the workers run small batches and accept lower per-GPU efficiency in exchange for low end-to-end delay. A mature system runs both, a periodic bulk re-embed that rebuilds the index when the encoder itself is upgraded, and a continuous streaming embed that keeps it current between rebuilds. The decision of which mode to use is the same throughput-versus-latency tension that separates offline batch inference from online serving throughout Part V, and recognizing that an embedding pipeline sits on both sides of it is the key to running it economically.

5. Writing the Vectors to the Store Intermediate

A vector that is computed but not written is wasted GPU-hours, so the write path deserves the same care as the forward pass. Each worker emits, for every chunk, a tuple of an identifier, the dense vector, and a small bundle of metadata (the source document, the chunk offset, any filterable attributes the retriever will use). These tuples flow into the distributed vector store of Chapter 25, which is itself sharded across machines so that no single node holds the whole index. The natural and efficient arrangement aligns the two shardings: a worker writes to the index shard that co-locates with it, so the vectors land where they will be searched without a network shuffle, the same partition-alignment principle that makes a co-partitioned join cheap in Chapter 7.

The write must be idempotent, because at corpus scale some worker will crash and its shard will be retried, and a retried chunk must overwrite rather than duplicate its earlier vector. Keying every write on the chunk identifier makes a re-run safe: the second write of a chunk simply replaces the first, so the job can be restarted from a checkpoint of completed shards without producing a single duplicate vector. Vectors are typically written in large batched commits rather than one at a time, both to amortize the index-update cost and because the approximate-nearest-neighbor structures that Chapter 25 builds are far cheaper to construct from a bulk load than from a stream of singletons. When the last shard commits, the index is complete and the corpus has become geometry, ready for the retriever that Section 36.6 builds on top of it.