import logging from typing import ( TYPE_CHECKING, Any, Dict, Iterable, Iterator, List, Optional, Tuple, Union, ) import numpy as np if TYPE_CHECKING: import jax logger = logging.getLogger(__name__) _GLOBAL_MESH_1D_AXIS = "data" NumpyBatch = Union[np.ndarray, Dict[str, np.ndarray]] JaxBatch = Union["jax.Array", Dict[str, "jax.Array"]] DTypeLikeSpec = Union["jax.typing.DTypeLike", Dict[str, "jax.typing.DTypeLike"]] Scalar = Union[int, float, bool] PaddingsSpec = Union[Scalar, Dict[str, Scalar]] def _get_column_value(mapping_or_value: Any, key: str) -> Any: """Get the value for a specific column from a mapping or a single value.""" if isinstance(mapping_or_value, dict): return mapping_or_value[key] return mapping_or_value def _unwrap_single_column_value(mapping_or_value: Any, name: str) -> Any: """Unwrap a single value from a mapping if it's a dictionary.""" if isinstance(mapping_or_value, dict): if len(mapping_or_value) != 1: raise ValueError( f"When constructing a single-tensor batch, only a single {name} " f"should be given, instead got: {mapping_or_value}" ) return next(iter(mapping_or_value.values())) return mapping_or_value def _create_sharding_1d(axis_name: str) -> "jax.sharding.Sharding": """Create a 1D JAX sharding, preferably using topology-aware mesh_utils.""" import jax from jax.sharding import Mesh, NamedSharding, PartitionSpec devices = None try: from jax.experimental import mesh_utils # Attempt to create a topology-aware mesh (e.g. for TPU/GPU interconnects) devices = mesh_utils.create_device_mesh((jax.device_count(),)) except Exception as e: logger.debug(f"Failed to use jax.experimental.mesh_utils: {e}") if devices is None: # Fallback to default device ordering if mesh_utils fails or is unavailable devices = np.array(jax.devices()) physical_mesh = Mesh(devices, (axis_name,)) return NamedSharding(physical_mesh, PartitionSpec(axis_name)) def _convert_ndarray_to_jax_array( ndarray: np.ndarray, sharding: "jax.sharding.Sharding", # noqa: F821 dtype: Optional["jax.typing.DTypeLike"] = None, ) -> "jax.Array": # noqa: F821 import jax local_batch_size = ndarray.shape[0] host_count = jax.process_count() # Global shape assumes each host gets the exact same local batch size. global_shape = (local_batch_size * host_count,) + ndarray.shape[1:] # Construct the globally aware 1D array from process-local data. # This automatically shards the local ndarray across the local devices # assigned to this process by the provided sharding. physical_array = jax.make_array_from_process_local_data( sharding, ndarray, global_shape ) if dtype is not None: physical_array = physical_array.astype(dtype) return physical_array def _convert_batch( ndarrays: NumpyBatch, sharding: "jax.sharding.Sharding", dtypes: Optional[DTypeLikeSpec] = None, ) -> JaxBatch: """Convert a NumPy ndarray batch to a globally sharded JAX Array batch. Args: ndarrays: A single NumPy ndarray or dictionary of NumPy ndarrays. sharding: The JAX sharding to use for the conversion. dtypes: A single JAX dtype or dictionary of JAX dtypes for the created arrays. Returns: A globally sharded JAX Array (or dictionary of arrays) residing in TPU/GPU memory. """ if isinstance(ndarrays, np.ndarray): dtype = _unwrap_single_column_value(dtypes, "dtype") jax_batch = _convert_ndarray_to_jax_array(ndarrays, sharding, dtype=dtype) else: jax_batch = {} for col_name, col_ndarray in ndarrays.items(): dtype = _get_column_value(dtypes, col_name) try: jax_batch[col_name] = _convert_ndarray_to_jax_array( col_ndarray, sharding, dtype=dtype ) except ValueError as e: raise ValueError( f"JAX Array Conversion Error for column '{col_name}'" ) from e return jax_batch def _get_batch_size(batch: NumpyBatch) -> int: """Get the batch size of a NumPy ndarray or dictionary of NumPy ndarrays.""" if isinstance(batch, dict): # Use the first column to determine the batch size try: return len(next(iter(batch.values()))) except StopIteration: return 0 return len(batch) def _pad_array(arr: np.ndarray, target_size: int, pad_value: Scalar) -> np.ndarray: """Pad a single array to target_size using pad_value.""" current_size = len(arr) if current_size == target_size: return arr padding_shape = (target_size - current_size,) + arr.shape[1:] padding = np.full(padding_shape, pad_value, dtype=arr.dtype) return np.concatenate([arr, padding], axis=0) def _dummy_array(arr: np.ndarray, target_size: int, pad_value: Scalar) -> np.ndarray: """Create a dummy array of target_size filled with pad_value.""" shape = (target_size,) + arr.shape[1:] return np.full(shape, pad_value, dtype=arr.dtype) def _pad_batch( batch: NumpyBatch, target_size: int, paddings: PaddingsSpec, ) -> NumpyBatch: """Pad a batch to target_size using paddings.""" if isinstance(batch, dict): return { k: _pad_array(v, target_size, _get_column_value(paddings, k)) for k, v in batch.items() } return _pad_array( batch, target_size, _unwrap_single_column_value(paddings, "padding"), ) def _create_dummy_batch( template_batch: NumpyBatch, target_size: int, paddings: PaddingsSpec, ) -> NumpyBatch: """Create a dummy batch of target_size filled with paddings.""" if isinstance(template_batch, dict): return { k: _dummy_array(v, target_size, _get_column_value(paddings, k)) for k, v in template_batch.items() } return _dummy_array( template_batch, target_size, _unwrap_single_column_value(paddings, "padding"), ) def _yield_batches_no_sync( iterator: Iterator[NumpyBatch], sharding: "jax.sharding.Sharding", num_local_devices: int, batch_size: int, paddings: Optional[PaddingsSpec], dtypes: Optional[DTypeLikeSpec] = None, ) -> Iterator[JaxBatch]: """Yield batches without multi-host synchronization.""" for batch in iterator: local_batch_size = _get_batch_size(batch) if local_batch_size == 0: continue if paddings is not None: if local_batch_size < batch_size: batch = _pad_batch(batch, batch_size, paddings) elif local_batch_size % num_local_devices != 0: # Without padding, batch size must be divisible by num_local_devices raise ValueError( f"The local batch size ({local_batch_size}) must be evenly " f"divisible by the number of local JAX devices " f"({num_local_devices}) on this host. " f"To safely truncate or pad the batch, " f"set `drop_last=True` or provide a `paddings` in `iter_jax_batches()`." ) yield _convert_batch(batch, sharding, dtypes=dtypes) def _fetch_lookahead_batches( iterator: Iterator[NumpyBatch], lookahead: int, ) -> Tuple[List[Optional[NumpyBatch]], List[int], Optional[NumpyBatch]]: """Fetch a window of batches and prepare synchronization info.""" local_batches = [] local_infos = [] template_batch: Optional[NumpyBatch] = None for _ in range(lookahead): try: batch = next(iterator) has_batch = True local_batch_size = _get_batch_size(batch) if template_batch is None: template_batch = batch except StopIteration: batch = None has_batch = False local_batch_size = 0 local_batches.append(batch) local_infos.extend([int(has_batch), local_batch_size]) if not has_batch: break return local_batches, local_infos, template_batch def _yield_batches_with_sync( iterator: Iterator[NumpyBatch], sharding: "jax.sharding.Sharding", num_local_devices: int, drop_last: bool, batch_size: int, paddings: Optional[PaddingsSpec], synchronize_lookahead: int, dtypes: Optional[DTypeLikeSpec] = None, ) -> Iterator[JaxBatch]: """Yield batches with multi-host synchronization.""" import jax.numpy as jnp from jax.experimental.multihost_utils import process_allgather template_batch: Optional[NumpyBatch] = None while True: local_batches, local_infos, window_template = _fetch_lookahead_batches( iterator, synchronize_lookahead ) if template_batch is None: template_batch = window_template # Pad local_infos to 2 * synchronize_lookahead padding_needed = 2 * synchronize_lookahead - len(local_infos) if padding_needed > 0: local_infos.extend([0] * padding_needed) gathered = process_allgather(jnp.array(local_infos, dtype=jnp.int32)) for i in range(synchronize_lookahead): h = gathered[:, 2 * i] s = gathered[:, 2 * i + 1] all_have_batch = bool(h.all()) any_have_batch = bool(h.any()) min_batch_size = int(s.min()) max_batch_size = int(s.max()) if not any_have_batch: return if not all_have_batch: # Some workers have exhausted their data while others have more. if drop_last: # If drop_last=True, we stop as soon as any worker is exhausted. return elif paddings is not None: # If paddings is set, we continue until all workers are exhausted. # Workers that are already exhausted will yield dummy batches. pass else: raise ValueError( "Uneven number of batches detected across JAX workers. " "To safely drop orphaned batches without hanging, " "set `drop_last=True` or provide a `paddings` in `iter_jax_batches()`." ) if paddings is not None: batch = local_batches[i] if batch is None: if template_batch is None: raise ValueError( "Cannot create dummy batches for synchronization because this " "JAX host has not received any data batches to use as a " "template. This usually happens if one JAX host's dataset " "shard is completely empty while others have data. " "Ensure that all JAX hosts have at least one batch of data, " "or use `drop_last=True` to avoid yielding dummy batches." ) batch = _create_dummy_batch(template_batch, batch_size, paddings) else: local_batch_size = _get_batch_size(batch) if local_batch_size < batch_size: batch = _pad_batch(batch, batch_size, paddings) assert batch is not None yield _convert_batch(batch, sharding, dtypes=dtypes) else: if max_batch_size > min_batch_size: raise ValueError( "Uneven batch sizes detected across JAX workers. " f"Host batch sizes range from {min_batch_size} to {max_batch_size}. " "To handle uneven batch sizes, provide a `paddings` in `iter_jax_batches()`." ) if min_batch_size % num_local_devices != 0: raise ValueError( f"The globally minimum batch size ({min_batch_size}) must be evenly " f"divisible by the number of local JAX devices " f"({num_local_devices}) on this host. " f"To safely truncate or pad the batch, " f"set `drop_last=True` or provide a `paddings` in `iter_jax_batches()`." ) batch = local_batches[i] assert batch is not None yield _convert_batch(batch, sharding, dtypes=dtypes) def jax_sync_generator( batch_iterable: Iterable[NumpyBatch], drop_last: bool, batch_size: int = 256, paddings: Optional[PaddingsSpec] = None, dtypes: Optional[DTypeLikeSpec] = None, synchronize_batches: bool = False, synchronize_lookahead: int = 10, ) -> Iterator[JaxBatch]: """A generator that synchronizes and shards batches across JAX workers. This generator wraps a locally yielded batch iterable and ensures that all JAX workers within a multi-host training setup receive the exact same number of batches and identical batch shapes, which is required for JAX's SPMD execution. Args: batch_iterable: An iterable yielding local data batches (either a NumPy ndarray or a dictionary of NumPy ndarrays). drop_last: Whether to drop partial or uneven batches. batch_size: The target batch size for each host. paddings: The value to use for padding uneven batches to `batch_size`. If a dictionary is provided, it must map column names to padding values. If None, padding is disabled. dtypes: A single JAX dtype or dictionary of JAX dtypes for the created arrays. synchronize_batches: Whether to synchronize batch shapes across all hosts. Setting this to False can improve performance if you guarantee that all hosts produce identical batch shapes and counts beforehand. synchronize_lookahead: The number of batches to look ahead and synchronize at once. Increasing this value reduces synchronization overhead but may increase memory usage as more batches are buffered locally. Yields: JaxBatch: Globally sharded batches. """ import jax # Physical Sharding (1D across the _GLOBAL_MESH_1D_AXIS dimension) # The sharding is created once for the lifetime of this generator and reused # across all batches. sharding = _create_sharding_1d(_GLOBAL_MESH_1D_AXIS) num_local_devices = jax.local_device_count() iterator = iter(batch_iterable) if not synchronize_batches or jax.process_count() == 1: yield from _yield_batches_no_sync( iterator, sharding, num_local_devices, batch_size, paddings, dtypes=dtypes, ) else: yield from _yield_batches_with_sync( iterator, sharding, num_local_devices, drop_last, batch_size, paddings, synchronize_lookahead, dtypes=dtypes, )