JAX-LM: Language Modelling and Distributed Training in JAX
Chuyi Shang
Introduction In this blog, we’ll implement a language model from scratch in JAX, then scale it up with distributed training across multiple GPUs/TPUs. We draw heavily upon 2 great resources. The first is Assignment 1 from Stanford’s CS336 course, which goes over how to implement language models...