GCP Professional ML Engineer Practice Question 14: Accelerator Choice Custom Ops
Google Cloud Professional Machine Learning Engineer · Test 2 · question 4 of 10
A video-understanding model is being moved from a research laptop onto Agent Platform custom training. Its main training loop calls two hand-written CUDA kernels for a custom optical-flow layer, and part of each step must fall back to the host CPU. A colleague suggests TPU slices because the model is large and trains for days.
Which accelerator configuration should you choose for the custom training job, and why?
Show answer and explanation
Answer:
Google's own hardware guidance lists GPUs for "models with a significant number of custom PyTorch/JAX operations that must run at least partially on CPUs" and for models using ops that are not available on Cloud TPU. The same guidance states that Cloud TPUs are not suited to neural network workloads that contain custom operations in the main training loop, because the XLA compiler has to lower the whole graph to TPU machine code. Hand-written CUDA kernels with a host fallback are precisely that case, so GPUs are the right choice regardless of how long the model trains (option A).
- B is wrong because size and duration alone do not decide the accelerator. Custom ops in the main training loop are a disqualifier for TPUs, and XLA does not compile arbitrary hand-written kernels for TPU hardware. - C is wrong because only part of each step falls back to the host. The convolutional bulk of the model is still heavily parallel work that a GPU accelerates enormously over CPUs. - D is wrong because a worker pool is configured with one accelerator type. You cannot split a single model's layers across TPU and GPU accelerators attached to the same machine spec.
Reference: Introduction to Cloud TPU