Merge pull request #5033 from IvyZX:shard-msg

PiperOrigin-RevId: 820437084

Author: Flax Authors

docs_nnx/flip/4844-var-eager-sharding.md

@@ -8,6 +8,8 @@
 
 Simplify the creation of sharded NNX models. When a sharding annotation is provided, all `nnx.Variable` creation will **require a mesh context** and automatically be sharded as annotated.
 
+See [GSPMD Guide](https://flax.readthedocs.io/en/latest/guides/flax_gspmd.html) for a comprehensive guide on how to make sharded NNX models.
+
 # Motivation
 
 To create a sharded model, user should only need to do this:
@@ -44,6 +46,19 @@ User can turn off this feature in two ways:
 * **Variable-specific flag**: Create a specific variable with metadata `eager_sharding=False`, such as: `nnx.Param(..., eager_sharding=False)`.
 
 
+# Flexibility options
+
+For debugging in a CPU environment, make a dummy mesh to run the model:
+
+```python
+mesh = jax.make_mesh(((1, 1, 1)), ('your', 'axes', 'names'))
+with jax.set_mesh(mesh):
+  ...
+```
+
+For JAX explicit mode, remove the `sharding_names=` annotation on the `nnx.Variable`.
+
+
 # Implementation
 [implementation]: #implementation
 

flax/core/spmd.py

@@ -40,9 +40,7 @@ def shard_value(value, sharding_names, sharding_rules, mesh):
     raise ValueError(
       'An auto mesh context or metadata is required if creating a variable'
       f' with annotation {sharding_names=}. '
-      'If running this on CPU for debugging, make a'
-      ' dummy mesh like `jax.make_mesh(((1, 1)), (<your axis names>))`. '
-      'If running on explicit mode, remove `sharding_names=` annotation.')
+      'For more guidance, see https://flax.readthedocs.io/en/latest/flip/4844-var-eager-sharding.html.')
   pspec = get_pspec(sharding_names, sharding_rules)
   if mesh is not None:
     jax.lax.with_sharding_constraint(value, NamedSharding(mesh, pspec))