# Copyright 2021 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import collections
from collections import OrderedDict
from collections.abc import Mapping, Sequence
import dataclasses
import functools
import itertools
import math
from typing import Any, NamedTuple, Union, cast
from jax._src import core
from jax._src import config
from jax._src import mesh as mesh_lib
from jax._src import sharding as jsharding
from jax._src import sharding_specs
from jax._src import tree_util
from jax._src import util
from jax._src import source_info_util
from jax._src import xla_bridge
from jax._src import mesh_utils
from jax._src.lib import xla_client as xc
from jax._src.lib.mlir.dialects import sdy
from jax._src.op_shardings import (
are_op_shardings_equal, get_num_ways_dim_sharded, is_op_sharding_replicated)
from jax._src.partition_spec import PartitionSpec
from jax._src.util import safe_map, safe_zip, use_cpp_class, use_cpp_method
import numpy as np
Shape = tuple[int, ...]
Device = xc.Device
Index = tuple[slice, ...]
XLADeviceAssignment = tuple[Device, ...]
# TODO(yashkatariya): Remove this after 3 months of deprecation.
XLACompatibleSharding = jsharding.Sharding
@dataclasses.dataclass(frozen=True)
class TransferToMemoryKind:
memory_kind: str
@util.cache(max_size=128, trace_context_in_key=False)
def _check_mesh_resource_axis(mesh, parsed_pspec, _manual_axes):
for p in parsed_pspec:
if p is not None:
for r in p:
if r not in mesh.shape:
raise ValueError(
f"Resource axis: {r} of {parsed_pspec.get_partition_spec()} "
f"is not found in mesh: {tuple(mesh.shape.keys())}.")
if r in _manual_axes:
raise ValueError(
f"Axis: {r} of {parsed_pspec.get_partition_spec()} "
f"is also found in manual_axes: {_manual_axes}.") from None
@util.cache(max_size=128, trace_context_in_key=False)
def _check_axis_type_consistency(mesh, parsed_pspec):
for p in parsed_pspec:
if p is not None:
if not all(mesh._name_to_type[p[0]] == mesh._name_to_type[r] for r in p):
raise ValueError(
'AxisTypes should be the same in a tuple subset of PartitionSpec:'
f' {parsed_pspec.get_partition_spec()}. Got subset {p} with axis'
f' types: ({", ".join(str(mesh._name_to_type[r]) for r in p)})')
if mesh_lib.AxisTypes.Hidden not in mesh.axis_types and None in parsed_pspec:
raise ValueError(
f'PartitionSpec {parsed_pspec.get_partition_spec()} cannot contain'
' `P.UNCONSTRAINED` when no mesh axis_types are `Hidden`. Got mesh'
f' axis_types: {mesh.axis_types}')
def hashed_index(x) -> int:
# This works for both `pjit` indices and `pmap` indices (which might
# have an integer instead of a slice).
assert all(v.step is None for v in x if isinstance(v, slice))
return hash(tuple((v.start, v.stop) if isinstance(v, slice) else v for v in x))
@util.cache(max_size=4096, trace_context_in_key=False)
def device_replica_id_map(sharding, global_shape: Shape) -> Mapping[Device, int]:
try:
device_indices_map_fn = sharding.devices_indices_map
except AttributeError:
raise ValueError(
f'Cannot calculate replica ids from sharding: {sharding}. Please '
'create a device to index mapping for your sharding from which replica '
'ids will be calculated.') from None
index_to_replica: dict[int, int] = collections.Counter()
out = {}
for device, index in device_indices_map_fn(global_shape).items():
h_index = hashed_index(index)
replica_id = index_to_replica[h_index]
index_to_replica[h_index] += 1
out[device] = replica_id
return out
@dataclasses.dataclass
class SdyDimSharding:
axes: Sequence[str]
is_closed: bool
priority: int | None = None
def build(self) -> sdy.DimensionShardingAttr:
return sdy.DimensionShardingAttr.get(
[sdy.AxisRefAttr.get(axis) for axis in self.axes],
is_closed=self.is_closed,
priority=self.priority)
def __repr__(self):
return f'SdyDimSharding({self._custom_repr()})'
def _custom_repr(self):
axes_repr = ', '.join(f"'{a}'" for a in self.axes)
open_repr = ''
if not self.is_closed:
open_repr = ', ?' if self.axes else '?'
priority_repr = '' if self.priority is None else f'p{self.priority}'
return f'{{{axes_repr}{open_repr}}}{priority_repr}'
@dataclasses.dataclass
class SdyArraySharding:
mesh_shape: tuple[tuple[str, int], ...] | None
dimension_shardings: Sequence[SdyDimSharding]
logical_device_ids: tuple[int, ...] | None = None
replicated_axes: tuple[str, ...] = ()
def build(self) -> sdy.TensorShardingAttr:
if self.mesh_shape is None:
mesh_attr = sdy.MeshAttr.get([])
else:
ldi = ([] if self.logical_device_ids is None else
list(self.logical_device_ids))
mesh_attr = sdy.MeshAttr.get(
[sdy.MeshAxisAttr.get(name, size) for name, size in self.mesh_shape],
ldi)
return sdy.TensorShardingAttr.get(
mesh_attr,
[dim_sharding.build() for dim_sharding in self.dimension_shardings],
replicated_axes=[sdy.AxisRefAttr.get(axis) for axis in self.replicated_axes])
def __repr__(self):
dim_sharding_repr = ', '.join(
d._custom_repr() for d in self.dimension_shardings)
device_id_repr = (f', device_ids={self.logical_device_ids}'
if self.logical_device_ids is not None else '')
rar = (f', replicated_axes={self.replicated_axes}'
if self.replicated_axes else '')
return f"SdyArraySharding([{dim_sharding_repr}]{device_id_repr}{rar})"
@dataclasses.dataclass
class SdyArrayShardingList:
shardings: Sequence[SdyArraySharding]
def build(self) -> sdy.TensorShardingPerValueAttr:
return sdy.TensorShardingPerValueAttr.get(
[sharding.build() for sharding in self.shardings])
@util.cache(max_size=4096, trace_context_in_key=False)
def named_sharding_to_xla_hlo_sharding(
self, num_dimensions: int) -> xc.HloSharding:
mesh_shape = self.mesh.shape
array_mapping = get_array_mapping(self._parsed_pspec)
mesh_axis_pos = {name: i for i, name in enumerate(self.mesh.axis_names)}
special_axes = {}
mesh_manual_axes = {n for n, t in self.mesh._name_to_type.items()
if t == mesh_lib.AxisTypes.Collective}
manual_axes = self._manual_axes.union(mesh_manual_axes)
if manual_axes:
axis_names = self.mesh.axis_names
for manual_axis in manual_axes:
special_axes[axis_names.index(manual_axis)] = xc.OpSharding.Type.MANUAL
replicated_mesh_axes = []
for i, (axis_name, axis_val) in enumerate(mesh_shape.items()):
if axis_name not in array_mapping: # type: ignore
replicated_mesh_axes.append((i, axis_val))
if len(replicated_mesh_axes) == len(mesh_shape) and not special_axes:
return xc.HloSharding.replicate()
mesh_permutation = []
new_mesh_shape = [1] * num_dimensions
for name, pos in sorted(array_mapping.items(), key=lambda x: x[1]): # type: ignore
new_mesh_shape[pos] *= mesh_shape[name]
mesh_permutation.append(mesh_axis_pos[name])
last_tile_dims = []
if replicated_mesh_axes:
axes_by_type = collections.defaultdict(list)
size_by_type = collections.defaultdict(lambda: 1) # type: ignore
assert {x[0] for x in replicated_mesh_axes}.issuperset(set(special_axes.keys()))
for i, size in replicated_mesh_axes:
ty = special_axes.get(i, xc.OpSharding.Type.REPLICATED)
axes_by_type[ty].append(i)
size_by_type[ty] *= size
for ty, axes in sorted(axes_by_type.items(), key=lambda x: x[0].value):
last_tile_dims.append(ty)
new_mesh_shape.append(size_by_type[ty])
mesh_permutation.extend(axes)
# Explanation of the parameters of `HloSharding.iota_tile`.
# This is the HloShardingV2 format:
# * dims: How many ways each dimension is sharded.
# Replicated/Manual dims are added added at the end
# * reshape_dims: This is the just the shape of the mesh.
# * transpose_perm: This is the order in which mesh axes in PartitionSpec
# appear relative to mesh.axis_names order.
# * subgroup_types: List of type of OpSharding. Type can be REPLICATED and MANUAL.
# Let's see an example:
# Consider input_shape=(8, 4, 2, 2), mesh={'a': 2, 'b': 2, 'c': 2, 'd': 2}
# and partition_spec=P(None, ('d', 'b'), 'c').
# Arguments to iota_tile will be:
# dims = [1, 4, 2, 1, 2] # 'a' is replicated hence `2` is at the end.
# reshape_dims = [2, 2, 2, 2]
# transpose_perm = [3, 1, 2, 0] # 'a' is replicated hence 0 is at the end
# subgroup_types = [xc.OpSharding.Type.REPLICATED]
dims = new_mesh_shape
reshape_dims = self.mesh.axis_sizes
if self._logical_device_ids is None:
return xc.HloSharding.iota_tile(
dims=dims, reshape_dims=reshape_dims, transpose_perm=mesh_permutation,
subgroup_types=last_tile_dims)
else:
return xc.HloSharding.subgroup_with_device_ordering(
np.asarray(self._logical_device_ids)
.reshape(dims).reshape(reshape_dims).transpose(mesh_permutation)
.reshape(dims), subgroup_types=last_tile_dims)
@use_cpp_class(xc.NamedSharding)
class NamedSharding(jsharding.Sharding):
r"""A :class:`NamedSharding` expresses sharding using named axes.
A :class:`NamedSharding` is a pair of a :class:`Mesh` of devices and
:class:`PartitionSpec` which describes how to shard an array across that
mesh.
A :class:`Mesh` is a multidimensional NumPy array of JAX devices,
where each axis of the mesh has a name, e.g. ``'x'`` or ``'y'``.
A :class:`PartitionSpec` is a tuple, whose elements can be a ``None``,
a mesh axis, or a tuple of mesh axes. Each element describes how an input
dimension is partitioned across zero or more mesh dimensions. For example,
``PartitionSpec('x', 'y')`` says that the first dimension of data
is sharded across ``x`` axis of the mesh, and the second dimension is sharded
across ``y`` axis of the mesh.
The Distributed arrays and automatic parallelization
(https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html#namedsharding-gives-a-way-to-express-shardings-with-names)
tutorial has more details and diagrams that explain how
:class:`Mesh` and :class:`PartitionSpec` are used.
Args:
mesh: A :class:`jax.sharding.Mesh` object.
spec: A :class:`jax.sharding.PartitionSpec` object.
Examples:
>>> from jax.sharding import Mesh
>>> from jax.sharding import PartitionSpec as P
>>> mesh = Mesh(np.array(jax.devices()).reshape(2, 4), ('x', 'y'))
>>> spec = P('x', 'y')
>>> named_sharding = jax.sharding.NamedSharding(mesh, spec)
"""
mesh: mesh_lib.Mesh | mesh_lib.AbstractMesh
spec: PartitionSpec
_memory_kind: str | None
_parsed_pspec: ParsedPartitionSpec
_manual_axes: frozenset[MeshAxisName]
_logical_device_ids: tuple[int, ...] | None
@use_cpp_method()
def __init__(
self, mesh: mesh_lib.Mesh | mesh_lib.AbstractMesh, spec: PartitionSpec, *,
memory_kind: str | None = None, _parsed_pspec=None,
_manual_axes=frozenset(), _logical_device_ids=None):
self.mesh = mesh
self.spec = spec
self._memory_kind = memory_kind
self._manual_axes = _manual_axes
self._logical_device_ids = _logical_device_ids
self._parsed_pspec = preprocess(self.mesh, self.spec, _parsed_pspec)
def __repr__(self):
mem = '' if self.memory_kind is None else f', memory_kind={self.memory_kind}'
ldi = ('' if self._logical_device_ids is None else
f', logical_device_ids={self._logical_device_ids}')
if isinstance(self.mesh, mesh_lib.AbstractMesh):
mesh_repr = f"{self.mesh}"
else:
nv_str = ", ".join(f"'{n}': {v}" for n, v in self.mesh.shape.items())
mesh_repr = f"Mesh({nv_str})"
return f'NamedSharding(mesh={mesh_repr}, spec={self.spec}{mem}{ldi})'
def __reduce__(self):
return (type(self), (self.mesh, self.spec),
{'memory_kind': self.memory_kind,
'_manual_axes': self._manual_axes,
'_logical_device_ids': self._logical_device_ids})
@property
def memory_kind(self) -> str | None:
return self._memory_kind
def __hash__(self):
if not hasattr(self, '_hash'):
self._hash = hash(
(self.mesh, self.memory_kind, self._parsed_pspec, self._manual_axes,
self._logical_device_ids))
return self._hash
def __eq__(self, other):
if not isinstance(other, NamedSharding):
return False
if self is other:
return True
if (self._parsed_pspec != other._parsed_pspec
or self.memory_kind != other.memory_kind
or self._manual_axes != other._manual_axes
or self._logical_device_ids != other._logical_device_ids):
return False
return self.mesh is other.mesh or self.mesh == other.mesh
[docs]
def check_compatible_aval(self, aval_shape: Shape) -> None:
assert self._parsed_pspec is not None
if len(aval_shape) < len(self._parsed_pspec):
extra_msg = (' For scalars the PartitionSpec should be P()'
if len(aval_shape) == 0 else '')
raise ValueError(
f"Sharding {self} is only valid for values of rank at least "
f"{len(self._parsed_pspec)}, but was applied to a value of rank "
f"{len(aval_shape)}.{extra_msg}")
@classmethod
def _from_parsed_pspec(
cls, mesh, parsed_pspec, *, memory_kind=None, _manual_axes=frozenset(),
_logical_device_ids=None,
):
return cls(mesh, parsed_pspec.get_partition_spec(),
memory_kind=memory_kind, _parsed_pspec=parsed_pspec,
_manual_axes=_manual_axes,
_logical_device_ids=_logical_device_ids)
@property
def num_devices(self) -> int:
return self.mesh.size
@property
def device_set(self) -> set[Device]:
if isinstance(self.mesh, mesh_lib.AbstractMesh):
raise ValueError(
'device_set is not implemented for `jax.sharding.AbstractMesh`.')
return self.mesh._flat_devices_set
@property
def _device_assignment(self) -> XLADeviceAssignment:
if isinstance(self.mesh, mesh_lib.AbstractMesh):
raise ValueError('_device_assignment is not implemented for'
' `jax.sharding.AbstractMesh`.')
return self.mesh._flat_devices_tuple
@property
def is_fully_addressable(self) -> bool:
if isinstance(self.mesh, mesh_lib.AbstractMesh):
raise ValueError('is_fully_addressable is not implemented for '
'`jax.sharding.AbstractMesh`.')
# Speed up `is_fully_addressable` since there is a high chance that the
# mesh across multiple NamedSharding objects will be the same.
return not self.mesh.is_multi_process
@property
def addressable_devices(self) -> set[Device]:
if isinstance(self.mesh, mesh_lib.AbstractMesh):
raise ValueError('addressable_devices is not implemented for '
'`jax.sharding.AbstractMesh`.')
# Override addressable devices because there is a high chance that the mesh
# across multiple NamedSharding objects will be the same.
return self.mesh._local_devices_set
@functools.cached_property
def is_fully_replicated(self) -> bool:
if self.mesh.size == 1:
return True
array_mapping = cast(ParsedPartitionSpec, get_array_mapping(self._parsed_pspec))
mesh_shape = self.mesh.shape
num_partitions = 1
for name in array_mapping:
num_partitions *= mesh_shape[name]
return num_partitions == 1
[docs]
def with_memory_kind(self, kind: str) -> NamedSharding:
return NamedSharding(self.mesh, self.spec, memory_kind=kind)
[docs]
def with_spec(self, spec: PartitionSpec | Sequence[Any]) -> NamedSharding:
if not isinstance(spec, PartitionSpec):
spec = PartitionSpec(*spec)
return NamedSharding(self.mesh, spec, memory_kind=self.memory_kind)
def _to_xla_hlo_sharding(self, num_dimensions: int) -> xc.HloSharding:
return named_sharding_to_xla_hlo_sharding(self, num_dimensions)
def _to_sdy_sharding(self, num_dimensions: int) -> SdyArraySharding:
dim_shardings = [SdyDimSharding(axes=[], is_closed=True)
for _ in range(num_dimensions)]
for i, dim_spec in enumerate(self._parsed_pspec):
if dim_spec is None:
dim_shardings[i].is_closed = False
elif not dim_spec:
# Already empty and closed sharding.
pass
else:
dim_shardings[i].axes = dim_spec
return SdyArraySharding(self.mesh.shape_tuple, dim_shardings,
self._logical_device_ids)
# TODO(yashkatariya): Upstream this into `_to_sdy_sharding` maybe with an extra
# parameter to it `_to_sdy_sharding(self, ndim, modify_wrt_axis_types=False)`
def modify_sdy_sharding_wrt_axis_types(sdy_sharding: SdyArraySharding, mesh):
if mesh._any_axis_hidden:
dim_shardings, used_axes = [], [] # type: ignore
for d in sdy_sharding.dimension_shardings:
# TODO(yashkatariya): Maybe if any mesh axis is auto, mark all axes as open?
dim_shardings.append(SdyDimSharding(axes=[], is_closed=False)
if not d.axes and d.is_closed else d)
used_axes.extend(d.axes)
remaining_axes = set(mesh.axis_names) - set(used_axes)
replicated_axes = tuple(r for r in remaining_axes
if mesh._name_to_type[r] == mesh_lib.AxisTypes.Visible)
return SdyArraySharding(sdy_sharding.mesh_shape, dim_shardings,
sdy_sharding.logical_device_ids, replicated_axes)
return sdy_sharding
@util.cache(max_size=128, trace_context_in_key=False)
def get_replicated_hlo_sharding():
return xc.HloSharding.replicate()
@use_cpp_class(xc.SingleDeviceSharding)
class SingleDeviceSharding(jsharding.Sharding):
"""A :class:`Sharding` that places its data on a single device.
Args:
device: A single :py:class:`Device`.
Examples:
>>> single_device_sharding = jax.sharding.SingleDeviceSharding(
... jax.devices()[0])
"""
_device: Device
_memory_kind: str | None
@use_cpp_method()
def __init__(self, device: Device, *, memory_kind: str | None = None):
self._device = device
self._memory_kind = memory_kind
def __reduce__(self):
return type(self), (self._device,), {'memory_kind': self._memory_kind}
def __repr__(self):
mem = '' if self._memory_kind is None else f', memory_kind={self._memory_kind}'
return f"SingleDeviceSharding(device={self._device!r}{mem})"
def __hash__(self):
if not hasattr(self, '_hash'):
self._hash = hash((self._device, self.memory_kind))
return self._hash
def __eq__(self, other):
if not isinstance(other, SingleDeviceSharding):
return False
if self is other:
return True
return (self._device == other._device and
self.memory_kind == other.memory_kind)
@property
def num_devices(self) -> int:
return len(self.device_set)
@property
def device_set(self) -> set[Device]:
return {self._device}
@property
def memory_kind(self) -> str | None:
return self._memory_kind
def with_memory_kind(self, kind: str) -> SingleDeviceSharding:
return SingleDeviceSharding(self._device, memory_kind=kind)
def devices_indices_map(self, global_shape: Shape) -> Mapping[Device, Index]: # type: ignore
return {self._device: (slice(None),) * len(global_shape)}
@property
def _device_assignment(self) -> XLADeviceAssignment:
return (self._device,)
def _to_xla_hlo_sharding(self, num_dimensions: int) -> xc.HloSharding:
return get_replicated_hlo_sharding()
def _to_sdy_sharding(self, num_dimensions: int) -> SdyArraySharding:
sdy_dim_sharding = [SdyDimSharding(axes=[], is_closed=True)
for _ in range(num_dimensions)]
return SdyArraySharding(None, sdy_dim_sharding)
@property
def is_fully_replicated(self) -> bool:
return True
@property
def is_fully_addressable(self) -> bool:
return True
@util.cache(max_size=4096, trace_context_in_key=False)
def pmap_sharding_devices_indices_map(
self, global_shape: Shape) -> Mapping[Device, Index]:
self.shard_shape(global_shape) # raises a good error message
indices = sharding_specs.spec_to_indices(global_shape, self.sharding_spec)
return dict(safe_zip(self.devices.flat, indices)) # type: ignore[arg-type]
@use_cpp_class(xc.PmapSharding)
class PmapSharding(jsharding.Sharding):
"""Describes a sharding used by :func:`jax.pmap`."""
devices: np.ndarray
sharding_spec: sharding_specs.ShardingSpec
_internal_device_list: xc.DeviceList
@use_cpp_method()
def __init__(self, devices: Sequence[Device] | np.ndarray,
sharding_spec: sharding_specs.ShardingSpec):
self.devices = np.asarray(devices)
# The sharding spec should be pmap's sharding spec.
self.sharding_spec = sharding_spec
def __reduce__(self):
return (type(self), (self.devices, self.sharding_spec),
{'memory_kind': self.memory_kind})
def __eq__(self, other):
if not isinstance(other, PmapSharding):
return False
if self is other:
return True
return (self.sharding_spec == other.sharding_spec and
self.devices.shape == other.devices.shape and
self._internal_device_list == other._internal_device_list)
def __hash__(self):
if not hasattr(self, '_hash'):
self._hash = hash((self._internal_device_list, self.sharding_spec))
return self._hash
def __str__(self):
device_ids = [d.id for d in self.devices.flat]
return (f'PmapSharding(sharding_spec={self.sharding_spec}, '
f'{device_ids=}, '
f'device_platform={self.devices.flat[0].platform.upper()}, '
f'device_shape={self.devices.shape})')
def __repr__(self):
return (f'PmapSharding(sharding_spec={self.sharding_spec}, '
f'devices={self.devices})')
def is_equivalent_to(self: PmapSharding, other: PmapSharding, # type: ignore
ndim: int) -> bool:
return self == other
# TODO(yashkatariya): Expose `sharded_dim_size` in the API if required.
@classmethod
def default(cls, shape: Shape, sharded_dim: int | None = 0,
devices: Sequence[xc.Device] | None = None) -> PmapSharding:
"""Creates a :class:`PmapSharding` which matches the default placement
used by :func:`jax.pmap`.
Args:
shape: The shape of the input array.
sharded_dim: Dimension the input array is sharded on. Defaults to 0.
devices: Optional sequence of devices to use. If omitted, the implicit
device order used by pmap is used, which is the order of
:func:`jax.local_devices`.
"""
if sharded_dim is None:
if devices is None:
raise ValueError("One of sharded_dim or devices must be set.")
nrep = len(devices)
return cls(np.array(devices),
sharding_specs.pmap_sharding_spec(nrep, nrep, shape, None))
# The dtype doesn't matter here. Its only used for creating the
# sharding_spec.
sharding_spec = sharding_specs.create_pmap_sharding_spec(
tuple(shape), sharded_dim)
num_ways_sharded = None
for s in sharding_spec.sharding:
if isinstance(s, sharding_specs.Unstacked):
assert num_ways_sharded is None
num_ways_sharded = s.size
elif isinstance(s, sharding_specs.Chunked):
assert num_ways_sharded is None
if len(s.chunks) == 1:
num_ways_sharded = s.chunks[0]
else:
raise NotImplementedError(
'Multiple chunks in Chunked dimension not supported.')
if devices is None:
pmap_devices: np.ndarray = np.array(
xla_bridge.local_devices()[:num_ways_sharded])
else:
pmap_devices = np.array(devices)
return cls(pmap_devices, sharding_spec)
@property
def num_devices(self) -> int:
return len(self.device_set)
@functools.cached_property
def device_set(self) -> set[Device]:
return set(self.devices.flat)
def devices_indices_map(self, global_shape: Shape) -> Mapping[Device, Index]:
return pmap_sharding_devices_indices_map(self, global_shape)
@functools.cached_property
def _device_assignment(self) -> XLADeviceAssignment:
return tuple(self.devices.flat)
@property
def memory_kind(self) -> str | None:
try:
return self._internal_device_list.default_memory_kind
except:
return None
def with_memory_kind(self, kind: str):
raise NotImplementedError("pmap does not support memories.")
def _to_xla_hlo_sharding(self, num_dimensions: int) -> xc.HloSharding:
raise NotImplementedError("pmap doesn't use OpSharding.")
def _to_sdy_sharding(self, num_dimensions: int) -> SdyArraySharding:
raise NotImplementedError("pmap doesn't use SdyArraySharding.")
@functools.cached_property
def is_fully_replicated(self) -> bool:
for s in self.sharding_spec.sharding:
if isinstance(s, (sharding_specs.Unstacked, sharding_specs.Chunked)):
return False
return True
@functools.cached_property
def is_fully_addressable(self) -> bool:
return self._internal_device_list.is_fully_addressable
def shard_shape(self, global_shape: Shape) -> Shape:
sharded_dim = None
sharded_dim_size = None
for i, s in enumerate(self.sharding_spec.sharding):
if isinstance(s, sharding_specs.Unstacked):
sharded_dim = i
sharded_dim_size = s.size
sharded_shape = util.tuple_delete(global_shape, sharded_dim)
break
elif isinstance(s, sharding_specs.Chunked):
sharded_dim = i
assert len(s.chunks) == 1, s.chunks
sharded_dim_size = s.chunks[0]
sharded_shape = util.tuple_update(global_shape, sharded_dim, 1)
break
if sharded_dim is None:
return global_shape
if global_shape[sharded_dim] != sharded_dim_size:
raise ValueError(
f'The sharded dimension must be equal to the number of '
f'devices passed to PmapSharding. Got sharded dimension {sharded_dim} '
f'with value {global_shape[sharded_dim]} in shape {global_shape} and '
f'the number of devices={len(self._device_assignment)}')
return sharded_shape
def _op_sharding_to_pos_sharding(
op_sharding: xc.OpSharding | xc.HloSharding,
device_assignment: Sequence[xc.Device],
memory_kind: str | None = None) -> PositionalSharding:
if isinstance(op_sharding, xc.OpSharding):
op_sharding = xc.HloSharding.from_proto(op_sharding)
if op_sharding.is_replicated():
return PositionalSharding(
device_assignment, memory_kind=memory_kind).replicate()
if len(op_sharding.subgroup_types()) > 1:
raise NotImplementedError(
'Unhandled HloSharding type. Please open a bug report!'
)
name = device_assignment[0].platform.upper()
ids = np.array(
[DeviceIdSet(name, i) for i in op_sharding.tile_assignment_devices()]
)
p = PositionalSharding._remake(tuple(device_assignment), ids,
memory_kind=memory_kind)
p = p.reshape(op_sharding.tile_assignment_dimensions())
if op_sharding.replicate_on_last_tile_dim():
p = p.replicate(-1, keepdims=False)
return p
@util.cache(max_size=4096, trace_context_in_key=False)
def _positional_sharding_to_xla_hlo_sharding(
self, num_dimensions: int) -> xc.HloSharding:
if self.shape == (1,) * self.ndim:
return get_replicated_hlo_sharding()
pbuf = xc.OpSharding()
shape = self.shape[self.ndim - num_dimensions:] # 'rank promotion' of val
set_size, = {len(device_set) for device_set in self._ids.flat}
pbuf.type = xc.OpSharding.Type.OTHER
if set_size > 1:
pbuf.last_tile_dims = [xc.OpSharding.Type.REPLICATED]
pbuf.tile_assignment_dimensions = (*shape, set_size)
else:
pbuf.tile_assignment_dimensions = shape
pbuf.tile_assignment_devices = [i for ids in self._ids.flat for i in ids]
product_of_dims = math.prod(pbuf.tile_assignment_dimensions)
num_devices = len(pbuf.tile_assignment_devices)
assert product_of_dims == num_devices, (product_of_dims, num_devices)
return xc.HloSharding.from_proto(pbuf)
class PositionalSharding(jsharding.Sharding):
_devices: tuple[xc.Device, ...]
_memory_kind: str | None
_ids: np.ndarray # dtype DeviceIdSet
def __init__(self, devices: Sequence[xc.Device] | np.ndarray,
*, memory_kind: str | None = None):
super().__init__()
if not isinstance(devices, np.ndarray):
devices = np.array(devices, dtype='object')
if not devices.size:
raise ValueError(f"{self.__class__.__name__}.__init__ requires at least "
f"one device, got {devices}")
self._devices = tuple(devices.flat)
self._memory_kind = memory_kind
name = self._devices[0].platform.upper()
self._ids = np.array([DeviceIdSet(name, i) for i in range(devices.size)],
dtype='object').reshape(devices.shape)
self._internal_device_list = xc.DeviceList(self._devices)
self._memory_kind = xc.check_and_canonicalize_memory_kind(
self._memory_kind, self._internal_device_list)
@property
def shape(self):
return self._ids.shape
@property
def ndim(self):
return self._ids.ndim
def __repr__(self) -> str:
cls_name = self.__class__.__name__
ids = self._ids.copy()
platform_name = self._devices[0].platform.upper()
for idx, x in np.ndenumerate(ids):
ids[idx] = DeviceIdSet(platform_name, *(self._devices[i].id for i in x)) # type: ignore # numpy 2.2
body = np.array2string(ids, prefix=cls_name + '(', suffix=')',
max_line_width=100)
mem = '' if self._memory_kind is None else f', memory_kind={self._memory_kind}'
return f'{cls_name}({body}{mem}, shape={self.shape})'
def reshape(self, *shape) -> PositionalSharding:
return self._remake(self._devices, self._ids.reshape(*shape),
memory_kind=self.memory_kind)
def transpose(self, *axes) -> PositionalSharding:
return self._remake(self._devices, self._ids.transpose(*axes),
memory_kind=self.memory_kind)
T = property(transpose)
def replicate(self, axis=None, keepdims=True) -> PositionalSharding:
new_ids = self._ids.sum(axis=axis, keepdims=keepdims) # union
return self._remake(self._devices, new_ids,
memory_kind=self.memory_kind)
def check_compatible_aval(self, aval_shape: Shape) -> None:
if len(aval_shape) != len(self.shape) and not self.is_fully_replicated:
raise ValueError(
f"Sharding {self} is only valid for values of rank "
f"{len(self.shape)}, but was applied to a value of rank "
f"{len(aval_shape)}")
@classmethod
def _remake(
cls, devices: tuple[xc.Device, ...], ids: np.ndarray,
*, memory_kind: str | None = None) -> PositionalSharding:
sharding = cls(devices, memory_kind=memory_kind)
sharding._ids = ids
return sharding
# Hashable
def __hash__(self) -> int:
if not hasattr(self, '_hash'):
self._hash = hash((self._internal_device_list, self.memory_kind))
return self._hash
def __eq__(self, other) -> bool:
if not isinstance(other, PositionalSharding):
return False
if self is other:
return True
all_ids_equal = np.array_equal(self._ids,other._ids)
mem_kind_equal = self.memory_kind == other.memory_kind
if self._devices is other._devices and mem_kind_equal and all_ids_equal:
return True
return (mem_kind_equal and all_ids_equal and
self._internal_device_list == other._internal_device_list)
# Sharding interface
@property
def num_devices(self) -> int:
return len(self.device_set)
@functools.cached_property
def device_set(self) -> set[xc.Device]:
return set(self._devices)
@property
def memory_kind(self) -> str | None:
return self._memory_kind
def with_memory_kind(self, kind: str) -> PositionalSharding:
return PositionalSharding(self._devices, memory_kind=kind)
@functools.cached_property
def is_fully_replicated(self) -> bool:
return self.shape == (1,) * self.ndim
# jsharding.Sharding interface
@property
def _device_assignment(self) -> XLADeviceAssignment:
return self._devices
def _to_xla_hlo_sharding(self, num_dimensions: int) -> xc.HloSharding:
return _positional_sharding_to_xla_hlo_sharding(self, num_dimensions)
def _to_sdy_sharding(self, num_dimensions: int) -> SdyArraySharding:
raise NotImplementedError(
"PositionalSharding can't be converted to an SdyArraySharding.")
@functools.cached_property
def is_fully_addressable(self) -> bool:
return self._internal_device_list.is_fully_addressable
class DeviceIdSet:
_name: str
_ids: frozenset[int]
def __init__(self, name, *ids):
self._name = name
self._ids = frozenset(ids)
def __iter__(self):
return iter(sorted(self._ids))
def __add__(self, other) -> DeviceIdSet:
assert isinstance(other, DeviceIdSet)
return DeviceIdSet(self._name, *(self._ids | other._ids))
def __len__(self) -> int:
return len(self._ids)
def __repr__(self) -> str:
ids = ', '.join(safe_map(str, sorted(self._ids)))
return f'{{{self._name} {ids}}}'
def __hash__(self) -> int:
return hash((self._name, self._ids))
def __eq__(self, other) -> bool:
return (isinstance(other, DeviceIdSet) and self._name == other._name and
self._ids == other._ids)
@use_cpp_class(xc.GSPMDSharding)
class GSPMDSharding(jsharding.Sharding):
_devices: tuple[Device, ...]
_hlo_sharding: xc.HloSharding
_memory_kind: str | None
_device_list: xc.DeviceList | None
_internal_device_list: xc.DeviceList
@use_cpp_method()
def __init__(self, devices: Sequence[Device],
op_sharding: xc.OpSharding | xc.HloSharding,
*, memory_kind: str | None = None,
_device_list: xc.DeviceList | None = None):
self._devices = tuple(devices)
if isinstance(op_sharding, xc.OpSharding):
self._hlo_sharding = xc.HloSharding.from_proto(op_sharding)
else:
self._hlo_sharding = op_sharding
self._memory_kind = memory_kind
def __reduce__(self):
return (type(self), (self._devices, self._hlo_sharding.to_proto()),
{'memory_kind': self._memory_kind})
@functools.cached_property
def _hlo_sharding_hash(self):
if self.is_fully_replicated:
return hash(get_replicated_hlo_sharding())
return hash(self._hlo_sharding)
def __eq__(self, other):
if not isinstance(other, GSPMDSharding):
return False
if self is other:
return True
return (are_op_shardings_equal(self._hlo_sharding, other._hlo_sharding)
and self.memory_kind == other.memory_kind
and self._internal_device_list == other._internal_device_list)
def __hash__(self):
if not hasattr(self, '_hash'):
self._hash = hash((self._internal_device_list, self._hlo_sharding_hash,
self.memory_kind))
return self._hash
def __repr__(self):
mem = '' if self._memory_kind is None else f', memory_kind={self._memory_kind}'
return f'GSPMDSharding({self._hlo_sharding!r}{mem})'
def check_compatible_aval(self, aval_shape: Shape) -> None:
num_ways_dim_sharded, _ = get_num_ways_dim_sharded(self._hlo_sharding)
if len(aval_shape) < len(num_ways_dim_sharded):
raise ValueError(
f"Sharding {self} is only valid for values of rank at least "
f"{len(num_ways_dim_sharded)}, but was applied to a value of rank "
f"{len(aval_shape)}")
@property
def num_devices(self) -> int:
return len(self.device_set)
@functools.cached_property
def device_set(self) -> set[Device]:
return set(self._devices)
@property
def memory_kind(self) -> str | None:
return self._memory_kind
def with_memory_kind(self, kind: str) -> GSPMDSharding:
return GSPMDSharding(self._devices, self._hlo_sharding, memory_kind=kind)
@property
def _device_assignment(self) -> XLADeviceAssignment:
return self._devices
def _to_xla_hlo_sharding(self, num_dimensions: int) -> xc.HloSharding:
return self._hlo_sharding
def _to_sdy_sharding(self, num_dimensions: int) -> SdyArraySharding:
raise NotImplementedError(
"GSPMDSharding can't be converted to SdyArraySharding.")
@functools.cached_property
def is_fully_replicated(self) -> bool:
return is_op_sharding_replicated(self._hlo_sharding)
@functools.cached_property
def is_fully_addressable(self) -> bool:
return self._internal_device_list.is_fully_addressable
@classmethod
def get_replicated(cls, device_assignment, *, memory_kind: str | None = None):
return cls(tuple(device_assignment), get_replicated_hlo_sharding(),
memory_kind=memory_kind)
class AUTO:
def __init__(self, mesh: mesh_lib.Mesh):
self.mesh = mesh
def _to_sdy_sharding(self, ndim: int) -> SdyArraySharding:
dim_shardings = [SdyDimSharding(axes=[], is_closed=False)
for _ in range(ndim)]
return SdyArraySharding(self.mesh.shape_tuple, dim_shardings)
class UnspecifiedValue:
def __repr__(self):
return "UnspecifiedValue"
UNSPECIFIED = UnspecifiedValue()
MeshAxisName = Any
"""
ArrayMapping specifies how an ndarray should map to mesh axes.
Note that the ordering is crucial for the cases when this mapping is non-injective
(i.e. when multiple mesh axes map to the same positional axis). Then, the
order of entries of the mapping determines a major-to-minor order on mesh axes,
according to which chunks of the value along the repeated dimension will be assigned.
For example, consider a mapping {'x': 1, 'y': 1} and a mesh with shape {'x': 2, 'y': 3}.
The second dimension of the value would get chunked into 6 pieces, and assigned to the
mesh in a way that treats 'y' as the fastest changing (minor) dimension. In this case,
that would mean that a flat list of chunks would get assigned to a flattened list of
mesh devices without any modifications. If the mapping was {'y': 1, 'x': 1}, then the
mesh devices ndarray would have to be transposed before flattening and assignment.
"""
ArrayMapping = OrderedDict[MeshAxisName, int]
ArrayMappingOrAutoOrUnspecified = Union[ArrayMapping, AUTO, UnspecifiedValue]
def array_mapping_to_axis_resources(array_mapping: ArrayMapping):
if not array_mapping:
return PartitionSpec()
max_index = -1
reverse_map = collections.defaultdict(list)
for axis, index in array_mapping.items():
reverse_map[index].append(axis)
if index > max_index:
max_index = index
partitions = []
for i in range(max_index + 1):
axis = reverse_map[i]
if axis:
partitions.append(axis[0] if len(axis) == 1 else tuple(axis))
else:
partitions.append(None)
return PartitionSpec(*partitions)
def get_array_mapping(
axis_resources: ParsedPartitionSpec | AUTO | UnspecifiedValue
) -> ArrayMappingOrAutoOrUnspecified:
if isinstance(axis_resources, (AUTO, UnspecifiedValue)):
return axis_resources
return OrderedDict((axis, i)
for i, axes in enumerate(axis_resources)
if axes is not None for axis in axes)
get_single_pspec = lambda p: array_mapping_to_axis_resources(
cast(ArrayMapping, get_array_mapping(p)))
class ParsedPartitionSpec:
__slots__ = ('_user_spec', 'partitions')
def __init__(self, user_spec, partitions):
self._user_spec = user_spec
# None in partitions represents unconstrained dim.
# TODO(yashkatariya): May use a sentinel value.
self.partitions = tuple(partitions)
def get_partition_spec(self) -> PartitionSpec:
if isinstance(self._user_spec, PartitionSpec):
return self._user_spec
else:
return get_single_pspec(self)
def insert_axis_partitions(self, dim, val):
parts = self.partitions
too_short = dim - len(parts)
if too_short > 0:
parts += ((),) * too_short
new_partitions = util.tuple_insert(parts, dim, val)
return ParsedPartitionSpec(None, new_partitions)
@classmethod
def from_user_input(cls, entry, arg_name, allow_unconstrained_dims=False):
if entry is None:
return cls(entry, ())
if not isinstance(entry, PartitionSpec):
raise TypeError(f"{arg_name} are expected to be "
f"PartitionSpec instances or None, but got {entry}")
axis_specs = []
for axis_spec in entry:
if axis_spec is None:
axis_spec = ()
elif isinstance(axis_spec, (list, tuple)):
axis_spec = tuple(axis_spec)
elif axis_spec == PartitionSpec.UNCONSTRAINED:
if not allow_unconstrained_dims:
raise ValueError(f"Unconstrained dims are not allowed: {entry}")
axis_spec = None
else:
axis_spec = (axis_spec,)
axis_specs.append(axis_spec)
new_entry = PartitionSpec(
*[tuple(e) if isinstance(e, (list, tuple)) else e for e in entry])
return cls(new_entry, axis_specs)
def __hash__(self):
return hash(self.partitions)
def __eq__(self, other):
if not isinstance(other, ParsedPartitionSpec):
return False
return self.partitions == other.partitions
def __len__(self):
return len(self.partitions)
def __getitem__(self, i):
return self.partitions[i]
def __iter__(self):
return iter(self.partitions)
def __repr__(self):
return f"ParsedPartitionSpec(partitions={self.partitions})"
def preprocess(mesh, spec, parsed_pspec, _manual_axes=frozenset()):
if parsed_pspec is None:
parsed_pspec = prepare_axis_resources(
PartitionSpec() if spec is None else spec,
"NamedSharding spec", allow_unconstrained_dims=True)
_check_mesh_resource_axis(mesh, parsed_pspec, _manual_axes)
_check_axis_type_consistency(mesh, parsed_pspec)
return parsed_pspec
def prepare_axis_resources(axis_resources, arg_name,
allow_unconstrained_dims=False):
# PyTrees don't treat None values as leaves, so we use an is_leaf function.
entries, treedef = tree_util.tree_flatten(
axis_resources, is_leaf=lambda x: x is None)
what = f"{arg_name} leaf specifications"
new_entries = []
for entry in entries:
if isinstance(entry, (UnspecifiedValue, AUTO)) or entry is None:
new_entries.append(entry)
elif isinstance(entry, jsharding.Sharding):
if isinstance(entry, PmapSharding):
raise ValueError(f'One of {what} got sharding {entry} which is not '
'allowed.')
new_entries.append(entry)
else:
new_entries.append(ParsedPartitionSpec.from_user_input(
entry, what, allow_unconstrained_dims=allow_unconstrained_dims))
_check_unique_resources(new_entries, arg_name)
return tree_util.tree_unflatten(treedef, new_entries)
def _check_unique_resources(axis_resources, arg_name):
for arg_axis_resources in axis_resources:
if not arg_axis_resources: continue
if isinstance(arg_axis_resources, (UnspecifiedValue, AUTO, jsharding.Sharding)):
continue
constrained_dims = [d for d in arg_axis_resources if d is not None]
resource_counts = collections.Counter(
itertools.chain.from_iterable(constrained_dims))
if not resource_counts: continue
if resource_counts.most_common(1)[0][1] > 1:
multiple_uses = [r for r, c in resource_counts.items() if c > 1]
if multiple_uses:
raise ValueError(
f'A single {arg_name} specification can map every mesh axis to at'
' most one positional dimension, but'
f' {arg_axis_resources.get_partition_spec()} has duplicate entries'
f' for {mesh_lib.show_axes(multiple_uses)}')
# Axis environments
class AxisEnv(NamedTuple):
"""Represents a pmap mesh (only along the replica axes)."""
nreps: int
names: tuple[Any, ...]
sizes: tuple[int, ...]
@dataclasses.dataclass(frozen=True)
class SPMDAxisContext:
"""A hardware axis context for parallel computations that use the GSPMD partitioner.
This includes the mesh that will later by used to execute this computation,
as well as a set of mesh axes that are currently lowered in the MANUAL
sharding mode.
"""
mesh: mesh_lib.Mesh
manual_axes: frozenset[MeshAxisName] = frozenset()
@property
def axis_env(self):
# All collectives that touch axis_env should remember to set use_global_device_ids
# when this context is enabled!
return self.unsafe_axis_env
@property
def unsafe_axis_env(self):
return AxisEnv(
nreps=self.mesh.size,
names=self.mesh.axis_names,
sizes=tuple(self.mesh.shape.values()))
def extend_manual(self, axes: frozenset[MeshAxisName]) -> SPMDAxisContext:
return SPMDAxisContext(self.mesh, self.manual_axes | axes)
@dataclasses.dataclass(frozen=True)
class ReplicaAxisContext:
"""A hardware axis context for parallel computations that are partitioned by JAX.
Unlike in the SPMDAxisContext, this means that JAX might need to emit calls to
explicit collectives.
"""
axis_env: AxisEnv
@dataclasses.dataclass(frozen=True)
class ShardingContext:
"""A hardware axis context for parallel computations that use the sharding
interface.
This context also uses the GSPMD partitioner.
"""
num_devices: int
device_assignment: tuple[xc.Device, ...] | None = None
abstract_mesh: mesh_lib.AbstractMesh | None = None
def __post_init__(self):
if self.device_assignment is not None:
assert isinstance(self.device_assignment, tuple)
assert self.num_devices == len(self.device_assignment)
# Similar to SPMDContext as ShardingContext also uses the GSPMD partitioner.
@property
def axis_env(self):
return AxisEnv(nreps=1, names=(), sizes=())
# -------------------- XLA OpSharding to PartitionSpec --------------------
# Note that OpSharding is more expressive than PartitionSpecs, so it's not
# always possible to convert them, but the code below should at least
# support handle all cases when this is possible.
def strides_for_sizes(sizes):
"""Returns an array of strides for major-to-minor sizes."""
return np.cumprod(sizes[::-1])[::-1] // np.asarray(sizes)
def unflatten_array(named_sizes, assignment):
"""Recovers the ordering of axis names based on a device assignment.
The device assignments that this function can convert into axis orders
are of the form::
np.arange(np.prod(named_sizes.values())).transpose(...).flatten()
for some transposition ``...``. This is satisfied by all OpSharding assignments
generated from partition specs.
Arguments:
named_sizes: A dictionary mapping axis names to their sizes.
assignment: A permutation of integers between 0 and the product of all
named sizes.
Returns:
A major-to-minor list of axis names that corresponds to the given assignment.
"""
named_sizes = {name: size for name, size in named_sizes.items() if size != 1}
sizes = np.fromiter(named_sizes.values(), dtype=np.int64)
strides = strides_for_sizes(sizes)
dims = explode_superdims(sizes, unflatten_superdims(assignment))
dim_to_name = {(size, stride): name for size, stride, name in zip(sizes, strides, named_sizes)}
return [dim_to_name[d] for d in dims]
def unflatten_superdims(assignment):
"""Unflatten a list of dimension sizes and their strides that generates assignment.
If this function succeeds for a given ``assignment``, then the following property
should be satisfied::
dims_with_strides = unflatten_superdims(assignment)
base_array = np.arange(map(fst, sorted(dims_with_strides, key=snd, reverse=True)))
assignment == base_array.transpose(argsort(dims_with_strides, key=snd, reverse=True)).flatten()
That is, the returned dimensions list all sizes of the base array (with strides
indicating their initial order). The order of dimensions in the list corresponds
to the permutation that applied to the base array generates the assignment.
"""
def check(cond):
if cond: return
raise NotImplementedError("Failed to convert OpSharding into a ShardingSpec. "
"Please open a bug report!")
flat_assignment = np.asarray(assignment, dtype=np.int64)
check(flat_assignment[0] == 0)
dims = []
while flat_assignment.size > 1:
stride = flat_assignment[1]
for i in range(len(flat_assignment)):
if flat_assignment[i] != i * stride: break
else:
# After this loop i should point to an "element after the sequence", so
# we have to increment it if the whole array is a strided sequence.
i += 1
size = i
dims.append((size, stride))
assert size > 1 # Ensure progress
flat_assignment = flat_assignment[::size]
return dims
def explode_superdims(sizes, dims):
"""Explode superdims to fit a known shape.
The unflattening process might mistakenly generate too few too large dimensions.
For example, ``unflatten_superdims(np.arange(n))`` always returns ``[(n, 1)]``.
This function takes a list of such contiguous super-dimensions and splits them
into smaller dimensions such that::
set(map(fst, explode_superdims(sizes, dims))) == set(sizes)
"""
strides_to_sizes = {stride: size for size, stride in zip(sizes, strides_for_sizes(sizes))}
dims = list(reversed(dims))
final_dims = []
for size, stride in dims:
target_size = strides_to_sizes[stride]
new_dims = []
while size > target_size:
assert target_size > 1 # Ensure progress
assert size % target_size == 0
new_dims.append((target_size, stride))
size //= target_size
stride *= target_size
target_size = strides_to_sizes[stride]
assert size == target_size
new_dims.append((size, stride))
final_dims += reversed(new_dims)
return final_dims
def parse_flatten_op_sharding(hlo_sharding: xc.OpSharding | xc.HloSharding,
mesh: mesh_lib.Mesh) -> Sequence[ParsedPartitionSpec]:
if isinstance(hlo_sharding, xc.OpSharding):
hlo_sharding = xc.HloSharding.from_proto(hlo_sharding)
if hlo_sharding.tuple_elements():
out: list[ParsedPartitionSpec] = []
for s in hlo_sharding.tuple_elements():
out.extend(parse_flatten_op_sharding(s, mesh))
return out
elif hlo_sharding.is_replicated():
return [ParsedPartitionSpec(PartitionSpec(), ())]
elif hlo_sharding.is_tiled():
mesh_shape = mesh.shape
mesh_axis_order = unflatten_array(
mesh.shape, hlo_sharding.tile_assignment_devices()
)
mesh_axis = iter(mesh_axis_order)
shape = hlo_sharding.tile_assignment_dimensions()
partitions = []
for dim_size in shape:
dim_partitions = []
while dim_size > 1:
axis = next(mesh_axis)
axis_size = mesh_shape[axis]
assert dim_size % axis_size == 0
dim_size //= axis_size
dim_partitions.append(axis)
partitions.append(tuple(dim_partitions))
if len(hlo_sharding.subgroup_types()) > 1:
raise NotImplementedError(
'Unhandled HloSharding type. Please open a bug report!'
)
if hlo_sharding.replicate_on_last_tile_dim():
partitions = partitions[:-1]
while partitions and partitions[-1] == ():
partitions.pop()
return [ParsedPartitionSpec(None, partitions)]
else:
raise AssertionError("Unhandled OpSharding type. Please open a bug report!")
def _slice_as_tuple(s: slice):
assert s.step is None
return (s.start, s.stop)
class NonUniformShardingError(ValueError):
"""Raised when sharding is not uniform across processes."""
def get_process_index_and_count(
tensor_sharding: jsharding.Sharding, dim: int, ndims: int) -> tuple[int, int]:
"""Get current process index and number of unique processes for given dimension.
This function facilitates mapping of process-level data to individual
devices. Each process can use its index to obtain the data corresponding
to that index. If process level data is sharded on multiple dimensions
this function can be used to build the cross product of indices in
each sharded axis. Processes that need to load the same data will have
the same index. For shardings whose per-process data is not distributed
on a grid, the number of distinct shards will be such that it is possible to
build the target shape while maintaining a "cube" shape of local-process data.
For example, in case of 4 hosts with sharding distributed like so:
1234
2143
For dim 0 (rows): all processes need to access all rows, so we return (0, 1)
For dim 1 (cols):
process 1 and 2 returns index 0 out of 2 (need cols 0 and 1),
process 3 and 4 returns index 1 out of 2 (need cols 2 and 3).
On the other hand, for a sharding like:
1212
3434
Dim 0 (rows): process 1 and 2 returns (0, 2), process 3 and 4 returns (1, 2)
Dim 1 (cols): process 1 and 3 returns (0, 2), process 2 and 4 returns (1, 2)
Note: This function requires sharding to be process uniform in dimension
`dim`:
each process has the same number of addressable indices in that
dimension and all index sets across processes are either disjoint or the same.
For sharding to be process uniform the addressable shards doesn't need to
form contiguous subtensor, or even a sparse grid and in case of
interleaved high-dimensional tensor it is possible for sharding to be
process uniform only in some dimensions but not others.
For example:
1111 and 12 and 1212 and 1212
2222 21 2121 1212
are all sharding uniform, in both dimensions. However
1122
2121
1121
1222
is uniform in dimension 0 (both hosts access all rows), but
is not uniform in dimension 1 (host 1 accesses columns: 0, 1, and 3),
while host 2 accesses (0, 1, 2, 3).
Returns:
A tuple of (index, num_distinct_shards) for the given dimension.
It is guaranteed that `index` will cover 0 to `num_distinct_shards - 1`,
across all processes.
Raises:
NonUniformShardingError: if the sharding is not process uniform in dimension
`dim`.
"""
# TODO(sandler, yashkatariya): Consider making this function public.
if (tensor_sharding.is_fully_addressable or
tensor_sharding.is_fully_replicated):
return (0, 1)
# Get device to indices map, we don't care about the concrete
# global shape here, only to get the distribution of shards across the tensor
# using (num_devices, num_devices, ...) This is a universal shape that is
# compatible with any mesh with num_devices.
device_map = tensor_sharding.devices_indices_map(
(tensor_sharding.num_devices,) * ndims)
# Get the slices for 'dim' for all devices.
global_slice = {k: v[dim] for k, v in device_map.items()}
# Contains mapping from process_index to a set of slices for that process.
process_to_slice = collections.defaultdict(set)
# Contains global set of slices across all processes.
all_slices = set()
# Compute the set of slices for each process and the global set of slices.
for d, v in global_slice.items():
key = (v.start, v.stop)
process_to_slice[d.process_index].add(key)
all_slices.add(key)
# Get the set of slices for the current process which we will use to compute
# the index of the current process.
current_pid = next(iter(tensor_sharding.addressable_devices)).process_index
addressable_slices = frozenset(process_to_slice[current_pid])
# Verify that all processes have the same number of slices.
slices_per_process = len(addressable_slices)
if any(len(x) != slices_per_process for x in process_to_slice.values()):
raise NonUniformShardingError(
f'{tensor_sharding=} is non-uniform on {dim=} as some processes have '
'different number of slices.'
)
unique_processes = list({frozenset(x) for x in process_to_slice.values()})
# After removing duplicate processes all unique slices should
# cover the dimension exactly once. If they don' it means that
# the sharding is not uniform.
if sum(len(h) for h in unique_processes) != len(all_slices):
raise NonUniformShardingError(
f'{tensor_sharding=} is non-uniform on {dim=}'
)
return (unique_processes.index(addressable_slices), len(unique_processes))
def local_to_global_shape(
sharding: jsharding.Sharding, local_shape: Shape) -> tuple[int | None, ...]:
"""Computes the global shape given the per process if possible.
The returned shape will have the size of the global tensor in that dimension
or None, if it is not computable. The latter can happen when sharding
is not uniform along that dimension, e.g. different hosts require
different shapes, or if different processes have partial data overlap.
If at most one dimension is sharded the shape is always computable.
Generally, global shape is computable for most practical meshes (including
topology aware such as meshes returned by mesh_utils.create_device_mesh)
Some examples: Suppose mesh is {'a': 2, 'b': 2, 'c': 2} with 2 devices
per host, 4 hosts total. For different specs we get:
- P():
global_shape = local_shape
- P(('a', 'b', 'c'), None):
global_shape = (4 * local_shape[0], local_shape[1])
Note: per device shape is (local_shape[0] / 2, local_shape[1])
- P(('a', 'b'), None)
global_shape = (4 * local_shape[0], local_shape[1])
# NB: the same global shape as above, since sharding along 'c' dimension
# happens to be within process, and thus doesn't affect the global shape.
# The underlying difference will be in the per *device* shape, which
# would be (local_shape[0], local_shape[1]) in this case.
- P(None, ('a', 'c'))
global_shape = (local_shape[0], 2 * local_shape[1])
# Per device shape is (local_shape[0], local_shape[1] / 2)
- P(('a', 'c'), 'b'):
global_shape = (2 * local_shape[0], 2 * local_shape[1])
# Per device shape is (local_shape[0] / 2, local_shape[1])
- If devices in the Mesh are randomly permuted: For any partition spec
which shards more than 1 axis: e.g. P('a', ('b', 'c')):
global_shape = (None, None)
Args:
local_shape: global shape of the tensor.
Returns:
global_shape with Nones in non-uniform dimensions.
"""
global_shape : list[int | None] = [None] * len(local_shape)
for i, local_dim in enumerate(local_shape):
try:
_, shard_count = get_process_index_and_count(
sharding, i, ndims=len(local_shape))
global_shape[i] = local_dim * shard_count
except NonUniformShardingError:
global_shape[i] = None
continue
return tuple(global_shape)
def num_addressable_indices(
tensor_sharding: jsharding.Sharding, dim: int, global_shape: Shape) -> int:
"""Returns the number of indices for given dimension this host has access to.
Each host can have multiple number of devices that are spanning
possibly discontiguous slices of data. This function computes the
total number of unique indices for dimension `dim` that any of its
addressable devices hold.
In most cases the addressable indices form a sparse grid (and in some
cases a subcube), and thus each host will hold the same of number of
indices for each dimension. However, it is possible to design a mesh that
addressable shards form a complicated pattern. In that case, the returned
value is the number of indices that are addressable by at least one device.
For example, suppose the sharding looks like this: (number indicates
the host index)
1221
1221
0000
Then on host 1 and 2, both dim 0 (rows), and dim=1 (cols) will have size 2,
while on host 0, dim 0 will have size 1, and dim 1 will have size 4.
Args:
tensor_sharding: Sharding of the tensor.
dim: dimension along which to compute the number of addressable indices.
global_shape: global shape of the tensor.
Returns:
The number of indices for dimension `dim` that this host holds.
"""
# TODO(sandler, yashkatariya): Consider making this function public.
addressables = tensor_sharding.addressable_devices_indices_map(global_shape)
addressables = cast(Mapping[jsharding.Device, Index], addressables)
num_unique_slices = len({
_slice_as_tuple(addressable[dim]) for addressable in addressables.values()
})
shard_size = tensor_sharding.shard_shape(global_shape)[dim]
return shard_size * num_unique_slices
def physical_hlo_sharding(aval, hlo_sharding: xc.HloSharding) -> xc.HloSharding:
elt_aval = core.physical_element_aval(aval.dtype)
new_op_sharding = hlo_sharding.to_proto().clone()
partitions, num_replicas = get_num_ways_dim_sharded(hlo_sharding)
suffix = [] if num_replicas == 1 else [num_replicas]
tad = partitions + [1] * elt_aval.ndim + suffix
new_op_sharding.tile_assignment_dimensions = tad
return xc.HloSharding.from_proto(new_op_sharding)
def is_single_device_sharding(sharding: jsharding.Sharding) -> bool:
# Special case PmapSharding here because PmapSharding maps away an axis
# and needs to be handled separately.test_pjit_single_device_sharding_add
return sharding.num_devices == 1 and not isinstance(sharding, PmapSharding)
def make_key_array_phys_sharding(aval, sharding):
if is_single_device_sharding(sharding):
return sharding
elif isinstance(sharding, PmapSharding):
elt_aval = core.physical_element_aval(aval.dtype)
trailing_sharding = [sharding_specs.NoSharding()] * elt_aval.ndim
phys_sharding_spec = sharding_specs.ShardingSpec(
sharding=(*sharding.sharding_spec.sharding, *trailing_sharding),
mesh_mapping=sharding.sharding_spec.mesh_mapping)
return PmapSharding(devices=sharding.devices,
sharding_spec=phys_sharding_spec)
elif isinstance(sharding, NamedSharding):
elt_aval = core.physical_element_aval(aval.dtype)
trailing_spec = [None] * elt_aval.ndim
return NamedSharding(
sharding.mesh,
PartitionSpec(*sharding.spec, *trailing_spec))
else:
hlos = sharding._to_xla_hlo_sharding(aval.ndim)
return GSPMDSharding(
sharding._device_assignment, physical_hlo_sharding(aval, hlos))
def physical_sharding(
aval, sharding: jsharding.Sharding) -> jsharding.Sharding:
return make_key_array_phys_sharding(aval, sharding)
def get_logical_gspmd_sharding(aval, phys_sharding):
elt_aval = core.physical_element_aval(aval.dtype)
phys_hlo_sharding = phys_sharding._to_xla_hlo_sharding(
aval.ndim + elt_aval.ndim)
partitions, num_replicas = get_num_ways_dim_sharded(phys_hlo_sharding)
suffix = [] if num_replicas == 1 else [num_replicas]
# Create logical sharding by cutting off the replicated trailing dims.
logical_op_sharding = phys_hlo_sharding.to_proto().clone()
tad = partitions[:-elt_aval.ndim] + suffix
logical_op_sharding.tile_assignment_dimensions = tad
return GSPMDSharding(phys_sharding._device_assignment,
xc.HloSharding.from_proto(logical_op_sharding))
def check_replicated_trailing_dims(sharding: jsharding.Sharding, aval):
if isinstance(sharding, PmapSharding):
return
phys_aval = core.physical_aval(aval)
hlo_s = sharding._to_xla_hlo_sharding(phys_aval.ndim)
partitions, _ = get_num_ways_dim_sharded(hlo_s)
num_trailing_dims = phys_aval.ndim - aval.ndim
if not all(i == 1 for i in partitions[-num_trailing_dims:]):
raise AssertionError(
"The trailing dims of extended dtypes should be replicated. Got"
f" sharding: {sharding}, partitions: {partitions}, "
f"num_trailing_dims: {num_trailing_dims}")
def logical_sharding(aval, phys_sharding) -> jsharding.Sharding:
# The trailing dims should always be replicated.
check_replicated_trailing_dims(phys_sharding, aval)
if is_single_device_sharding(phys_sharding):
return phys_sharding
elif isinstance(phys_sharding, PmapSharding):
elt_aval = core.physical_element_aval(aval.dtype)
logical_sharding_spec = sharding_specs.ShardingSpec(
sharding=phys_sharding.sharding_spec.sharding[:-elt_aval.ndim],
mesh_mapping=phys_sharding.sharding_spec.mesh_mapping)
return PmapSharding(devices=phys_sharding.devices,
sharding_spec=logical_sharding_spec)
elif isinstance(phys_sharding, NamedSharding):
logical_gs = get_logical_gspmd_sharding(aval, phys_sharding)
assert isinstance(phys_sharding.mesh, mesh_lib.Mesh)
return _gspmd_to_named_sharding_via_mesh(
logical_gs, phys_sharding.mesh)
else:
return get_logical_gspmd_sharding(aval, phys_sharding)
@util.cache()
def create_mesh_pspec_sharding(
mesh: mesh_lib.Mesh, pspec: PartitionSpec | None, parsed_pspec=None,
memory_kind: str | None = None) -> NamedSharding:
if pspec is None:
pspec, parsed_pspec = PartitionSpec(), None
return NamedSharding(mesh, pspec, _parsed_pspec=parsed_pspec,
memory_kind=memory_kind)
def _gspmd_to_named_sharding_via_mesh(
out_s: GSPMDSharding, mesh: mesh_lib.Mesh) -> NamedSharding:
parsed_pspec = parse_flatten_op_sharding(
out_s._hlo_sharding, mesh)[0]
return create_mesh_pspec_sharding(
mesh, parsed_pspec.get_partition_spec(), parsed_pspec,
out_s.memory_kind)
def flatten_spec(spec):
out = []
for s in spec:
if s is None:
continue
if isinstance(s, tuple):
out.extend(s)
else:
out.append(s)
return out
def canonicalize_sharding(sharding: NamedSharding | PartitionSpec | None,
check_mesh_consistency: bool = True
) -> NamedSharding | None:
if not config.sharding_in_types.value:
return sharding # type: ignore
if sharding is None:
return sharding
if isinstance(sharding, PartitionSpec):
sharding = NamedSharding(mesh_lib.get_abstract_mesh(), sharding) # type: ignore
else:
if (check_mesh_consistency and
sharding.mesh != mesh_lib.get_abstract_mesh()):
raise ValueError(
f'Context mesh {mesh_lib.get_abstract_mesh()} should match the mesh'
f' of sharding {sharding.mesh}. This error occurs at source: '
f' {source_info_util.summarize(source_info_util.current())}')
for s in flatten_spec(sharding.spec):
if sharding.mesh._name_to_type[s] in {
mesh_lib.AxisTypes.Hidden, mesh_lib.AxisTypes.Collective}:
raise ValueError(
'PartitionSpec cannot contain axis names that are of type Hidden or'
f' Collective. Got PartitionSpec: {sharding.spec} with axis name:'
f' {s} or type: {sharding.mesh._name_to_type[s]}')
return sharding
def make_mesh(axis_shapes: Sequence[int], axis_names: Sequence[str],
*, devices: Sequence[xc.Device] | None = None,
axis_types: mesh_lib.MeshAxisType | None = None) -> mesh_lib.Mesh:
"""Creates an efficient mesh with the shape and axis names specified.
This function attempts to automatically compute a good mapping from a set of
logical axes to a physical mesh. For example, on a TPU v3 with 8 devices:
>>> mesh = jax.make_mesh((8,), ('x')) # doctest: +SKIP
>>> [d.id for d in mesh.devices.flat] # doctest: +SKIP
[0, 1, 2, 3, 6, 7, 4, 5]
The above ordering takes into account the physical topology of TPU v3.
It orders the devices into a ring, which yields efficient all-reduces on a
TPU v3.
Now, let's see another example with 16 devices of TPU v3:
>>> mesh = jax.make_mesh((2, 8), ('x', 'y')) # doctest: +SKIP
>>> [d.id for d in mesh.devices.flat] # doctest: +SKIP
[0, 1, 2, 3, 6, 7, 4, 5, 8, 9, 10, 11, 14, 15, 12, 13]
>>> mesh = jax.make_mesh((4, 4), ('x', 'y')) # doctest: +SKIP
>>> [d.id for d in mesh.devices.flat] # doctest: +SKIP
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
As you can see, logical axes (`axis_shapes`) affect the ordering of the
devices.
You can use `jax.experimental.mesh_utils.create_device_mesh` if you want to
use the extra arguments it provides like `contiguous_submeshes` and
`allow_split_physical_axes`.
Args:
axis_shapes: Shape of the mesh. For example, axis_shape=(4, 2)
axis_names: Names of the mesh axes. For example, axis_names=('x', 'y')
devices: Optional keyword only argument, that allows you to specify the
devices you want to create a mesh with.
Returns:
A `jax.sharding.Mesh` object.
"""
if devices is None:
devices = xla_bridge.devices()
new_axis_shapes = mesh_utils._canonicalize_axis_sizes(axis_shapes)
if new_axis_shapes is None:
raise ValueError(
'`axis_shapes` passed to `make_mesh` should be a sequence of ints.'
f' Got {axis_shapes}')
del axis_shapes
axis_size = math.prod(new_axis_shapes)
if axis_size > len(devices):
raise ValueError(
f'Number of devices {len(devices)} must be >= the product '
f'of mesh_shape {new_axis_shapes}')
elif axis_size < len(devices):
devices = devices[:axis_size]
if devices[0].device_kind in (mesh_utils._TPU_V5_LITE, mesh_utils._TPU_V5E):
allow_split_physical_axes = True
else:
allow_split_physical_axes = False
mesh_devices = mesh_utils.create_device_mesh(
new_axis_shapes, devices,
allow_split_physical_axes=allow_split_physical_axes)
return mesh_lib.Mesh(mesh_devices, axis_names, axis_types=axis_types)