add logical axis name indirection to new flax partitioning API

PiperOrigin-RevId: 495211165

Author: levskaya

CHANGELOG.md

@@ -6,7 +6,7 @@ vNext
 (Add your change to a random empty line to avoid merge conflicts)
 -
 -
--
+- Added logical partitioning helpers for using pjit with Flax.
 -
 -
 -

flax/linen/__init__.py

@@ -15,7 +15,7 @@
 """The Flax Module system."""
 
 
-# pylint: disable=g-multiple-import
+# pylint: disable=g-multiple-import,useless-import-alias
 # re-export commonly used modules and functions
 from .activation import (
   PReLU as PReLU,
@@ -67,6 +67,16 @@
     unbox as unbox,
     PARTITION_NAME as PARTITION_NAME,
 )
+from .spmd import (
+    logical_axis_rules as logical_axis_rules,
+    set_logical_axis_rules as set_logical_axis_rules,
+    get_logical_axis_rules as get_logical_axis_rules,
+    logical_to_mesh_axes,
+    logical_to_mesh,
+    with_logical_constraint,
+    LogicallyPartitioned as LogicallyPartitioned,
+    with_logical_partitioning as with_logical_partitioning,
+)
 from .initializers import (
   ones as ones,
   zeros as zeros

flax/linen/partitioning.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Utilities for working with pjit and partitioned models.
+"""Legacy utilities for working with pjit and partitioned models.
 
 **Experimental: please give feedback, and expect changes.**
 
@@ -28,247 +28,57 @@
 logical axis metadata to the underlying Lifted transformations.
 """
 
-import collections
-import contextlib
-import enum
 import functools
 import re
-import threading
-from typing import (Any, Callable, List, Mapping, Optional, Sequence, Tuple,
-                    Union)
+from typing import (Any, Callable, Mapping, Optional, Tuple)
+
 import flax
 from flax import linen as nn
+from flax import struct
 from flax.core.frozen_dict import freeze
 from flax.core.frozen_dict import unfreeze
 from flax.core.lift import In as ScanIn  # pylint: disable=unused-import
 from flax.core.lift import Out as ScanOut  # pylint: disable=unused-import
-import flax.struct
+from flax.linen.spmd import _axis_rules  # pylint: disable=unused-import
+from flax.linen.spmd import _AxisRules  # pylint: disable=unused-import
+from flax.linen.spmd import _is_logical_spec
+from flax.linen.spmd import _with_sharding_constraint  # pylint: disable=unused-import
+from flax.linen.spmd import Array  # pylint: disable=unused-import
+from flax.linen.spmd import ArrayPytree  # pylint: disable=unused-import
+from flax.linen.spmd import get_logical_axis_rules as get_axis_rules  # pylint: disable=unused-import
+from flax.linen.spmd import logical_axis_rules as axis_rules  # pylint: disable=unused-import
+from flax.linen.spmd import logical_to_mesh  # pylint: disable=unused-import
+from flax.linen.spmd import logical_to_mesh_axes  # pylint: disable=unused-import
+from flax.linen.spmd import LogicalPartitionSpec  # pylint: disable=unused-import
+from flax.linen.spmd import LogicalPartitionSpecPytree
+from flax.linen.spmd import LogicalRules  # pylint: disable=unused-import
+from flax.linen.spmd import PartitionSpecPytree  # pylint: disable=unused-import
+from flax.linen.spmd import RulesFallback
+from flax.linen.spmd import set_logical_axis_rules as set_axis_rules  # pylint: disable=unused-import
+from flax.linen.spmd import with_logical_constraint as with_sharding_constraint
 from flax.traverse_util import flatten_dict
 from flax.traverse_util import unflatten_dict
 import jax
-from jax.experimental import maps
 from jax.experimental import pjit
 
-# Real types and dummy aliases for documentation
-LogicalRules = Sequence[Tuple[str, Union[str, Tuple[str], None]]]
-Array = Any  # pylint: disable=invalid-name
-ArrayPytree = Any  # pylint: disable=invalid-name
-LogicalPartitionSpec = Any  # pylint: disable=invalid-name
-LogicalPartitionSpecPytree = Any  # pylint: disable=invalid-name
-PartitionSpecPytree = Any  # pylint: disable=invalid-name
 
-# Dynamic Axis Mapping Context
 # ------------------------------------------------------------------------------
-
-
-class _AxisRules:
-  """Dynamic logical axis to mesh axis binding context."""
-
-  def __init__(self):
-    self._thread_data = threading.local()
-
-  @property
-  def rules(self) -> LogicalRules:
-    if not hasattr(self._thread_data, 'rules'):
-      self._thread_data.rules = ()
-    return self._thread_data.rules
-
-  @rules.setter
-  def rules(self, value: LogicalRules):
-    self._thread_data.rules = value
-
-
-# Global axis binding context.
-_axis_rules = _AxisRules()
-
-
-def set_axis_rules(rules: LogicalRules):
-  """Sets the global logical axis to mesh axis binding."""
-  _axis_rules.rules = rules
-
-
-def get_axis_rules() -> LogicalRules:
-  """Returns the global logical axis to mesh axis binding."""
-  return _axis_rules.rules
-
-
-@contextlib.contextmanager
-def axis_rules(rules: LogicalRules):
-  """Context manager for setting the logical to mesh axis bindings."""
-  old_rules = _axis_rules.rules
-  try:
-    _axis_rules.rules = rules
-    yield
-  finally:
-    _axis_rules.rules = old_rules
-
-
-class _UnassignedAxis:
-  """Sentinel class for unassigned logical axis name."""
-
-  def __repr__(self):
-    return 'UnassignedAxis'
-
-  def __bool__(self):
-    return False
-
-
-_unassigned_axis = _UnassignedAxis()
-
-
-def _mesh_assignment_free(new_assignment, existing_assignments):
-  """Determines if a given mesh axis has already been assigned."""
-  new = set(jax.tree_util.tree_leaves(new_assignment))
-  existing = set(jax.tree_util.tree_leaves(existing_assignments))
-  if existing.intersection(new):
-    return False
-  return True
-
-
-def _logical_to_mesh_axes(
-    array_dim_names: Optional[Sequence[Optional[str]]],
-    rules: Optional[LogicalRules] = None,
-) -> Optional[List[Union[_UnassignedAxis, None, str, Tuple[str]]]]:
-  """Same as logical_to_mesh_axes, but doesn't fill in _unassigned_axis."""
-  if array_dim_names is None:
-    return None
-  if rules is None:
-    rules = _axis_rules.rules
-  axis_name_counts = collections.Counter(array_dim_names)
-  dups = tuple(
-      k for k, v in axis_name_counts.items() if v > 1 and k is not None)
-  if dups:
-    raise ValueError(
-        f'Unsupported: Dimensions {dups} occur more than once in array names.')
-  if not isinstance(rules, (tuple, list)):
-    raise ValueError('Unknown axis rule specification type.')
-  # We assign mesh axes using a priority based ruleset over logical axis names.
-  result: List[Union[_UnassignedAxis, None, str, Tuple[str]]]
-  result = [_unassigned_axis] * len(array_dim_names)
-  for rule_model_name, rule_mesh_names in rules:
-    if rule_model_name in array_dim_names:
-      pos = array_dim_names.index(rule_model_name)
-      if (_mesh_assignment_free(rule_mesh_names, result) and
-          result[pos] == _unassigned_axis):
-        result[pos] = rule_mesh_names
-  return result
-
-
-def logical_to_mesh_axes(
-    array_dim_names: Optional[Sequence[Optional[str]]],
-    rules: Optional[LogicalRules] = None,
-) -> Optional[pjit.PartitionSpec]:
-  """Compute layout for an array.
-
-  The rules are in order of precedence, and consist of pairs:
-  (ArrayDimensionName, MeshDimensionName), meaning that the given array
-  dimension (if present and unused) should be sharded across the given
-  mesh dimension (if present and unused).
-
-  A Layout of an Array is expressed as a tuple with one element for each
-  dimension in the Array. The element is either None, or is the name of a
-  mesh-dimension, meaning that this dimension of the array is sharded across
-  this dimension of the mesh.
-
-  For example, given an array with
-    array_dim_names = ('batch', 'length', 'heads', 'features')
-  and the layout rules are:
-    rules = (('batch', 'X'),
-             ('features', 'X'),
-             ('heads', 'Y'),
-             ('batch', 'Z'))
-
-  then this function will return
-
-    PartitionSpec('X', None, 'Y', None)
-
-  Args:
-    array_dim_names: Tuple of array dimension names or None.
-    rules: Optional logical to mesh rules override.  Defaults to using the
-      rules defined in the dynamic context set from the `axis_rules` function.
-
-  Returns:
-    PartitionSpec for the parameter.
-  """
-  result = _logical_to_mesh_axes(array_dim_names, rules)
-  if result is None:
-    return None
-  # We default to None - ie unsharded along the dimension.
-  result = [None if x is _unassigned_axis else x for x in result]
-  return pjit.PartitionSpec(*result)
-
-
-def _global_mesh_defined() -> bool:
-  """Checks if global xmap/pjit mesh resource environment is defined."""
-  maps_env = maps.thread_resources.env
-  return maps_env.physical_mesh.devices.shape != ()  # pylint: disable=g-explicit-bool-comparison
-
-
-class RulesFallback(enum.Enum):
-  """How a sharding constraint should behave when no matching rule is found."""
-  AXIS_IS_UNSHARDED = 'axis_is_unsharded'
-  RAISE_ERROR = 'raise_error'
-  NO_CONSTRAINT = 'no_constraint'
-
-
-def _with_sharding_constraint(x: Array, axis_resources: Optional[pjit.PartitionSpec]):
-  """Wrapper for pjit with_sharding_constraint, no-op on cpu or outside pjit."""
-  if jax.devices()[0].platform == 'cpu' or not _global_mesh_defined():
-    return x
-  else:
-    return pjit.with_sharding_constraint(x, axis_resources)
-
-
-def _with_sharding_constraint_one_fallback(
-    axis_resources: LogicalPartitionSpec,
-    x: Array,
-    fallback: RulesFallback = RulesFallback.AXIS_IS_UNSHARDED):
-  """Either imposes a sharding constraint or applies fallback."""
-  mesh_axes = _logical_to_mesh_axes(axis_resources)
-  if mesh_axes is None:
-    return _with_sharding_constraint(x, None)
-
-  if fallback == RulesFallback.AXIS_IS_UNSHARDED:
-    mesh_axes = [None if x is _unassigned_axis else x for x in mesh_axes]
-  else:
-    if any(x is _unassigned_axis for x in mesh_axes):
-      if fallback == RulesFallback.RAISE_ERROR:
-        raise ValueError(f'Axis names {axis_resources} did not match a rule')
-      else:
-        return x
-  return _with_sharding_constraint(x, pjit.PartitionSpec(*mesh_axes))
-
-
-def _is_logical_spec(x):
-  return x is None or (
-      isinstance(x, tuple) and all(isinstance(e, str) or e is None for e in x))
-
-
-def with_sharding_constraint(
-    x: ArrayPytree,
-    logical_axis_resources: LogicalPartitionSpecPytree,
-    fallback: RulesFallback = RulesFallback.AXIS_IS_UNSHARDED):
-  """Version of pjit's with_sharding_constraint that uses logical axis names."""
-  # If no axis binding is set, this is a no-op.
-  if not _axis_rules.rules or logical_axis_resources is None:
-    return x
-  # Translate logical names to mesh assignments.
-  return jax.tree_util.tree_map(
-      functools.partial(
-          _with_sharding_constraint_one_fallback, fallback=fallback),
-      logical_axis_resources,
-      x,
-      is_leaf=_is_logical_spec)
+# NOTICE:  This experimental partitioning utility API is deprecated
+#
+# We intend to continue supporting it indefinitely for those using it, but
+# we encourage new users to adopt the simpler metadata handling system found
+# in "spmd.py".
+# ------------------------------------------------------------------------------
 
 
 # Annotated parameters and Module axis metadata handling.
 # ------------------------------------------------------------------------------
 
 
-@flax.struct.dataclass
+@struct.dataclass
 class AxisMetadata:
   """Contains a tuple of axis names, which is passed through FLAX."""
-  names: LogicalPartitionSpecPytree = flax.struct.field(pytree_node=False)
+  names: LogicalPartitionSpecPytree = struct.field(pytree_node=False)
 
 
 def _param_with_axes_sow_reduce_fn(x, y):
@@ -306,7 +116,7 @@ def param_with_axes(
     init_fn,
     *init_args,
     axes: Optional[Tuple[str, ...]] = None,
-    module: Optional[nn.Module] = None):
+    module: Optional['nn.Module'] = None):
   """Declares and returns a parameter with logical axes in the current Module.
 
   See :mod:`flax.linen.module.param` for original docstring.
@@ -340,7 +150,7 @@ def param_with_axes(
                                             pjit.PartitionSpec(*axes))
     # record logical axis constraint for global axis metadata
     module.sow(
-        'params_axes', f'{name}_axes', AxisMetadata(axes), # type: ignore
+        'params_axes', f'{name}_axes', AxisMetadata(axes),  # type: ignore
         reduce_fn=_param_with_axes_sow_reduce_fn)
   return module_param
 
@@ -418,7 +228,7 @@ def variable_with_axes(
     init_fn,
     *init_args,
     axes: Optional[Tuple[str, ...]] = None,
-    module: Optional[nn.Module] = None,
+    module: Optional['nn.Module'] = None,
     fallback: RulesFallback = RulesFallback.AXIS_IS_UNSHARDED):
   """Declares and returns a variable with logical axes in the current Module.
 
@@ -459,7 +269,7 @@ def variable_with_axes(
   if axes is not None:
     # record logical axis constraint for global axis metadata
     module.sow(
-        f'{collection}_axes', f'{name}_axes', AxisMetadata(axes), # type: ignore
+        f'{collection}_axes', f'{name}_axes', AxisMetadata(axes),  # type: ignore
         reduce_fn=_param_with_axes_sow_reduce_fn)
   return module_var
 
@@ -567,7 +377,7 @@ def remove_fn(x):
 
 # pylint: disable=dangerous-default-value
 def scan_with_axes(
-    target: flax.linen.transforms.Target,
+    target: 'flax.linen.transforms.Target',
     variable_axes: Mapping[flax.core.lift.CollectionFilter,
                            flax.core.lift.InOutScanAxis] = {},
     variable_broadcast: flax.core.lift.CollectionFilter = False,
@@ -581,7 +391,7 @@ def scan_with_axes(
     axis_name: str = 'layers',
     axes_collections: Tuple[str, ...] = ('params',),
     data_transform: Optional[Callable[..., Any]] = None,
-    methods=None) -> flax.linen.transforms.Target:
+    methods=None) -> 'flax.linen.transforms.Target':
   """Wrapped version of nn.scan that handles logical axis metadata."""
 
   # we broadcast the static metadata collections.
@@ -616,7 +426,7 @@ def scan_with_axes(
 
 
 # pylint: disable=dangerous-default-value
-def vmap_with_axes(target: flax.linen.transforms.Target,
+def vmap_with_axes(target: 'flax.linen.transforms.Target',
                    variable_axes: Mapping[flax.core.lift.CollectionFilter,
                                           flax.core.lift.InOutAxis],
                    split_rngs: Mapping[flax.core.lift.PRNGSequenceFilter,
@@ -627,7 +437,7 @@ def vmap_with_axes(target: flax.linen.transforms.Target,
                    axis_name: Optional[str] = None,
                    partitioning_axis_names: Mapping[Any, str] = {},
                    spmd_axis_name: Optional[str] = None,
-                   methods=None) -> flax.linen.transforms.Target:
+                   methods=None) -> 'flax.linen.transforms.Target':
   """Wrapped version of nn.vmap that handles logical axis metadata."""
 
   # tell normal vmap to broadcast axis metadata.