nnx.Optimizer doesn't respect extra sharding axis added by nnx.scan/nnx.vmap

[qGentry]

Hi folks, me again.

I keep playing around with nnx and seems like nnx.Optimizer, when creating optimizer state for models intended to be used with 'scan', use sharding information for original, non-stacked tensor, without taking into the account extra dimension added by vmap.

Repro script:

import jax
import flax.nnx as nnx
import jax
import jax.numpy as jnp
import flax.nnx as nnx
import optax


mesh1 = jax.make_mesh((2, 4), ("a", "b"))
rules1 = (("A", "a"), ("B", "b"))


class Model(nnx.Module):
    def __init__(self, num_layers, rngs: nnx.Rngs):
        @nnx.split_rngs(splits=num_layers)
        @nnx.vmap(in_axes=(0,), out_axes=0)
        def create_linear(rngs: nnx.Rngs):
            return nnx.Param(
                jnp.ones((16, 16)), 
                sharding=("A", "B"), 
                mesh=mesh1,
                sharding_rules=rules1,
            )
        self.linears = create_linear(rngs=rngs)


@nnx.jit
def init():
    model = Model(num_layers=1, rngs=nnx.Rngs(params=0))
    optimizer = nnx.Optimizer(
        model,
        optax.adam(learning_rate=0.001),
        wrt=nnx.Param,
    )
    return model, optimizer

model, optimizer = init()

Output:

Traceback (most recent call last):
  File "/papyrax/test_scan_axis.py", line 37, in <module>
    model, optimizer = init()
                       ^^^^^^
  File "/usr/local/lib/python3.11/dist-packages/flax/nnx/transforms/compilation.py", line 474, in __call__
    pure_args_out, pure_kwargs_out, pure_out = self.jitted_fn(
                                               ^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/dist-packages/flax/nnx/transforms/compilation.py", line 135, in __call__
    out = self.f(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^
  File "/papyrax/test_scan_axis.py", line 30, in init
    optimizer = nnx.Optimizer(
                ^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/dist-packages/flax/nnx/pytreelib.py", line 400, in __call__
    return _graph_node_meta_call(cls, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/dist-packages/flax/nnx/pytreelib.py", line 412, in _graph_node_meta_call
    cls._pytree_meta_construct(node, *args, **kwargs)
  File "/usr/local/lib/python3.11/dist-packages/flax/nnx/pytreelib.py", line 403, in _pytree_meta_construct
    self.__init__(*args, **kwargs)
  File "/usr/local/lib/python3.11/dist-packages/flax/nnx/training/optimizer.py", line 88, in _check_wrt_wrapper
    return f(*args, wrt=wrt, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/dist-packages/flax/nnx/training/optimizer.py", line 160, in __init__
    to_opt_state(tx.init(nnx.state(model, wrt)))
  File "/usr/local/lib/python3.11/dist-packages/flax/nnx/training/optimizer.py", line 57, in to_opt_state
    tree = jax.tree.map(
           ^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/dist-packages/jax/_src/tree.py", line 155, in map
    return tree_util.tree_map(f, tree, *rest, is_leaf=is_leaf)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/dist-packages/flax/nnx/training/optimizer.py", line 52, in _to_opt_state
    opt_state = OptVariable(x.get_value(), **x.get_metadata())  # type: ignore
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/dist-packages/flax/nnx/variablelib.py", line 904, in __call__
    return cls._variable_meta_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/dist-packages/flax/nnx/variablelib.py", line 907, in _variable_meta_call
    variable = super().__call__(*args, **kwargs)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/dist-packages/flax/nnx/variablelib.py", line 1108, in __init__
    value = core_spmd.shard_value(
            ^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/dist-packages/flax/core/spmd.py", line 49, in shard_value
    return _apply_sharding(value, NamedSharding(mesh, pspec))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/dist-packages/flax/core/spmd.py", line 37, in _apply_sharding
    return jax.jit(lambda x: x, out_shardings=sharding)(value)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: One of pjit outputs with pytree key path result was given the sharding of NamedSharding(mesh=Mesh('a': 2, 'b': 4, axis_types=(Auto, Auto)), spec=PartitionSpec('a', 'b'), memory_kind=device), which implies that the global size of its dimension 0 should be divisible by 2, but it is equal to 1 (full shape: (1, 16, 16))
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

[qGentry]

The issue seems to be with the following function in flax:

flax/flax/nnx/training/optimizer.py

Line 49 in 697f4e5

def to_opt_state(tree):

It tries to apply sharding of the model's tensor which contains sharding metadata for original, non-vmapped Variable.

Generally, I think that for vmapped arrays, we should automatically add additional axis to its sharding metadata, shouldn't we? Because rn we have this discrepancy between sharding metadata VariableMetadata and actual sharding of the array which lead to issue like this one.

[vfdev-5]

Thanks for stress testing flax in these cases, @qGentry !
Let me see what happens inside.

[qGentry]

Seems like extending nnx.vmap args with transform_metadata solves this issue.

        @nnx.vmap(
            in_axes=(0,), 
            out_axes=0,
            transform_metadata={
                nnx.spmd.PARTITION_NAME: None,
            }
        )

I wonder if it should happen automatically or mentioned in documentation here
https://flax.readthedocs.io/en/latest/nnx_basics.html#scan-over-layers

[mohsinm-dev]

I was able to reproduce the issue, and I think the problem comes from how sharding metadata is handled when NNX transforms add new axes. I think vmap, scan, and pmap should always insert a default sharding metadata entry ({PARTITION_NAME: None}) so the extra axis they introduce is explicitly marked as unsharded.

Also, probably the sharding add/remove helper functions should be idempotent and tolerant of missing axes, so sharding_names stays aligned with the actual array rank after a lift. So, when nnx.Optimizer rewraps parameters and reapplies eager sharding; it would no longer try to shard the new leading axis, which should avoid the dimension-0 divisibility error.

@vfdev-5 Any thoughts or suggestions?

[samanklesaria]

To summarize what's been said so far, the core issue is that when we nnx.vmap over a function that returns a Param, we currently get back a Param that has the same sharding metadata rather than one with an extra None in it. You have to explicitly add the nnx.PARTITION_NAME key in transform_metadata if you want the returned Param to have the proper sharding. This is already documented in the nnx guides https://flax.readthedocs.io/en/stable/guides/transforms.html#axis-metadata. But it is pretty unintuitive.

One way to get around this is to use Jax's explicit sharding instead:

mesh = jax.make_mesh((2, 4), ("a", "b"), axis_types=(AxisType.Explicit, AxisType.Explicit))
jax.set_mesh(mesh)

class Model(nnx.Module):
    def __init__(self, num_layers, rngs: nnx.Rngs):
        @nnx.split_rngs(splits=num_layers)
        @nnx.vmap(in_axes=(0,), out_axes=0)
        def create_linear(rngs: nnx.Rngs):
            return nnx.Param(
                jnp.ones((16, 16), out_sharding=P("a", "b"))
            )
        self.linears = create_linear(rngs=rngs)

This works just fine with the Optimizer case above.

[mohsinm-dev]

@samanklesaria Good point on the metadata gap. I’ve put together a fix that makes vmap/scan/pmap insert a default {PARTITION_NAME: None} so sharding metadata tracks the extra axis automatically, and made the axis add/remove helpers idempotent. This avoids having to set transform_metadata manually or use explicit sharding for the Optimizer case. I’ll open a PR could you take a look?

[samanklesaria]

@mohsinm-dev that seems reasonable, but I'll have to check with @cgarciae , who will be on break until next week.

[mohsinm-dev]

@samanklesaria, cool. let me know after you will discuss with @cgarciae then we can discuss and work on it.

[qGentry]

I've also see quite unexpected results even with the use of transform_metadata:

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


spmd_mesh = jax.make_mesh((2, 1, 4), ("layers", "fsdp", "tensor"))

jax.set_mesh(spmd_mesh)

class Block(nnx.Module):
    def __init__(self, in_features, out_features, rngs):
        self.dense1 = nnx.Linear(
            in_features=in_features,
            out_features=out_features,
            kernel_init=nnx.with_partitioning(
                nnx.initializers.lecun_normal(), 
                ("fsdp", "tensor"),
            ),
            rngs=rngs,
        )
        self.dense2 = nnx.Linear(
            in_features=out_features,
            out_features=in_features,
            kernel_init=nnx.with_partitioning(
                nnx.initializers.lecun_normal(), 
                ("tensor", "fsdp"),
            ),
            rngs=rngs,
        )
        self.relu = nnx.relu

    def __call__(self, x):
        x = self.dense1(x)
        x = self.relu(x)
        x = self.dense2(x)
        return x


class Model(nnx.Module):
    def __init__(self, in_features, out_features, n_layers, rngs):
        @nnx.split_rngs(splits=n_layers)
        @nnx.vmap(in_axes=0, out_axes=0, transform_metadata={nnx.PARTITION_NAME: "layers"})
        def get_blocks(rngs):
            return Block(in_features=in_features, out_features=out_features, rngs=rngs)
        self.blocks = get_blocks(rngs)

@jax.jit
def get_model():
    model = Model(in_features=16, out_features=16, n_layers=8, rngs=nnx.Rngs(params=0))
    return model

model = get_model()
print(f"dense1 sharding_names: {model.blocks.dense1.kernel.sharding_names}")
print(f"dense1 sharding: {model.blocks.dense1.kernel.sharding}")

output:

dense1 sharding_names: ('layers', 'fsdp', 'tensor')
dense1 sharding: NamedSharding(mesh=Mesh('layers': 2, 'fsdp': 1, 'tensor': 4, axis_types=(Auto, Auto, Auto)), spec=PartitionSpec(None, None, 'tensor'), memory_kind=device)

For some reason, dense1 is not sharded over the leading axis even though "layers" is in sharding_names

[samanklesaria]

@qGentry The transform_metadata={nnx.PARTITION_NAME: "layers"} keyword argument is about transforming Variable metadata. In the model you present, there is no variable metadata: nnx.with_partitioning(nnx.initializers.lecun_normal(), ("tensor", "fsdp")) shards your kernels, but doesn't encode that sharding within flax as metadata. Here's how you would write the code using flax metadata:

import jax
jax.config.update('jax_num_cpu_devices', 8)
from flax import nnx
import jax.numpy as jnp

spmd_mesh = jax.make_mesh((2, 2, 2), ("layers", "fsdp", "tensor"), axis_types=(jax.sharding.AxisType.Auto,) * 3)
jax.set_mesh(spmd_mesh)

def make_model(rngs):
  return nnx.Linear(
      in_features=16,
      out_features=16,
      kernel_metadata={'sharding_names': ('fsdp', 'tensor')},
      rngs=rngs)

def test_problem():
      @nnx.split_rngs(splits=2)
      @nnx.vmap(in_axes=0, out_axes=0, transform_metadata={nnx.PARTITION_NAME: "layers"})
      def get_blocks(rngs):
          return make_model(rngs)

      result = get_blocks(nnx.Rngs(0))
      print(result.kernel.sharding)

The transform_metadata argument doesn't do any sharding itself. It's just there to notify flax about how to update the metadata about variables that are sharded when you transform them.

[cgarciae]

@samanklesaria nnx.with_partitioning does result in metadata being added to the Variable.

@qGentry I'm starting to incline to a yes after thinking about this, feel free to send a PR.