Gemma's training infrastructure is a highly optimized, massively parallel system designed for efficient large language model training, built on Google's proprietary hardware and software stack. The models are trained entirely on TPUv5e Tensor Processing Units, with the Gemma 2B and 7B versions using 512 and 4,096 chips respectively, organized into pods configured as 16×16 2D torus networks. On the software side, training uses the JAX and ML Pathways frameworks, leveraging their "single controller" programming paradigm — a single Python process coordinates the entire distributed training run, dramatically simplifying development.
To harness such massive computational resources, training employs a combination of data parallelism, model sharding (e.g., 16-way sharding within a pod for the 7B model), and ZeRO-3-like optimizer state sharding, with the GSPMD partitioner and MegaScale XLA compiler performing extreme optimization of the computation graph. Additionally, Google estimated the carbon footprint of this training (approximately 131 tonnes of CO2 equivalent) and ensured carbon neutrality for its data center operations and this training run through improved energy efficiency, renewable energy purchases, and carbon offsets — reflecting a commitment to environmental responsibility alongside the pursuit of cutting-edge AI performance.