Generative AI Leader Fundamentals of Generative AI Practice Question
A team is training a custom foundation model using JAX on TPUs on Google Cloud. They encounter frequent Out of Memory (OOM) errors. Which action is most effective in resolving the OOM error?
⚠ Common exam trap
Google Cloud often tests the misconception that mixed precision (bfloat16) alone is sufficient to resolve OOM errors, when in fact for very large models the memory bottleneck is the model size itself, not just the precision, and model parallelism is required.
Answer choices
Why each option matters
Answer the question above first, then reveal the full breakdown to understand why each option is right or wrong.
Correct answer & explanation
✓
Enable model parallelism using GSPMD to distribute the model across TPU cores.
OOM errors when training large foundation models on TPUs often stem from the model exceeding the memory of a single TPU core. GSPMD (Generalized SPMD) enables automatic model parallelism, sharding the model's parameters, gradients, and optimizer states across multiple TPU cores, thereby reducing per-core memory pressure without altering the model architecture or precision.
Answer analysis
Option-by-option breakdown
For each option: why learners choose it and why it is or isn't the right answer here.
- ✗
Reduce the model size by decreasing the number of layers.
Why it's wrong here
Shrinking the model reduces quality; it may resolve OOM but is not optimal.
- ✗
Increase the batch size to maximize TPU utilization.
Why it's wrong here
Larger batch size increases memory usage, worsening OOM.
- ✗
Use mixed precision training (bfloat16) to reduce memory footprint.
Why it's wrong here
Mixed precision reduces memory but may not be sufficient for very large models causing OOM.
- ✓
Enable model parallelism using GSPMD to distribute the model across TPU cores.
Why this is correct
Model parallelism directly addresses memory constraints by partitioning the model.
Visual reference
Go deeper
Related to this question
About these practice questions
One of 683 original Generative AI Leader practice questions on Courseiva, each with a full explanation and wrong-answer analysis — not exam dumps or protected exam content. Learn why practice questions differ from exam dumps →
JA
Written by Johnson Ajibi, MSc IT Security
Senior Network & Security Engineer · founder of Courseiva
This Generative AI Leader 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 Generative AI Leader exam.