Scalable Training of Language Models using JAX pjit and TPUv4. - Joanna Yoo; Kuba Perlin; Siddhartha Rao Kamalakara; João G. M. Araújo (CoRR, 2022)

Interlocking Backpropagation: Improving depthwise model-parallelism. - Aidan N. Gomez; Oscar Key; Kuba Perlin; Stephen Gou; Nick Frosst; Jeff Dean; Yarin Gal (J. Mach. Learn. Res., 2022)