# 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 from collections import defaultdict from collections.abc import Sequence import enum import functools import math import operator as op from typing import Any, Callable, TYPE_CHECKING, cast from jax._src import abstract_arrays from jax._src import api from jax._src import api_util from jax._src import basearray from jax._src import config from jax._src import core from jax._src import dispatch from jax._src import dtypes from jax._src import errors from jax._src import profiler from jax._src import tree_util from jax._src import xla_bridge from jax._src.interpreters import mlir from jax._src.interpreters import pxla from jax._src.interpreters import xla from jax._src.layout import AutoLayout, DeviceLocalLayout, Layout from jax._src.lib import xla_client as xc from jax._src.lib import xla_extension as xe from jax._src.sharding import Sharding from jax._src.sharding_impls import ( PmapSharding, SingleDeviceSharding, device_replica_id_map, hashed_index, num_addressable_indices) # pyformat: disable from jax._src.typing import ArrayLike, DLDeviceType from jax._src.util import safe_zip, unzip3, use_cpp_class, use_cpp_method import numpy as np Shape = tuple[int, ...] Device = xc.Device Index = tuple[slice, ...] PRNGKeyArray = Any # TODO(jakevdp): fix cycles and import this. def _get_device(a: ArrayImpl) -> Device: devices = a.sharding._internal_device_list # pytype: disable=attribute-error assert len(devices) == 1 return devices[0] class Shard: """A single data shard of an Array. Attributes: device : Which device this shard resides on. index : The index into the global array of this shard. replica_id : Integer id indicating which replica of the global array this shard is part of. Always 0 for fully sharded data (i.e. when there’s only 1 replica). data : The data of this shard. None if ``device`` is non-local. """ def __init__(self, device: Device, sharding: Sharding, global_shape: Shape, data: None | ArrayImpl | PRNGKeyArray = None): self._device = device self._sharding = sharding self._global_shape = global_shape self._data = data def __repr__(self): try: return (f'Shard(device={self.device!r}, index={self.index}, ' f'replica_id={self.replica_id}, data={self.data})') except ValueError: return f'Shard(device={self.device!r}, data={self.data})' @functools.cached_property def index(self) -> Index: try: device_indices_map_fn = self._sharding.devices_indices_map except AttributeError: raise ValueError('Cannot calculate indices from sharding: ' f'{self._sharding}. Please create a device to index ' 'mapping for your sharding.') from None index = device_indices_map_fn(self._global_shape)[self.device] assert index is not None return index @functools.cached_property def replica_id(self) -> int: return device_replica_id_map(self._sharding, self._global_shape)[self.device] @property def device(self): return self._device @property def data(self): return self._data def _reconstruct_array(fun, args, arr_state, aval_state): """Method to reconstruct a device array from a serialized state.""" np_value = fun(*args) np_value.__setstate__(arr_state) jnp_value = api.device_put(np_value) jnp_value.aval = jnp_value.aval.update(**aval_state) return jnp_value @functools.lru_cache(maxsize=4096) def _cached_index_calc(s, shape): map_ = s.addressable_devices_indices_map(shape) seen_h_indices = set() l = [] for array_index, index in enumerate(map_.values()): h_index = hashed_index(index) if h_index not in seen_h_indices: seen_h_indices.add(h_index) l.append((array_index, index)) return l @functools.lru_cache(maxsize=4096) def _process_has_full_value_in_mcjax(s, shape): # Return False for single host as a fast path. if xla_bridge.process_count() == 1: return False num_unique_indices = len( {hashed_index(v) for v in s.devices_indices_map(shape).values()}) num_addressable_unique_indices = len( {hashed_index(v) for v in s.addressable_devices_indices_map(shape).values()}) return num_unique_indices == num_addressable_unique_indices class ArrayImpl(basearray.Array): # TODO(yashkatariya): Add __slots__ here. aval: core.ShapedArray _sharding: Sharding _arrays: list[ArrayImpl] _committed: bool _skip_checks: bool _npy_value: np.ndarray | None @use_cpp_method() def __init__(self, aval: core.ShapedArray, sharding: Sharding, arrays: Sequence[ArrayImpl], committed: bool, _skip_checks: bool = False): # NOTE: the actual implementation of the constructor is moved to C++. self.aval = aval self._sharding = sharding self._arrays = [a._arrays[0] for a in arrays] self._committed = committed self._npy_value = None # Don't rearrange if skip_checks is enabled because this assumes that the # input buffers are already arranged properly. This usually happens when # Array's are created as output of a JAX transformation # (like pjit, xmap, etc). if not _skip_checks or config.enable_checks.value: self._check_and_rearrange() def _check_and_rearrange(self): for db in self._arrays: if db.dtype != self.dtype: raise ValueError( "Input buffers to `Array` must have matching dtypes. " f"Got {db.dtype}, expected {self.dtype} for buffer: {db}") device_id_to_buffer = {_get_device(db).id: db for db in self._arrays} addressable_dev = self.sharding.addressable_devices if len(self._arrays) != len(addressable_dev): raise ValueError( f"Expected {len(addressable_dev)} per-device arrays " "(this is how many devices are addressable by the sharding), but " f"got {len(self._arrays)}") array_device_ids = set(device_id_to_buffer.keys()) addressable_device_ids = {d.id for d in addressable_dev} # Calculate a symmetric difference because the device ids between sharding # and _arrays should match. diff = array_device_ids ^ addressable_device_ids if diff: dev_in_sharding_not_in_arrays = addressable_device_ids - array_device_ids dev_in_arrays_not_in_sharding = array_device_ids - addressable_device_ids err_msg = ( "Addressable devices and per-device arrays devices do not match.") if dev_in_sharding_not_in_arrays: err_msg += (f" Sharding contains devices {dev_in_sharding_not_in_arrays} " "that are not present in per-device arrays.") if dev_in_arrays_not_in_sharding: err_msg += (f" Per-device arrays contain devices {dev_in_arrays_not_in_sharding} " "that are not present in the sharding.") raise ValueError(err_msg) ss = self.sharding.shard_shape(self.shape) for db in self._arrays: if db.shape != ss: raise ValueError( f"Expected shard shape {ss} doesn't match the single device array " f"shape {db.shape}. Shape of Array is " f"{self.aval.str_short()} with sharding {self.sharding}") # Rearrange arrays based on the device assignment. addressable_da = self.sharding._addressable_device_assignment self._arrays = [device_id_to_buffer[device.id] for device in addressable_da] @property def shape(self) -> Shape: return self.aval.shape @property def dtype(self): return self.aval.dtype @property def ndim(self): return len(self.shape) @property def size(self): return math.prod(self.shape) @property def sharding(self): return self._sharding @property def weak_type(self): return self.aval.weak_type def __str__(self): return str(self._value) def __len__(self): try: return self.shape[0] except IndexError as err: raise TypeError("len() of unsized object") from err # same as numpy error def __bool__(self): core.check_bool_conversion(self) return bool(self._value) def __float__(self): core.check_scalar_conversion(self) return self._value.__float__() def __int__(self): core.check_scalar_conversion(self) return self._value.__int__() def __complex__(self): core.check_scalar_conversion(self) return self._value.__complex__() def __hex__(self): core.check_integer_conversion(self) return hex(self._value) def __oct__(self): core.check_integer_conversion(self) return oct(self._value) def __index__(self): core.check_integer_conversion(self) return op.index(self._value) def tobytes(self, order="C"): return self._value.tobytes(order) def tolist(self): return self._value.tolist() def __format__(self, format_spec): # Simulates behavior of https://github.com/numpy/numpy/pull/9883 if self.ndim == 0: return format(self._value[()], format_spec) else: return format(self._value, format_spec) def __getitem__(self, idx): from jax._src.lax import lax from jax._src.numpy import lax_numpy self._check_if_deleted() if isinstance(self.sharding, PmapSharding): if config.pmap_no_rank_reduction.value: cidx = idx if isinstance(idx, tuple) else (idx,) padded_cidx = tuple( slice(i, i + 1, None) if isinstance(i, int) else i for i in cidx ) + (slice(None),) * (len(self.shape) - len(cidx)) else: if not isinstance(idx, tuple): padded_cidx = (idx,) + (slice(None),) * (len(self.shape) - 1) else: padded_cidx = idx + (slice(None),) * (len(self.shape) - len(idx)) indices = tuple(self.sharding.devices_indices_map(self.shape).values()) try: arr_idx = indices.index(padded_cidx) except ValueError: arr_idx = None if arr_idx is not None: a = self._arrays[arr_idx] out = ArrayImpl( a.aval, SingleDeviceSharding(_get_device(a)), [a], committed=False, _skip_checks=True) if config.pmap_no_rank_reduction.value: # If cidx was the index of a single shard, then it corresponds to one # shard of the chunked dimension. dims = tuple(i for i, x in enumerate(cidx) if isinstance(x, int)) return lax.squeeze(out, dimensions=dims) else: return out return lax_numpy._rewriting_take(self, idx) def __iter__(self): if self.ndim == 0: raise TypeError("iteration over a 0-d array") # same as numpy error else: assert self.is_fully_replicated or self.is_fully_addressable if dispatch.is_single_device_sharding(self.sharding) or self.is_fully_replicated: return (sl for chunk in self._chunk_iter(100) for sl in chunk._unstack()) elif isinstance(self.sharding, PmapSharding): return (self[i] for i in range(self.shape[0])) else: # TODO(yashkatariya): Don't bounce to host and use `_chunk_iter` path # here after uneven partitioning support is added. return (api.device_put(self._value[i]) for i in range(self.shape[0])) @property def is_fully_replicated(self) -> bool: return self.sharding.is_fully_replicated def __repr__(self): prefix = 'Array(' if self.aval is not None and self.aval.weak_type: dtype_str = f'dtype={self.dtype.name}, weak_type=True)' else: dtype_str = f'dtype={self.dtype.name})' if self.is_fully_addressable or self.is_fully_replicated: line_width = np.get_printoptions()["linewidth"] if self.size == 0: s = f"[], shape={self.shape}" else: s = np.array2string(self._value, prefix=prefix, suffix=',', separator=', ', max_line_width=line_width) last_line_len = len(s) - s.rfind('\n') + 1 sep = ' ' if last_line_len + len(dtype_str) + 1 > line_width: sep = ' ' * len(prefix) return f"{prefix}{s},{sep}{dtype_str}" else: return f"{prefix}{self.shape}, {dtype_str}" @property def is_fully_addressable(self) -> bool: """Is this Array fully addressable? A jax.Array is fully addressable if the current process can address all of the devices named in the :class:`Sharding`. ``is_fully_addressable`` is equivalent to "is_local" in multi-process JAX. Note that fully replicated is not equal to fully addressable i.e. a jax.Array which is fully replicated can span across multiple hosts and is not fully addressable. """ return self.sharding.is_fully_addressable def __array__(self, dtype=None, context=None, copy=None): # copy argument is supported by np.asarray starting in numpy 2.0 kwds = {} if copy is None else {'copy': copy} return np.asarray(self._value, dtype=dtype, **kwds) def __dlpack__(self, *, stream: int | Any | None = None, max_version: tuple[int, int] | None = None, dl_device: tuple[DLDeviceType, int] | None = None, copy: bool | None = None): from jax._src.dlpack import to_dlpack # pylint: disable=g-import-not-at-top device_set = self.sharding.device_set if len(device_set) > 1: raise BufferError( "to_dlpack can only pack a dlpack tensor from an array on a singular " f"device, but an array with a Sharding over {len(device_set)} devices " "was provided." ) device, = device_set return to_dlpack(self, stream=stream, max_version=max_version, src_device=device, dl_device=dl_device, copy=copy) def __dlpack_device__(self) -> tuple[enum.Enum, int]: if len(self._arrays) != 1: raise BufferError("__dlpack__ only supported for unsharded arrays.") from jax._src.dlpack import DLDeviceType # pylint: disable=g-import-not-at-top if self.platform() == "cpu": return DLDeviceType.kDLCPU, 0 elif self.platform() == "gpu": platform_version = _get_device(self).client.platform_version if "cuda" in platform_version: dl_device_type = DLDeviceType.kDLCUDA elif "rocm" in platform_version: dl_device_type = DLDeviceType.kDLROCM else: raise BufferError("Unknown GPU platform for __dlpack__: " f"{platform_version}") local_hardware_id = _get_device(self).local_hardware_id if local_hardware_id is None: raise BufferError("Couldn't get local_hardware_id for __dlpack__") return dl_device_type, local_hardware_id else: raise BufferError( "__dlpack__ device only supported for CPU and GPU, got platform: " f"{self.platform()}" ) def __reduce__(self): fun, args, arr_state = self._value.__reduce__() aval_state = {'weak_type': self.aval.weak_type, 'named_shape': self.aval.named_shape} return (_reconstruct_array, (fun, args, arr_state, aval_state)) @use_cpp_method() def unsafe_buffer_pointer(self): if len(self._arrays) != 1: raise ValueError("unsafe_buffer_pointer() is supported only for unsharded" " arrays.") return self._arrays[0].unsafe_buffer_pointer() @property @use_cpp_method() def __cuda_array_interface__(self): if len(self._arrays) != 1: raise ValueError("__cuda_array_interface__() is supported only for " "unsharded arrays.") return self._arrays[0].__cuda_array_interface__ # pytype: disable=attribute-error # bind-properties @use_cpp_method() def on_device_size_in_bytes(self): """Returns the total global on-device size of the array in bytes.""" arr = self._arrays[0] per_shard_size = arr.on_device_size_in_bytes() return per_shard_size * len(self.sharding.device_set) def devices(self) -> set[Device]: self._check_if_deleted() return self.sharding.device_set @property def device_buffer(self): raise AttributeError( "arr.device_buffer has been deprecated. Use arr.addressable_data(0)") @property def device_buffers(self): raise AttributeError( "arr.device_buffers has been deprecated. Use [x.data for x in arr.addressable_shards]") def addressable_data(self, index: int) -> ArrayImpl: self._check_if_deleted() if self.is_fully_replicated: return self._fully_replicated_shard() return self._arrays[index] @functools.cached_property def addressable_shards(self) -> Sequence[Shard]: self._check_if_deleted() out = [] for a in self._arrays: out.append(Shard(_get_device(a), self.sharding, self.shape, a)) return out @property def layout(self): # TODO(yashkatariya): Remove the deleted check from here. if self.is_deleted(): return Layout(None, self.sharding) try: return Layout(DeviceLocalLayout(self._pjrt_layout), self.sharding) except xe.XlaRuntimeError as e: msg, *_ = e.args if type(msg) is str and msg.startswith("UNIMPLEMENTED"): return Layout(None, self.sharding) else: raise @property def global_shards(self) -> Sequence[Shard]: """Returns list of all `Shard`s of the Array across all devices. The result includes shards that are not addressable by the current process. If a `Shard` is not addressable, then its `data` will be `None`. """ self._check_if_deleted() if self.is_fully_addressable: # pylint: disable=using-constant-test return self.addressable_shards out = [] device_id_to_buffer = {_get_device(a).id: a for a in self._arrays} for global_d in self.sharding.device_set: if device_id_to_buffer.get(global_d.id, None) is not None: array = device_id_to_buffer[global_d.id] else: array = None out.append(Shard(global_d, self.sharding, self.shape, array)) return out @use_cpp_method() def delete(self): if self._arrays is None: return for buf in self._arrays: buf.delete() self._arrays = None self._npy_value = None @use_cpp_method() def is_deleted(self): if self._arrays is None: return True # This path is taken when a view of `Array` is created and the original # Array is deleted. In that case, the buffers the view represents also get # deleted. return any(buf.is_deleted() for buf in self._arrays) def _check_if_deleted(self): if self.is_deleted(): raise RuntimeError( f"Array has been deleted with shape={self.aval.str_short()}.") @use_cpp_method() def block_until_ready(self): self._check_if_deleted() for db in self._arrays: db.block_until_ready() return self @use_cpp_method() def _single_device_array_to_np_array(self): return np.asarray(self._arrays[0]) @use_cpp_method() def _copy_single_device_array_to_host_async(self): self._arrays[0].copy_to_host_async() @profiler.annotate_function def copy_to_host_async(self): self._check_if_deleted() if self._npy_value is None: if self.is_fully_replicated: self._copy_single_device_array_to_host_async() return for i, _ in _cached_index_calc(self.sharding, self.shape): self._arrays[i]._copy_single_device_array_to_host_async() @property @functools.partial(profiler.annotate_function, name="np.asarray(jax.Array)") def _value(self) -> np.ndarray: self._check_if_deleted() if self._npy_value is None: if self.is_fully_replicated: self._npy_value = self._single_device_array_to_np_array() self._npy_value.flags.writeable = False return cast(np.ndarray, self._npy_value) # TODO(yashkatariya): Merge `_process_has_full_value_in_mcjax` with # is_fully_addressable. if (not self.is_fully_addressable and not _process_has_full_value_in_mcjax(self.sharding, self.shape)): raise RuntimeError( "Fetching value for `jax.Array` that spans non-addressable" " (non process local) devices is not possible. You can use" " `jax.experimental.multihost_utils.process_allgather` to print the" " global array or use `.addressable_shards` method of jax.Array to" " inspect the addressable (process local) shards." ) for i, _ in _cached_index_calc(self.sharding, self.shape): self._arrays[i]._copy_single_device_array_to_host_async() npy_value = np.empty(self.shape, self.dtype) for i, ind in _cached_index_calc(self.sharding, self.shape): npy_value[ind] = self._arrays[i]._single_device_array_to_np_array() self._npy_value = npy_value self._npy_value.flags.writeable = False # https://docs.python.org/3/library/typing.html#typing.cast return cast(np.ndarray, self._npy_value) # TODO(b/273265390): ideally we would write this as a decorator on the ArrayImpl # class, however this triggers a pytype bug. Workaround: apply the decorator # after the fact. if not TYPE_CHECKING: ArrayImpl = use_cpp_class(xc.ArrayImpl)(ArrayImpl) # explicitly set to be unhashable. setattr(ArrayImpl, "__hash__", None) setattr(ArrayImpl, "__array_priority__", 100) # TODO(yashkatariya): Remove None from callback input type. def make_array_from_callback( shape: Shape, sharding: Sharding | Layout, data_callback: Callable[[Index | None], ArrayLike]) -> ArrayImpl: # pyformat: disable """Returns a ``jax.Array`` via data fetched from ``data_callback``. ``data_callback`` is used to fetch the data for each addressable shard of the returned ``jax.Array``. This function must return concrete arrays, meaning that ``make_array_from_callback`` has limited compatibility with JAX transformations like :func:`jit` or :func:`vmap`. Args: shape : Shape of the ``jax.Array``. sharding: A ``Sharding`` instance which describes how the ``jax.Array`` is laid out across devices. data_callback : Callback that takes indices into the global array value as input and returns the corresponding data of the global array value. The data can be returned as any array-like object, e.g. a ``numpy.ndarray``. Returns: A ``jax.Array`` via data fetched from ``data_callback``. Example: >>> import math >>> from jax.sharding import Mesh >>> from jax.sharding import PartitionSpec as P >>> import numpy as np ... >>> input_shape = (8, 8) >>> global_input_data = np.arange(math.prod(input_shape)).reshape(input_shape) >>> global_mesh = Mesh(np.array(jax.devices()).reshape(2, 4), ('x', 'y')) >>> inp_sharding = jax.sharding.NamedSharding(global_mesh, P('x', 'y')) ... >>> def cb(index): ... return global_input_data[index] ... >>> arr = jax.make_array_from_callback(input_shape, inp_sharding, cb) >>> arr.addressable_data(0).shape (4, 2) """ # pyformat: enable dll = sharding.device_local_layout if isinstance(sharding, Layout) else None if isinstance(dll, AutoLayout): raise TypeError( "`DeviceLocalLayout.AUTO` cannot be used in place of a device-local" f" layout when calling `jax.make_array_from_callback`. Got {sharding}") sharding = sharding.sharding if isinstance(sharding, Layout) else sharding # type: ignore if not isinstance(sharding, Sharding): raise TypeError( f"sharding should be an instance of `jax.sharding`. Got {sharding} of" f" type {type(sharding)}") if sharding.is_fully_replicated: devices = list(sharding._internal_device_list.addressable_device_list) # type: ignore per_device_values = [data_callback((slice(None),) * len(shape))] * len(devices) else: device_to_index_map = sharding.addressable_devices_indices_map(shape) devices = list(device_to_index_map.keys()) per_device_values = [data_callback(device_to_index_map[device]) for device in devices] if isinstance(per_device_values[0], core.Tracer): raise errors.UnexpectedTracerError( "jax.make_array_from_callback cannot be called within a traced context.") first_value = xla.canonicalize_dtype(per_device_values[0]) aval = core.ShapedArray(shape, first_value.dtype, weak_type=False) # first value can be numpy array, python scalar, etc. if (sharding.is_fully_replicated and not isinstance(first_value, ArrayImpl) and not dtypes.issubdtype(aval.dtype, dtypes.extended) and dll is None): # Do this check outside because `batched_device_put` won't do these checks # like ArrayImpl. if shape != first_value.shape: raise ValueError( f"Expected shard shape {shape} doesn't match the single device " f"array shape {first_value.shape}. Shape of Array is " f"{aval.str_short()} with sharding {sharding}") return pxla.batched_device_put( aval, sharding, per_device_values, devices, committed=True) if (isinstance(first_value, ArrayImpl) and first_value._committed and sharding.is_fully_replicated and first_value.is_fully_replicated and first_value.sharding._device_assignment == tuple(devices) and (first_value.layout.device_local_layout == pxla._maybe_get_default_layout(Layout(dll, sharding), None, sharding, aval))): return first_value if dll is not None: devices = [Layout(dll, SingleDeviceSharding(d)) for d in devices] arrays = api.device_put(per_device_values, devices) if dtypes.issubdtype(aval.dtype, dtypes.extended): return aval.dtype._rules.make_sharded_array(aval, sharding, arrays, committed=True) return ArrayImpl(aval, sharding, arrays, committed=True) def make_array_from_process_local_data( sharding: Sharding, local_data: np.ndarray, global_shape: tuple[int, ...], ) -> ArrayImpl: # pyformat: disable """Creates distributed tensor using the data available in process. This function is a common special case of `make_array_from_callback`. It assumes that the data is available in the process and takes care of the index wrangling. Note, if the two hosts are replicas, host_local_data should be identical as well. Each dimension of the shape of host_local_data should either match global_shape or the # indices the devices on this process need to address. For example if dimension $i$ is fully sharded then this size would be `per_device_shape[i] * jax.local_device_count()`. If the shape matches global shape, each device slice will just lookup the slice in the local_data. In the latter case the global slice of each device will be mapped into local slice of `local_data` array. For example, if given process only addresses slices (8, 12) and (24, 28), then these slices will be mapped into (0, 4) and (4, 8) of the `local_data`. This function can be used to create tensors from dataset feeding pipelines. The most common case is when the sharding is fully sharded across the batch dimension and each host just loads its corresponding sub-batch. This function supports more general case as well, such as multi-host replication but you would need to compute the size and the contents of process-local data correctly to satisfy the replication constraints. Examples: >>> from jax.sharding import PartitionSpec as P >>> mesh_rows = 2 >>> mesh_cols = jax.device_count() // 2 ... >>> mesh = jax.sharding.Mesh(np.array(jax.devices()).reshape(mesh_rows, mesh_cols), ('x', 'y')) >>> sharding = jax.sharding.NamedSharding(mesh, P(('x', 'y'),)) >>> rows_per_device = 2 >>> feature_length = 32 >>> per_device_shape = (rows_per_device, feature_length) >>> per_host_shape = (rows_per_device * len(mesh.local_devices), feature_length) >>> per_host_generator = lambda : np.arange(np.prod(per_host_shape)).reshape(per_host_shape) >>> per_host_data = per_host_generator() # replace with your own per-host data pipeline that outputs numpy arrays >>> global_shape = (rows_per_device * len(sharding.device_set), ) + per_device_shape[1:] >>> output_global_array = jax.make_array_from_process_local_data(sharding, per_host_data, global_shape) ... >>> assert output_global_array.addressable_data(0).shape == per_device_shape >>> assert output_global_array.shape == global_shape Args: sharding: sharding of the global tensor. host_local_data: data on the host to be placed on local devices. Each dimension should either match global_shape, or match num_addressable_indices(dim). global_shape: the target shape of the global tensor. In some cases this parameter can be inferred from sharding and host_local_data, however it is useful to catch common sharding errors. Returns: Tensor that will have sharding=sharding. """ # pyformat: enable shard_shape = sharding.shard_shape(global_shape) full_dim = [] for i, (data_dim, global_dim) in enumerate( zip(local_data.shape, global_shape) ): full_dim.append(data_dim == global_dim) if data_dim != global_dim: process_slice = num_addressable_indices(sharding, i, global_shape) if process_slice != data_dim: raise ValueError( "Invalid host data, each dimension should match either global or " f"process shape. In dimension {i=}, the process data has {data_dim}" f"elements. Process addresses {process_slice} elements and " f"{global_shape=}." ) addressable_shards = sharding.addressable_devices_indices_map(global_shape) slices_for_each_dim: list[list[int]] = [[] for _ in global_shape] for shard_index in addressable_shards.values(): assert shard_index is not None for i, slc in enumerate(shard_index): slices_for_each_dim[i].append(slc.start or 0) for i in range(len(global_shape)): slices_for_each_dim[i] = sorted(set(slices_for_each_dim[i])) def local_slice(i, slc): # Looks up the index of this slice in the list of slices for this dimension. # This will determine the slice in host_local_data start = slices_for_each_dim[i].index(slc.start or 0) * shard_shape[i] end = start + shard_shape[i] return slice(start, end) def cb(index: Index | None) -> ArrayLike: assert index is not None data_slice = [ slc if full_dim[i] else local_slice(i, slc) for i, slc in enumerate(index) ] return local_data[tuple(data_slice)] return make_array_from_callback(global_shape, sharding, cb) def make_array_from_single_device_arrays( shape: Shape, sharding: Sharding, arrays: Sequence[basearray.Array] ) -> ArrayImpl: r"""Returns a ``jax.Array`` from a sequence of ``jax.Array``\s each on a single device. Every device in input ``sharding``\'s mesh must have an array in ``arrays``\s. Args: shape : Shape of the output ``jax.Array``. This conveys information already included with ``sharding`` and ``arrays`` and serves as a double check. sharding: Sharding: A global Sharding instance which describes how the output jax.Array is laid out across devices. arrays: Sequence of ``jax.Array``\s that are each single device addressable. ``len(arrays)`` must equal ``len(sharding.addressable_devices)`` and the shape of each array must be the same. For multiprocess code, each process will call with a different ``arrays`` argument that corresponds to that processes' data. These arrays are commonly created via ``jax.device_put``. Returns: A global ``jax.Array``, sharded as ``sharding``, with shape equal to ``shape``, and with per-device contents matching ``arrays``. Examples: In this single-process example, we use ``make_array_from_single_device_arrays`` to create an a global array. >>> import math >>> from jax.sharding import Mesh >>> from jax.sharding import PartitionSpec as P >>> import numpy as np ... >>> mesh_rows = 2 >>> mesh_cols = jax.device_count() // 2 ... >>> global_shape = (8, 8) >>> mesh = Mesh(np.array(jax.devices()).reshape(mesh_rows, mesh_cols), ('x', 'y')) >>> sharding = jax.sharding.NamedSharding(mesh, P('x', 'y')) >>> inp_data = np.arange(math.prod(global_shape)).reshape(global_shape) ... >>> arrays = [ ... jax.device_put(inp_data[index], d) ... for d, index in sharding.addressable_devices_indices_map(global_shape).items()] ... >>> arr = jax.make_array_from_single_device_arrays(global_shape, sharding, arrays) >>> assert arr.shape == (8,8) # arr.shape is (8,8) regardless of jax.device_count() When using multiple processes, a common data pipeline is to have data parallelism across devices, with each device receiving at least one example. In this case, the following recipe will use `make_array_from_single_device_arrays` to create a global jax.Array. First, we create the per host data as Numpy arrays. >>> sharding = jax.sharding.NamedSharding(mesh, P(('x', 'y'),)) >>> rows_per_device = 2 >>> feature_length = 32 >>> per_device_shape = (rows_per_device, feature_length) >>> per_host_shape = (rows_per_device * len(mesh.local_devices), feature_length) >>> per_host_generator = lambda : np.arange(np.prod(per_host_shape)).reshape(per_host_shape) >>> per_host_data = per_host_generator() # replace with your own per-host data pipeline that outputs numpy arrays Second, we put the Numpy data onto the local devices as single device Jax Arrays. Then we call make_array_from_single_device_arrays to make the global Array. >>> global_shape = (rows_per_device * len(sharding.device_set), ) + per_device_shape[1:] >>> per_device_data = np.split(per_host_data, len(mesh.local_devices), axis = 0) # per device data, but on host >>> per_device_data_on_device = jax.device_put(per_device_data, mesh.local_devices) # per device data, now on device >>> output_global_array = jax.make_array_from_single_device_arrays(global_shape, sharding, per_device_data_on_device) ... >>> assert output_global_array.addressable_data(0).shape == per_device_shape >>> assert output_global_array.shape == global_shape When using tensor parallelism (equivalent to sharding across both rows and columns in the above example), the above example doesn't generate the data in the sharding that you plan to consume it with. The most common fix is to simply load the data in this data parallel sharding and have the reshard happen automatically within the downstream jitted function. Depending on your use case, you might prefer to directly load sharded data, something that ``make_array_from_single_device_arrays`` can do but will depend on your data loading pipeline also loading in the matching sharding. Loading in a data parallel format is typically fully satisfactory for data loading for LLM use cases. """ # All input arrays should be committed. Checking it is expensive on # single-controller systems. if any(isinstance(arr, core.Tracer) for arr in arrays): raise ValueError( "jax.make_array_from_single_device_arrays requires a list of concrete" f" arrays as input. got types {set(map(type, arrays))}") aval = core.ShapedArray(shape, arrays[0].dtype, weak_type=False) if dtypes.issubdtype(aval.dtype, dtypes.extended): return aval.dtype._rules.make_sharded_array(aval, sharding, arrays, committed=True) # TODO(phawkins): ideally the cast() could be checked. return ArrayImpl(aval, sharding, cast(Sequence[ArrayImpl], arrays), committed=True) core.pytype_aval_mappings[ArrayImpl] = abstract_arrays.canonical_concrete_aval xla.pytype_aval_mappings[ArrayImpl] = op.attrgetter('aval') xla.canonicalize_dtype_handlers[ArrayImpl] = pxla.identity api_util._shaped_abstractify_handlers[ArrayImpl] = op.attrgetter('aval') # TODO(jakevdp) replace this with true inheritance at the C++ level. basearray.Array.register(ArrayImpl) def _array_mlir_constant_handler(val): return mlir.ir_constants(val._value) mlir.register_constant_handler(ArrayImpl, _array_mlir_constant_handler) # NOTE(skye): we could refactor to generate _multi_slice parameters directly # from the input ShardingSpec, rather than the indices. However, this would # require duplicating the ordering logic of spec_to_indices, which is more # subtle and more likely to change than the index logic we have to support here. def as_slice_indices(arr: Any, idx: Index) -> tuple[ tuple[int, ...], tuple[int, ...], tuple[int, ...]]: """Returns start_indices, limit_indices, removed_dims""" start_indices = [0] * arr.ndim limit_indices = list(arr.shape) removed_dims: list[int] = [] tuple_idx = idx if isinstance(idx, tuple) else (idx,) for dim, sub_idx in enumerate(tuple_idx): if isinstance(sub_idx, int): start_indices[dim] = sub_idx limit_indices[dim] = sub_idx + 1 removed_dims.append(dim) elif sub_idx == slice(None): continue else: assert isinstance(sub_idx, slice), sub_idx assert isinstance(sub_idx.start, int), sub_idx assert isinstance(sub_idx.stop, int), sub_idx start_indices[dim] = sub_idx.start limit_indices[dim] = sub_idx.stop return tuple(start_indices), tuple(limit_indices), tuple(removed_dims) def shard_device_array(x, devices, indices, sharding): start_indices, limit_indices, removed_dims = unzip3( as_slice_indices(x, idx) for idx in indices) if sharding.is_fully_replicated: shards = [x] * len(devices) else: shards = x._multi_slice(start_indices, limit_indices, removed_dims) aval = api_util.shaped_abstractify(x) return pxla.batched_device_put(aval, sharding, shards, devices) def _hashable_index(idx): return tree_util.tree_map( lambda x: (x.start, x.stop) if type(x) == slice else x, idx) def shard_sharded_device_array_slow_path(x, devices, indices, sharding): candidates = defaultdict(list) bufs = [buf.data for buf in x.addressable_shards] arr_indices = tuple(x.sharding.devices_indices_map(x.shape).values()) for buf, idx in safe_zip(bufs, arr_indices): candidates[_hashable_index(idx)].append(buf) bufs = [] for idx, device in safe_zip(indices, devices): # Look up all buffers that contain the correct slice of the logical array. candidates_list = candidates[_hashable_index(idx)] if not candidates_list: # This array isn't sharded correctly. Reshard it via host roundtrip. # TODO(skye): more efficient reshard? return pxla.shard_arg(x._value, sharding, canonicalize=False) # Try to find a candidate buffer already on the correct device, # otherwise copy one of them. for buf in candidates_list: if buf.devices() == {device}: bufs.append(buf) break else: bufs.append(buf) return pxla.batched_device_put(x.aval, sharding, bufs, devices) @functools.lru_cache(maxsize=4096) def _sharding_indices_and_eq(src_sharding, shape, dst_sharding): src_indices = src_sharding.addressable_devices_indices_map(shape).values() dst_indices = dst_sharding.addressable_devices_indices_map(shape).values() return dst_indices, tuple(src_indices) == tuple(dst_indices) def _array_shard_arg(x, sharding): x._check_if_deleted() indices, same_indices = _sharding_indices_and_eq(x.sharding, x.shape, sharding) if not x.is_fully_addressable: if same_indices: return x else: raise NotImplementedError( "Cannot reshard an input that is not fully addressable") else: devices = sharding._addressable_device_assignment if same_indices: return xc.copy_array_to_devices_with_sharding(x, list(devices), sharding) # Resharding starts here: if dispatch.is_single_device_sharding(x.sharding): return shard_device_array(x, devices, indices, sharding) else: return shard_sharded_device_array_slow_path(x, devices, indices, sharding) pxla.shard_arg_handlers[ArrayImpl] = _array_shard_arg def _array_global_result_handler(global_aval, out_sharding, committed): if global_aval.dtype == dtypes.float0: return lambda _: np.zeros(global_aval.shape, dtypes.float0) if dtypes.issubdtype(global_aval.dtype, dtypes.extended): return global_aval.dtype._rules.global_sharded_result_handler( global_aval, out_sharding, committed) return xc.array_result_handler( global_aval, out_sharding, committed=committed, _skip_checks=True ) pxla.global_result_handlers[core.ShapedArray] = _array_global_result_handler pxla.global_result_handlers[core.ConcreteArray] = _array_global_result_handler # Only used for Arrays that come out of pmap. def _array_local_result_handler(aval, sharding, indices): if aval.dtype == dtypes.float0: return lambda _: np.zeros(aval.shape, dtypes.float0) if dtypes.issubdtype(aval.dtype, dtypes.extended): return aval.dtype._rules.local_sharded_result_handler( aval, sharding, indices) return xc.array_result_handler( aval, sharding, committed=True, _skip_checks=True ) pxla.local_result_handlers[core.ShapedArray] = _array_local_result_handler pxla.local_result_handlers[core.ConcreteArray] = _array_local_result_handler # Token handlers def _token_shard_arg(x, sharding): return _array_shard_arg(x._buf, sharding) pxla.shard_arg_handlers[core.Token] = _token_shard_arg def _token_global_result_handler(global_aval, out_sharding, committed): array_handler = _array_global_result_handler( core.token_shaped_array, out_sharding, committed) def wrapper(*args, **kwargs): out_buf = array_handler(*args, **kwargs) return core.Token(out_buf) return wrapper pxla.global_result_handlers[core.AbstractToken] = _token_global_result_handler