Merge pull request #5079 from samanklesaria:eager_sharding_context

PiperOrigin-RevId: 831520025

Author: Flax Authors

docs_nnx/guides/flax_gspmd.ipynb

@@ -43,6 +43,7 @@
     "from jax import numpy as jnp\n",
     "from jax.sharding import PartitionSpec as P, NamedSharding, AxisType\n",
     "import optax\n",
+    "import flax\n",
     "from flax import nnx\n",
     "\n",
     "# Ignore this if you are already running on a TPU or GPU\n",
@@ -56,7 +57,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Set up a `2x4` device mesh as the [JAX data sharding tutorial](https://docs.jax.dev/en/latest/sharded-computation.html#key-concept-data-sharding) instructs. \n",
+    "Set up a `2x4` device mesh as the [JAX data sharding tutorial](https://docs.jax.dev/en/latest/sharded-computation.html#key-concept-data-sharding) instructs.\n",
     "\n",
     "In this guide we use a standard FSDP layout and shard our devices on two axes - `data` and `model`, for doing batch data parallelism and tensor parallelism."
    ]
@@ -75,7 +76,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "> Compatibility Note: This guide covers the [eager sharding feature](https://flax.readthedocs.io/en/latest/flip/4844-var-eager-sharding.html) that greatly simplifies creating sharded model. If your project already used Flax GSPMD API on version `flax<0.12`, you might have turned the feature off to keep your code working. Check the flag and read on to learn how to use the feature."
+    "> Compatibility Note: This guide covers the [eager sharding feature](https://flax.readthedocs.io/en/latest/flip/4844-var-eager-sharding.html) that greatly simplifies creating sharded model. If your project already used Flax GSPMD API on version `flax<0.12`, you might have turned the feature off to keep your code working. Users can toggle this feature using the `nnx.use_eager_sharding` function."
    ]
   },
   {
@@ -84,8 +85,45 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import flax\n",
-    "assert flax.config.flax_always_shard_variable is True"
+    "nnx.use_eager_sharding(True)\n",
+    "assert nnx.using_eager_sharding()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c24144d8",
+   "metadata": {},
+   "source": [
+    "The `nnx.use_eager_sharding` function can also be used as a context manager to toggle the eager sharding feature within a specific scope."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2d849e2e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with nnx.use_eager_sharding(False):\n",
+    "  assert not nnx.using_eager_sharding()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c9f808ec",
+   "metadata": {},
+   "source": [
+    "You can also enable eager sharding on a per-variable basis by passing `eager_sharding=False` during variable initialization. The mesh can also be passed this way."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "67bbd440",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "nnx.Param(jnp.ones(4,4), sharding_names=(None, 'model'), eager_sharding=True, mesh=auto_mesh)"
    ]
   },
   {
@@ -256,7 +294,7 @@
     "with jax.set_mesh(auto_mesh):\n",
     "  # Create your input data, sharded along `data` dimension, as in data parallelism\n",
     "  x = jax.device_put(jnp.ones((16, 4)), P('data', None))\n",
-    "  \n",
+    "\n",
     "  # Run the model forward function, jitted\n",
     "  y = jax.jit(lambda m, x: m(x))(linear, x)\n",
     "  print(y.sharding.spec)                       # sharded: ('data', 'model')\n",
@@ -313,7 +351,7 @@
     "    def create_sublayers(r):\n",
     "      return DotReluDot(depth, r)\n",
     "    self.layers = create_sublayers(rngs.fork(split=num_layers))\n",
-    "  \n",
+    "\n",
     "  def __call__(self, x):\n",
     "    def scan_over_layers(x, layer):\n",
     "      return layer(x), None\n",
@@ -364,7 +402,7 @@
     "  # Model and optimizer\n",
     "  model = MultiDotReluDot(1024, 2, rngs=nnx.Rngs(0))\n",
     "  optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)\n",
-    "  \n",
+    "\n",
     "  # The loop\n",
     "  for i in range(5):\n",
     "    model, loss = train_step(model, optimizer, input, label)\n",
@@ -496,7 +534,7 @@
     "    def create_sublayers(r):\n",
     "      return LogicalDotReluDot(depth, r)\n",
     "    self.layers = create_sublayers(rngs.fork(split=num_layers))\n",
-    "  \n",
+    "\n",
     "  def __call__(self, x):\n",
     "    def scan_over_layers(x, layer):\n",
     "      return layer(x), None\n",
@@ -617,7 +655,7 @@
     "    def create_sublayers(r):\n",
     "      return ExplicitDotReluDot(depth, r)\n",
     "    self.layers = create_sublayers(rngs.fork(split=num_layers))\n",
-    "  \n",
+    "\n",
     "  def __call__(self, x):\n",
     "    def scan_over_layers(x, layer):\n",
     "      return layer(x), None\n",

docs_nnx/guides/flax_gspmd.md

@@ -29,6 +29,7 @@ import jax
 from jax import numpy as jnp
 from jax.sharding import PartitionSpec as P, NamedSharding, AxisType
 import optax
+import flax
 from flax import nnx
 
 # Ignore this if you are already running on a TPU or GPU
@@ -37,7 +38,7 @@ if not jax._src.xla_bridge.backends_are_initialized():
 print(f'You have 8 “fake” JAX devices now: {jax.devices()}')
 ```
 
-Set up a `2x4` device mesh as the [JAX data sharding tutorial](https://docs.jax.dev/en/latest/sharded-computation.html#key-concept-data-sharding) instructs. 
+Set up a `2x4` device mesh as the [JAX data sharding tutorial](https://docs.jax.dev/en/latest/sharded-computation.html#key-concept-data-sharding) instructs.
 
 In this guide we use a standard FSDP layout and shard our devices on two axes - `data` and `model`, for doing batch data parallelism and tensor parallelism.
 
@@ -46,11 +47,24 @@ In this guide we use a standard FSDP layout and shard our devices on two axes -
 auto_mesh = jax.make_mesh((2, 4), ('data', 'model'))
 ```
 
-> Compatibility Note: This guide covers the [eager sharding feature](https://flax.readthedocs.io/en/latest/flip/4844-var-eager-sharding.html) that greatly simplifies creating sharded model. If your project already used Flax GSPMD API on version `flax<0.12`, you might have turned the feature off to keep your code working. Check the flag and read on to learn how to use the feature.
+> Compatibility Note: This guide covers the [eager sharding feature](https://flax.readthedocs.io/en/latest/flip/4844-var-eager-sharding.html) that greatly simplifies creating sharded model. If your project already used Flax GSPMD API on version `flax<0.12`, you might have turned the feature off to keep your code working. Users can toggle this feature using the `nnx.use_eager_sharding` function.
 
 ```{code-cell} ipython3
-import flax
-assert flax.config.flax_always_shard_variable is True
+nnx.use_eager_sharding(True)
+assert nnx.using_eager_sharding()
+```
+
+The `nnx.use_eager_sharding` function can also be used as a context manager to toggle the eager sharding feature within a specific scope.
+
+```{code-cell} ipython3
+with nnx.use_eager_sharding(False):
+  assert not nnx.using_eager_sharding()
+```
+
+You can also enable eager sharding on a per-variable basis by passing `eager_sharding=False` during variable initialization. The mesh can also be passed this way.
+
+```{code-cell} ipython3
+nnx.Param(jnp.ones(4,4), sharding_names=(None, 'model'), eager_sharding=True, mesh=auto_mesh)
 ```
 
 ## Shard a single-array model
@@ -107,7 +121,7 @@ You should still make sure to `jax.jit` for maximum performance, and also to exp
 with jax.set_mesh(auto_mesh):
   # Create your input data, sharded along `data` dimension, as in data parallelism
   x = jax.device_put(jnp.ones((16, 4)), P('data', None))
-  
+
   # Run the model forward function, jitted
   y = jax.jit(lambda m, x: m(x))(linear, x)
   print(y.sharding.spec)                       # sharded: ('data', 'model')
@@ -153,7 +167,7 @@ class MultiDotReluDot(nnx.Module):
     def create_sublayers(r):
       return DotReluDot(depth, r)
     self.layers = create_sublayers(rngs.fork(split=num_layers))
-  
+
   def __call__(self, x):
     def scan_over_layers(x, layer):
       return layer(x), None
@@ -182,7 +196,7 @@ with jax.set_mesh(auto_mesh):
   # Model and optimizer
   model = MultiDotReluDot(1024, 2, rngs=nnx.Rngs(0))
   optimizer = nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)
-  
+
   # The loop
   for i in range(5):
     model, loss = train_step(model, optimizer, input, label)
@@ -266,7 +280,7 @@ class LogicalMultiDotReluDot(nnx.Module):
     def create_sublayers(r):
       return LogicalDotReluDot(depth, r)
     self.layers = create_sublayers(rngs.fork(split=num_layers))
-  
+
   def __call__(self, x):
     def scan_over_layers(x, layer):
       return layer(x), None
@@ -354,7 +368,7 @@ class ExplicitMultiDotReluDot(nnx.Module):
     def create_sublayers(r):
       return ExplicitDotReluDot(depth, r)
     self.layers = create_sublayers(rngs.fork(split=num_layers))
-  
+
   def __call__(self, x):
     def scan_over_layers(x, layer):
       return layer(x), None

flax/nnx/__init__.py

@@ -203,6 +203,8 @@
 from .variablelib import register_variable_name as register_variable_name
 from .variablelib import use_refs as use_refs
 from .variablelib import using_refs as using_refs
+from .variablelib import use_eager_sharding as use_eager_sharding
+from .variablelib import using_eager_sharding as using_eager_sharding
 from .visualization import display as display
 from .extract import to_tree as to_tree
 from .extract import from_tree as from_tree

flax/nnx/variablelib.py

@@ -64,10 +64,103 @@
 @dataclasses.dataclass
 class VariableContext(threading.local):
   mutable_variable_stack: list[bool] = dataclasses.field(default_factory=list)
+  eager_shard_stack: list[bool] = dataclasses.field(default_factory=list)
+
 
 
 VARIABLE_CONTEXT = VariableContext()
 
+class UseEagerShardContext:
+  def __init__(self, prev_value: bool | None, new_value: bool):
+    self.prev_value: bool | None = prev_value
+    self.new_value: bool = new_value
+
+  def __enter__(self):
+    if self.prev_value is not None:
+      VARIABLE_CONTEXT.eager_shard_stack.insert(-1, self.prev_value)
+
+  def __exit__(self, exc_type, exc_value, traceback):
+    VARIABLE_CONTEXT.eager_shard_stack.pop()
+
+  def __call__(self, f: F) -> F:
+    # undo eager stack change
+    VARIABLE_CONTEXT.eager_shard_stack.pop()
+    if self.prev_value is not None:
+      VARIABLE_CONTEXT.eager_shard_stack.append(self.prev_value)
+
+    @functools.wraps(f)
+    def use_eager_sharding_wrapper(*args, **kwargs):
+      VARIABLE_CONTEXT.eager_shard_stack.append(self.new_value)
+      try:
+        return f(*args, **kwargs)
+      finally:
+        VARIABLE_CONTEXT.eager_shard_stack.pop()
+
+    return use_eager_sharding_wrapper  # type: ignore[return-value]
+
+def using_eager_sharding() -> bool:
+  """Returns whether Variables are using eager sharding by default.
+
+  Example::
+
+    >>> from flax import nnx
+    >>> nnx.use_eager_sharding(True)
+    <...>
+    >>> nnx.using_eager_sharding()
+    True
+    >>> nnx.use_eager_sharding(False)
+    <...>
+    >>> nnx.using_eager_sharding()
+    False
+
+
+  Returns:
+    A boolean indicating if Variables are using eager sharding by default.
+  """
+  do_eager_sharding = config.flax_always_shard_variable
+  if VARIABLE_CONTEXT.eager_shard_stack:
+    do_eager_sharding = VARIABLE_CONTEXT.eager_shard_stack[-1]
+  return do_eager_sharding
+
+def use_eager_sharding(value: bool, /):
+  """Sets whether Variables should use eager sharding by default or not.
+
+  Example usage::
+
+    >>> from flax import nnx
+    >>> # Use eager sharding by default
+    >>> nnx.use_eager_sharding(True)
+    <...>
+    >>> # Variable will now use eager sharding
+    >>> nnx.using_eager_sharding()
+    True
+
+  It can also be used as a context manager to temporarily
+  change the default behavior for a block of code::
+
+    >>> nnx.use_eager_sharding(False)
+    <...>
+    >>> with nnx.use_eager_sharding(True):
+    ...   nnx.using_eager_sharding()
+    True
+    >>> # it will reset outside
+    >>> v = nnx.Variable(jax.numpy.ones((2, 3)))
+    >>> nnx.using_eager_sharding()
+    False
+
+  Args:
+    value: A boolean indicating if Variables should use eager sharding by default.
+
+  Returns:
+    A context manager that resets the context to the previous value.
+  """
+  if VARIABLE_CONTEXT.eager_shard_stack:
+    prev_value = VARIABLE_CONTEXT.eager_shard_stack[-1]
+    VARIABLE_CONTEXT.eager_shard_stack[-1] = value
+  else:
+    prev_value = None
+    VARIABLE_CONTEXT.eager_shard_stack.append(value)
+  return UseEagerShardContext(prev_value, value)
 
 def using_refs() -> bool:
   """Returns whether Variables are using ArrayRefs by default.
@@ -289,10 +382,7 @@ def __init__(
     value = self.create_value(self.raw_value)
 
     # shard the value if applicable
-    do_eager_sharding = config.flax_always_shard_variable
-    if 'eager_sharding' in metadata:
-      do_eager_sharding = metadata['eager_sharding']
-    if do_eager_sharding and 'sharding_names' in metadata:
+    if metadata.get('eager_sharding', using_eager_sharding()) and 'sharding_names' in metadata:
       value = core_spmd.shard_value(
         value, metadata['sharding_names'], metadata.get('sharding_rules', None),
         metadata.get('mesh', None))

tests/nnx/spmd_test.py

@@ -192,6 +192,21 @@ def __call__(self, x: jax.Array):
     self.assertEqual(badds, [(0, 'layers'), (0, 'layers')])
     self.assertEqual(bremoves, [(0, 'layers')])
 
+
+  @parameterized.product(use_eager_sharding=[True, False])
+  def test_eager_sharding_context(self, use_eager_sharding):
+    rngs = nnx.Rngs(0)
+    with nnx.use_eager_sharding(use_eager_sharding):
+      mesh = jax.make_mesh(((2, 2)), ("data", "model"))
+      with jax.set_mesh(mesh):
+        w = nnx.Param(
+          rngs.lecun_normal()((4, 8)),
+          sharding_names=(None, 'model'))
+        if use_eager_sharding:
+          assert has_sharding_spec(w)
+        else:
+          assert not has_sharding_spec(w)
+
   @parameterized.product(use_ref=[True, False])
   def test_logical_rules(self, use_ref):
     self.enter_context(nnx.use_refs(use_ref))
@@ -302,6 +317,14 @@ def test_explicit_sharding_mesh_context(self):
       P('row', 'col'),
     )
 
+def has_sharding_spec(array):
+    sharding = array.sharding
+    if hasattr(sharding, 'spec'):
+        # For NamedSharding or PositionalSharding
+        return sharding.spec is not None and any(
+            s is not None for s in sharding.spec
+        )
+    return False
 
 if __name__ == '__main__':
   absltest.main()