arXiv:2605.25645cs.DCcs.AI2026-05

首次实现在TPU上微调和部署Gemma 4 31B模型,性能优于同成本GPU。

Fine-Tuning and Serving Gemma 4 31B on Google Cloud TPU: A Technical Comparison with GPU Baselines

论文配图:Fine-Tuning and Serving Gemma 4 31B on Google Cloud TPU: A Technical Comparison with GPU Baselines
图 1 · 摘自论文原文
  • 用LoRA在TPU上实现模型微调,适配了JAX与Tunix/Qwix框架
  • 训练速度比2xH100快1.61倍,成本低2.12倍;长文本推理吞吐高66%
  • 适合想在TPU上部署大模型的工程师,提供完整代码迁移方案

本文首次在TPU硬件上完整演示了Google Gemma 4 31B模型的微调与服务,对比了TPU与GPU平台在大语言模型适配中的表现。采用LoRA方法,在Google TPU v5p-8上训练,v6e-8(Trillium)上推理,将原基于PyTorch、HuggingFace TRL和FSDP的GPU训练流程,迁移到JAX + Tunix/Qwix栈,涉及网格配置、LoRA命名、分片标注、梯度检查点、数据流水线重构及自定义Orbax到safetensor合并流程。推理方面,构建vLLM-TPU Docker环境,评估从512到16k输入令牌的QPS范围下延迟与吞吐。与相同成本的2xH100 GPU基线相比,TPU训练快1.61倍,成本低2.12倍;长上下文(4096令牌)时,吞吐高66%,首次响应时间(TTFT)快23.6倍(61毫秒 vs 1,443毫秒,QPS=4)。本工作填补了开源工具链空白,为开发者提供在TPU上部署Gemma 4 Dense 31B的完整方案。

原文摘要 · Abstract (English)

We present the first end-to-end demonstration of fine-tuning and serving Google's Gemma 4 31B model on TPU hardware, providing an empirical comparison of TPU and GPU platforms for large language model adaptation. Using LoRA on a Google TPU v5p-8 for training and TPU v6e-8 (Trillium) for inference, we document the full set of code-level adaptations required to port a GPU-native training recipe - built on PyTorch, HuggingFace TRL, and FSDP - to the JAX + Tunix/Qwix stack. These adaptations span mesh configuration, LoRA module naming conventions, sharding annotation corrections, gradient checkpoint, data pipeline restructuring, and a custom Orbax-to-safetensor checkpoint merging procedure. For inference, we detail the vLLM-TPU Docker setup necessary to serve Gemma 4 on v6e-8 and characterize the resulting latency and throughput profile. Compared with a similar-costing 2xH100 GPU baseline under identical hyperparameters, TPU training completes 1.61x faster at 2.12x lower cost. For inference, we cover the vLLM-TPU Docker setup required to serve Gemma 4 on v6e-8 and explain the observed latency and throughput characteristics across a QPS sweep spanning 512 to 16k input tokens. Across both workloads we compare performance and cost against a 2xH100 GPU baseline running identical hyperparameters. The TPU completes training 1.61x faster at 2.12x lower cost. For inference, TPU v6e-8 matches GPU at short context (<=2048 tokens) and decisively outperforms at long context: 66% higher throughput and 23.6x faster TTFT at 4096-token inputs (61 ms vs 1,443 ms at QPS=4). Our work removes a critical gap in the open tooling ecosystem and provides practitioners with a recipe for Gemma 4 Dense 31B deployment on the TPU infrastructure.

TPU大模型部署微调推理优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。