Scaling prototypes into ML models · Topic 4 of 4
Training Hardware and Distributed Training
Choose CPUs, GPUs or TPUs for training, and use data parallelism, model parallelism and worker pools to train large models across many accelerators fast.
- 5 min
- Includes a diagram
- CPU for classical ML and small models; GPU for most deep learning; TPU for large matrix-heavy training.
- Data parallelism copies the model and splits the data; it speeds up training.
- Model parallelism splits the model; it is needed when the model will not fit.
What it is
Training hardware decides how fast and how cheaply a model trains. CPUs suit small and classical models, GPUs suit most deep learning, and TPUs are Google's accelerators for large matrix-heavy models. When one device is too slow or too small, distributed training spreads the work across many devices using data or model parallelism.
How it works
Comparing compute options
| Option | Best for | Watch out for |
|---|---|---|
| CPU | scikit-learn, XGBoost on moderate data, small neural networks, preprocessing | Slow for deep learning on large data |
| GPU (NVIDIA, several generations) | Most deep learning in PyTorch, TensorFlow or JAX; custom operations; flexible workloads | Cost per hour; memory per GPU limits model and batch size |
| TPU | Large models dominated by matrix multiplication, large batches, long training; TensorFlow, JAX and PyTorch through XLA | Custom operations that XLA cannot compile run poorly; best with static shapes |
Rules of thumb:
- Classical ML on tables: CPU first; XGBoost can also use GPUs for large datasets.
- Deep learning prototype or custom layers: GPU.
- Very large transformers or embeddings trained for days: consider TPU slices.
- Lower cost for fault-tolerant training: Spot VMs with checkpoints.
- Faster, lower-memory training on modern accelerators: mixed precision (bfloat16 on TPUs).
Data parallelism
Every device holds a full copy of the model and trains on a different slice of each batch. Gradients are combined so all copies stay in sync.
| Approach | How | Examples |
|---|---|---|
| Synchronous all-reduce | Devices exchange gradients each step | TensorFlow MirroredStrategy (one machine), MultiWorkerMirroredStrategy (many machines), PyTorch DistributedDataParallel |
| Parameter servers | Workers send gradients to parameter servers that hold the weights; can be asynchronous | TensorFlow ParameterServerStrategy |
| TPU | Data parallel across TPU cores | TensorFlow TPUStrategy, JAX sharding |
Data parallelism is the answer when the model fits on one device but training is too slow.
import tensorflow as tf
strategy = tf.distribute.MultiWorkerMirroredStrategy()
with strategy.scope():
model = build_model()
model.compile(optimizer="adam", loss="binary_crossentropy",
metrics=["AUC"])
model.fit(train_ds, epochs=10)Model parallelism
When the model does not fit in one device's memory, its parts are split across devices.
| Strategy | What is split |
|---|---|
| Tensor parallelism | Individual layers (large matrix multiplications) across devices |
| Pipeline parallelism | Groups of layers on different devices, with micro-batches flowing through |
| Sharded data parallelism (for example PyTorch FSDP) | Parameters, gradients and optimizer state sharded across data-parallel workers |
Large LLM training combines these. Fast interconnects matter: GPUs in one machine and TPU slices are linked much faster than separate machines over the network.
Distributed training on Agent Platform custom training
A custom training job defines worker pools:
| Pool | Role |
|---|---|
| First pool | The primary (chief) replica, which also coordinates and saves the model |
| Second pool | Additional workers |
| Third pool | Parameter servers, or Reduction Server replicas that speed up all-reduce for GPU training |
| Fourth pool | Evaluators |
The service sets the cluster information each replica needs (for TensorFlow, the TF_CONFIG variable), so the same code runs on every replica.
"Multi-node GPU training is limited by gradient exchange bandwidth" points to Reduction Server. "The model has more parameters than one accelerator's memory" points to model parallelism or sharding, not more workers.
A worked example
A team trains a 1-billion-parameter vision transformer on one GPU; an epoch takes 20 hours. The model fits in memory, so the engineer uses data parallelism: a custom training job with 4 machines of 8 GPUs, PyTorch DistributedDataParallel and a larger global batch with an adjusted learning rate.
Later they move to a much larger model that no longer fits on one GPU. They switch to FSDP to shard parameters and optimizer state, and compare cost against a TPU slice.
How the exam asks about this
- "Train XGBoost on 5 GB of tabular data at the lowest cost." CPU machines; accelerators are not needed.
- "A model fits on one GPU but training takes a week." Data parallelism across more GPUs, for example multi-worker mirrored or DistributedDataParallel.
- "A transformer is too large for a single accelerator's memory." Model parallelism (tensor or pipeline) or sharded training.
What to remember for the exam
- CPU for classical ML and small models; GPU for most deep learning; TPU for large matrix-heavy training.
- Data parallelism copies the model and splits the data; it speeds up training.
- Model parallelism splits the model; it is needed when the model will not fit.
- Worker pools: chief, workers, parameter servers or Reduction Server, evaluators.
- Spot VMs plus checkpoints cut cost; mixed precision cuts memory and time.
Adding workers with data parallelism does not help a model that is too big for one device's memory. Every replica still needs the whole model.
Common mistakes
Choosing TPUs for models with many custom operations. Operations XLA cannot compile run slowly or fail; GPUs are more flexible.
Scaling workers without changing the batch and learning rate. The global batch grows with workers; adjust the learning rate or training becomes unstable.
Paying for GPUs for tree models on small data. CPUs are cheaper and often just as fast.
Check yourself
What is the difference between data and model parallelism?
Data parallelism puts a full model copy on each device and splits the data; model parallelism splits the model itself across devices because it does not fit on one.
What does Reduction Server speed up?
The all-reduce gradient exchange between GPU workers in multi-node synchronous data-parallel training.
When are TPUs a good choice?
For large models dominated by matrix operations, with large batches and long training, written in TensorFlow, JAX or PyTorch through XLA.
Topic 4 of 4 in Scaling prototypes into ML models
Test yourself on Scaling prototypes into ML models
The full Google Cloud Professional Machine Learning Engineer question bank has 749 questions, with timed practice exams and an explanation for every answer.
This study guide is independent and is not affiliated with, endorsed by, or sponsored by Google Cloud. Vendor names, certification names and exam codes are trademarks of their respective owners and are used for identification only.