11.1 Choosing CPUs, GPUs & TPUs for Training

Google Cloud Professional Machine Learning Engineer PMLE

Track your progress as you read

Finished sections get a check mark.

1 Chapter 1: Exam Overview & the Agent Platform Map

1.1 PMLE Exam Facts, Policies & the June 2026 Blueprint 1.2 From Vertex AI to Gemini Enterprise Agent Platform: The Product Name Map

2 Chapter 2: Low-Code ML with BigQuery ML

2.1 Choosing, Training & Predicting with BigQuery ML Models 2.2 Feature Engineering & Feature Selection in BigQuery ML 2.3 Fine-Tuning & Evaluating Gemini Models from BigQuery

3 Chapter 3: AutoML & Google Cloud AI APIs

3.1 Training Models with Agent Platform AutoML 3.2 Building with Document AI, Vision, Translation & Other AI APIs

4 Chapter 4: Foundation Models in Model Garden

4.1 Evaluating & Selecting Models in Model Garden 4.2 Building & Tuning Solutions with Gemini, Imagen, Veo & MaaS 4.3 Optimizing Gemini Applications for Cost, Latency & Availability

5 Chapter 5: Exploring, Preprocessing & Protecting Data

5.1 Organizing & Exploring Tabular, Text & Image Data 5.2 Choosing a Preprocessing Tool: BigQuery, Dataflow, Spark or Python 5.3 Creating & Consolidating Features in Agent Platform Feature Store 5.4 Protecting PII & Sensitive Data in ML Workflows

6 Chapter 6: Notebooks, Prototyping & Experiment Tracking

6.1 Workbench vs. Colab Enterprise: Secure, Collaborative Notebooks 6.2 Prototyping with PyTorch, scikit-learn, JAX & Model Garden in Notebooks 6.3 Experiment Environments, Tracking & Lineage with Experiments and ML Metadata

7 Chapter 7: Evaluating Predictive & Generative AI Solutions

7.1 Evaluation Metrics for Predictive Models 7.2 Evaluating Gen AI: Computation Metrics, Rubrics & LLM-as-a-Judge

8 Chapter 8: Choosing the Modeling Approach

8.1 Choosing the Model Type: ARIMA, Trees, DNNs & LLMs 8.2 Choosing the Product & Deployment Strategy 8.3 Modeling Techniques for Interpretability Requirements

9 Chapter 9: Training Data & Training SDKs

9.1 Organizing & Ingesting Training Data from Cloud Storage & BigQuery 9.2 Custom Training on Agent Platform: Containers, Jobs & Worker Pools 9.3 Kubeflow on GKE, AutoML & Tabular Workflows as Training Choices

10 Chapter 10: Troubleshooting, Hyperparameter Tuning & Fine-Tuning

10.1 Troubleshooting ML Training Failures 10.2 Hyperparameter Tuning on Agent Platform 10.3 Fine-Tuning Foundation Models & When to Tune

11 Chapter 11: Training Hardware & Distributed Training

11.1 Choosing CPUs, GPUs & TPUs for Training 11.2 Distributed Training: Data & Model Parallelism

12 Chapter 12: Deploying Models for Inference

12.1 Batch vs. Online Inference on Agent Platform 12.2 Serving Beyond Agent Platform Endpoints: Model Garden, Cloud Run & GKE 12.3 Prebuilt & Custom Serving Containers

13 Chapter 13: Model Versions, Rollouts & Inference Logic

13.1 Organizing & Versioning Models in Model Registry 13.2 Rollout Strategies: A/B Tests, Canaries & Traffic Splits 13.3 Inference Preprocessing & Postprocessing

14 Chapter 14: Scaling Online Serving

14.1 Serving Features Online with Feature Store 14.2 Public, Dedicated & Private Endpoints 14.3 Serving Hardware & Autoscaling for Throughput 14.4 Tuning Models for Production Training & Serving

15 Chapter 15: End-to-End ML Pipelines

15.1 Validating Data & Models in Pipelines 15.2 Agent Platform Pipelines with KFP: Components, Templates & Caching 15.3 Ensuring Consistent Preprocessing Between Training & Serving

16 Chapter 16: Orchestration with Managed Airflow & Ray

16.1 Orchestrating ML Workflows with Managed Service for Apache Airflow 16.2 Ray on Agent Platform for Distributed ML Workloads

17 Chapter 17: Automating Retraining & CI/CD/CT

17.1 Choosing a Retraining Policy 17.2 CI/CD/CT for ML with Cloud Build

18 Chapter 18: Securing AI Systems, Responsible AI & Explainability

18.1 Protecting AI Systems from Data & Model Exfiltration 18.2 Defending Gen AI Apps: Safety Filters, Model Armor & Regex Guardrails 18.3 Responsible AI: Detecting & Monitoring Bias 18.4 Model Explainability on Agent Platform

19 Chapter 19: Monitoring AI Solutions

19.1 Configuring Model Monitoring on Agent Platform 19.2 Training-Serving Skew, Data Drift, Concept Drift & Attribution Drift 19.3 Monitoring, Testing & Evaluating Gen AI Solutions in Production

11.1 Choosing CPUs, GPUs & TPUs for Training

Key Takeaways Explain This Topic AI

Google recommends CPUs for quick prototyping, simple models, and models with many custom C++ operations or I/O limits.

GPUs suit medium-to-large models with larger batch sizes and models with many custom PyTorch or JAX operations that must partly run on CPUs.

TPUs suit large matrix-heavy models with large batches, weeks-long training, and ultra-large embeddings, but not models with custom operations in the main training loop.

Code that runs on TPUs is compiled by the XLA compiler, and TPU v6e no longer supports TensorFlow.

Agent Platform serverless training supports NVIDIA GPUs from T4 and L4 through A100, H100, H200, B200, and GB200, plus TPU VMs.

Last updated: September 2026

Section 3.3 of the exam guide begins with evaluating compute and accelerator options (for example, CPU, GPU, and TPU) . The right answer depends on model type, framework, custom operations, model and batch size, and how long training runs.

Google's Guidance: CPU vs. GPU vs. TPU

Hardware Best for Not ideal for

CPU Quick prototyping. Simple models that train fast. Small models with small batches. Models with many custom C++ operations . Workloads limited by I/O or host networking. Classical ML (scikit-learn, many XGBoost jobs) Large deep networks

GPU Medium-to-large deep learning models with larger batches. Models with many custom PyTorch or JAX operations that must run at least partly on CPUs. TensorFlow ops not available on TPUs Tiny models, where host overhead dominates and GPUs sit idle

TPU Models dominated by matrix computations . No custom operations inside the main training loop. Training that runs weeks or months . Large models with large batches . Ultra-large embeddings in ranking and recommendation Frequent branching, many element-wise operations, high-precision arithmetic, custom ops in the training loop

Tensor Processing Units (TPUs) are Google-designed ASICs with on-chip high-bandwidth memory. They connect into slices that scale with little code change. TPU code is compiled by XLA (Accelerated Linear Algebra). JAX, PyTorch/XLA, and TensorFlow can target TPUs, and starting with TPU v6e, TensorFlow is no longer supported . On Agent Platform, TPU training runs on TPU VMs (for example, v2, v3, v5e, and v6e generations).

GPU Options on Agent Platform Training

• GPU Memory Typical training role
• NVIDIA T4 16 GB Low-cost small models, experimentation

NVIDIA L4 24 GB Cost-efficient small-to-medium models and fine-tuning small LLMs with parameter-efficient methods

NVIDIA A100 40 GB or 80 GB Mainstream large deep learning training

NVIDIA H100 80 GB High-performance large-model training. Mega variants add high-bandwidth GPU networking (GPUDirect-TCPXO)

• NVIDIA H200 141 GB Very large models needing more memory per GPU. Includes GPUDirect-RDMA
• NVIDIA B200 / GB200 Latest generation Frontier-scale training. Includes GPUDirect-RDMA

Older GPUs (P4, P100, V100) are also listed for training. Each accelerator is available only in certain regions and machine families, such as a2 for A100, a3 for H100/H200, a4 for B200, and g2 for L4.

Sizing Memory

Training memory holds weights + gradients + optimizer states + activations . With the Adam optimizer in mixed precision, a rough rule is about 16 bytes per parameter before activations:

• Model size Approximate weight/gradient/optimizer memory Implication
• 350 M parameters ~5.6 GB Fits on one T4 or L4 with room for activations

7 B parameters ~112 GB Doesn't fit on one 80 GB GPU for full fine-tuning. Use sharding (FSDP/ZeRO) across GPUs, or parameter-efficient tuning (LoRA or QLoRA)

70 B parameters ~1.1 TB Needs multi-node model parallelism and sharding

Ways to reduce memory before buying bigger hardware:

• Mixed precision (bfloat16 on A100, H100, and TPUs, or float16 with loss scaling)
• Gradient checkpointing (recompute activations to save memory)
• Gradient accumulation (smaller micro-batches with the same effective batch)
• Parameter-efficient fine-tuning (LoRA, QLoRA)
• Sharded optimizer states (Section 11.2)
• Matching Hardware to Scenarios

Scenario Recommended hardware Reason

XGBoost churn model on 5 GB of tabular data High-memory CPU machine (or GPU if using XGBoost GPU training) Tree methods are fast on CPU. GPUs help mainly at larger scale

Fine-tuning a vision transformer on 2 million images in PyTorch with custom augmentation ops A100 or H100 GPUs Large batches plus custom PyTorch ops

Training a large recommendation model with huge embedding tables in JAX, for weeks TPU slice Matrix-heavy with ultra-large embeddings and long training

• Prototype of a new loss function on a sample CPU or a single small GPU Flexibility and low cost matter most
• LoRA tuning of a 7B open model on a tight budget L4 or A100 with QLoRA Memory-efficient tuning on cheaper GPUs
• Cost and Capacity Strategies

• Strategy Effect
• Spot VMs Large discounts for fault-tolerant jobs with checkpoints. Not supported with TPU Pods

Dynamic Workload Scheduler (flex-start) Waits for GPUs (L4, A100, H100, H200, B200) to become available, then starts all nodes together

• Reservations / Managed Training Clusters Guaranteed capacity for critical or very large jobs
• Right-sizing Measure GPU utilization. If it's low, fix the input pipeline or use fewer or smaller accelerators

Faster hardware is often cheaper overall. A job that costs 3× more per hour but finishes 5× faster costs less in total. Compare cost per completed training run , not hourly price.

Checking Utilization Before Scaling Up

Before choosing bigger or more accelerators, look at how the current ones are used:

Observation Meaning Action

GPU utilization low, CPU high Input pipeline or preprocessing bottleneck Optimize data loading (Chapter 9), then reassess

GPU memory near full, utilization high Hardware is the limit Mixed precision, then more or bigger accelerators

GPU utilization spiky in multi-GPU jobs Communication or synchronization waits Better interconnect, Reduction Server, larger per-device batch

• Many short CPU-bound trials Accelerators unnecessary CPU machines or BigQuery ML
• Exam Traps

Choosing TPUs for a PyTorch model full of custom CUDA operations in the training loop. GPUs are the better fit.

Choosing GPUs for a small scikit-learn model. CPUs are simpler and cheaper.

Recommending TensorFlow on the newest TPU generation where it isn't supported.

Adding accelerators when utilization shows the input pipeline is the bottleneck.