CHAPTER 26
Data, Features and Training Infrastructure
Pipelines, feature stores, distributed training and experiment management
Source: PDF pp. 68–70
PDF page 68Anatomy of an ML platform p. 68
Feature store p. 68
Data pipelines p. 68
- Batch (Spark, dbt, SQL on a warehouse; orchestrated by Airflow / Dagster / Prefect) for history and heavy joins; streaming (Kafka + Flink / Spark Structured Streaming) for fresh signals such as "clicks in the last 5 minutes" (Ch. 8).
- Data contracts and validation: producers publish schemas; consumers assert expectations (nullability, ranges, uniqueness, freshness). Fail the pipeline, do not silently train on bad data.
- Lineage and reproducibility: track which raw data, code and parameters produced each dataset and model.
Feature engineering essentials p. 68
| Feature type | Techniques |
|---|---|
| Numeric | Scaling/normalisation, log transforms for skew, bucketing, clipping outliers, missing-value indicators |
| Categorical | One-hot (low cardinality), target/mean encoding (with leakage-safe folds), hashing trick, learned embeddings (high cardinality IDs) |
| Text | TF-IDF / n-grams, pretrained embeddings (sentence encoders), tokens for transformers |
| Time / sequence | Cyclical encodings (hour, weekday), lags, rolling-window aggregates (counts, sums, ratios over 1 h / 1 d / 30 d), recency, session features |
| Interaction / cross | User x item affinity, feature crosses, similarity between embeddings |
| Graph / location | Neighbour aggregates, PageRank, geohash/H3 cell statistics |
Feature store p. 68
A feature store is a system that defines features once and serves them consistently to training and inference. It solves training-serving skew, duplicate work across teams, and point-in-time leakage.
| Component | Role |
|---|---|
| Registry / definitions | Feature names, owners, entities, transformation code, freshness SLAs, documentation |
| Offline store | Warehouse / lake tables with full history; point-in-time joins attach each training row to feature values as of its event timestamp |
| Online store | Low-latency key-value store (Redis, DynamoDB, Cassandra) with the latest values per entity (user_id), populated by streaming and batch materialisation; read in a few ms |
| Transformation engines | Batch (Spark/SQL), streaming (Flink) and on-demand (request-time) computation, one definition for all |
Training infrastructure p. 69
Memory arithmetic: why one GPU is not enough Training with Adam in mixed precision needs roughly 16 bytes per parameter: 2 (fp16/bf16 weights) + 2 (gradients) + 4 (fp32 master weights) + 4 + 4 (Adam moments), plus activations. A 7-billion-parameter model needs about 7 B x 16 = ~112 GB before activations, more than an 80 GB GPU. A 70 B model needs >1 TB. Inference needs only the weights: ~2 bytes/param in fp16 (7 B is 14 GB; 70 B is 140 GB), or ~0.5 byte/param with 4-bit quantisation.
- Data parallelism with all-reduce (PyTorch DDP, Horovod) is the default. ZeRO / FSDP shards optimizer states, gradients and parameters across data-parallel workers, cutting memory per GPU almost linearly.
- Mixed precision (bf16/fp16 or fp8), gradient accumulation, and activation checkpointing (recompute instead of store) trade compute for memory.
- Interconnect matters: NVLink/NVSwitch inside a node (hundreds of GB/s) vs InfiniBand/RoCE between nodes. Communication-to-compute ratio decides scaling efficiency.
- Fault tolerance: checkpoint frequently to object storage; jobs must resume after node failure (large clusters fail often). Use elastic training or auto-restart; spot/preemptible instances need cheap checkpoints.
- Scheduling and orchestration: Kubernetes + Kueue / Volcano, Slurm, Ray, or managed services. Gang scheduling, topology-aware placement and fair-share queues keep expensive GPUs busy.
- Data loading must not starve GPUs: sharded sequential formats (Parquet, WebDataset, TFRecord), prefetch, caching on local NVMe.
Experiments, tuning and the registry p. 69
- Experiment tracking (MLflow, Weights & Biases): log code version, data version, hyperparameters, metrics, artifacts, and environment so any result is reproducible.
- Hyperparameter search: random search beats grid; Bayesian optimisation and successive-halving (Hyperband/ASHA) cut cost by killing poor trials early.
- Model registry: immutable versions with metadata (training data, metrics, owner, lineage) and stage transitions (staging, production, archived); deployment pulls from the registry, never from a notebook.
- Cost control: right-size GPUs, use spot capacity with checkpointing, profile utilisation (MFU), stop unpromising runs, cache preprocessed datasets.
Review questions for this chapter have not been written and reviewed yet. The chapter text above is the full source content.
Try it
References
System Design & AI System Design: The Complete Guide — 2026 Edition, pp. 68–70; figures 47, 48.
Tip: use ← and → keys to move between chapters.