The ‘single controller’ programming model of Jax (Bradbury et al., 2018) and Pathways (Barham
et al., 2022) allows a single Python process to orchestrate the entire training run, dramatically
simplifying the development workflow. The GSPMD partitioner (Xu et al., 2021) in the XLA compiler
partitions the training step computation, and the MegaScale XLA compiler (XLA, 2019) pass statically
schedules appropriate collectives so that they maximally overlap with the computation with very little
variation in step time.
gemini_1_report