Choosing CPUs, GPUs and TPUs for ML

Designing scalable and cost-optimized training architectures means combining distributed orchestration with the right capacity and pricing model to balance training speed, model accuracy, and budget. On Google Cloud, Google Kubernetes Engine (GKE) (opens in a new tab) , Slurm, Cluster Director, Cluster Toolkit, and Ray form the orchestration layer, while committed use discounts (CUDs), Spot VMs, and Dynamic Workload Scheduler (DWS) modes control cost and availability. Training traffic flows from job queues and autoscalers to colocated GPU or TPU hosts, and the choice of machine type, reservation mode, and management model determines throughput, stability, and price.

1.1Optimize Availability and Cost for Training Capacity

Optimizing availability and cost for training capacity is the practice of matching training duration and fault tolerance to a pricing and provisioning model. For stable, long-running training with fixed or predictable capacity, use committed use discounts (CUDs) (opens in a new tab) for Compute Engine VMs and accelerators to lower cost. For fault-tolerant evaluation and experimentation, use Spot VMs on GKE nodes and Compute Engine VMs, because Spot VMs carry the deepest discount but can be preempted at any time.

To improve obtainability of high-demand accelerators, use DWS (opens in a new tab) to queue and provision GPUs and TPUs when capacity appears. For complex batch training that needs high-end GPUs for up to seven days, use DWS Flex-Start mode, where GKE persists the GPU request and automatically provisions Flex-Start VMs when capacity becomes available. For longer runs up to three months, use DWS Calendar mode to reserve specific GPUs such as H100 and H200 and TPUs such as Trillium for a future date and time. Flex-Start is billed pay as you go with up to a 53% discount for vCPUs, GPUs, TPUs, and attached Local SSD disks, and it requires no complex quota management.

1.2Compare Consumption Options for Clustered GPU Training

Consumption options for clustered GPUs define how dense multi-host capacity is requested, how long it lasts, and what discount applies. Clustered GPUs support dense allocation with colocation for multi-GPU training, and the request flows from Compute Engine, Cluster Director, Cluster Toolkit, or GKE to the available accelerator pool.

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 only when checkpoints and retries can absorb preemption.

1.3Orchestrate Distributed Training at Scale

Distributed training orchestration automates job queuing, resource allocation, colocation, and scaling across many accelerator hosts so developers do not manage each instance. Cluster Director (opens in a new tab) 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. Ray on Gemini Enterprise Agent Platform abstracts distributed infrastructure and lets applications request resources for training and serving without directly managing VMs and containers.

The level of control determines the deployment path. Highly managed deployment automates setup of compute, networking, and storage and reduces operational overhead:

• Cluster Director automates cluster setup and configuration of compute, networking, and storage to maximize performance and minimize downtime.
• 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.

For more granular control, create a Compute Engine cluster by using managed instance groups (MIGs) or by creating instances in bulk, then manually install required software. The orchestrator then governs daily operations: Slurm handles job queuing and resource allocation for high-performance computing, AI, or ML workloads; GKE handles container orchestration, autoscaling, and hardware-aware scheduling; and a bring-your-own orchestrator on Compute Engine leaves setup, maintenance, and updates to the operator.

1.4Select Accelerators, Machine Types, and Images for Training Nodes

Accelerator and image selection determines what training nodes can execute and how fast multi-host jobs communicate. Graphical processing units (GPUs) accelerate machine learning and data processing, and GKE offers nodes with NVIDIA H100, A100, and L4 GPUs plus clustered options including A3 Ultra H200, A4 B200, and A3 Edge machine types. All clustered GPU types except A4X Max and A4X support Flex-Start, Calendar mode, and Spot, while Future reservations in AI Hypercomputer also cover A3 Edge. TPUs including Trillium are obtainable through the same DWS Flex-Start and Calendar mechanisms with matching vCPU, GPU, TPU, and Local SSD discount treatment.

The operating system image depends on GPU infrastructure and deployment target. For GKE clusters, use the default Container-Optimized OS images for Standard and Autopilot modes, with Ubuntu available as an alternative in Standard mode, and Cluster Toolkit blueprints use container OS images. For Compute Engine and Slurm clusters, select an accelerator-optimized OS such as Rocky Linux 9 accelerator, Rocky Linux 8 accelerator, Ubuntu 24.04 LTS accelerator, or Ubuntu 22.04 LTS accelerator, where Cluster Toolkit extends the Ubuntu LTS Accelerator images. For a ready-to-use deep learning environment, use Deep Learning Software Layer (DLSL) container images that install NVIDIA CUDA, NCCL, and ML frameworks such as PyTorch and are tested for GKE clusters.

Analyzing workload characteristics for hardware selection means mapping the computational demands of a machine learning (ML) workload to the processor that executes it most efficiently. On Google Cloud, the choice centers on central processing units (CPUs), graphics processing units (GPUs), and Tensor Processing Units (TPUs) (opens in a new tab), Google's custom application-specific integrated circuits (ASICs) for ML. The optimal compute profile follows from four signals: the dominance of matrix computation in the model architecture, the effective batch size and memory footprint, the need for data-parallel scale, and the precision and shape constraints the code can satisfy.

2.1Compute Patterns by Model Architecture

Compute patterns by model architecture describe whether an accelerator can stay busy with the math it was built to accelerate. TPUs are purpose-built for large-scale matrix multiplication, which is the foundation of deep learning, and they efficiently train models by executing large matrix operations in hardware. Models dominated by matrix computations, models that train for weeks or months, and models with ultra-large embeddings common in advanced ranking and recommendation workloads fit TPUs well. By contrast, a program dominated by non-matrix operations such as add, reshape, or concatenate will likely not achieve high matrix unit (MXU) utilization, and linear algebra programs that require frequent branching or many element-wise operations are not suited to TPUs.

Efficient models are often co-designed for a specific platform, so a model tuned for one generation or vendor may not fully utilize another without configuration or software changes. Sparse models, for example, can improve quality and reduce computation by 3-10 times compared with dense models, while a distilled model such as DistilBERT can deliver similar performance to BERT with far less overhead. Frontier models spanning hundreds of billions to trillions of parameters, such as large language models (LLMs) like Gemini, and large models spanning tens to hundreds of billions of parameters therefore demand a hardware decision that considers both architecture and scale, not advertised peak specifications alone.

2.2Batch Size, Memory, and Data Parallelism

Batch size, memory, and data parallelism decide how much parallelism the hardware can exploit and how large a model can fit. TPUs provide on-chip high-bandwidth memory (HBM) that lets teams use larger models and larger batch sizes, and TPUs can be connected in groups called slices that scale workloads with little to no code changes. Pre-training foundation models illustrates the extreme: large clusters of accelerators continuously read large volumes of data, adjust weights through forward and backward passes, and run for weeks or even months. Fine-tuning the same large models uses smaller clusters and moderate data volumes to adapt weights to specific tasks, while inference splits by whether the load fits on one host or must be shared across multiple hosts.

The practical selection rule follows effective batch size and training duration:

• CPUs map to quick prototyping that needs maximum flexibility, simple models that train quickly, and small models with small effective batch sizes, especially when limited by host input/output (I/O) or networking bandwidth.
• GPUs map to medium-to-large models with larger effective batch sizes, especially when the 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 models with large effective batch sizes that train for weeks or months and contain no custom operations inside the main training loop.

Microbenchmarking of floating point operations per second (FLOPS), HBM, and networking bandwidth, combined with roofline analysis of operational intensity, shows what a platform can theoretically achieve before model benchmarking measures real tokens per second per chip across training and inference.

2.3Precision, Shapes, and Compilation Constraints

Precision, shapes, and compilation constraints determine whether code can compile and stay efficient on TPUs. Code that runs on TPUs must be compiled by the accelerated linear algebra (XLA) compiler, a just-in-time compiler that takes the graph emitted by the ML framework and compiles the 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 that contain 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 that reduces parameters from 32-bit floating-point to 8-bit integers can significantly decrease memory footprint and power consumption without a noticeable loss in accuracy.

TPU efficiency also depends on static, tile-friendly shapes because the MXU hardware 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 the graph 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 a 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 MXU 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; the op_profile tool can reveal how much padding occurred.

2.4Matching Training and Inference Profiles to Accelerators

The training or serving profile converts workload size and duration into a concrete accelerator and deployment shape. For pre-training foundation models on AI Hypercomputer (opens in a new tab), Google Cloud recommends A4X Max (NVIDIA GB300), A4X (NVIDIA GB200), A4 (NVIDIA B200), A3 Ultra (NVIDIA H200 141GB), A3 Mega (NVIDIA H100 80GB), or A3 High (NVIDIA H100 80GB) machine types deployed as orchestrated clusters through Google Kubernetes Engine (GKE) 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. For sustained multi-week runs, a reservation long enough to earn committed use discounts provides assurance of large-scale capacity.

The three processor families divide work by flexibility, operation support, and scale:

TPUs are available through Compute Engine, Google Kubernetes Engine, and Gemini Enterprise Agent Platform, and parameter-efficient fine-tuning such as low-rank adaptation (LoRA), which freezes base weights and trains only small new layers, further reduces the compute profile for adaptation work.

Comparing performance and cost metrics across Google Cloud accelerators means weighing training throughput and inference latency against memory capacity, interconnect scale, and cost per training hour. AI Hypercomputer (opens in a new tab) is the unified portfolio that packages performance-optimized infrastructure, open software, and flexible consumption models to maximize goodput, the measure of actual machine learning productivity. The choice centers on three families: high-end NVIDIA GPUs in A-series machines for pre-training and large-model work, cost-efficient G-series and N1 GPUs for inference and graphics, and Cloud Tensor Processing Units (TPUs) for large matrix operations with high-bandwidth memory.

3.1Accelerator Families and Machine Series for Training and Inference

Accelerator families on Google Cloud map directly to workload intensity, from large-scale pre-training to small-model serving. A-series machines including A4X Max, A4X, A4, A3 Ultra, A3 Mega, A3 High, and A2 (opens in a new tab) cover pre-training, while fine-tuning adds G4 and serving inference adds A3 Edge and G4 for lower-cost deployment. Graphics-intensive workloads use G4, G2, and N1 with T4, and high performance computing can use any accelerator-optimized series depending on how much computation must be offloaded to the GPU. 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.

The table shows that A100 and H100 class hardware justifies its higher hourly cost when model size, batch size, or distributed training demands exceed what L4, T4, or RTX PRO 6000 can serve efficiently. Choose A-series clustered capacity when training must scale across many accelerators as one system, and choose G-series or N1+T4 when inference latency targets can be met with smaller memory and lower hourly rates.

3.2Memory, Interconnect, and Scaling Architecture

Memory, interconnect, and scaling architecture determines whether additional accelerators produce linear speedup or stall on communication and I/O. Cloud TPUs use hardware designed for large matrix operations and include on-chip high-bandwidth memory (HBM) to hold larger models and batch sizes, and TPUs connect in groups called slices that scale workloads with little to no code changes. Clustered GPUs provide automated hardware health monitoring and proactive node replacement for tightly coupled multi-GPU training, while General GPUs use standard node auto-repair to maintain pod-level uptime for serving fleets. Dynamic Workload Scheduler (DWS) (opens in a new tab) provides atomic provisioning of entire blocks of interconnected nodes for large-scale training, so either all requested nodes start together or none start.

Cluster Director builds on this foundation by deploying and managing accelerator groups with colocation and scheduling for multi-GPU training on A3 Ultra with H200 and A4 with B200, supporting both Google Kubernetes Engine (GKE) and Slurm clusters. Storage completes the data path with Google Cloud Managed Lustre, a parallel file system 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. Together, colocated compute, synchronized maintenance, atomic scheduling, and parallel storage keep accelerators spending time on productive computation rather than waiting for data or failed nodes.

3.3Inference Latency and Throughput Metrics

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 and the model observability dashboard for queries per second, token throughput, and first-token latency.

• Time to First Token (TTFT) measures time in milliseconds to generate the first token for a request.
• Time Per Output Token (TPOT) measures time per token as request latency minus TTFT divided by total output tokens minus one.
• Inter-token latency (ITL) measures time between two output token generations, aggregated as mean, median, and percentiles such as p90.
• Normalized Time Per Output Token (NTPOT) and request latency measure efficiency and total time, where NTPOT is request latency divided by total output tokens and request latency is end-to-end time to complete a request.
• Output tokens per second, input tokens per second, tokens per second, and requests per second measure sustained throughput, where output and input rates divide total tokens generated by elapsed seconds, total tokens per second combines both to compare prefill-heavy versus decode-heavy workloads, and requests per second can vary widely with context length.

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.

3.4Cost per Training Hour and Consumption Options

Cost per training hour on Google Cloud is the sum of attached GPUs plus predefined vCPU, memory, and bundled Local SSD, with only one discount type applying to a resource at any time. For stable 24x7 inference or training with fixed capacity, committed use discounts (CUDs) lower the hourly rate, while elastic or fault-tolerant work uses discounted preemptible or scheduled capacity. GKE Inference Gateway (opens in a new tab) and custom compute classes further control cost by routing compute-intensive generative AI tasks to reserved GPU and TPU instances and falling back to Spot VMs when reserved capacity is unavailable.

• On-demand VMs receive compute flexible CUDs only for eligible N1, G2, and G4 spend, where N1 with attached GPUs covers only vCPUs, memory, and Local SSD, while G2 and G4 cover vCPUs, memory, GPUs, and Local SSD.
• Spot VMs automatically receive Spot discounts and suit fault-tolerant evaluation and experimentation, with the tradeoff that Spot VMs can be preempted.
• Flex-start with DWS instances automatically receive DWS pricing discounts and support complex batch training on high-end GPUs for up to seven days with improved obtainability.
• Reservations enable resource-based CUDs for all accelerator-optimized series, including A3 and A4, but resource-based commitments for all GPUs and most Local SSD require attached reservations.

For longer horizons, DWS Calendar mode reserves specific GPUs such as H100 and H200 and TPUs such as Trillium for up to three months. A vLLM engine with GPU and TPU fungibility can shift dynamically between accelerator types to balance fluctuating capacity, performance, and cost.

Key Facts

• Dynamic Workload Scheduler Flex-Start supports batch training up to seven days with queued provisioning, while Calendar mode reserves H100, H200, and Trillium capacity up to 90 days or three months.
• Spot VMs provide the deepest discount between 61% and 90% but can be preempted at any time, so they fit only fault-tolerant training with checkpoints and retries.
• Tensor Processing Units (TPUs) require XLA-compiled graphs, static shapes, and no custom operations in the main training loop, with batch or feature dimensions sized to multiples of 128 to avoid zero-padding of the 128x128 MXU.
• A100 and H100 GPUs in A-series machines fit pre-training and large-model inference, while L4, T4, and RTX PRO 6000 in G-series and N1 machines fit cost-efficient medium and small model inference.
• Time to First Token measures milliseconds to the first token and Time Per Output Token measures streaming speed after that, separating prefill and weight-loading bottlenecks from decode bottlenecks.