Study guide
11.2 Distributed Training: Data & Model Parallelism
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.2 Distributed Training: Data & Model Parallelism
Key Takeaways Explain This Topic AI
Data parallelism copies the full model to each device, splits each batch across devices, and synchronizes gradients, usually with all-reduce.
Model parallelism splits a model that doesn't fit on one device, either within layers (tensor parallelism) or across layers (pipeline parallelism).
Sharded data parallelism such as FSDP or ZeRO divides parameters, gradients, and optimizer states across workers to train models larger than one device's memory.
Reduction Server is an all-reduce algorithm on Agent Platform that speeds multi-node GPU data-parallel training using NCCL; it runs in worker pool 2.
When scaling the global batch size across more devices, the learning rate usually needs adjusting, often with warmup, to keep training stable.
Last updated: September 2026
The exam guide asks you to understand the options for distributed training on GPUs and TPUs using data and model parallelism strategies . The core question is always: does the model fit on one device? If it does and you need speed, use data parallelism . If it doesn't, add model parallelism or sharding .
Data Parallelism
Every device holds a full copy of the model . Each global batch is split into per-device micro-batches, each device computes gradients, and the gradients are combined before the weights update.
Variant How gradients combine Pros Cons
Synchronous all-reduce All workers average gradients each step (ring or tree all-reduce over NCCL) Consistent model and predictable convergence The slowest worker sets the pace. Communication overhead grows with model size
Asynchronous parameter server Workers push gradients to parameter servers that hold the weights. Workers don't wait for each other Tolerates slow or preempted workers. Scales to many CPU workers Stale gradients can hurt convergence
Framework strategies
- Framework Single machine, multiple GPUs Multiple machines
- TensorFlow MirroredStrategy MultiWorkerMirroredStrategy (sync), ParameterServerStrategy (async), TPUStrategy
- PyTorch DistributedDataParallel (DDP) DDP across nodes (launched with torchrun ), FSDP
- JAX pmap / sharding APIs Multi-host sharding across TPU or GPU slices
- Scaling the batch
With N devices at the same per-device batch, the global batch grows N times . Larger batches usually need a higher learning rate (a common starting heuristic is to scale it proportionally) plus warmup , and sometimes a different optimizer (LAMB or LARS for very large batches). If accuracy drops as you add devices, check the learning rate schedule before blaming the hardware.
Model Parallelism
Technique What gets split When to use
Tensor (intra-layer) parallelism Individual weight matrices across devices Very wide layers, such as large transformer attention and MLP blocks. Needs fast interconnect
Pipeline (inter-layer) parallelism Consecutive layers placed on different devices, with micro-batches flowing through stages Very deep models. Watch for idle "bubbles" between stages
Sharded data parallelism (PyTorch FSDP, DeepSpeed ZeRO) Parameters, gradients, and optimizer states sharded across data-parallel workers Models too large for one GPU's memory, with simpler code than full tensor or pipeline parallelism
Expert parallelism Mixture-of-experts experts on different devices MoE architectures
Large-model training often combines them. For example: tensor parallelism within a node over fast GPU links, pipeline parallelism across nodes, and data parallelism across replicas.
- Distributed Training on Agent Platform
- Cluster structure
- Worker pool Role
- workerPoolSpecs[0] Primary replica (exactly 1)
- workerPoolSpecs[1] Workers
- workerPoolSpecs[2] Parameter servers or Reduction Server
- workerPoolSpecs[3] Evaluators
Code reads CLUSTER_SPEC (general) or TF_CONFIG (TensorFlow) to find its role and peers. Use the same container image in every pool to avoid version mismatches.
Reduction Server
Gradient communication between nodes can dominate step time. Reduction Server is an all-reduce algorithm that Agent Platform provides as a container image for worker pool 2 . It can raise throughput and cut latency for multi-node GPU data-parallel training. Requirements:
GPU workers, with TensorFlow or PyTorch configured for multi-host data-parallel training using NCCL all-reduce .
Primary and worker containers that support it: prebuilt TensorFlow 2.3+ or PyTorch 1.4+ training containers, or a custom container with NCCL 2.7+ and the google-reduction-server package.
Interconnects
Machine families for H100, H200, B200, and GB200 include high-bandwidth GPU networking ( GPUDirect-TCPXO on H100 Mega, GPUDirect-RDMA on H200, B200, and GB200). Tensor parallelism across nodes without fast interconnects scales poorly.
Other orchestration choices
Managed Training Clusters for reserved large-scale clusters, including Slurm-based workflows.
Ray on Agent Platform with Ray Train for Python-native distributed training.
Kubeflow Trainer on GKE for Kubernetes-native jobs (Chapter 9).
Fault Tolerance at Scale
The more machines a job uses, the more likely one fails during a long run.
Checkpoint frequently to Cloud Storage, and make restarts resume from the latest checkpoint.
Synchronous jobs stop when any worker fails, so enable automatic restarts and make startup idempotent.
Asynchronous parameter-server training tolerates losing individual workers better, at the cost of gradient staleness.
For multi-week runs, prefer reserved capacity (Managed Training Clusters or reservations) over Spot VMs.
Choosing a Strategy
Situation Strategy
Model fits on one GPU, and training on one GPU takes too long Synchronous data parallelism (DDP or MirroredStrategy), multiple GPUs in one node first
Scaling across many GPU nodes, with communication as the bottleneck Multi-node data parallelism + Reduction Server
- Many CPU workers, some preemptible, sparse models Asynchronous parameter server strategy
- Model doesn't fit on one GPU (for example, full fine-tuning at 7-13B) FSDP/ZeRO sharding
Very large transformer (tens to hundreds of billions of parameters) Tensor + pipeline + data parallelism on fast-interconnect clusters or TPU slices
- Large matrix-heavy model on TPUs TPU slice with data parallelism, and model sharding via JAX or PyTorch/XLA
- Worked Scenario
A company trains a 1-billion-parameter vision model. One A100 takes 10 days. The goal is under 2 days.
The model fits in 80 GB, so start with data parallelism : 8 A100s in one node with DDP, which is roughly 7× faster after communication overhead.
Scale the learning rate with the larger global batch, and add warmup.
To go further, use 2 nodes × 8 GPUs, and add Reduction Server in worker pool 2 to cut the gradient-synchronization bottleneck.
Checkpoint to Cloud Storage and use Dynamic Workload Scheduler so all 16 GPUs start together.