Scaling Up (Part 2)

Key Takeaways

[music] Hi, I'm Robert Crowe. Uh, welcome back to part two of our three-part series on sharding and parallelism. Let's explore how to scale your flax enx models using Jax's powerful distributed computing capabilities, specifically its SPMD paradigm. If you're coming from PyTorch and have started using Jax and Flax NNX, you know that modern models often outgrow single accelerators. We'll discuss Jax's approach to parallelism and how NX integrates with it seamlessly. This episode focuses on the main workflow, how to integrate Jax's sharding primitives with Flaxen and X with a special focus on the critical sharded initialization pattern. So let's get started. In our last episode, we laid the foundation for scaling our models. We explored why distributed training is essential and covered the core strategies like data and model parallelism. We then introduced the fundamental JAX primitives, the mesh to define our hardware topology, the partition spec to describe how our tensors should be laid out, and JAX.jit JIT which triggers the SPMD compilation. Now we're going to build on that foundation and connect it directly to Flax NX. Okay, let's quickly recap Flax NX. Unlike Pure Jacks, NNX modules are stateful Python objects much like PyTorch modules. They hold their parameters and other state directly as attributes usually initialized eagerly. A crucial feature for us today is that these state variables can store metadata. We'll use this to attach our sharding hints. However, this statefulness traditionally clashes with Jax's functional nature. The big news is that in NNx uh version 0.11 and and after is that NNX modules are now native Jack PI trees. This closes a major conceptual gap that we had in earlier versions. While the module object itself is still a mutable Python object, it can now be understood directly by Jack's transformations. And NNX provides the bridge to handle this correctly. So how do we reconcile NX's mutable objects with Jax's functional transformations? Well, NX provides two ways. First, the functional API. You explicitly split your module into its static definition and its dynamic state which is a JAX compatible PI tree. You pass this state through jacks.jit or jacks.grad to get the result and then merge or update the original module. Alternatively, NNX provides its own wrappers like nx.jit and nx.grad. These handle the split merge process under the hood, giving you a more seamless experience, although potentially with some slight overhead compared to the manual functional API. Now, let's connect NNX and SPMD. The Flax NX SPMD module lets us attach those partition spec hints directly to our parameters as metadata. When we define the module, we can use helper functions like nxxwith metadata around your initializer passing a sharding argument with the partition spec tupil. Sometimes you can even pass the sharding argument directly to nx.param. This stores a partition spec tupil in ashsharding attribute on the variable state. Crucially, this doesn't shard the parameter yet. It's just metadata, a hint for later. Here's what those two options look like in code. First, the preferred method using nx.withth metadata. Second, annotating the variable directly. In this case, nx.param. Okay, this is the core workflow and addresses a common pitfall. If you just instantiate a very large NNx model, even with sharding annotations, Jax might try to create all parameters on a single default device first, leading to a anom or out of memory error. The solution is to perform the initialization and the application of sharding inside a JIT compiled function. We'll use nx.jit for convenience. Inside this function, the first steps are instantiate the module. The parameters are created but likely still on device zero and then extract its functional state pi tree using nx.state. Continuing inside our jitted initialization function, next we use nx.spmd.get partition spec. This utility function walks through the state pi tree we just extracted and pulls out all the dotsharding metadata we added earlier creating a parallel pi tree of partitionspec objects. Then comes the crucial step. We use jacks.laxwithsharting constraint passing in the original state and the partition specs we just extracted. This tells the Jack's compiler, when you produce the final result for this state, make sure it conforms to these sharding specifications. Here's what that last part looks like in code. We use nx.spmd.getartition spec to pull out all thesharting metadata we added earlier, creating a parallel pie of partition spec objects. Then we use jacks.laxwish lax with sharding constraint passing in the original state and the partition specs we just extracted. The final step inside the jitted function is to use nx.update. This takes the sharted state which Jackson XLA will ensure is correctly distributed upon execution and merges it back into our original mutable model object. The function then returns this model. The key is to call this entire create sharded model function while inside a width mesh block or a mesh context. This provides the context Jax needs to actually fulfill the sharding constraints across the specified devices. The result sharted model now holds parameters physically distributed across the mesh avoiding the single device om or out of memory. So inside the jitted function we get our sharted state and merge it back into the model with nx update. Then finally we execute that jitted function in a with mesh context. For better abstraction instead of hard- coding data or model inside your module definitions you can use logical access names like batch sequence or embed or hidden. You then define a separate mapping called sharding rules that translates these logical names to your actual physical mesh access names. You can provide these rules when initializing variables or attach them later. This makes your model code more portable and easier to adapt to different hardware setups or sharding strategies as you only need to change the rules and not the model definition itself. Here's an example of using sharding rules to map logical axis names to physical mesh axes. We define sharding rules as a tupole of pairs where each pair maps a logical name like batch or hidden to a physical mesh axis like data or model. For example, batch data means that anything annotated with batch will be sharted along the data axis of the mesh. You then use these logical names in your NNX module definition when initializing parameters, providing the sharding rules through nx.with metadata. And this approach allows you to define your model independently of the specific hardware layout, making your code more flexible and adaptable to different sharding strategies. You can learn more about Jax and the entire Jax AI stack with these coding exercises, quick reference docs, and slides. To watch the whole Learning Jack series, check out our playlist on YouTube. There's a growing community on Discord for Jax. And here's an invite link and here are links to the docs for Jax Flax and the Jax AI stack. We've just walked through the most critical part of this entire process. how to safely initialize a massive sharded NNx model. By annotating our parameters with sharding metadata and wrapping the instantiation inside a jitted function with the wish sharding constraint, we've solved the single device om or out of memory problem and now have a model with its parameters physically distributed across our accelerator mesh. We've started our model, but what's next? How do we feed it data, run the training step, manage checkpoints, and see all of this in action with a real world example? Coming up in our next episode, we'll put it all together. We'll build the distributed training loop, discuss efficient data loading with grain and sharded checkpointing with Orbbacks, and then walk through a complete practical example of sharding a GPT2 style transformer block from start to finish. See you there and don't forget to hit like and subscribe.