TPU vs GPU: Differences for AI Training and Serving

What a GPU is

A GPU is a processor built for parallel work. A modern GPU contains thousands of arithmetic logic units that run multiplications and additions at the same time. Neural networks are mostly matrix arithmetic, so this parallelism is why GPUs became the default accelerator for deep learning. Google's own TPU documentation says that on a typical deep learning training workload a GPU can deliver an order of magnitude higher throughput than a CPU.

A GPU is still a general-purpose processor. It has to run many kinds of software, so each calculation reads operands from registers or shared memory and writes results back. The software layer most GPU programs use is NVIDIA CUDA, which NVIDIA describes as the platform for accelerated computing and the foundation for GPU computing.

What a TPU is

A Tensor Processing Unit is an application-specific integrated circuit (ASIC) that Google designs to accelerate machine learning. Google describes it as a matrix processor for neural network workloads. It cannot run a word processor, but it handles large matrix operations quickly.

The core of a TPU is a systolic array. Multiply-accumulate units are wired directly to each other, so each partial result passes to the next unit instead of going back to memory. Google's architecture page says a matrix-multiply unit (MXU) is a 128 x 128 array of multiply-accumulators on TPU versions before v6e, and 256 x 256 on v6e and TPU7x. Inputs are bfloat16 and accumulation is in FP32. Chips carry on-chip high-bandwidth memory (HBM) and can be connected into slices.

Google lists its generations as TPU7x (Ironwood), v6e, v5p, v5e, v4, v3 and v2. Each version has its own chip and slice configurations, so check the page for the one you plan to rent.

Architecture differences

The biggest difference is what each chip is for. A GPU is built to run many kinds of parallel code. A TPU spends its silicon on matrix multiplication and relies on the compiler to map the whole model onto it.

Software: CUDA and PyTorch vs XLA, JAX and PyTorch/XLA

The software gap decides most switching projects. GPU code usually runs through CUDA on NVIDIA hardware, and most ML libraries were built on that path first. TPU code must go through XLA, a just-in-time compiler that takes the graph emitted by the framework and compiles the linear algebra, loss and gradient parts into TPU machine code. The rest of the program runs on the TPU host.

• Google names PyTorch and JAX as frameworks that run on TPUs. Plan for the XLA compile step in both.
• TensorFlow ops that are not on the Cloud TPU list keep some models on GPU.
• Custom PyTorch or JAX operations that must run inside the training loop are a common reason teams keep a model on GPU.
• Moving from CUDA-specific code to XLA is a code change, not a configuration change. Budget engineering time for it.

A practical way to measure the switch: take a checkpoint of your model, run the same training step and the same serving request on each target, and compare step time, memory use and the bill at the same batch size and precision. Keep the input pipeline identical, because a slow data loader can hide the accelerator difference entirely.

Training vs inference

Google's guidance says TPUs suit models dominated by matrix computation, with no custom operations in the main training loop, large effective batch sizes, and training runs that last weeks or months. It also lists ultra-large embeddings in ranking and recommendation workloads as a fit. It says TPUs are a poor fit for linear algebra with frequent branching, many element-wise operations, or workloads that need high-precision arithmetic. Its GPU guidance covers models with many custom PyTorch or JAX operations, models that use TensorFlow ops not on TPU, and medium-to-large models.

Serving is judged differently. The limits are latency per request, whether the model and its cache fit in accelerator memory, and the cost per request at your real traffic. AWS positions Inferentia2 (Inf2) for deep learning inference and the NVIDIA L4 based G6 for inference and graphics. Those are vendor positions, not measurements of your model, so test on your own traffic before you choose.

Availability and lock-in

TPUs are available only from Google Cloud, through Compute Engine, Google Kubernetes Engine and Google's managed AI platform. A TPU workload that moves to another cloud has to move its hardware target as well as its compute. NVIDIA GPUs run on most major clouds and in your own data center, so the hardware is portable even though the CUDA stack still ties you to NVIDIA. Trainium and Inferentia are AWS chips and run only on AWS.

In practice the lock-in sits in the compiler and runtime more than the chip: XLA on TPU, CUDA on NVIDIA, and the AWS Neuron SDK on Trainium and Inferentia. Each one is a toolchain you build on.

If you expect to move between clouds, or to burst to another environment at peak, a GPU-based design keeps the hardware target fixed. A TPU workload stays tied to Google Cloud for as long as it runs there, so decide that before the first long training run.

Cost model: compare cost to finish, not hourly price

A cheaper hourly rate can cost more if the job takes longer or needs weeks of porting. Compare the total cost to finish your training run or to serve your traffic for a month. This article quotes no prices, because prices change by region, commitment and spot or on-demand choice. Check the provider pricing pages for the exact configuration you would use.

• Hourly price of the instance or slice in the region and purchase option you would really use.
• Utilization. An accelerator waiting on data bills the same as one doing work.
• Time to finish, measured on your model, with the same batch size and precision on each side.
• Engineering cost of porting to XLA or Neuron, and of debugging the port.
• Storage, data transfer and the cluster network around the accelerator.
• Capacity. Accelerator instances can be scarce in some regions, and reservations or capacity blocks change the price model.

Side by side: which one fits which job

Use this table as a starting point. The decision rests on your model, your framework and the cloud you already pay for.

Decision checklist

• Is your code PyTorch or JAX on NVIDIA today, with custom CUDA kernels? Stay on GPU.
• Do you already run on Google Cloud, with a matrix-heavy model and no custom ops in the training loop? Benchmark a TPU run against your GPU run.
• Do you run on AWS and want a lower-cost option? Check that your operations are supported by Neuron, then test Trainium for training or Inferentia for inference.
• Do you need portability across clouds or on-premises? Choose GPU.
• Do you need to serve an LLM without managing accelerators at all? Consider a managed model API before you buy any accelerator.
• Did you measure time to finish and total cost on your own model? If not, you have an opinion, not a decision.

On AWS: GPUs, Trainium and Inferentia

AWS does not offer Google TPUs. Its AI accelerators on EC2 are NVIDIA GPUs and its own chips. The EC2 pages name these families:

• P5 instances, powered by NVIDIA H100 GPUs, and P5e and P5en, powered by NVIDIA H200 GPUs, for deep learning and HPC applications.
• G6 instances, powered by NVIDIA L4 GPUs, for deep learning inference and graphics.
• Trn2 instances, powered by 16 AWS Trainium2 chips, for training and deploying generative AI models, with native PyTorch and JAX support.
• Inf2 instances, powered by AWS Inferentia2 chips, for deep learning inference, including large language models.
• AWS Neuron SDK, which compiles and runs models on Trainium and Inferentia. It supports PyTorch and JAX, and vLLM Neuron for serving LLMs.

Neuron changes the toolchain, so check that your model and its custom operations are supported before you plan a move. The price and performance case for Trainium or Inferentia has to be measured on your workload, not taken from the marketing pages.

If your team already runs PyTorch on NVIDIA GPU instances, the lowest-change first step on AWS is to keep the GPU instances and measure cost per run. Port to Trainium or Inferentia only when the measured saving covers the engineering time.

Where SeaGit fits

SeaGit deploys apps to Amazon EKS in your own AWS account. It runs on AWS only and does not offer TPUs or Google Cloud and Azure targets.

What the product code shows about accelerators: the node group form offers instance types from the P3, P4d, P5, G5, G4dn, Trn1, Inf1 and Inf2 families, with labels for GPU or Trainium or Inferentia hardware. Some of these are marked as having limited spot capacity. The RAG stack and the LLM gateway call Amazon Bedrock models, so their model inference does not run on a node you manage.

What to check before you schedule GPU pods

The node group picker chooses the instance. I found no NVIDIA driver, GPU AMI, Neuron device plugin or GPU scheduling setup in the SeaGit platform code. Confirm the driver, device plugin and container image for your accelerator in your own cluster before you rely on a GPU node.

Frequently asked questions