Log In to continue learning
Build and Train an LLM with JAX
Intermediate
Topics
GenAI Applications
In this lesson, you'll get a high-level introduction to JAX and see why it has become a core tool for modern ML research. You'll also learn how JAX compares to NumPy, PyTorch, and TensorFlow, and where it fits into the broader JAX AI stack. Let's dive in. To build our mini LLM, we'll be using JAX, which is an open source Python library built at Google. JAX's core is designed to be a small and focused Python library. At the root, there is a NumPy-like programming interface which allows you to do complex matrix math. On top of that, JAX adds automatic differentiation, which is a key component of back propagation in neural networks. Because JAX uses functional programming, all the gradient calculations in JAX are pure functions, which means you can chain them together to do things like double differentiation, which is hard or even impossible with other AI libraries. Another unique part of JAX is the just-in-time compiler, or "jit". Just-in-time compilation allows JAX to be compiled to run on the accelerator of your choice, CPU, GPU, or TPU. using the XLA compiler. Once a JAX function is compiled, it will then be cached, which means subsequent calls to that function are super fast. The other thing this enables is vectorized and parallelized execution. which means the compiler can vectorize your functions to accept batches of data and can split up your workload across multiple accelerators. Just by adding these function calls in your pipeline. Together, this is why JAX can run across CPUs, GPUs, and TPUs all without code rewrites. This means you can rapidly experiment on a small scale, and then when it comes time to scale your training or inference, you can split your workload across as many chips as you want. Let's take a quick look at the core NumPy-like coding style of JAX. Here you can see that we can import jax.numpy which replicates many of the NumPy functions, as functional jax. You can then use those familiar functions like tanh or sum. But these JAX versions will be able to be composed with other JAX functions later. One of the most powerful things that you can do with these functions is get the gradient. by importing grad from jax, then calling that grad function on your other operations. JAX keeps track of all those gradients for you. and you can get the result at any time, which is the key component of training a neural network. Batching is accomplished with another function call called vmap, or vectorized map. vmap, like grad, is just another function in JAX, and it lets you operate on batches of data at a time. Then, you can compile your code with the jit function. or, as we'll see later, with the jit decorator around functions. This step will compile your code just once, the first time it runs, and then we'll reuse that compiled code later. At this point, the XLA compiler, which is used by JAX, Now it knows everything about your function and data. So finally, we can spread out all that work across different accelerators by defining a sharding, which is how JAX will split up your workload. All of that work happens in the core JAX library, but there's an entire ecosystem of other libraries that we can use to build upon that foundation. The JAX ecosystem includes JAX, of course. which is compiled by the XLA compiler and can run on any of the supported hardware, which includes CPUs, GPUs, and TPUs. That core loop of JAX to XLA to hardware can be used to accelerate any matrix operation. which could be a neural network, but it can also be used for scientific computing or weather modeling or price prediction or any number of other things. To build an LLM specifically, in this course we'll also dig into the JAX AI stack, which includes a group of libraries designed to make training a neural network fast and easy. After this course, if you're interested in training even larger LLMs, you can take a look at MaxText. Or if you're interested in diffusion models, take a look at MaxDiffusion. These are a collection of libraries and examples from Google that make it easy to get started with large and powerful language or image models. Then once you're done training your LLM, you'll want to use an LLM inference engine like vLLM. which implements all of the latest tips and tricks to get the very fastest results with the most efficient use of compute possible. Altogether, these libraries are part of what we call the JAX ecosystem. And this ecosystem of libraries is one of the reasons so many people like to use JAX. From a programming perspective though, how does JAX compare to NumPy, PyTorch, or TensorFlow? This table contains a high-level overview of the differences. Compared to NumPy, JAX has a very similar set of core functionality, but JAX also has automatic differentiation and a built-in way to accelerate on TPUs or GPUs. Compared to PyTorch, JAX has a functional paradigm, whereas PyTorch is more imperative. It is mostly just a user preference, but for some researchers and workloads, it's absolutely vital that JAX is functional because researchers can compose those functions and accomplish things like taking the gradient of a gradient, which is something that PyTorch has trouble with. And finally, compared to TensorFlow, JAX is a more flexible programming paradigm because you don't have to define your compute graph for yourself. Instead, you just define your function calls and then call the just-in-time compiler which builds the graph for you. All of those together is what makes JAX so powerful for research and production. Here we can see some examples of PyTorch code side by side along with some JAX code. using the Flax.NNX library for neural networks. You can see it looks almost identical. One major difference is the forward method of PyTorch versus the call method of the functional Flax.NNX code. But if you've used PyTorch before, you should find switching to JAX with Flax.NNX easy to pick up. So, for this course, we'll be using the JAX core library, but we'll be interacting with it using the other JAX stack libraries. First, we'll use Grain to load our data. Grain is a powerful data loader that is especially useful for loading huge data sets, like the type you need if you're going to train your own LLM from scratch. We'll also use Flax/NNX to define our model architecture. NNX lets you write models that look very similar to PyTorch. So if you already know PyTorch, this part will seem very familiar. Next, we'll use Optax during training, which contains all of the optimizers and gradient processing utilities that you'll need to train a neural network. So you don't have to implement these yourself. And finally, we'll use Orbax for saving and loading our model. Like Grain, Orbax is specifically designed to be very useful at scale, and it can handle model checkpoints for you across huge datasets and large clusters of GPUs or TPUs. All of these features, from the flexible functional coding paradigm, to all the utilities which handle training and serving at scale, is why we at Google use JAX for all of our latest AI models. In fact, Gemini, Gemma, Nano Banana, Veo, and many, many more models are all built on top of JAX. You can read about these and a lot more models and how they use JAX by searching for Google DeepMind papers about the models. In this lesson, you learned about JAX and how it fits into a larger AI stack ecosystem. You learned how it compares to other libraries like PyTorch and TensorFlow, and you learned how we use it at Google to build our latest AI models. In the next lesson, you'll build your first model architecture using JAX.
