Distributed Preprocessing

"I normalized my shard by its own mean and felt magnificent, right up until validation, which had never heard of my mean, returned numbers from a different planet."

A Worker That Fit On Itself

By this point in the chapter the data already lives in shards on distributed storage (building on the aggregation patterns of Chapter 6) and a sharded loader feeds it into training, as Section 8.5 set up. What we have not yet asked is what happens to each example between the raw bytes on disk and the tensor the model consumes. Numerical features get normalized, categorical features get encoded against a vocabulary, text gets tokenized, images get resized and scaled. That transformation pipeline is preprocessing, and the part of it that genuinely scales out, as opposed to the generic feature engineering you would do on a laptop, is the subject here. The distinction this section keeps insisting on is between work that is per-example and therefore trivially parallel, and work that depends on the whole dataset and therefore demands coordination.

1. Two Places Preprocessing Runs: Offline and Online Beginner

The first design decision is when the transform happens relative to training. In offline preprocessing, a batch job reads the raw dataset, applies the transform, and writes preprocessed shards back to storage; training then reads those finished shards directly. This is the natural home of a distributed DataFrame engine, so a normalization-and-encoding pass over a multi-terabyte table is exactly the kind of wide transformation that Spark and distributed DataFrames (Chapter 7) exist to run, with the fit expressed as a grouped aggregation and the transform as a column map. In online preprocessing, the transform runs inside the training data loader (the per-worker pipeline of Section 8.5), turning raw examples into tensors on the fly as each batch is assembled, often overlapped with computation on the accelerator.

Neither is universally correct, and the right answer depends on how often the transform changes and how expensive it is. A fixed, cheap transform (a known normalization, a frozen tokenizer) is a fine candidate for the loader, because recomputing it costs little and keeps the storage footprint small. An expensive, stable transform (heavy text cleaning, audio feature extraction, a resize-and-decode that dominates the loader budget) is better materialized once offline, so that no epoch pays for it twice. The cost model in Section 5 of this section makes that trade-off concrete. What both modes share, and what the rest of this section is about, is that the constants the transform depends on must be learned from the whole training set before either mode can run a single example correctly.

2. The Fit/Transform Split, Done Right Across Machines Intermediate

Consider the most ordinary preprocessing step of all, standardizing a numeric feature to zero mean and unit variance. The transform is $z = (x - \mu)/\sigma$, applied to each value independently. The fit is the pair $(\mu, \sigma)$, and the entire correctness question is: which examples does that pair summarize? The answer must be the training distribution, the same $\mu$ and $\sigma$ for the first shard, the last shard, the validation set, and every production request, because the model learned weights that assume inputs were centered and scaled by those exact numbers. If each worker instead standardizes its own shard with that shard's local mean, every shard arrives at the model on a slightly different scale, the validation set uses yet another, and serving uses a fourth. Section 8.8 treats this family of mistakes as data leakage and correctness failures; here we focus on computing the shared statistic right.

Computing $(\mu, \sigma)$ over data spread across $K$ machines without shipping the raw data anywhere is a mergeable aggregation, the same idea that made distributed aggregation work in Chapter 6 (aggregation and combiners). The trick is to choose per-shard partial summaries that combine by simple addition. Mean and variance both fall out of three additive quantities, the count, the sum, and the sum of squares,

from which $\mu = S_1/n$ and $\sigma^2 = S_2/n - \mu^2$. Each worker computes its own triple $(n_k, S_{1,k}, S_{2,k})$ over its shard alone, the triples are summed in a single all-reduce, and the merged triple yields exactly the statistic a single machine would have computed over all the data at once. Nothing but three numbers per worker crosses the network, no matter how many billions of examples each shard holds. The streaming mean and variance formula above is the foundation, and a numerically stabler version (Welford's update) is used in libraries when $S_2$ would overflow, but the additive, mergeable structure is identical.

The code below makes the equivalence concrete on a skewed feature column. It computes the mean and standard deviation the single-pass way over the whole column, then again as $K$ shard-local triples merged by addition, and finally applies the shared transform to confirm the standardized data has zero mean and unit variance everywhere.

import numpy as np

rng = np.random.default_rng(7)
N, K = 1_000_000, 16            # examples, shards (one per worker)
x = rng.gamma(shape=2.0, scale=3.0, size=N) + 5.0   # skewed feature column
shards = np.array_split(x, K)   # K disjoint shards, deliberately unequal-ish

# --- Single-pass "ground truth" statistic on the whole training column ---
true_mean = x.mean()
true_std  = x.std()             # population std (ddof=0)

# --- Distributed fit: each worker emits a MERGEABLE triple (count, sum, sumsq) ---
def local_partial(col):
    return np.array([col.size, col.sum(), np.square(col).sum()])

partials = [local_partial(s) for s in shards]        # K independent map outputs

# --- Reduce step: triples are additive, so merging is a plain elementwise sum ---
n, s1, s2 = np.sum(partials, axis=0)                  # the all-reduce
merged_mean = s1 / n
merged_var  = s2 / n - merged_mean * merged_mean      # E[x^2] - E[x]^2
merged_std  = np.sqrt(merged_var)

print("shards (workers)        :", K)
print("single-pass mean        :", f"{true_mean:.10f}")
print("merged   mean           :", f"{merged_mean:.10f}")
print("single-pass std         :", f"{true_std:.10f}")
print("merged   std            :", f"{merged_std:.10f}")
print("max abs diff (mean,std) :",
      f"{max(abs(merged_mean-true_mean), abs(merged_std-true_std)):.2e}")

# --- Transform: apply the SHARED stats identically on every shard ---
def transform(col, mu, sd):
    return (col - mu) / sd

z_shards = [transform(s, merged_mean, merged_std) for s in shards]
z_all = np.concatenate(z_shards)
print("post-transform mean     :", f"{z_all.mean():+.2e}")   # ~0
print("post-transform std      :", f"{z_all.std():.6f}")     # ~1

shards (workers)        : 16
single-pass mean        : 10.9935636126
merged   mean           : 10.9935636126
single-pass std         : 4.2374883765
merged   std            : 4.2374883765
max abs diff (mean,std) : 3.55e-15
post-transform mean     : -2.14e-16
post-transform std      : 1.000000

The agreement is not approximate; it is exact up to the rounding of floating-point addition, exactly as the data-parallel gradient was in Section 1.1. This is the same lesson in a new costume: the operation that looks like it should require seeing all the data at once actually decomposes into additive partials plus one all-reduce. A vocabulary is the same story with a different summary, where each shard emits a partial token-count dictionary and the reduce merges them by summing counts, after which a global frequency threshold selects the shared vocabulary. Quantile boundaries for bucketization merge mergeable sketches the same way. In every case the fit is a distributed aggregation and the transform is a per-example map.

3. Tokenization and Encoding at Scale Intermediate

Text and categorical features sharpen the fit/transform split because the "constant" being learned is a whole vocabulary or a subword merge table. Training a tokenizer (learning byte-pair-encoding merges, or counting token frequencies for a fixed vocabulary) is a fit: it must see a representative sample of the corpus, and at web scale that sample is itself sharded, so the counting is a distributed aggregation that merges per-shard count tables exactly as the mean did. Crucially, the tokenizer is fit once and then frozen. Applying it, mapping a string to a sequence of token ids, is a pure per-example transform with no cross-machine dependency, which is why tokenization parallelizes cleanly across every loader worker and every serving replica.

The discipline that matters at scale is that the frozen vocabulary travels as a small artifact (a file of merges or a token-to-id map) and is loaded identically everywhere. A token id is only meaningful relative to the vocabulary that defined it; if training and serving disagree on the mapping by even one entry, every downstream id shifts and the model receives gibberish. The same holds for categorical encoders: the learned category-to-index map is the shared artifact, and unseen categories at serving time must be routed to a reserved out-of-vocabulary slot rather than silently assigned a fresh index that the model never trained on. The fit produces a portable constant; the transform is stateless given that constant.

from pyspark.sql import functions as F

# FIT: one distributed aggregation returns the shared constants (one all-reduce internally)
mu, sd = df.select(F.mean("x"), F.stddev_pop("x")).first()

# TRANSFORM: one per-row expression, applied across every partition, no extra shuffle
df_z = df.withColumn("x_z", (F.col("x") - F.lit(mu)) / F.lit(sd))

4. Materialize Once or Recompute Every Epoch Intermediate

Once the constants are fixed, the transform still has to be applied, and there are two moments to apply it: once, writing preprocessed shards to storage that every epoch then reads directly, or every epoch, recomputing the transform inside the loader as the raw data streams by. Materializing once turns an $E$-epoch training run from $E$ applications of the transform into one, at the cost of extra storage for the preprocessed copy and a loss of flexibility, since changing the transform means regenerating the whole materialized set. Recomputing keeps storage lean and the pipeline flexible, but pays the transform cost $E$ times and can make the loader the bottleneck if the transform is heavy, the exact loader-starvation failure mode of Section 8.5.

The decision reduces to comparing two costs over the life of the run. Let $C_t$ be the per-example transform cost, $E$ the number of epochs, $N$ the dataset size, and $C_{io}$ the extra per-example cost of reading a (typically larger) materialized example versus a raw one. Materializing once costs roughly $N \cdot C_t + E \cdot N \cdot C_{io}$ of added work; recomputing costs $E \cdot N \cdot C_t$. Materializing wins when

which says: materialize when the transform is expensive relative to the extra input/output it creates and you run more than a couple of epochs, since the saved recomputation $C_t(1 - 1/E)$ then exceeds the added read cost $C_{io}$. Cheap transforms over few epochs favor recomputation; expensive transforms over many epochs favor materialization. Many production pipelines split the difference: materialize the expensive, stable steps (decode, resize, heavy text cleaning) offline once, and leave the cheap, varying steps (a final normalization, random augmentation) in the loader, getting the storage profile of recomputation for most of the work and the compute profile of materialization for the heavy part.

5. Beyond Training: Feature Stores and Serving Consistency Advanced

The fit/transform discipline does not end at training. A model in production receives raw features at request time and must transform them with the same constants it trained on, or it silently sees inputs from a different distribution than the one it learned. The system that enforces this is a feature store: a service that holds feature definitions and their learned constants, computes features offline in batch for training, and serves the identical computation online at low latency for inference, so the offline and online paths cannot drift apart. This online/offline consistency guarantee is the production-grade form of "fit once, apply everywhere", and it is developed in depth where it matters most, in the distributed recommendation case study (Chapter 38), whose feature pipelines fan out across hundreds of features and machines.

The thread that connects this section to that one is exactly the artifact we have been tracking: the shared constant produced by a distributed fit. Whether it is a mean, a vocabulary, or a learned bucket boundary, it is computed once by a mergeable aggregation across the training shards and then must be applied byte-for-byte identically in every later context. The feature store is the machinery that makes "byte-for-byte identically" a system guarantee rather than a hope.

We have separated preprocessing into a distributed fit and a parallel transform, computed the fit as an exact mergeable aggregation, weighed materializing against recomputing, and followed the shared constant all the way to the serving boundary. The one thing we have treated only by warning is what goes wrong when these boundaries blur, when the fit accidentally peeks at validation data, or when training and serving transforms drift apart. Those are correctness failures with a name, and the next section gives them one. The treatment of data leakage and correctness in distributed pipelines begins in Section 8.8.