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

Amit Singh; Jatin Kishnani; Mayank Goel; Pulkit Agrawal; Sairanjan Mishra

arxiv: 2605.25645 · v3 · pith:XIIRTNSMnew · submitted 2026-05-25 · 💻 cs.DC · cs.AI

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

Jatin Kishnani , Mayank Goel , Amit Singh , Pulkit Agrawal , Sairanjan Mishra This is my paper

Pith reviewed 2026-06-29 20:37 UTC · model grok-4.3

classification 💻 cs.DC cs.AI

keywords fine-tuningTPUGPU baselineinference servingLoRAthroughputcost comparisonmodel porting

0 comments

The pith

TPU hardware completes fine-tuning of the 31B model 1.61 times faster and at 2.12 times lower cost than a matched GPU baseline.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper establishes that a GPU-native fine-tuning recipe for a 31 billion parameter model can be ported to TPU hardware through targeted code changes in the JAX ecosystem. It then measures training and inference performance against a comparable GPU configuration under matched conditions. A reader would care because the reported outcomes indicate a concrete path to lower-cost model adaptation and serving when TPU resources are available. The work supplies the exact adaptations needed so others can replicate the setup rather than starting from scratch.

Core claim

By carrying out the full set of required changes to mesh layout, LoRA naming, sharding annotations, gradient checkpointing, data pipelines, and checkpoint merging, the TPU v5p-8 run finishes training 1.61 times faster at 2.12 times lower cost than the 2xH100 GPU baseline. For inference the TPU v6e-8 platform matches GPU results at short contexts and produces 66 percent higher throughput together with 23.6 times faster time-to-first-token at 4096-token inputs.

What carries the argument

The sequence of code-level adaptations that translate a PyTorch FSDP plus Hugging Face recipe into the JAX plus Tunix/Qwix stack while preserving identical hyperparameters and model behavior.

If this is right

Training time and monetary cost both drop when the adapted recipe runs on TPU rather than the GPU baseline.
Inference throughput rises 66 percent on TPU once input lengths reach 4096 tokens.
Time-to-first-token latency falls by more than an order of magnitude on TPU at the same long-context operating point.
The documented vLLM-TPU Docker configuration and checkpoint procedure make the model immediately servable on the target hardware.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

Hardware choice for inference workloads may shift toward TPU when average context length exceeds a few thousand tokens.
The porting steps could be reused as a template for other dense models of similar scale once the same JAX stack is targeted.
Direct energy or power measurements on both platforms would allow a fuller cost comparison that includes electricity rather than only billed dollars.

Load-bearing premise

The GPU baseline and TPU runs truly share identical hyperparameters, effective batch sizes, and model outputs after all described adaptations.

What would settle it

A controlled side-by-side execution in which loss curves or generated token distributions diverge between the two platforms despite the claimed identical settings.

Figures

Figures reproduced from arXiv: 2605.25645 by Amit Singh, Jatin Kishnani, Mayank Goel, Pulkit Agrawal, Sairanjan Mishra.

**Figure 2.** Figure 2: Gradient norm over training (log scale). Both runs are stable after warmup. The [PITH_FULL_IMAGE:figures/full_fig_p009_2.png] view at source ↗

**Figure 3.** Figure 3: pass@1 score per problem for both models. Most problems are solved at 100% by [PITH_FULL_IMAGE:figures/full_fig_p010_3.png] view at source ↗

**Figure 4.** Figure 4: Distribution of pass@1 scores across 156 problems. GPU-trained model: 101 problems [PITH_FULL_IMAGE:figures/full_fig_p011_4.png] view at source ↗

**Figure 5.** Figure 5: Inference comparison across throughput, latency, and cost dimensions at the 128-token [PITH_FULL_IMAGE:figures/full_fig_p013_5.png] view at source ↗

**Figure 6.** Figure 6: TPU vs GPU relative advantage across all dimensions (ratio [PITH_FULL_IMAGE:figures/full_fig_p015_6.png] view at source ↗

read the original abstract

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.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Desk Editor's Note private letter to a colleague

Practical porting recipe and TPU-vs-GPU numbers for Gemma 4 31B, but the claimed speedups rest on an unverified assumption that the JAX version matches the PyTorch baseline exactly.

read the letter

This paper walks through the full set of changes needed to move a PyTorch + HuggingFace TRL + FSDP recipe for Gemma 4 31B over to JAX on TPU v5p-8 for training and v6e-8 for inference. It reports 1.61x faster training and 2.12x lower cost than a 2xH100 setup, plus inference wins at longer contexts (66% higher throughput and much lower TTFT at 4096 tokens). The value is in the concrete list of adaptations: mesh config, LoRA naming fixes, sharding annotations, gradient checkpointing, data pipeline changes, and the Orbax-to-safetensors merge step. Those details are the part that actually helps someone else repeat the work.

The numbers are presented as direct measurements under identical hyperparameters. The soft spot is that the text does not show any side-by-side check that effective batch size, per-step token count, loss curves, or final model outputs match between the two stacks. Small differences in how gradient accumulation or checkpointing is implemented after the port could shift the real comparison. No error bars or run counts are mentioned either.

The work is aimed at practitioners who already have access to Google Cloud TPU capacity and need a working recipe for this model size. It fills a narrow tooling gap rather than introducing new methods. The porting details are specific enough to be useful, and the empirical comparison is at least a data point worth having.

I would send it to peer review. The central claim is testable once the equivalence check is added, and the engineering documentation stands on its own.

Referee Report

2 major / 1 minor

Summary. The manuscript presents the first end-to-end demonstration of fine-tuning Gemma 4 31B (LoRA) on TPU v5p-8 using JAX/Tunix/Qwix and serving on TPU v6e-8 with vLLM-TPU, detailing porting adaptations from a PyTorch/FSDP/HuggingFace TRL baseline. It reports that, under identical hyperparameters, TPU training completes 1.61× faster at 2.12× lower cost than a comparable 2×H100 GPU baseline; for inference, TPU v6e-8 matches GPU at short contexts and outperforms at long contexts (66% higher throughput, 23.6× faster TTFT at 4096-token inputs and QPS=4).

Significance. If the JAX and PyTorch implementations are shown to be numerically and dynamically equivalent and the measurements are reproducible with statistical controls, the work would supply practitioners with concrete, open recipes for 31B-scale model deployment on TPU hardware and close a documented tooling gap. The focus on real hardware cost/performance trade-offs and full porting details has practical utility in distributed systems and cloud ML.

major comments (2)

[Abstract] Abstract: the headline training ratios (1.61× speed, 2.12× cost) and inference metrics (throughput, TTFT) are computed under the explicit premise of 'identical hyperparameters' after porting, yet no verification is supplied that effective batch size, tokens per update, gradient accumulation, loss curves, or final model quality match between the PyTorch and JAX stacks; the listed adaptations (mesh config, LoRA naming, sharding, checkpointing, Orbax merging) can alter these quantities.
[Abstract] Abstract: none of the reported performance numbers include error bars, number of runs, run-to-run variance, data-exclusion rules, or statistical tests, so the reliability and generalizability of the direct ratio claims cannot be assessed from the given text.

minor comments (1)

[Abstract] Abstract contains duplicated phrasing in the inference-setup sentence.

Simulated Author's Rebuttal

2 responses · 0 unresolved

We thank the referee for the careful reading and constructive feedback on the equivalence of the two implementations and the statistical robustness of the reported metrics. We address each major comment below and will revise the manuscript to improve clarity on these points.

read point-by-point responses

Referee: [Abstract] Abstract: the headline training ratios (1.61× speed, 2.12× cost) and inference metrics (throughput, TTFT) are computed under the explicit premise of 'identical hyperparameters' after porting, yet no verification is supplied that effective batch size, tokens per update, gradient accumulation, loss curves, or final model quality match between the PyTorch and JAX stacks; the listed adaptations (mesh config, LoRA naming, sharding, checkpointing, Orbax merging) can alter these quantities.

Authors: We agree that the manuscript should explicitly demonstrate that the port preserves the intended training dynamics. The adaptations were chosen to maintain identical per-device batch size, gradient accumulation steps, and optimizer state; however, we did not include a direct side-by-side comparison of effective batch size, tokens-per-update, or loss curves in the original submission. In the revision we will add a dedicated subsection (Section 3.3) that (1) tabulates the hyperparameter mapping between the two stacks, (2) confirms that the effective global batch size and tokens per update are identical, and (3) presents training loss curves from both implementations on the same data subset to show comparable convergence behavior and final loss values. We will also note that model quality was not re-evaluated on downstream benchmarks because the focus was on systems-level porting cost, but the loss curves provide the requested dynamic equivalence check. revision: yes
Referee: [Abstract] Abstract: none of the reported performance numbers include error bars, number of runs, run-to-run variance, data-exclusion rules, or statistical tests, so the reliability and generalizability of the direct ratio claims cannot be assessed from the given text.

Authors: The reported ratios derive from single training and inference runs; repeating 31 B-scale experiments multiple times was not feasible within the available cloud budget. We will revise the abstract and add a short “Limitations” paragraph that states the single-run nature of the measurements, the fixed random seeds used where applicable, and the absence of run-to-run variance statistics. No data-exclusion rules were applied beyond standard checkpointing. While we cannot retroactively supply error bars without new experiments, the revised text will make the statistical limitations explicit so readers can interpret the ratios accordingly. revision: partial

Circularity Check

0 steps flagged

No circularity: empirical hardware measurements only

full rationale

The paper contains no derivations, equations, fitted parameters, or self-citations that could reduce claims to inputs by construction. All reported ratios (1.61x training speedup, 2.12x cost reduction, inference throughput/TTFT numbers) are presented as direct empirical measurements from separate TPU and GPU runs. The central premise of 'identical hyperparameters' after porting is an unverified empirical assumption, not a mathematical reduction or self-referential fit. No patterns from the enumerated list (self-definitional, fitted-input-called-prediction, self-citation load-bearing, etc.) are present. This is a standard benchmark report whose claims rest on external hardware execution rather than internal definitional closure.

Axiom & Free-Parameter Ledger

0 free parameters · 0 axioms · 0 invented entities

Purely empirical benchmarking study; contains no mathematical model, free parameters, or new theoretical entities.

pith-pipeline@v0.9.1-grok · 5937 in / 1297 out tokens · 33223 ms · 2026-06-29T20:37:14.992097+00:00 · methodology

discussion (0)

Reference graph

Works this paper leans on

5 extracted references · 1 canonical work pages

[1]

Codev: Empowering llms for verilog generation through multi-level summariza- tion.arXiv preprint arXiv:2407.10424, 2024

Y. Zhu et al., “CodeV: Empowering LLMs for Verilog Generation through Multi-Level Summarization,”arXiv:2407.10424, 2024

work page arXiv 2024
[2]

VerilogEval: Evaluating LLMs for Verilog Code Generation,

M. Liu et al., “VerilogEval: Evaluating LLMs for Verilog Code Generation,”ICCAD, 2023

2023
[3]

Gemma 4 Technical Report,

Google DeepMind, “Gemma 4 Technical Report,” 2025

2025
[4]

LoRA: Low-Rank Adaptation of Large Language Models,

E. Hu et al., “LoRA: Low-Rank Adaptation of Large Language Models,”ICLR, 2022

2022
[5]

Efficient Memory Management for Large Language Model Serving with PagedAttention,

W. Kwon et al., “Efficient Memory Management for Large Language Model Serving with PagedAttention,”SOSP, 2023. 16

2023

[1] [1]

Codev: Empowering llms for verilog generation through multi-level summariza- tion.arXiv preprint arXiv:2407.10424, 2024

Y. Zhu et al., “CodeV: Empowering LLMs for Verilog Generation through Multi-Level Summarization,”arXiv:2407.10424, 2024

work page arXiv 2024

[2] [2]

VerilogEval: Evaluating LLMs for Verilog Code Generation,

M. Liu et al., “VerilogEval: Evaluating LLMs for Verilog Code Generation,”ICCAD, 2023

2023

[3] [3]

Gemma 4 Technical Report,

Google DeepMind, “Gemma 4 Technical Report,” 2025

2025

[4] [4]

LoRA: Low-Rank Adaptation of Large Language Models,

E. Hu et al., “LoRA: Low-Rank Adaptation of Large Language Models,”ICLR, 2022

2022

[5] [5]

Efficient Memory Management for Large Language Model Serving with PagedAttention,

W. Kwon et al., “Efficient Memory Management for Large Language Model Serving with PagedAttention,”SOSP, 2023. 16

2023