Prepare and test your skills
Prepare and test your skills
Evaluation of compute and accelerator options is the process of matching a machine learning workload to central processing units (CPUs), graphics processing units (GPUs), or Tensor Processing Units (TPUs) based on math patterns, batch size, cost, and operational constraints. On Google Cloud (opens in a new tab), training traffic flows from job queues and autoscalers to colocated accelerator hosts, and the choice of processor, machine series, orchestration system, and pricing model determines throughput, stability, and price per training hour.
Processor selection starts from whether the model is dominated by large matrix multiplications and how large an effective batch the hardware can exploit. Cloud Tensor Processing Units (TPUs) are Google's custom application-specific integrated circuits for machine learning that are purpose-built for large-scale matrix multiplication, while CPUs offer maximum flexibility for prototyping and GPUs offer broad operation coverage for medium-to-large models. Models that train for weeks or months, models dominated by matrix math, and models with ultra-large embeddings for ranking and recommendation fit TPUs well, while programs dominated by add, reshape, concatenate, frequent branching, or many element-wise operations achieve low matrix unit utilization and do not suit TPUs.
The practical rule follows batch size, duration, and operation support. CPUs map to quick prototyping, simple models that train quickly, and small models with small effective batch sizes, especially when limited by host input/output or networking bandwidth. GPUs (opens in a new tab) map to medium-to-large models with larger effective batch sizes, especially when code depends on TensorFlow operations not available on Cloud TPU or on custom PyTorch or JAX operations that must run at least partially on CPUs. TPUs map to large matrix-dominated models with large effective batch sizes that train for weeks or months and contain no custom operations inside the main training loop. Cloud TPU (opens in a new tab) provides on-chip high-bandwidth memory (HBM) for larger models and batches, and TPUs connect in groups called slices that scale data-parallel work with little to no code changes.
| Option | Best-fit workload pattern | Key constraint |
|---|---|---|
| CPUs | Prototyping and simple models, including models with many custom operations | Limited throughput for large matrix-heavy training |
| GPUs | Medium-to-large models needing broader operation coverage | Unavailable TensorFlow operations or custom PyTorch or JAX operations force partial CPU execution |
| TPUs | Matrix-dominated large models, long runs, and ultra-large embedding workloads | Requires compilable graphs, static shapes, and no custom operations in the main loop |
TPU efficiency depends on whether code can compile through the accelerated linear algebra (XLA) compiler and keep the matrix hardware full. The XLA compiler is a just-in-time compiler that takes the graph emitted by the machine learning framework and compiles linear algebra, loss, and gradient components into TPU machine code, while the rest of the program runs on the TPU host machine. Neural network workloads with custom operations in the main training loop are not suited to TPUs, and workloads that require high-precision arithmetic are also excluded. When lower precision is acceptable, quantization from 32-bit floating-point to 8-bit integers can significantly decrease memory footprint and power consumption without noticeable accuracy loss.
Shape rules follow directly from the hardware. The matrix unit (MXU) is a 128x128 systolic array, and the memory subsystem prefers dimensions that are multiples of 8. The XLA compiler tiles matrix multiplies into smaller blocks and compiles just in time for the first batch, so tensors with dynamic shapes are not well suited because recompiling on every shape change is too slow. When computation cannot fill an entire MXU, the compiler pads tensors with zeros, which underutilizes the core and increases on-chip memory that can lead to out-of-memory errors in the extreme case. To maximize use, either the batch size or a feature dimension must be a multiple of 128 or the compiler will pad it to 128, and ideally both batch and feature dimensions are multiples of 8.
Accelerator families on Google Cloud map directly to workload intensity, from foundation-model pre-training to small-model serving. Accelerator-optimized machines (opens in a new tab) in the A-series including A4X Max with NVIDIA GB300, A4X with NVIDIA GB200, A4 with NVIDIA B200, A3 Ultra with NVIDIA H200 141GB, A3 Mega and A3 High with NVIDIA H100 80GB, and A2 cover pre-training, while fine-tuning adds G4 and inference adds A3 Edge and G4 for lower-cost deployment. Graphics-intensive work uses G4, G2, and N1 with T4, and high performance computing can use any accelerator-optimized series depending on how much computation must be offloaded. Within GPUs, General GPUs are managed as independent compute units with asynchronous maintenance, while Clustered GPUs are managed as a single tightly coupled system with synchronized maintenance for multi-host training.
| Workload | High-end choice | Cost-efficient choice |
|---|---|---|
| Model fine-tuning | A100 and H100 GPUs in A4 and A3 series | G4 for smaller fine-tuning jobs |
| Large model inference | A100 and H100 GPUs for largest models | RTX PRO 6000 and L4 GPUs |
| Medium and small model inference | High-end GPUs typically not required | RTX PRO 6000, L4, and T4 GPUs in G4, G2, and N1 machines |
For pre-training on AI Hypercomputer, Google Cloud recommends A4X Max, A4X, A4, A3 Ultra, A3 Mega, or A3 High deployed as orchestrated clusters through Google Kubernetes Engine (GKE) (opens in a new tab) or Slurm with Cluster Director or Cluster Toolkit. Data flows from high-throughput storage into the accelerator cluster, gradients flow between accelerators during forward and backward passes, and checkpoints flow back to multi-regional or dual-region Cloud Storage buckets for redundancy.
Cost per training hour is the sum of attached GPUs plus predefined virtual CPUs (vCPUs), memory, and bundled Local SSD, with only one discount type applying to a resource at any time. For stable, long-running training with predictable capacity, committed use discounts lower the hourly rate, while elastic or fault-tolerant work uses discounted preemptible or scheduled capacity. The Dynamic Workload Scheduler (DWS) improves obtainability of high-demand accelerators by queuing GPU and TPU requests and provisioning them when capacity appears.
| Consumption option | Best for | Duration and discount |
|---|---|---|
| Future reservations in AI Hypercomputer | Pre-training foundation models needing extended stability | Extended period, dense allocation, up to 53% discount for vCPUs and GPUs |
| Future reservations in calendar mode | Pre-training or fine-tuning needing stability up to 90 days | Up to 90 days, dense allocation, up to 53% discount |
| Flex-Start | Short-lived, dense clusters without a fixed start date | Up to seven days, dense allocation, up to 53% discount on supported types |
| Spot VMs | Fault-tolerant, batch, or short-lived training | Immediate creation, deepest discount between 61% and 90%, but preemptible at any time |
Choose Future reservations or Calendar mode when training cannot tolerate interruption and needs guaranteed topology for weeks to months. Choose Flex-Start when batch training can wait for capacity and finish within seven days, and choose Spot VMs only when checkpoints and retries can absorb preemption. Compute flexible discounts cover GPUs only on G2 and G4, while A3, A4, and A2 GPU discounts require reservations with resource-based commitments. DWS Calendar mode reserves specific GPUs such as H100 and H200 and TPUs such as Trillium for up to three months.
Exam tip: Flex-Start lasts up to seven days with queued provisioning when capacity appears, while Calendar mode reserves H100, H200, and Trillium capacity for up to 90 days or three months.
Distributed training orchestration automates job queuing, resource allocation, colocation, and scaling across many accelerator hosts. Cluster Director deploys and manages accelerator groups with colocation and scheduling for multi-GPU training on A3 Ultra H200 and A4 B200, and Cluster Director supports GKE and Slurm clusters. Cluster Toolkit deploys common configurations from predefined blueprints, such as A4 machine types with Slurm, for GKE or Compute Engine, where blueprints can be modified to customize deployments. GKE provides autoscaling, high availability, container orchestration, specialized hardware support, and integration with the Google Cloud ecosystem.
The level of control determines the deployment path. Highly managed deployment automates setup of compute, networking, and storage to maximize performance and minimize downtime, while granular control uses Compute Engine managed instance groups or bulk-created instances with manually installed software. Slurm then handles job queuing and resource allocation for high-performance computing, GKE handles container orchestration and hardware-aware scheduling, and Ray abstracts distributed infrastructure so applications request resources without directly managing virtual machines and containers. DWS provides atomic provisioning of entire blocks of interconnected nodes, so either all requested nodes start together or none start. Storage completes the path with Google Cloud Managed Lustre that can attach to 10,000 or more Pods, Cloud Storage Rapid and Cloud Storage FUSE for high-throughput weight access, and Hyperdisk ML block storage that can attach to up to 2,500 Pods.
Inference performance metrics separate how fast the first result appears from how fast the full response streams and how many requests the system sustains. Latency indicators focus on token timing for generative models, while throughput indicators focus on aggregate token and request rates under load. These metrics flow from the model server, such as vLLM or Triton, into Cloud Monitoring for queries per second, token throughput, and first-token latency.
High TTFT points to prefill, weight-loading, or accelerator bottlenecks, while high TPOT or ITL points to decode bottlenecks during streaming. Choose higher-memory accelerators and parallel storage when TTFT dominates startup, and choose scaled-out GPU or TPU replicas with model-aware routing when tokens-per-second throughput must grow under concurrent load.
Understanding distributed training options means choosing how to replicate models, shard weights, and place computation across GPU and TPU workers so communication scales with model size. On Google Cloud, Cloud TPU (opens in a new tab) provides the systolic hardware substrate for this pattern, with single-host and multi-slice topologies acting as worker pools. Effective design matches worker topology to model size, keeps every replica fed with data, and separates high-bandwidth intra-slice work from data center network-limited inter-slice work.
The TPU execution foundation is a matrix processor specialized for neural-network multiply-and-accumulate operations, unlike general-purpose CPUs and GPUs. A CPU based on von Neumann architecture loads values from memory, performs one calculation, and stores the result back for every operation, which creates the von Neumann bottleneck where slow memory access limits throughput. A GPU increases throughput with 2,500 to 5,000 Arithmetic Logic Units that execute thousands of multiplications and additions at once, but each calculation must still read operands from registers or shared memory and store intermediate results. Cloud TPU v3 contains two systolic arrays of 128x128 units on a single processor, and as each multiplication executes the result passes directly to the next multiply-accumulator with no memory access during matrix multiplication.
| Processor | Execution model | Memory behavior per calculation |
|---|---|---|
| CPU | General-purpose von Neumann processor | Loads operands from memory and stores result back for every calculation |
| GPU | General-purpose processor with thousands of units for parallelism | Must access registers or shared memory for every calculation |
| Cloud TPU | Matrix processor with systolic arrays | Passes results directly between connected multiply-accumulators with no memory access during multiplication |
Data flow through the TPU worker is the fixed host-to-TPU-to-host path that moves training data from the host CPU through on-chip memory and compute and back to the host. Throughput depends on the host keeping the infeed queue full and draining the outfeed queue promptly, because the matrix unit depends on high-bandwidth memory staging and staging depends on the host.
Frameworks such as PyTorch (opens in a new tab) and JAX run workloads on TPUs through Compute Engine, Google Kubernetes Engine, and Vertex AI using this same path.
The TPU7x worker pool building block is a full-host, 4-chip virtual machine that scales from single-host to multi-slice topologies. Each TPU7x virtual machine contains 4 chips with 224 vCPUs, 960 GB of RAM, and access to 2 Non-Uniform Memory Access (NUMA) nodes, and all TPU7x slices use this 4-chip shape. The virtual machine boot disk is Hyperdisk Balanced by default, with Hyperdisk Balanced and Hyperdisk ML supported for additional attached storage. Because the unit of scale is the 4-chip virtual machine connected by high-bandwidth interconnects, adding workers means adding virtual machines or slices rather than subdividing a virtual machine.
Pod-scale interconnects determine whether distributed training scales linearly or stalls on communication. TPU slices connect many chips with reconfigurable high-speed links, and large slices use Inter-Chip Interconnect to form a mesh across hosts. A TPU v2 slice can scale to 512 chips interconnected in one fabric and is requested by accelerator type such as v2-8 for a single device with 8 cores or v2-128 for 128 TensorCores. TPU7x continues this design with high-bandwidth interconnects, high-bandwidth memory, and mixed-precision compute in FP8 and BF16 to raise throughput per chip. Training-optimized models on TPU7x are tuned for dual-chiplet architecture, interconnects, memory, and mixed precision without code changes, with recipes that configure data, tensor, pipeline, and context parallelism.
Capacity mode controls visibility, isolation, and responsibility for failures at large scale. Managed TPU capacity holds back part of the reservation so Google can automatically replace faulty machines and restart slices, while All Capacity mode exposes the full reservation with no holdbacks and makes the user responsible for managing failures and planned maintenance. All Capacity mode provides dedicated, physically colocated capacity with ultra-low-latency networking and topology-aware scheduling based on full views of health and utilization. Its physical topology has four levels where a cluster is a deployment unit of Pod multiples in one zone, and blocks, sub-blocks, and hosts map to physical components below the cluster.
Data parallelism scales training by replicating the same model on every worker and splitting the training data among the replicas. Data parallelism across TPU Multislice replicates the full model on each slice and shards only the data batch, synchronizing gradients with all-reduce over the data center network. Fully sharded data parallelism shards both the data and the weights and optimizer state, using all-gather for weights and reduce-scatter for gradients. In data parallelism, every step performs all-reduce over the data center network to synchronize gradients, using 32 bits of bandwidth per parameter for two slices and up to 64 bits per parameter when more than two slices trigger a full-shuffle all-reduce. In fully sharded mode, each operation all-gathers the needed weights so each slice has the weights it needs, then reduce-scatters gradients as they are produced so each slice receives only the gradients for the weights it owns.
Fully sharded data parallelism therefore decreases weights and optimizer-state memory per slice as slice count grows, but it increases network traffic and the risk of blocking on a delayed collective. Both pure data parallelism and fully sharded data parallelism require scaling the global batch size linearly with the number of slices, and strong scaling requires the backward-pass time to exceed the all-reduce time. For Transformer language models, operational intensity equals 4 * B / 32 for two slices and 4 * B / 64 for many slices, where B is batch size in tokens, which implies a theoretical minimum of 176k to 352k tokens and a practical margin of at least 350k tokens per Pod for two Pods to 700k tokens for many Pods. Choose a training method (opens in a new tab) that keeps the host input pipeline feeding the infeed queue, because parallel workers stall when host memory cannot sustain staging.
Model parallelism partitions a large model that exceeds single-device memory across multiple accelerators so each device holds and computes only part of the model. On TPU Multislice, which connects multiple TPU slices over a data center network, model parallelism means deciding how weights, activations, and pipeline stages are placed within a slice versus across slices. Tensor parallelism shards individual layers or activations across model dimensions so multiple devices cooperate on the same operation, while pipeline parallelism places consecutive groups of layers on different slices and streams micro-batches through those stages in order.
Tensor parallelism across TPU slices is not recommended because sharding activations across model dimensions requires too much inter-slice bandwidth. The data center network between slices cannot sustain the fine-grained, high-frequency communication that tensor-parallel matrix multiplications need, so the workload becomes communication-bound. Tensor parallelism remains useful within a slice where interconnect bandwidth is high, but it should not be used as the cross-slice strategy. Pipeline parallelism moves forward activations forward and gradients backward across the data center network boundary from the first stage slice to the next stage slice and then to the final stage slice. Pipeline parallelism requires micro-batching to keep multiple stages busy, and micro-batching decreases effective batch size, arithmetic intensity, and ultimately model FLOPS utilization. Pipeline parallelism also incurs the pipeline bubble, where chips sit idle because they are waiting for data from an upstream or downstream stage.
Choosing a parallelism strategy is a batch-size decision driven by convergence limits and data center network costs. Pure data parallelism and fully sharded data parallelism both need very large batches for strong scaling, while pipeline parallelism allows pipeline slices to share one global batch instead of scaling it. Start with pure data parallelism, use fully sharded data parallelism to relieve memory pressure or grow parameters and activations, and use pipeline parallelism only when those approaches require too large a global batch.
| Strategy | Memory effect | Communication and batch tradeoff |
|---|---|---|
| Data parallelism | Each slice holds full weights and optimizer state | All-reduce each step; needs very large batch for strong scaling |
| Fully sharded data parallelism | Weights and optimizer state per slice shrink as slices grow | All-gather plus reduce-scatter; still needs linearly scaled batch but stores more parameters |
| Pipeline parallelism | Each slice holds only its stages | Stage-to-stage transfers plus pipeline bubble; shares one batch but lowers utilization through micro-batching |
Fully sharded data parallelism tends to achieve higher model FLOPS utilization, so pipeline parallelism is the better choice only when convergence per sample slows at the batch size needed for efficient sharded training. The open-source MaxText framework is a high-performance large language model training framework written in JAX and optimized for Cloud TPU, with official training recipes that include precision configurations, step-time benchmarks, and throughput metrics. Inference-optimized models use vLLM with the tpu-inference plugin, which provides hardware-optimized kernels and tensor parallelism with optimized PagedAttention kernels for serving throughput.
Exam tip: Do not use tensor parallelism across TPU slices; use pipeline parallelism across slices only when data parallelism or fully sharded data parallelism would require a global batch larger than the convergence limit.
CPUs fit quick prototyping, simple models that train quickly, and small models with small effective batch sizes. GPUs fit medium-to-large models with larger effective batch sizes that need broader operation coverage, especially when TensorFlow operations are unavailable on Cloud TPU or custom PyTorch or JAX operations must run at least partially on CPUs. Cloud Tensor Processing Units (TPUs) fit large matrix-dominated models with large effective batch sizes that train for weeks or months and contain no custom operations inside the main training loop.
Training must compile through the accelerated linear algebra (XLA) compiler with static shapes and no custom operations in the main training loop, and workloads requiring high-precision arithmetic are excluded. The matrix unit (MXU) is a 128x128 systolic array that pads tensors with zeros when computation cannot fill it, so either the batch size or a feature dimension must be a multiple of 128 and ideally both batch and feature dimensions are multiples of 8. When lower precision is acceptable, quantization from 32-bit floating-point to 8-bit integers can significantly decrease memory footprint and power consumption without noticeable accuracy loss.
Spot VMs provide the deepest discount between 61% and 90% but can be preempted at any time, so they fit only fault-tolerant, batch, or short-lived training with checkpoints and retries. Flex-Start supports short-lived, dense clusters up to seven days with queued provisioning when capacity appears for batch training that can wait. Calendar mode and Future reservations provide guaranteed stability and dense allocation for up to 90 days with up to 53% discount, including H100 and H200 GPUs and Trillium TPUs, for pre-training or fine-tuning that cannot tolerate interruption.
Start with pure data parallelism where each slice holds full weights and optimizer state, use fully sharded data parallelism to shrink weights and optimizer-state memory per slice as slices grow, and use pipeline parallelism only when those approaches would require a global batch larger than the convergence limit. Both data parallelism and fully sharded data parallelism require scaling the global batch size linearly with the number of slices, while pipeline parallelism lets pipeline slices share one global batch but lowers utilization through micro-batching and the pipeline bubble. Tensor parallelism across Tensor Processing Unit (TPU) slices is not recommended because the data center network cannot sustain its fine-grained communication, though it remains useful within a slice where interconnect bandwidth is high.
Last updated