Technical Summary: TPU vs. GPU for Gemma 4 31B LoRA SFT

Problem Statement

The fine-tuning and serving of large language models (LLMs) are increasingly constrained by compute costs and iteration speeds rather than data availability. While Google Cloud TPUs offer high theoretical performance, the ecosystem is often perceived as requiring significant engineering effort compared to the mature PyTorch-based GPU ecosystem. This report addresses the gap in open tooling by documenting the end-to-end process of porting a standard GPU-native training recipe (PyTorch, HuggingFace TRL, FSDP) to the JAX-based TPU stack (Tunix, Qwix, Orbax) for Gemma 4 31B. The study specifically targets LoRA Supervised Fine-Tuning (SFT) for Verilog code generation and subsequent inference serving, providing an empirical comparison between TPU and GPU hardware under identical hyperparameters.

Methodology

The authors executed a side-by-side comparison of fine-tuning and serving Gemma 4 31B on two distinct hardware configurations:

• Training: TPU v5p-8 (4 chips) vs. 2× NVIDIA H100 80GB GPUs.
• Inference: TPU v6e-8 (Trillium, 8 chips) vs. 2× NVIDIA H100 80GB GPUs.
• Dataset: CodeV-R1 (10k natural-language-to-Verilog pairs) for training; NVlabs/verilog-eval (156 spec-to-RTL problems) for evaluation.

Engineering Adaptations (GPU to TPU)

The core of the methodology involved a comprehensive porting effort to adapt the training and inference pipelines from PyTorch to JAX:

1. Framework & Mesh Configuration: Transitioned from PyTorch FSDP to JAX with explicit device mesh configuration (fsdp=1, tp=4 for v5p-8). The authors noted that tp=8 failed due to Gemma 4's attention head count not being divisible by 8.
2. LoRA Injection & Naming: Addressed discrepancies between HuggingFace (PyTorch) and Tunix (JAX) module naming. Notably, Tunix fuses K and V projections into a single kv_einsum tensor, requiring custom handling during checkpoint merging.
3. Sharding & Optimization: Implemented fixes for LoRA sharding annotations where Qwix-injected parameters inherited mismatched PartitionSpec ranks from base weights. The optimizer was switched to Optax AdamW, requiring custom state sharding logic.
4. Data Pipeline: Replaced HuggingFace DataLoader with Google's Grain library, implementing custom logic for loss masking (stripping system prompts and reasoning chains) and sequence handling.
5. Checkpointing & Conversion: Developed a custom pipeline to merge LoRA adapters from Orbax (JAX format) back into Safetensors (PyTorch/HF format) for inference, accounting for tensor shape transpositions and fused kv_einsum splits.
6. Inference Setup: Deployed vLLM-TPU on v6e-8 using a specific Docker image (vllm/vllm-tpu:gemma4), navigating constraints such as XLA compilation at startup, HBM allocation, and the disabling of chunked multi-modal inputs due to Gemma 4's heterogeneous attention dimensions.

Key Results

Training Performance

• Speed: The TPU v5p-8 completed training 1.61× faster (3.34 hours) than the 2×H100 baseline (5.39 hours).
• Throughput: TPU achieved 763 tokens/sec vs. 486 tokens/sec on GPU.
• Cost: Due to lower hourly rates ($16.80 vs. $22.12) and faster execution, TPU training was 2.12× cheaper ($56.11 vs. $119.23).
• Loss Convergence: Both models followed similar convergence trajectories. The absolute loss difference (0.072 for TPU vs. 2.258 for GPU) was attributed to different loss computation scopes (TPU used assistant-only loss; GPU used full-sequence loss).

Inference Performance

• Short Context (≤2048 tokens): GPU showed a marginal advantage in peak output throughput (+6%) and slightly better P99 tail latency, attributed to higher per-chip HBM bandwidth and the absence of XLA bucket padding overhead.
• Long Context (≥4096 tokens): TPU v6e-8 demonstrated dominance:
  • Throughput: +66% higher output tokens/sec at 4096 tokens.
  • Time-to-First-Token (TTFT): 23.6× faster at 4096 tokens (61 ms vs. 1,443 ms) and 7.1× faster at 8192 tokens.
  • Scalability: TPU sustained up to 32 concurrent requests (QPS) at medium context, compared to GPU saturation at ~8, enabled by fp8 KV cache (doubling effective capacity) and higher aggregate compute.
• Cost: For long-context workloads, TPU was 41% cheaper per output token.

Evaluation Quality

On the Verilog generation benchmark, the GPU-trained model achieved a slightly higher pass@1 (0.697 vs. 0.641) and pass@5 (0.814 vs. 0.795). The authors attribute this marginal difference to training objective variations (loss masking) and stochastic sampling variance rather than hardware capabilities.

Significance and Claims

The paper claims to provide the first end-to-end demonstration of fine-tuning and serving Gemma 4 31B on TPU hardware, effectively removing a critical gap in the open tooling ecosystem.

• Production Readiness: The work delivers a reproducible, production-ready recipe for deploying Gemma 4 on TPU infrastructure, including specific Docker configurations, mesh settings, and checkpoint conversion scripts.
• Total Cost of Ownership (TCO): The study concludes that for a representative train-plus-serve workload, the TPU configuration is 1.82× cheaper than the GPU baseline. This advantage is driven primarily by training cost savings, with inference costs compounding significantly for long-context applications.
• Hardware Suitability: The authors position TPU as the preferred infrastructure for long-context workloads and cost-sensitive production pipelines, while acknowledging GPU advantages in short-context throughput and ecosystem maturity (simpler prototyping).

The report emphasizes that while the engineering overhead of porting to JAX/Tunix is non-trivial (estimated at roughly one week for mesh configuration, sharding fixes, and pipeline restructuring), this is a one-time cost amortized across many training runs, making TPU a viable and superior alternative for large-scale LLM adaptation.