Training Pipelines

An ML training pipeline is more than 'call model.fit()'. In production, training is a multi-stage workflow that must be reproducible, auditable, and resilient to failures. A typical pipeline includes: (1) data ingestion and validation -- pull training data, verify schema, check for distribution drift; (2) preprocessing and feature engineering -- normalize, tokenize, compute derived features; (3) training -- run distributed gradient descent across multiple GPUs; (4) evaluation -- measure metrics on hold-out sets, compare against the production baseline; (5) registration -- if the new model passes quality gates, register it in the model registry for deployment.

Distributed training is the compute-intensive core of the pipeline. Data parallelism -- the most common strategy -- replicates the model on every GPU and shards the training data across them. After each forward-backward pass, gradients are synchronized across GPUs using AllReduce (ring-reduce or tree-reduce). This scales well up to ~64 GPUs but communication overhead grows with model size: a 7B-parameter model produces 28 GB of gradients (in FP32) that must be exchanged every step. Gradient compression, mixed-precision training (FP16/BF16 forward pass, FP32 master weights), and gradient accumulation (local steps before sync) reduce communication pressure.

For models too large to fit on a single GPU (70B+ parameters), model parallelism splits the model itself across GPUs. Pipeline parallelism assigns different layers to different GPUs and pipelines micro-batches through them. Tensor parallelism splits individual layers (e.g., attention heads) across GPUs on the same node. ZeRO (DeepSpeed) partitions optimizer states, gradients, and parameters across GPUs to reduce per-GPU memory, enabling training of 100B+ models on clusters of commodity GPUs. FSDP (PyTorch) implements similar ideas natively.

Pipeline orchestration manages the DAG of stages, handles retries on transient failures (GPU errors, preemptions on spot instances), and tracks lineage. Modern orchestrators (Kubeflow Pipelines, Metaflow, Vertex AI Pipelines, Flyte) treat each stage as a containerized step with declared inputs, outputs, and resource requirements. Checkpointing is critical: a 24-hour training run that fails at hour 23 without checkpoints wastes $10,000+ in GPU compute. Best practice is to checkpoint every 15-30 minutes and implement automatic resume from the latest checkpoint.

An e-commerce company retrains its search ranking model nightly. At midnight, the pipeline (1) pulls the latest 30 days of click data from the data warehouse, (2) validates that the data schema is unchanged and volume is within 20% of the previous day, (3) computes features using the feature store, (4) trains for 3 epochs on 8 GPUs using data parallelism, checkpointing every 20 minutes, (5) evaluates NDCG@10 on a held-out set and compares against the production model, (6) if NDCG improves by >0.5%, registers the new model in the registry and triggers a canary deployment. The entire pipeline takes 4 hours and costs ~$200 in GPU compute.