Professional Machine Learning Engineer · Domain 3: Scaling prototypes into ML models · 20% of the exam
Task 3.3: Choosing appropriate hardware for training.
Choosing CPUs, GPUs or TPUs for a training job and planning distributed training with data and model parallelism.
Study it
CPUs, GPUs and TPUs; data and model parallelism
Lesson coming
Sample questions
Try each one before opening the answer. Every option is explained, with the Google Cloud documentation page that proves it.
Question 1 · choose 1
A JAX ranking model for product recommendations has very large embedding tables, and embedding lookups dominate its step time. Under the team's sharding plan, each TPU chip must hold about 60 GiB of embedding shards and optimizer state. Which of these options should the ML engineer choose?
ATPU v5p, with 95 GiB of usable HBM and four SparseCores per chip
BTPU v6e, with two SparseCores and 32 GB of HBM per chip
CTPU v5e, with 16 GB of HBM per chip, chosen for its lower cost
DCPU machines with large host memory, which hold the whole embedding table
Show the answer and why
ATPU v5p, with 95 GiB of usable HBM and four SparseCores per chip
Correct
v5p has room for the shards on each chip, and its SparseCores accelerate the embedding lookups that recommendation models rely on.
BTPU v6e, with two SparseCores and 32 GB of HBM per chip
Incorrect
v6e has SparseCores, but 32 GB of HBM per chip cannot hold 60 GiB of shards.
CTPU v5e, with 16 GB of HBM per chip, chosen for its lower cost
Incorrect
16 GB per chip is far below what the sharding plan needs.
DCPU machines with large host memory, which hold the whole embedding table
Incorrect
CPUs suit small models or I/O-bound work, not a large model whose step time is dominated by embedding lookups.
SparseCores are dataflow processors that accelerate models using sparse operations, with recommendation models that rely heavily on embeddings as a primary use case. TPU v5p and TPU7x have four SparseCores per chip and v6e has two. Per-chip HBM differs too: v5p lists 95 GiB of usable HBM, v6e 32 GB and v5e 16 GB, so the memory each chip must hold also narrows the choice.
A JAX model trains on a TPU v5e slice with a batch size of 100 per step and dense layers of 300 units. The memory viewer shows heavy padding and the matrix units are poorly used. Accuracy is fine, and the team wants better throughput on the same hardware. What should the ML engineer change?
AMove the job to TPU v6e, whose larger matrix units speed up any model
BUse a total batch size that is a multiple of 64 and layer widths that are multiples of 128
CAdd more TPU chips to the slice and keep the same batch and layer shapes
DSwitch the training job to GPUs with the same model
Show the answer and why
AMove the job to TPU v6e, whose larger matrix units speed up any model
Incorrect
A newer TPU does not fix tensors that tile badly; the padding moves with the model, and the team wants to keep its hardware.
BUse a total batch size that is a multiple of 64 and layer widths that are multiples of 128
Correct
These sizes tile into the TPU's chunks, which reduces padding and keeps the matrix units busy.
CAdd more TPU chips to the slice and keep the same batch and layer shapes
Incorrect
More chips split the same badly shaped tensors, so the padding stays.
DSwitch the training job to GPUs with the same model
Incorrect
This changes the hardware the team wants to keep and does not fix the tensor shapes.
The XLA compiler pads tensors so they tile into the TPU's matrix units, which wastes compute and memory. The Cloud TPU performance guide recommends a total batch size that is a multiple of 64 (8 per core) with feature dimensions that are multiples of 128, or batch sizes that are multiples of 1,024 with features that are multiples of 8.
A PyTorch job does multi-host data-parallel training on 16 GPU nodes with NCCL all-reduce on Agent Platform serverless training. Profiling shows the GPUs waiting on gradient synchronization for a large share of each step, and the model fits easily in each GPU's memory. Which change most directly targets the bottleneck?
AAdd a Reduction Server worker pool to the distributed training job
BMove the training data to a Filestore share mounted over NFS
CMove the job to Spot VMs with the same machine types
DAdd an evaluator worker pool to run validation in parallel
Show the answer and why
AAdd a Reduction Server worker pool to the distributed training job
Correct
Reduction Server is an all-reduce algorithm for GPU data-parallel training with NCCL that increases throughput and reduces latency.
BMove the training data to a Filestore share mounted over NFS
Incorrect
Faster storage helps an input-bound job, but the profile shows the GPUs waiting on gradient synchronization, not on data.
CMove the job to Spot VMs with the same machine types
Incorrect
Spot VMs lower cost but do not change how gradients are synchronized.
DAdd an evaluator worker pool to run validation in parallel
Incorrect
An evaluator pool runs evaluation; it does not speed up gradient all-reduce between workers.
Reduction Server is an all-reduce algorithm that Agent Platform provides as a container image for one of the worker pools. It applies to distributed training with GPU workers whose code uses TensorFlow or PyTorch with NCCL all-reduce for multi-host data parallelism.
A team runs LoRA tuning of a mid-sized open model that fits on a single host and wants the lowest-cost GPU machines that can do it. Large multi-node training is not planned. Which Compute Engine machine series should the ML engineer consider first?
AA4, recommended for large-scale ML training
BA3 Ultra, a top-end A3 machine type with eight GPUs per host
CN1 general-purpose machines with a few older GPUs attached for the tuning run
DG4, positioned as a low-cost option for single-host tuning and inference
Show the answer and why
AA4, recommended for large-scale ML training
Incorrect
A4 targets large-scale training, more than a single-host tuning job needs.
BA3 Ultra, a top-end A3 machine type with eight GPUs per host
Incorrect
A3 machines are A series machine types, which cost more than this single-host tuning job needs.
CN1 general-purpose machines with a few older GPUs attached for the tuning run
Incorrect
N1 attaches older GPUs, some of them approaching end of support, and is not among the series the page recommends for LoRA-style fine-tuning.
DG4, positioned as a low-cost option for single-host tuning and inference
Correct
G4 machines offer a low-cost solution for single-host inference and model tuning compared with A series machine types.
Compute Engine's accelerator-optimized families include A4 (recommended for ML training at large scale), A3 (network-bound ML training and HPC) and G4 and G2. G4 machines provide a low-cost solution for single-host inference and model tuning compared with A series machine types.
A team wants to run its existing TensorFlow 2 training code on TPUs in Agent Platform serverless training and asks which TPU version to request. It does not plan to rewrite the code in another framework. What should the ML engineer advise?
ATPU v6e, the newest version, for the best performance
BAny TPU version, because all TPUs run TensorFlow the same way
CTPU v5e, which supports TensorFlow 2.15 and later
DGPUs only, because serverless training does not offer TPUs
Show the answer and why
ATPU v6e, the newest version, for the best performance
Incorrect
Starting with TPU v6e, TensorFlow is no longer supported, so this code would not run.
BAny TPU version, because all TPUs run TensorFlow the same way
Incorrect
Framework support differs by TPU version.
CTPU v5e, which supports TensorFlow 2.15 and later
Correct
TPU v5e supports JAX, TensorFlow 2.15+ and PyTorch 2.1+ in serverless training.
DGPUs only, because serverless training does not offer TPUs
Incorrect
Serverless training does offer TPU VMs, including v5e.
Agent Platform serverless training supports TPU VMs. TPU v5e supports JAX 0.4.6+, TensorFlow 2.15+ and PyTorch 2.1+, while starting with TPU v6e TensorFlow is no longer supported, so TensorFlow code that is not being rewritten should target v5e.
A team scales a JAX model from one TPU slice to four slices with Multislice. Inside each slice it already shards the model with fully sharded data parallelism. Which kind of parallelism should it run across the four slices?
AData parallelism across the slices, with the number of shards equal to four
BTensor parallelism across the slices, splitting every layer into four shards over the DCN
CNo parallelism across slices; each slice trains its own model
DPipeline parallelism coded by hand for the DCN links
Show the answer and why
AData parallelism across the slices, with the number of shards equal to four
Correct
Slices talk over the data center network, which has higher latency and lower throughput than ICI, so data parallelism goes there.
BTensor parallelism across the slices, splitting every layer into four shards over the DCN
Incorrect
Tensor parallelism needs heavy, frequent communication, which suits the fast ICI inside a slice, not the slower DCN.
CNo parallelism across slices; each slice trains its own model
Incorrect
Independent models do not combine into one trained model.
DPipeline parallelism coded by hand for the DCN links
Incorrect
Developers do not write inter-slice communication; XLA generates it, and the usual choice across slices is data parallelism.
In a Multislice configuration, chips within a slice communicate over the fast inter-chip interconnect (ICI), and slices communicate over the data center network (DCN), which has higher latency and lower throughput. A common setup uses FSDP or tensor parallelism within a slice and sets DCN data parallelism to the number of slices; the product of the DCN parallelism settings must equal the slice count.
A scientific team trains a neural network whose loss computation needs double-precision floating-point arithmetic throughout. It asks whether Cloud TPUs are a good fit. What should the ML engineer answer?
AYes; TPUs suit every neural network dominated by matrix math, whatever its arithmetic needs
BNo; TPUs are not suited to workloads that require high-precision arithmetic
CYes, as long as the job uses a Multislice configuration
DNo; TPUs cannot run neural networks written in JAX
Show the answer and why
AYes; TPUs suit every neural network dominated by matrix math, whatever its arithmetic needs
Incorrect
Matrix-heavy models suit TPUs only when they do not need high-precision arithmetic.
BNo; TPUs are not suited to workloads that require high-precision arithmetic
Correct
The Cloud TPU guidance lists high-precision arithmetic among the workloads TPUs are not suited to.
CYes, as long as the job uses a Multislice configuration
Incorrect
Multislice adds scale, not higher numeric precision.
DNo; TPUs cannot run neural networks written in JAX
Incorrect
TPUs run JAX models well; precision is the issue here.
Cloud TPUs are optimized for matrix-dominated models, but the guidance lists workloads they are not suited to, among them linear algebra with frequent branching or many element-wise operations, and workloads that require high-precision arithmetic.