The answer is that the learning rate is too high for the batch size. When scaling hyperparameter scaling from GPU to TPU training on Vertex AI, the linear scaling rule dictates that the learning rate should increase proportionally with the batch size to maintain stable gradient updates. Here, the batch size jumped from 32 on a single GPU to 256 across 8 TPU cores, but the learning rate remained at 1e-4, causing the optimizer to overshoot minima and produce NaN loss. On the Google Professional Machine Learning Engineer exam, this scenario tests your understanding of distributed training dynamics and the common pitfall of forgetting to adjust hyperparameters when moving to TPUs. A frequent trap is assuming the same learning rate works for any batch size, but the rule is simple: double the batch size, double the learning rate (within reason). Memory tip: “Bigger batch, bigger step—scale your LR to keep the gradient in check.”
PMLE Scaling prototypes into ML models Practice Question
This PMLE practice question tests your understanding of scaling prototypes into ml models. Read the scenario carefully and evaluate each option against the stated constraints before committing to an answer. After answering, compare your reasoning against the explanation and wrong-answer breakdown below. Once you have made your selection, read the full explanation to reinforce the concept and understand why each distractor is designed to mislead on exam day.
A machine learning engineer is scaling a prototype natural language processing model that uses a transformer encoder. The prototype was trained on a small corpus on a single GPU. For production, they need to train on a much larger corpus using TPUs on Vertex AI. They convert the TensorFlow code to work with TPUStrategy. The training starts but after a few steps, the loss becomes NaN and training diverges. The learning rate scheduler uses a warm-up and then linear decay. The initial learning rate is 1e-4. The batch size per TPU core is 32, with 8 cores total (batch size 256). What is the most likely cause?
Clue words in this question
Noticing these words before you look at the options changes how you read each choice.
Clue: "most likely"
Why it matters: Probability qualifier — the question wants the most probable cause or outcome, not a guaranteed one. Eliminate low-probability options.
Answer the question above first, then reveal the full breakdown to understand why each option is right or wrong.
Correct answer & explanation
✓
The learning rate is too high for the batch size.
When scaling from a single GPU to 8 TPU cores, the global batch size increases from 32 to 256. The learning rate of 1e-4, which was appropriate for batch size 32, becomes too high for the larger batch size. This violates the linear scaling rule (learning rate should be scaled proportionally to batch size), causing gradient updates to overshoot minima and leading to NaN loss and divergence.
Key principle: Answer the scenario, not the keyword: identify the specific constraint before choosing the most familiar-sounding option.
Answer analysis
Option-by-option breakdown
For each option: why learners choose it and why it is or isn't the right answer here.
✗
The batch size is too small for TPU.
Why it's wrong here
TPUs can effectively use small per-core batch sizes.
✓
The learning rate is too high for the batch size.
Why this is correct
Larger batch size requires lower learning rate to maintain stability.
Clue confirmation
The clue word "most likely" in the question point toward this answer.
Related concept
Read the scenario before looking for a memorised answer.
✗
The learning rate schedule should be cosine instead of linear.
Why it's wrong here
Switching schedule does not address the fundamental scaling issue.
✗
The warm-up steps are insufficient.
Why it's wrong here
Warm-up helps but is not the most likely cause of immediate divergence.
Common exam traps
Common exam trap: answer the scenario, not the keyword
Google Cloud often tests the misconception that TPU-specific issues (like batch size or hardware compatibility) are the root cause, when in fact the problem is a fundamental hyperparameter scaling error that applies to any distributed training setup.
Detailed technical explanation
How to think about this question
The linear scaling rule states that when batch size is multiplied by k, the learning rate should also be multiplied by k to maintain the same gradient variance. Here, batch size increased 8x (32 to 256), so the learning rate should be around 8e-4, not 1e-4. However, using 1e-4 with a 256 batch size means the effective learning rate per sample is too high, causing gradient explosion. In practice, TPUs with bfloat16 precision are more sensitive to large gradients, making this mismatch especially problematic.
KKey Concepts to Remember
Read the scenario before looking for a memorised answer.
Find the constraint that changes the correct option.
Eliminate answers that are true in general but not in this case.
TExam Day Tips
→Watch for words such as best, first, most likely and least administrative effort.
→Review why wrong options are wrong, not only why the correct option is correct.
Key takeaway
Answer the scenario, not the keyword: identify the specific constraint before choosing the most familiar-sounding option.
Real-world example
How this comes up in practice
An e-commerce site experiences heavy traffic on Black Friday and near-zero traffic during off-peak weeks. Rather than provisioning permanent large VMs, the team uses auto-scaling groups that add capacity automatically under load and reduce it overnight. Questions like this test whether you understand elasticity, availability zones, and cloud compute scaling patterns.
What to study next
Got this wrong? Here's your next step.
Identify which exam domain this question belongs to, review the core concept, then practise similar questions from the same domain.
Scaling prototypes into ML models — This question tests Scaling prototypes into ML models — Read the scenario before looking for a memorised answer..
What is the correct answer to this question?
The correct answer is: The learning rate is too high for the batch size. — When scaling from a single GPU to 8 TPU cores, the global batch size increases from 32 to 256. The learning rate of 1e-4, which was appropriate for batch size 32, becomes too high for the larger batch size. This violates the linear scaling rule (learning rate should be scaled proportionally to batch size), causing gradient updates to overshoot minima and leading to NaN loss and divergence.
What should I do if I get this PMLE question wrong?
Identify which exam domain this question belongs to, review the core concept, then practise similar questions from the same domain.
Are there clue words in this question I should notice?
Yes — watch for: "most likely". Probability qualifier — the question wants the most probable cause or outcome, not a guaranteed one. Eliminate low-probability options.
What is the key concept behind this question?
Read the scenario before looking for a memorised answer.
About these practice questions
Courseiva creates original exam-style practice questions with explanations and wrong-answer analysis. It does not publish real exam questions, exam dumps, or protected exam content. Learn why practice questions differ from exam dumps →
Share a tip, memory trick, or ask about the reasoning behind this question. Do not post real exam questions, leaked content, braindumps, or copyrighted exam material. Comments are moderated and may be removed without notice.
This PMLE practice question is part of Courseiva's free Google Cloud certification practice question bank. Courseiva provides original exam-style practice questions with explanations, topic-based practice, mock exams, readiness tracking, and study analytics to help learners prepare for the PMLE exam.
Question Discussion
Share a tip, memory trick, or ask about the reasoning behind this question. Do not post real exam questions, leaked content, braindumps, or copyrighted exam material. Comments are moderated and may be removed without notice.
Sign in to join the discussion.