How do we perform gradient sync in nnx for data parallel training?

[jinyongyoo]

Context
I'm trying to perform data parallel training where model is replicated on each device and data is sharded along the batch axis. Typically in PyTorch / Tensorflow, we have the gradient synced across the devices during the backward pass (Pytorch) or right before optimizer update (Tensorflow) via an all_reduce op.

Question
I've been looking at some examples from for both NNX and Linen, and for Linen examples, we see explicit jax.lax.pmean calls to sync the gradients (example).. However, for NNX, I cannot find any resource or example that explains to me where exactly gradient synchronization happens. All examples simply call value_and_grad and passes the gradients to the optimizer. Where and how does gradient sync exactly happen with NNX?

[vladfi1]

NNX is just a wrapper around jax that takes care of "mutable" state -- it doesn't add any gradient synchronization beyond what jax already does. Whether you need to insert an explicit pmean will depend on how you wrote your code, e.g. are you using shard_map to explicitly control sharding or are you relying on jax.jit to do it automatically. In the shard_map case, you can either a) apply pmean to your loss so that the gradient automatically gets a pmean during the backward pass, or b) no pmean on the loss meaning the gradients are computed independently on each shard and you will then need to explictly pmean the gradients before applying your optimizer. In case (b) you will also need to explictly pcast your model parameters to varying before passing them to your loss function, otherwise jax will insert an unwanted psum on the backwards pass -- see https://docs.jax.dev/en/latest/notebooks/shard_map.html for more info.