Explicit sharding in NNX

[cgarciae]

This is a temporary hack as we add support for explict sharding:

import functools
from flax import nnx
import jax.numpy as jnp
import jax
from jax import P

jax.config.update('jax_num_cpu_devices', 8)

def nnx_auto_axes(f):
  @functools.wraps(f)
  def outer_fn(*args, out_sharding=None, **kwargs):
    args_tree, kwargs_tree = nnx.extract.to_tree((args, kwargs))
    args_sharding, kwargs_sharding = nnx.extract.clear_non_graph_nodes((args, kwargs))
    args_sharding, kwargs_sharding = nnx.extract.to_tree((args_sharding, kwargs_sharding))
    args_sharding, kwargs_sharding = jax.tree.map(lambda x: getattr(x, 'sharding', None), (args_sharding, kwargs_sharding))

    @jax.sharding.auto_axes(out_sharding=(args_sharding, kwargs_sharding, out_sharding))
    def inner_fn(*args_tree, **kwargs_tree):
      args, kwargs = nnx.extract.from_tree((args_tree, kwargs_tree))
      out = f(*args, **kwargs)
      args_out, kwargs_out = nnx.extract.clear_non_graph_nodes((args, kwargs))
      args_out, kwargs_out = nnx.extract.to_tree((args_out, kwargs_out))
      return args_out, kwargs_out, out

    args_out, kwargs_out, out = inner_fn(*args_tree, **kwargs_tree)
    nnx.update((args, kwargs), nnx.state((args_out, kwargs_out)))
    return out

  return outer_fn


nnx.Linear.__call__ = nnx_auto_axes(nnx.Linear.__call__)

mesh = jax.make_mesh((8,), ('data',), axis_types=(jax.sharding.AxisType.Explicit,))

def sharding(*names):
  return jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(*names))

def sharded_init(
  initializer: jax.nn.initializers.Initializer, out_sharding=None
):
  def init_fn(*args, **kwargs):
    return initializer(*args, **kwargs, out_sharding=out_sharding)

  return init_fn

with jax.set_mesh(mesh):
  model = nnx.Linear(
    32,
    32,
    rngs=nnx.Rngs(0),
    kernel_init=sharded_init(
      nnx.initializers.xavier_uniform(), out_sharding=P(None, 'data')
    ),
    use_bias=False,
  )
  x = jnp.ones((64, 32), out_sharding=P('data'))

  jax.debug.visualize_array_sharding(model.kernel.value)
  jax.debug.visualize_array_sharding(x)

  y = model(x, out_sharding=P('data'))

  jax.debug.visualize_array_sharding(y)

Update: fixed nnx_auto_axes.

[jackopenn]

@cgarciae Is this still the solution for explicit sharding in nnx? I see docs here https://flax.readthedocs.io/en/latest/guides/flax_gspmd.html#explicit-sharding. But this still doesn't allow to call something like intermediate = self.up_proj(x, out_sharding=P("data", None, None))

[vfdev-5]

Right, for explicit sharding you put out_sharding to jax initializers and in certain jax ops where it can be needed. Sorry, can you restate your question, please again? I'm not sure to get it correctly, thanks!

By the way, there is now new stable flax release v0.12.0 available, so you can have available the new sharding features.

[jackopenn]

Sure, lets look at an example - a SwiGLU block using FSDP and TP

https://colab.research.google.com/drive/1X2Dx6LkEs42za6yWfk-6tc3bYglnIK3s#scrollTo=900AkpNyERjp

• We can pass out_sharding to initializers 👍
• For nnx layers, we can not pass out_sharding ❌
• The last down_proj is ambiguous. I can't pass out_sharding to nnx.Linear. So I would have to wrap the whole call in auto axes which at the point i'm not doing explicit sharding.

So back to my original question, is this still the solution for explicit sharding in nnx? As mentioned in the post this is just a temporary hack. I would assume we want to add out_sharding to all nnx layers.

Just saw 0.12.0 - eager sharding is nice. Thanks team :)

[vfdev-5]

@jackopenn thanks for clarification! I understand your point, so I'm translating the question as whether we want to add out_sharding argument to all linear layers like nnx.Linear, nnx.Conv etc? I'll ask the team about this feature.

By the way, we can override dot_general method used in nnx.Linear:

    self.down_proj = nnx.Linear(
        intermediate_dim,
        hidden_dim,
        use_bias=False,
        dtype=jnp.bfloat16,
        kernel_init=partial(nnx.initializers.lecun_normal(), out_sharding=P("model", "fsdp")),
        rngs=rngs,
        dot_general=partial(jax.lax.dot_general, out_sharding=P("model", "fsdp")),
    )

so doing that, I could make run your code:

bfloat16[8@fsdp,4096,256@model]
bfloat16[8@fsdp,4096,256@model]
bfloat16[8@fsdp,4096,256@model]
bfloat16[8@model,4096@fsdp,128]

[jackopenn]

Yeah I think that would be a nice feature. But I didn't think about the dot product - that is a great idea :). But all these partials might become too verbose, I still think having the out_sharding in the call of the layers will be a better API. Thanks for helping on this though, much appreciated :)

[cgarciae]

@jackopenn we do plan to add out_sharding to all standard layers and maybe upstream the decorator above.