ML System Design · Module 15 of 47 · 11 concepts · about 14 min

Training Pipelines and Training at Scale

Turning training from a notebook into an automated, repeatable pipeline: reproducibility, experiment tracking, tuning, distributed training and its costs, checkpoints, how often to retrain, the gate a new model must pass, and the registry that stores every version.

At work: What lets a team retrain a model every night without anyone watching, and answer 'exactly how was the model in production built?' months later.

The training pipeline

What it is

In production, training is a pipeline that runs automatically, on a schedule or a trigger:

snapshot the training data at a recorded version;
validate it (module 8);
train with a recorded configuration;
evaluate on held-out data, overall and per slice;
compare with the model currently in production (the champion);
register the new model with its scores, data version and code version, if it passes.

Why it matters

A pipeline makes retraining routine and safe.

Nobody needs to remember the steps or run a notebook.
Every model in the registry can be traced to its data, code and settings.
A bad candidate is stopped by the gate, not discovered by users.

In the real world: Kirana Kart's ETA pipeline runs every night: it trains on the latest twelve months, and the new model is registered only if it beats the champion on last week's data without hurting any city.

Visual example

Pipeline

A failing check at any step stops the run and alerts the owner.

A nightly training pipeline

SnapshotFixed, recorded data version
ValidateSchema, counts, distributions
TrainRecorded config, commit, seed
EvaluateOverall, slices, latency
CompareAgainst the champion
RegisterOnly if it passes

Retraining by re-running a notebook by hand. It works until the person who knows the steps is on holiday, or until a cell is run out of order and nobody can reproduce the model.

Practise it: Design an ML Training Pipeline

Reproducibility

What it is

A training run is reproducible if it can be run again to get the same model. Pin down:

code: the exact commit;
data: the exact version or snapshot;
configuration: every setting, in a file, not typed by hand;
environment: library versions, in a container image or lock file;
randomness: seeds set, and known sources of non-determinism (some GPU operations) noted.

Why it matters

Without reproducibility you cannot debug, compare or audit.

'The new model is worse' cannot be investigated if the old one cannot be rebuilt.
A regulator asking how a decision was made needs the exact model and data.
A library upgrade that changes results can be spotted only if everything else is fixed.

In the real world: Kirana Kart builds one container image per training-code release, and the registry entry for each model records the image, the commit and the data version.

Installing 'the latest' libraries in every training run. A minor release of a dependency changes a default, and the model shifts with no change to your code.

Practise it: Design an ML Training Pipeline

Experiment tracking

What it is

An experiment tracker records every training run so runs can be compared and found later:

parameters and configuration;
metrics, overall and per slice, and learning curves;
the data version and code commit;
artefacts: the model file, plots, feature importances.

Common tools: MLflow, Weights & Biases, Neptune, and the clouds' built-in trackers.

Why it matters

Models are developed by trying many things.

Without a tracker, the best run of last Tuesday is lost in someone's notebook.
Teams can see what has already been tried.
The tracker's records become the evidence in a model review.

In the real world: Many teams' model-review template simply links to the tracked runs of the champion and the candidate; reviewers compare them side by side.

import mlflow

with mlflow.start_run(run_name="eta-lgbm-rain-feature"):
    mlflow.log_params(config)                         # every setting
    mlflow.set_tags({"data_version": "orders@v812", "git_commit": GIT_SHA})
    model = train(config, train_df)
    mlflow.log_metrics({"mae_min": mae, "p90_abs_err_min": p90})
    for city, err in per_city_mae.items():
        mlflow.log_metric(f"mae_{city}", err)
    mlflow.lightgbm.log_model(model, "model")

Logging only the final headline metric. When a model fails in one city in production, the run records cannot show whether that city was already worse offline.

Practise it: Design an ML Training Pipeline

Hyperparameter tuning

What it is

Hyperparameters are settings chosen before training (tree depth, learning rate, number of layers). Ways to search:

grid search: every combination; wasteful beyond two or three settings;
random search: random combinations; surprisingly effective;
Bayesian optimisation (Optuna, Hyperopt): uses past results to choose promising settings next;
early stopping: abandon runs that are clearly poor, to spend the budget on good ones.

Why it matters

Tuning helps, but with fast-diminishing returns.

The first few dozen trials usually capture most of the gain.
Better features and labels usually help more.
Tune on the validation set, never the test set, or the final score is too optimistic.

In the real world: Kirana Kart runs a 100-trial Optuna search for its ETA model once a quarter; nightly retraining uses the best settings from the last search.

Re-tuning from scratch every night. Tune occasionally, then retrain nightly with the chosen settings; the settings rarely need to change as fast as the data.

Distributed training

What it is

When data or models are too big for one machine, training is spread across many:

Data parallelism: each machine (or GPU) holds a full copy of the model and trains on a different slice of the data; their updates are averaged. The common case.
Model parallelism: the model itself is split across devices, because it does not fit on one. Needed for very large neural networks.
Sharded embeddings: huge embedding tables (billions of IDs) are split across machines; common in recommendation models.

Frameworks: PyTorch DDP and FSDP, Horovod, Ray Train, and for trees, distributed XGBoost, LightGBM or Spark.

Why it matters

It cuts the time per run, which speeds up iteration.

But it is never perfectly efficient: machines spend time exchanging updates.
So each extra GPU buys less speed, and the total GPU-hours (and bill) go up.
Use it when time-to-result matters, not by default.

In the real world: Large recommendation models at Meta and Google train across hundreds of GPUs mainly because their embedding tables do not fit on one; Kirana Kart's models all train on a single machine.

Visual example

Computed

A 24-hour single-GPU job, losing 15% efficiency at each doubling. Faster results cost more in total.

speed-up = GPUs × 0.85^log2(GPUs); illustrative ₹250 per GPU-hour

GPUs	speed-up	hours to finish	GPU-hours	cost
1	1×	24	24	₹6,000
2	1.7×	14.1	28.2	₹7,059
4	2.9×	8.3	33.2	₹8,304
8	4.9×	4.9	39.1	₹9,770
16	8.4×	2.9	46	₹11,494

Distributing a job that fits on one machine. Boosted trees on a few hundred million rows often train fastest on one large machine; spreading them across a cluster adds communication for little gain.

Checkpoints and failures

What it is

Long training runs fail: a machine dies, a cheap spot machine is reclaimed, memory runs out. Checkpointing saves progress so a run can resume:

save the model's weights, optimiser state and position in the data at regular intervals;
on restart, load the latest checkpoint and continue;
keep the best checkpoint by validation score, not only the last.

Why it matters

It makes long runs and cheap hardware practical.

Spot or pre-emptible machines can cost a fraction of the normal price, but can disappear at any time.
Without checkpoints, a failure at hour 20 of 24 wastes 20 hours.

In the real world: Teams training on spot GPUs typically checkpoint every 15 to 30 minutes, accepting the small overhead to save large amounts on hardware.

Saving only the weights. Resuming without the optimiser state and data position changes the training, and the result differs from an uninterrupted run.

How often to retrain

What it is

Three common policies:

On a schedule: nightly, weekly, monthly. Simple and predictable.
On a trigger: when monitoring shows drift or a performance drop (module 21).
Continuously: update the model on a stream of new data, every few minutes or hours (online learning).

And two ways to update: retrain from scratch on a recent window, or warm-start (fine-tune yesterday's model on new data).

Why it matters

The right cadence depends on how fast the world changes and how fast labels arrive.

Fraud patterns and trending items change daily; a churn model may be fine monthly.
Measure it: train a model, then test it on each following week; how fast does accuracy fall?
Warm starts are cheaper and faster but can accumulate errors; periodic full retrains reset them.

In the real world: Kirana Kart retrains recommendations nightly, ETA nightly, fraud weekly with a trigger on drift, and demand forecasts nightly with a full re-tune each quarter.

Visual example

Cadence

Most teams combine a schedule with a drift trigger.

Retraining policies compared

policy	good for	risk
Schedule	Predictable change; simple operations	Too slow for sudden shifts; wasteful if nothing changed
Trigger	Irregular change	Needs reliable monitoring; may react to noise
Continuous	Fast-moving signals (trends, ads)	Complex; errors spread quickly; harder to test

Retraining more often than labels arrive. Retraining a fraud model daily when chargebacks take 60 days mostly re-learns the same labels, at full cost.

The promotion gate

What it is

A new model is promoted only if it passes automatic checks against the champion, on the same recent data:

better (or not worse) on the main metric;
no important slice gets meaningfully worse;
calibration within limits, if probabilities are used;
latency and size within the serving budget;
sanity checks: predictions in a plausible range; no feature suddenly dominating.

Why it matters

The gate is what makes automated retraining safe.

A broken data day can produce a broken model; the gate stops it.
Slice checks catch a model that improves the average by sacrificing a small city.
Passing the gate earns an online test (module 17), not instant full release.

In the real world: Kirana Kart's gate requires the ETA candidate to beat the champion's average error on last week's orders, with no city more than 3% worse and p99 latency under 5 ms.

Comparing the candidate's score on new data with the champion's score from when it was trained. Data changes; both must be scored on the same recent set.

The model registry

What it is

A model registry stores every trained model with what is needed to use and trust it:

the model file and its version;
lineage: data version, code commit, configuration, training run;
scores: offline metrics, slices, gate results;
stage: candidate, staging, production, archived;
owner and approval history.

Tools: the MLflow Model Registry, and the registries built into SageMaker, Vertex AI and Databricks.

Why it matters

The registry is the hand-off between training and serving.

Serving loads 'the production version' from it, never a file someone copied.
Rolling back means pointing production at the previous version.
Audits start here: which model was live on a given date, and how was it built?

In the real world: At Kirana Kart, a model can only be served if it is in the registry with a passing gate result; the serving system refuses anything else.

Deploying model files from a shared folder. Two weeks later nobody can say which file is live or how it was trained.

Practise it: Design an ML Training Pipeline, Design a Model Serving Platform, Design Continuous ML Deployment

Interview questions: training pipelines

What it is

The questions interviewers ask about training pipelines, from the first screening call to the senior round. Cover the answer, say yours aloud, then compare.

Why it matters

Reading an idea and explaining it are different skills.

Each answer below is the two or three sentences an interviewer is listening for.
The reasoning behind each one is in the cards above. If an answer surprises you, go back to that card.

In the real world: MLOps interviews dwell here; product ML interviews usually ask one question: 'how does a new model get from training to production safely?'.

Q: What does a production training pipeline look like?

Snapshot versioned data, validate it, train with recorded config, evaluate overall and per slice, compare with the champion on the same data, and register the model with lineage if it passes.

Q: How do you make training reproducible?

Pin the code commit, data version, configuration file, environment (container or lock file) and random seeds, and record them all with the model.

Q: When do you need distributed training?

When data or the model is too large for one machine, or time-to-result matters enough to pay for more total GPU-hours. Data parallelism is common; model parallelism and sharded embeddings for very large models.

Q: How often should you retrain?

As often as the world changes and labels arrive: measure how fast accuracy decays on later weeks. Usually a schedule plus a drift trigger, with periodic full retrains if warm-starting.

Q: What checks must a new model pass before release?

Not worse than the champion on the main metric on the same recent data, no important slice regressing, calibration, latency and size within limits, and sanity checks on predictions.

Q: What is a model registry for?

Storing every version with its lineage, scores, stage and owner, so serving loads an approved version, rollback is a pointer change, and audits are possible.

Practise it: Design an ML Training Pipeline

Exercise: automate a training pipeline

What it is

A short build that uses this module's ideas. Write it yourself before looking anything up; the checks at the end tell you whether it works.

Why it matters

An idea you have only read fades in a week.

Building it once, even in fifty lines, is what makes it yours.
Interviewers can tell the difference between someone who has read about a loop and someone who has debugged one.

In the real world: This is close to what an ML engineer's first month at a company with a mature platform looks like.

Your task. Turn a notebook model into a pipeline that trains, compares with a champion and registers the winner.

Steps

Take any model you have trained in this course.
Move its settings into a configuration file and its steps into functions: load data at a version, validate, train, evaluate, compare, register.
Log every run to MLflow with parameters, per-slice metrics, data version and commit.
Write a gate: the candidate must beat the registered champion on the latest data, with no slice worse by more than a set margin.
Run it twice on different data snapshots and confirm only a passing model becomes the new champion.
Delete your local model files and rebuild the champion exactly from its registry record.

You are done when

A failing candidate is rejected with a clear reason.
The rebuilt champion gives identical predictions to the original.

Stretch. Package the pipeline in a container and schedule it with Airflow, Dagster or Prefect.