Scaling prototypes into ML models Study Guide — Professional Machine Learning Engineer

Scaling Prototypes into ML Models

This domain covers the end-to-end workflow of taking a working prototype and operationalizing it as a production-ready model using Vertex AI custom training. It tests your ability to select the right infrastructure (containers, accelerators, distributed strategies), tune models systematically, handle data quality issues like class imbalance, and manage the training lifecycle through checkpointing and retraining. At 18% of the exam, this is one of the heaviest-weighted domains — expect 8–12 questions that mix conceptual understanding with scenario-based reasoning.

Custom Training Jobs: Containers, Python Packages, GPUs/TPUs

Core Concepts

Vertex AI custom training lets you run ML workloads on managed infrastructure. You define a training application using one of two mechanisms:

Custom training job anatomy (via CustomJob or CustomTrainingJob in Vertex AI SDK):

• workerPoolSpecs: defines machine type, accelerator, replica count, and container/package
• baseOutputDirectory: Cloud Storage path for model artifacts
• serviceAccount: controls GCS/BigQuery access during training
• Machine types: n1-standard-*, n1-highmem-*, n1-highcpu-*, a2-highgpu-* (A100), c2-*

GPU/TPU Selection

TPU-specific rules:

• TPUs require TensorFlow (or JAX); PyTorch support is limited/experimental on Vertex AI managed TPUs
• TPU workers are specified in a separate worker pool with tpuTfVersion
• TPU VMs use TPU pods for scaling; a TPU v3-8 has 8 cores

Common Traps

• Using Container Registry instead of Artifact Registry — Google deprecated Container Registry; exam expects Artifact Registry
• Thinking Python packages eliminate the need for a container — they still run inside a pre-built container
• Selecting TPUs for a PyTorch workload — TPUs on Vertex AI are TensorFlow/JAX optimized
• Forgetting that baseOutputDirectory must be a Cloud Storage URI (gs://), not a local path

Worked Scenario

Requirement: A team has a PyTorch training script with a custom C++ extension that must be compiled at build time. What training mechanism should they use?

Reasoning: Pre-built containers cannot accommodate custom compiled extensions. A Python package won't work because the C++ compilation happens at build time. Custom container pushed to Artifact Registry is the correct answer. The container's Dockerfile handles the compilation, and the container URI is referenced in workerPoolSpecs.

Distributed Training Strategies

Decision tree for choosing distributed training strategy in Vertex AI

TensorFlow Distribution Strategies

Key exam fact: Vertex AI automatically injects TF_CONFIG when you use MultiWorkerMirroredStrategy with multiple replicas. You do NOT need to set it manually.

PyTorch Distributed Training

• Use torch.distributed with DistributedDataParallel (DDP)
• Vertex AI sets MASTER_ADDR, MASTER_PORT, WORLD_SIZE, RANK environment variables automatically
• Specify replicaCount > 1 in workerPoolSpec
• Use torchrun (formerly torch.distributed.launch) as the entry point

Reduction Server

Reduction Server is a Vertex AI feature that accelerates all-reduce operations in multi-node training:

• Acts as a relay to reduce cross-machine gradient communication latency
• Specified as an additional worker pool with machine type n1-highcpu-*
• Beneficial for large-batch, bandwidth-sensitive multi-node GPU training
• Works with both TF MultiWorkerMirroredStrategy and PyTorch DDP

Common Traps

• Assuming you must manually configure TF_CONFIG — Vertex AI injects it automatically
• Using MirroredStrategy across multiple machines — it only works on a single machine
• Forgetting that ParameterServerStrategy requires separate worker pools for PS and workers
• Confusing Reduction Server (gradient communication optimizer) with a parameter server

Worked Scenario

Requirement: A team trains a large TF model on 16 machines, each with 8 GPUs. Training is slow due to gradient synchronization overhead. What should they add?

Reasoning: They're already using MultiWorkerMirroredStrategy. The bottleneck is all-reduce communication. Adding a Reduction Server worker pool reduces cross-node bandwidth bottlenecks. This is distinct from changing the distribution strategy.