diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 5669e012ae..04f9d7f8f1 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -294,6 +294,22 @@ norm_topk_prob: false # boolean to enable the top-k probability normalization. q # when moe weight matrices are sharded on both fsdp and fsdp-transpose axes, use two separate all-gather calls moe_fsdp_use_two_stage_all_gather: false +# Comma-separated list of MoE ops to run on the TPU SparseCore instead of the TensorCore, freeing +# TensorCore cycles for the expert GEMMs. Empty (the default) keeps everything on the TensorCore. +# Supported targets, or "all" for every one of them: +# fsdp_all_gather - the MoE weight all-gather over the fsdp / fsdp_transpose axes +# ep_collectives - the expert-parallel activation all-gathers and ragged all-to-alls +# ragged_sort - the routing index math (argsorts, group sizes, offsets) in the ragged sort kernels +# Requires a TPU with a SparseCore (v5p, v6e, tpu7x or newer); MaxText raises otherwise. +# +# Offloading never changes results. For collectives it is not merely a hint though: XLA's +# SparseCore collective-offload pass reads the annotation as "force this one" and CHECK-fails, +# aborting the compile, on an annotated collective the chip cannot lower. MaxText therefore +# only annotates a collective its SparseCore is known to run (utils/sparsecore.py). Offloading +# an all-gather needs Ironwood, and a reduce-scatter transposes to one in the backward pass, so on +# v5p/v6e the fsdp_all_gather target logs a warning and is ignored, and ep_collectives is left with +# just its ragged all-to-alls. ragged_sort is not a collective and runs everywhere. +moe_sparse_core_offload_targets: "" # Shard the expert dimension of the MLP weights on the FSDP axis. # This configuration is recommended only when num_experts is a multiple of fsdp_parallelism shard_exp_on_fsdp: false diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 353656c665..f21e16b8fb 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -37,6 +37,7 @@ from maxtext.utils import elastic_utils from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT, HF_IDS from maxtext.utils import accelerator_to_spec_map +from maxtext.utils import sparsecore from pydantic.config import ConfigDict from pydantic.fields import Field from pydantic.functional_validators import field_validator, model_validator @@ -979,6 +980,13 @@ class MoEGeneral(BaseModel): False, description="Use two separate All-Gather calls for MoE weights sharded on both FSDP and FSDP-transpose.", ) + moe_sparse_core_offload_targets: str = Field( + "", + description="Comma-separated list of MoE ops to run on the TPU SparseCore instead of the TensorCore. " + f"Supported targets: {', '.join(sparsecore.OFFLOAD_TARGETS)}; 'all' enables every one of them. " + "Empty (the default) keeps everything on the TensorCore. Requires a TPU with a SparseCore; targets " + "whose collectives that SparseCore cannot run are warned about and ignored.", + ) shard_exp_on_fsdp: bool = Field( False, description="Shard the expert dimension of the MLP weights on the FSDP axis, " @@ -3283,6 +3291,23 @@ def _validate_check_vma_is_supported(self): f"Found other ICI axes enabled: {active}." ) + def _validate_sparse_core_offload(self): + """Validates moe_sparse_core_offload_targets against the target hardware.""" + # Raises on unrecognized target names. + targets = sparsecore.parse_offload_targets(self.moe_sparse_core_offload_targets) + if not targets: + return + if not sparsecore.has_sparse_core(self.compile_topology, self.hardware): + raise ValueError( + f"moe_sparse_core_offload_targets={self.moe_sparse_core_offload_targets!r} requires a TPU with a " + "SparseCore (v5p, v6e, tpu7x or newer), but the target hardware has none. Set it to '' to keep " + "these ops on the TensorCore." + ) + # Warns now, at startup, for any target this chip cannot serve, rather than + # leaving the first one to surface mid-trace. Unserviceable targets are + # dropped instead of rejected so the same config runs on every chip. + sparsecore.supported_offload_targets(self.moe_sparse_core_offload_targets, self.compile_topology, self.hardware) + def validate_ragged_buffer_factor(self): if self.ragged_buffer_factor <= 0: return # Not using a ragged buffer factor @@ -4807,6 +4832,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de ) self._validate_check_vma_is_supported() + self._validate_sparse_core_offload() # Final string-to-enum conversions if they haven't been coerced by pydantic yet. if isinstance(self.decoder_block, str): diff --git a/src/maxtext/kernels/ragged/ragged_sort.py b/src/maxtext/kernels/ragged/ragged_sort.py index 364c4a00b7..abe22373b8 100644 --- a/src/maxtext/kernels/ragged/ragged_sort.py +++ b/src/maxtext/kernels/ragged/ragged_sort.py @@ -12,12 +12,24 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Ragged token sorting operations with custom VJP.""" +"""Ragged token sorting operations with custom VJP. + +Every entry point here takes ``offload_index_math``. When set, the routing index +math around the Pallas kernels -- argsorts, one-hot histograms, group offsets, +permutations of 1D index/weight vectors and validity masks -- is annotated to run +on the TPU SparseCore, freeing the TensorCore for the expert GEMMs. Ops touching +the hidden dimension stay on the TensorCore: they are dense vector work the +SparseCore is not the right place for. The annotation only picks which core runs +an op, so it never changes results; nothing here is a collective, so unlike the +MoE's collective offload targets it also cannot fail to lower. See +:mod:`maxtext.utils.sparsecore`. +""" import jax import jax.numpy as jnp from maxtext.kernels.ragged.ragged_gather import ragged_gather from maxtext.kernels.ragged.ragged_gather_reduce_v2 import ragged_gather_reduce +from maxtext.utils import sparsecore def ring_ragged_sort( @@ -35,6 +47,7 @@ def ring_ragged_sort( gather_bytes_accessed_override=-1, gather_reduce_bytes_accessed_override=-1, use_single_sparsecore=False, + offload_index_math=False, ): """Ragged-gather variant for AG-RS Expert Parallelism token routing. @@ -63,6 +76,7 @@ def ring_ragged_sort( ep_name: ``str`` identifying the expert parallel axis name. ep_size: scalar ``int`` representing the expert parallel mesh size. buffer_size: optional scalar ``int`` representing the size of the local buffer. + offload_index_math: run the routing index math on the SparseCore. Returns: A tuple containing: @@ -83,24 +97,26 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local): """Sort and gather activations forward pass.""" num_tokens_local = hidden_states_local.shape[0] + shard_idx = jax.lax.axis_index(ep_name) + local_num_experts = num_experts // ep_size - topk_indices_flat = topk_indices_local.flatten() # num_tokens_local x topk - topk_argsort_indices = jnp.argsort(topk_indices_flat) # num_tokens_local x topk + with sparsecore.offload(offload_index_math): + topk_indices_flat = topk_indices_local.flatten() # num_tokens_local x topk + topk_argsort_indices = jnp.argsort(topk_indices_flat) # num_tokens_local x topk - token_indices = jnp.arange(num_tokens_local, dtype=jnp.int32).repeat(topk) # num_tokens_local x topk - token_indices_sorted = token_indices[topk_argsort_indices] # num_tokens_local x topk + token_indices = jnp.arange(num_tokens_local, dtype=jnp.int32).repeat(topk) # num_tokens_local x topk + token_indices_sorted = token_indices[topk_argsort_indices] # num_tokens_local x topk - group_sizes_local = jax.nn.one_hot(topk_indices_flat, num_experts, dtype=jnp.int32).sum(axis=0) # GLOBAL_NUM_EXPERTS + # GLOBAL_NUM_EXPERTS + group_sizes_local = jax.nn.one_hot(topk_indices_flat, num_experts, dtype=jnp.int32).sum(axis=0) - topk_argsort_revert_indices = jnp.argsort(topk_argsort_indices) # num_tokens_local x topk - shard_idx = jax.lax.axis_index(ep_name) + topk_argsort_revert_indices = jnp.argsort(topk_argsort_indices) # num_tokens_local x topk - local_num_experts = num_experts // ep_size - experts_start = shard_idx * local_num_experts - experts_end = experts_start + local_num_experts - group_offsets = jnp.cumulative_sum(group_sizes_local, include_initial=True) - shard_output_start = group_offsets[experts_start] - shard_output_end = group_offsets[experts_end] + experts_start = shard_idx * local_num_experts + experts_end = experts_start + local_num_experts + group_offsets = jnp.cumulative_sum(group_sizes_local, include_initial=True) + shard_output_start = group_offsets[experts_start] + shard_output_end = group_offsets[experts_end] if buffer_size is None or buffer_size >= num_tokens_local * topk: local_buffer_size = num_tokens_local * topk @@ -116,18 +132,19 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local): ) else: local_buffer_size = buffer_size - # We only gather up to the available buffer size or the actual number of - # tokens destined for this shard's experts, whichever is smaller. - gather_end = jnp.minimum(shard_output_end - shard_output_start, local_buffer_size) - # Pad the indices to ensure we can safely slice a block of size `local_buffer_size` - # starting at `shard_output_start` without going out-of-bounds during compilation. - padded_token_indices_sorted = jnp.pad(token_indices_sorted, (0, local_buffer_size)) - sliced_indices = jax.lax.dynamic_slice_in_dim( - padded_token_indices_sorted, - shard_output_start, - local_buffer_size, - axis=0, - ) + with sparsecore.offload(offload_index_math): + # We only gather up to the available buffer size or the actual number of + # tokens destined for this shard's experts, whichever is smaller. + gather_end = jnp.minimum(shard_output_end - shard_output_start, local_buffer_size) + # Pad the indices to ensure we can safely slice a block of size `local_buffer_size` + # starting at `shard_output_start` without going out-of-bounds during compilation. + padded_token_indices_sorted = jnp.pad(token_indices_sorted, (0, local_buffer_size)) + sliced_indices = jax.lax.dynamic_slice_in_dim( + padded_token_indices_sorted, + shard_output_start, + local_buffer_size, + axis=0, + ) x = ragged_gather( hidden_states_local, sliced_indices, @@ -175,9 +192,10 @@ def _ring_ragged_sort_bwd(res, g_out): n = topk_argsort_revert_indices.shape[0] if local_buffer_size >= n: - valid_rows_mask = (topk_argsort_revert_indices >= shard_output_start) & ( - topk_argsort_revert_indices < shard_output_end - ) + with sparsecore.offload(offload_index_math): + valid_rows_mask = (topk_argsort_revert_indices >= shard_output_start) & ( + topk_argsort_revert_indices < shard_output_end + ) # The forward scatter-add over `token_indices_sorted` is equivalent to a # gather-reduce: each input token has exactly `topk` contributions located # at sorted positions `topk_argsort_revert_indices[t*topk:(t+1)*topk]`. @@ -197,16 +215,17 @@ def _ring_ragged_sort_bwd(res, g_out): # Buffering: g_x has size `local_buffer_size` (packed). # The revert indices are global [0, n), but they must map to the local # packed g_x buffer. - shifted_indices = topk_argsort_revert_indices - shard_output_start - local_num_tokens = shard_output_end - shard_output_start - # We only reduce gradients from the valid portion of the local buffer. - limit = jnp.minimum(local_num_tokens, local_buffer_size) - # Mask out tokens that were not gathered (either because they belong to - # other shards, or they exceeded the local buffer size). - valid_rows_mask = (shifted_indices >= 0) & (shifted_indices < limit) - # Clamp invalid indices to 0 to prevent compile-time/run-time out-of-bounds - # in JAX. These clamped values will be ignored due to `valid_rows_mask`. - safe_indices = jnp.where(valid_rows_mask, shifted_indices, 0) + with sparsecore.offload(offload_index_math): + shifted_indices = topk_argsort_revert_indices - shard_output_start + local_num_tokens = shard_output_end - shard_output_start + # We only reduce gradients from the valid portion of the local buffer. + limit = jnp.minimum(local_num_tokens, local_buffer_size) + # Mask out tokens that were not gathered (either because they belong to + # other shards, or they exceeded the local buffer size). + valid_rows_mask = (shifted_indices >= 0) & (shifted_indices < limit) + # Clamp invalid indices to 0 to prevent compile-time/run-time out-of-bounds + # in JAX. These clamped values will be ignored due to `valid_rows_mask`. + safe_indices = jnp.where(valid_rows_mask, shifted_indices, 0) grad_hidden_states = ragged_gather_reduce( g_x, @@ -241,6 +260,7 @@ def ring_ragged_unsort( gather_bytes_accessed_override=-1, gather_reduce_bytes_accessed_override=-1, use_single_sparsecore=False, + offload_index_math=False, ): """Dual of :func:`ring_ragged_sort`. @@ -267,6 +287,7 @@ def ring_ragged_unsort( ep_name: ``str`` identifying the expert parallel axis name. topk_weights: ``[num_tokens_local * topk]`` tensor of per-slot routing weights. Differentiated: its gradient is what trains the router. + offload_index_math: run the routing index math on the SparseCore. Returns: A 2D ``[num_tokens_local, hidden]`` tensor with expert outputs scattered back @@ -296,14 +317,16 @@ def _ring_ragged_unsort_fwd( topk_weights_flat, ): """Executes unsorting sending tokens back.""" - group_offsets = jnp.cumulative_sum(group_sizes_local, include_initial=True) - shard_idx = jax.lax.axis_index(ep_name) - experts_start = shard_idx * local_num_experts - experts_end = experts_start + local_num_experts - shard_output_start = group_offsets[experts_start] - shard_output_end = group_offsets[experts_end] + with sparsecore.offload(offload_index_math): + group_offsets = jnp.cumulative_sum(group_sizes_local, include_initial=True) + + experts_start = shard_idx * local_num_experts + experts_end = experts_start + local_num_experts + + shard_output_start = group_offsets[experts_start] + shard_output_end = group_offsets[experts_end] buffer_size = sorted_tokens_local.shape[0] num_tokens = topk_argsort_revert_indices.shape[0] @@ -318,9 +341,10 @@ def _ring_ragged_unsort_fwd( # from sorted_tokens_local at position `topk_argsort_revert_indices[i]` if # that position is within this shard's [start, end) range, else zero. # The routing weights are applied per-row before the topk reduction. - valid_rows_mask = (topk_argsort_revert_indices >= shard_output_start) & ( - topk_argsort_revert_indices < shard_output_end - ) + with sparsecore.offload(offload_index_math): + valid_rows_mask = (topk_argsort_revert_indices >= shard_output_start) & ( + topk_argsort_revert_indices < shard_output_end + ) out = ragged_gather_reduce( sorted_tokens_local, topk_argsort_revert_indices, @@ -334,11 +358,12 @@ def _ring_ragged_unsort_fwd( ) else: # Shift indices so they map to the packed local buffer [0, local_num_tokens). - shifted_indices = topk_argsort_revert_indices - shard_output_start - local_num_tokens = shard_output_end - shard_output_start - limit = jnp.minimum(local_num_tokens, buffer_size) - valid_rows_mask = (shifted_indices >= 0) & (shifted_indices < limit) - safe_indices = jnp.where(valid_rows_mask, shifted_indices, 0) + with sparsecore.offload(offload_index_math): + shifted_indices = topk_argsort_revert_indices - shard_output_start + local_num_tokens = shard_output_end - shard_output_start + limit = jnp.minimum(local_num_tokens, buffer_size) + valid_rows_mask = (shifted_indices >= 0) & (shifted_indices < limit) + safe_indices = jnp.where(valid_rows_mask, shifted_indices, 0) out = ragged_gather_reduce( sorted_tokens_local, @@ -392,14 +417,16 @@ def _ring_ragged_unsort_bwd(res, g_out): n = topk_argsort_revert_indices.shape[0] # Build the inverse permutation idx_inv such that idx_inv[j] = i # where revert[i] = j. - idx_inv = jnp.argsort(topk_argsort_revert_indices) + with sparsecore.offload(offload_index_math): + idx_inv = jnp.argsort(topk_argsort_revert_indices) # Handle the same two buffering modes for backward pass. # ragged_gather does the fan-out, by indexing into the un-expanded # g_hidden_states_local via idx_inv // topk. It gathers unweighted so the same rows # feed both gradients: the activation one after scaling, the weight one after a dot. if buffer_size >= n: - weight_for_sorted = topk_weights_flat[idx_inv] + with sparsecore.offload(offload_index_math): + weight_for_sorted = topk_weights_flat[idx_inv] gathered = ragged_gather( g_hidden_states_local, idx_inv // topk, @@ -412,20 +439,23 @@ def _ring_ragged_unsort_bwd(res, g_out): ) # Mask out gradients that correspond to elements outside the valid shard # output range. - mask = (jnp.arange(n) >= shard_output_start) & (jnp.arange(n) < shard_output_end) + with sparsecore.offload(offload_index_math): + mask = (jnp.arange(n) >= shard_output_start) & (jnp.arange(n) < shard_output_end) gathered = jnp.where(mask[:, None], gathered, 0.0) grad_sorted_tokens = (gathered * weight_for_sorted[:, None]).astype(gathered.dtype) # Row-wise dot in sorted order, then permuted back to flat slot order. dot_sorted = jnp.sum(gathered.astype(jnp.float32) * sorted_tokens_local[:n].astype(jnp.float32), axis=-1) - grad_topk_weights = dot_sorted[topk_argsort_revert_indices] + with sparsecore.offload(offload_index_math): + grad_topk_weights = dot_sorted[topk_argsort_revert_indices] else: - # Slice the inverse permutation to match the packed local buffer. - padded_idx_inv = jnp.pad(idx_inv, (0, buffer_size)) - sliced_idx_inv = jax.lax.dynamic_slice_in_dim(padded_idx_inv, shard_output_start, buffer_size, axis=0) - gather_end = jnp.minimum(shard_output_end - shard_output_start, buffer_size) - # Slice the per-slot routing weights to match the packed local buffer. - padded_weights = jnp.pad(topk_weights_flat[idx_inv], (0, buffer_size)) - sliced_weights = jax.lax.dynamic_slice_in_dim(padded_weights, shard_output_start, buffer_size, axis=0) + with sparsecore.offload(offload_index_math): + # Slice the inverse permutation to match the packed local buffer. + padded_idx_inv = jnp.pad(idx_inv, (0, buffer_size)) + sliced_idx_inv = jax.lax.dynamic_slice_in_dim(padded_idx_inv, shard_output_start, buffer_size, axis=0) + gather_end = jnp.minimum(shard_output_end - shard_output_start, buffer_size) + # Slice the per-slot routing weights to match the packed local buffer. + padded_weights = jnp.pad(topk_weights_flat[idx_inv], (0, buffer_size)) + sliced_weights = jax.lax.dynamic_slice_in_dim(padded_weights, shard_output_start, buffer_size, axis=0) gathered = ragged_gather( g_hidden_states_local, sliced_idx_inv // topk, @@ -437,14 +467,16 @@ def _ring_ragged_unsort_bwd(res, g_out): use_single_sparsecore=use_single_sparsecore, ) # Mask out gradients for elements beyond the valid limit of the local buffer. - limit = jnp.minimum(shard_output_end - shard_output_start, buffer_size) - mask = jnp.arange(buffer_size) < limit + with sparsecore.offload(offload_index_math): + limit = jnp.minimum(shard_output_end - shard_output_start, buffer_size) + mask = jnp.arange(buffer_size) < limit gathered = jnp.where(mask[:, None], gathered, 0.0) grad_sorted_tokens = (gathered * sliced_weights[:, None]).astype(gathered.dtype) # Scatter the per-slot dot back to flat slot order; dropped slots stay zero. dot_local = jnp.sum(gathered.astype(jnp.float32) * sorted_tokens_local.astype(jnp.float32), axis=-1) - slots = jnp.where(mask, sliced_idx_inv, n) - grad_topk_weights = jnp.zeros((n,), jnp.float32).at[slots].set(dot_local, mode="drop") + with sparsecore.offload(offload_index_math): + slots = jnp.where(mask, sliced_idx_inv, n) + grad_topk_weights = jnp.zeros((n,), jnp.float32).at[slots].set(dot_local, mode="drop") return grad_sorted_tokens, None, None, grad_topk_weights _ring_ragged_unsort.defvjp(_ring_ragged_unsort_fwd, _ring_ragged_unsort_bwd) @@ -467,6 +499,7 @@ def a2a_ragged_sort( enforce_gather_fallback=False, enforce_gather_reduce_fallback=False, use_single_sparsecore=False, + offload_index_math=False, ): """Ragged-gather variant for ``local_permute``. @@ -493,6 +526,7 @@ def a2a_ragged_sort( ordering. Values at positions ``>= valid_end`` are ignored. valid_end: scalar ``int32`` indicating the exclusive end of the valid prefix. + offload_index_math: run the routing index math on the SparseCore. Returns: A 2D ``[num_tokens, hidden]`` tensor sorted by ``sort_indices`` over the @@ -515,7 +549,8 @@ def _a2a_ragged_sort_fwd(inputs, sort_indices, valid_end): use_single_sparsecore=use_single_sparsecore, ) n = sort_indices.shape[0] - valid_mask = jnp.arange(n) < end + with sparsecore.offload(offload_index_math): + valid_mask = jnp.arange(n) < end out = jnp.where(valid_mask[:, None], out, 0.0) res = (sort_indices, end, inputs.shape) return out, res @@ -524,17 +559,19 @@ def _a2a_ragged_sort_fwd(inputs, sort_indices, valid_end): def _a2a_ragged_sort_bwd(res, g_out): sort_indices, end, _ = res n = sort_indices.shape[0] - valid_rows_mask = jnp.arange(n) < end - # g_inputs[sort_indices[i]] += g_out[i], for i in [0, end). This is a - # ragged scatter-add, which we express as a gather-reduce along the inverse - # permutation: each input row j receives exactly one contribution from - # output row i where sort_indices[i] == j. - idx_inv = jnp.argsort(sort_indices) + with sparsecore.offload(offload_index_math): + valid_rows_mask = jnp.arange(n) < end + # g_inputs[sort_indices[i]] += g_out[i], for i in [0, end). This is a + # ragged scatter-add, which we express as a gather-reduce along the inverse + # permutation: each input row j receives exactly one contribution from + # output row i where sort_indices[i] == j. + idx_inv = jnp.argsort(sort_indices) + sorted_valid_rows_mask = valid_rows_mask[idx_inv] grad_inputs = ragged_gather_reduce( g_out, idx_inv, topk_weights=jnp.ones((n,), dtype=jnp.float32), - valid_rows_mask=valid_rows_mask[idx_inv], + valid_rows_mask=sorted_valid_rows_mask, reduce_group_size=1, enforce_fallback=enforce_gather_reduce_fallback, use_single_sparsecore=use_single_sparsecore, @@ -554,6 +591,7 @@ def a2a_ragged_unsort( enforce_gather_fallback=False, enforce_gather_reduce_fallback=False, use_single_sparsecore=False, + offload_index_math=False, ): """Dual of :func:`a2a_ragged_sort`. @@ -573,6 +611,7 @@ def a2a_ragged_unsort( revert_indices: 1D permutation of ``[0, num_tokens)``. valid_end: scalar ``int32`` indicating the exclusive end of the valid prefix. + offload_index_math: run the routing index math on the SparseCore. Returns: A 2D ``[num_tokens, hidden]`` tensor with rows reordered by @@ -588,7 +627,8 @@ def _a2a_ragged_unsort_fwd(sorted_tokens, revert_indices, valid_end): start = jnp.int32(0) end = valid_end.astype(jnp.int32) if hasattr(valid_end, "astype") else jnp.int32(valid_end) n = revert_indices.shape[0] - valid_rows_mask = jnp.arange(n) < end + with sparsecore.offload(offload_index_math): + valid_rows_mask = jnp.arange(n) < end out = ragged_gather_reduce( sorted_tokens, revert_indices, @@ -607,7 +647,8 @@ def _a2a_ragged_unsort_bwd(res, g_out): # g_sorted_tokens[revert_indices[i]] = g_out[i] for i in [0, end). # Because revert_indices is a permutation, build the inverse and use # ragged_gather to pull the per-row gradients to the right positions. - idx_inv = jnp.argsort(revert_indices) + with sparsecore.offload(offload_index_math): + idx_inv = jnp.argsort(revert_indices) grad_sorted = ragged_gather( g_out, idx_inv, @@ -616,8 +657,9 @@ def _a2a_ragged_unsort_bwd(res, g_out): use_single_sparsecore=use_single_sparsecore, ) num_rows = sorted_tokens_shape[0] - pos = jnp.arange(num_rows) - valid = pos < end + with sparsecore.offload(offload_index_math): + pos = jnp.arange(num_rows) + valid = pos < end grad_sorted = jnp.where(valid[:, None], grad_sorted, jnp.zeros_like(grad_sorted)) return grad_sorted, None, None diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 5ea638f426..10a426c0f1 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -43,12 +43,15 @@ from maxtext.utils import max_logging from maxtext.utils import max_utils from maxtext.utils import maxtext_utils +from maxtext.utils import sparsecore from maxtext.utils.sharding import ( + all_gather_axes_between_pspecs, create_sharding, get_logical_axis_rules, logical_to_mesh_axes, maybe_shard_with_logical, maybe_shard_with_pspec, + mesh_axes_in_pspec, remove_expert_from_partition_spec, remove_incompatible_mesh_axes_from_partition_spec, remove_mesh_axes_from_partition_spec, @@ -756,6 +759,59 @@ def _maybe_shard_with_pspec(self, inputs, pspec: jax.sharding.PartitionSpec | No logical_axes=logical_axes, ) + def _sparse_core_offload_targets(self): + """The MoE ops the config asks to run on the SparseCore, minus any this chip cannot serve.""" + return sparsecore.supported_offload_targets( + self.config.moe_sparse_core_offload_targets, self.config.compile_topology, self.config.hardware + ) + + def _offload_to_sparse_core(self, target, collective=None): + """Context manager running the ops traced inside it on the SparseCore, if `target` is enabled. + + Args: + target: the `moe_sparse_core_offload_targets` entry that owns these ops. + collective: the collective traced inside the block, if any. Passed through + so the annotation is skipped on a chip whose SparseCore cannot run it; + see `sparsecore.offload`. + """ + if target not in self._sparse_core_offload_targets(): + return sparsecore.offload(False) + return sparsecore.offload( + True, + collective=collective, + compile_topology=self.config.compile_topology, + hardware=self.config.hardware, + ) + + def _offload_ragged_sort_index_math(self): + """Whether the ragged sort/unsort helpers should offload their index math.""" + return sparsecore.RAGGED_SORT in self._sparse_core_offload_targets() + + def _fsdp_weight_all_gather_plan(self, source_logical_axes, target_pspec, ndim=3): + """Plan for gathering a MoE weight from `source_logical_axes` to `target_pspec`. + + The MoE has always let GSPMD synthesize this all-gather: the weights enter the + `sparse_matmul` shard_map with an `in_spec` that drops the FSDP mesh axes, and + the partitioner inserts the collective at that boundary. An implicit collective + has no op to annotate, so for the `fsdp_all_gather` offload target we instead + hand the weights to the shard_map still FSDP-sharded and perform the very same + gather inside the manual region, where it is a real `all_gather` that can carry + a SparseCore compute type. Making the gather explicit also improves the + backward pass on its own: the weight gradients come back as reduce-scatters + (the transpose of the manual all-gather) instead of all-reduces. + + The caller decides whether the target is on; this only works out the gathers. + + Returns: + A `(dim, axes)` list for `manual_all_gather_weights` -- empty when this + weight is already in its target layout -- or `None` to keep the implicit + boundary gather because the transition is not a pure all-gather. + """ + source_pspec = self._logical_to_mesh_axes(source_logical_axes) + if source_pspec is None or target_pspec is None: + return None + return all_gather_axes_between_pspecs(source_pspec, target_pspec, ndim) + def _maybe_shard_moe_dispatch(self, inputs, logical_axis, peel_expert): """Shard a MoE dispatch/MLP activation. When `peel_expert` is set, drop the 'expert' mesh axis from the batch dim (index 1) so the GEMM stays expert-parallel (AllToAll) @@ -1070,6 +1126,7 @@ def permute( gather_bytes_accessed_override=self.config.ragged_gather_cost_estimate_bytes_accessed, gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed, use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, + offload_index_math=self._offload_ragged_sort_index_math(), ) else: flatten_selected_experts = jnp.ravel(selected_experts) @@ -1175,6 +1232,7 @@ def unpermute( gather_bytes_accessed_override=self.config.ragged_gather_cost_estimate_bytes_accessed, gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed, use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, + offload_index_math=self._offload_ragged_sort_index_math(), ) else: unsort_intermediate = _sort_activations( @@ -1240,6 +1298,7 @@ def local_permute( use_ragged_sort=False, ragged_buffer_factor=-1.0, use_single_sparsecore=False, + offload_index_math=False, ): """Permutes tokens locally within an expert shard. @@ -1268,6 +1327,8 @@ def local_permute( (`a2a_ragged_sort`) to sort only the valid prefix of `inputs`. The ragged buffer can be much larger than the actually-routed token count, so this avoids touching the padded tail in both forward and backward. + offload_index_math: Run the local routing index math (group-size slicing, + expert-index construction and the argsort) on the TPU SparseCore. Returns: A tuple containing: @@ -1279,59 +1340,63 @@ def local_permute( inputs. """ - # Slice the count of local expert IDs in each batch shard. - # all_shard_local_sizes.shape: [expert_shard, local_expert_size] - all_shard_local_sizes = jax.lax.dynamic_slice_in_dim( - global_group_sizes, - shard_index * local_expert_size, - local_expert_size, - axis=1, - ) - local_sizes = all_shard_local_sizes.reshape(-1) - - # Total count of the local expert IDs is the sum of the counts across all - # batch shards, since all batch shards will send their contributions to the - # current expert shard. - local_group_size = RoutedMoE._maybe_truncate_local_group_size( - all_shard_local_sizes, inputs.shape[0], ragged_buffer_factor - ) - - # In this case, the data that needs to be processed by the local shard - # does not start from row 0 but actually starts at - # (jnp.concatenate((jnp.array([0]), - # jnp.cumsum(local_group_sizes[:-1]))[shard_id]). - # This happens if batches (`inputs`) are replicated across expert shards and - # pre-sorted by global Expert ID (via permute()). - if is_offset: - divided_assignments = jnp.floor_divide(global_sorted_experts, local_expert_size) - expert_indices = jnp.where( - divided_assignments == shard_index, - jnp.mod(global_sorted_experts, local_expert_size), + with sparsecore.offload(offload_index_math): + # Slice the count of local expert IDs in each batch shard. + # all_shard_local_sizes.shape: [expert_shard, local_expert_size] + all_shard_local_sizes = jax.lax.dynamic_slice_in_dim( + global_group_sizes, + shard_index * local_expert_size, local_expert_size, + axis=1, ) + local_sizes = all_shard_local_sizes.reshape(-1) - # In this case the `input` data has been received from the batch shards and - # needs to be reorganized in order of local Expert IDs. - else: - base_indices = jnp.mod(jnp.arange(local_sizes.shape[0]), local_expert_size) - expert_indices = jnp.repeat(base_indices, local_sizes, total_repeat_length=inputs.shape[0]) + # Total count of the local expert IDs is the sum of the counts across all + # batch shards, since all batch shards will send their contributions to the + # current expert shard. + local_group_size = RoutedMoE._maybe_truncate_local_group_size( + all_shard_local_sizes, inputs.shape[0], ragged_buffer_factor + ) + + # In this case, the data that needs to be processed by the local shard + # does not start from row 0 but actually starts at + # (jnp.concatenate((jnp.array([0]), + # jnp.cumsum(local_group_sizes[:-1]))[shard_id]). + # This happens if batches (`inputs`) are replicated across expert shards and + # pre-sorted by global Expert ID (via permute()). + if is_offset: + divided_assignments = jnp.floor_divide(global_sorted_experts, local_expert_size) + expert_indices = jnp.where( + divided_assignments == shard_index, + jnp.mod(global_sorted_experts, local_expert_size), + local_expert_size, + ) + + # In this case the `input` data has been received from the batch shards and + # needs to be reorganized in order of local Expert IDs. + else: + base_indices = jnp.mod(jnp.arange(local_sizes.shape[0]), local_expert_size) + expert_indices = jnp.repeat(base_indices, local_sizes, total_repeat_length=inputs.shape[0]) + + sorted_indices = jnp.argsort(expert_indices) + sorted_experts_ids = expert_indices[sorted_indices] - sorted_indices = jnp.argsort(expert_indices) if use_ragged_sort: # Only the first `valid_end` rows of `inputs` carry actual tokens for # this shard (`local_group_size.sum()`), the remainder is padding from # the worst-case ragged buffer. Restricting the gather to that prefix # makes both forward and backward proportional to the routed token count. - valid_end = jnp.sum(local_group_size).astype(jnp.int32) + with sparsecore.offload(offload_index_math): + valid_end = jnp.sum(local_group_size).astype(jnp.int32) sorted_inputs = a2a_ragged_sort( inputs, sorted_indices, valid_end, use_single_sparsecore=use_single_sparsecore, + offload_index_math=offload_index_math, ) else: sorted_inputs = _sort_activations(inputs, sorted_indices, use_custom_sort_vjp) - sorted_experts_ids = expert_indices[sorted_indices] return ( sorted_inputs, sorted_indices, @@ -1856,6 +1921,78 @@ def get_routed_moe_shardings(is_batch_sharded_by_expert, has_input_ids): decoder_tokens_pspec = maybe_replicate_incompatible_batch(decoder_tokens_pspec, input_ids) output_pspec = maybe_replicate_incompatible_batch(output_pspec, inputs) + # SparseCore offload of the FSDP weight all-gather. `w{0,1,o}_in_pspec` is the + # layout the weights are handed to the shard_map in; when a plan is present it + # keeps the FSDP axes so the gather happens inside the manual region (see + # `_fsdp_weight_all_gather_plan`), otherwise it is `w{0,1,o}_pspec` and the + # partitioner inserts the gather at the boundary exactly as before. + # Everything below is skipped unless the target is on, so with the default + # config this method reads and computes exactly what it did before. The + # target is also dropped on a chip whose SparseCore cannot offload an + # all-gather (`sparsecore.supported_offload_targets`), so the graph is never + # restructured for an annotation that would not survive. + weight_ag_plans = None + # Quantized weights carry their own scales and, under `explicitly_weight_ag()`, + # their own hand-written gather inside the shard_map; leave both alone. + if ( + sparsecore.FSDP_ALL_GATHER in self._sparse_core_offload_targets() + and not explicitly_weight_ag() + and not any(isinstance(k, aqt.QTensor) for k in (w0_kernel, w1_kernel, wo_kernel)) + ): + if self.config.moe_fsdp_use_two_stage_all_gather: + # The two-stage path already gathered the FSDP axes off above. + wi_source_axes = ("exp_with_fsdp", None, "mlp_no_fsdp") + wo_source_axes = ("exp_with_fsdp", "mlp_no_fsdp", None) + else: + wi_source_axes = self.wi_kernel_axes + wo_source_axes = self.wo_kernel_axes + plans = ( + self._fsdp_weight_all_gather_plan(wi_source_axes, w0_pspec), + self._fsdp_weight_all_gather_plan(wi_source_axes, w1_pspec), + self._fsdp_weight_all_gather_plan(wo_source_axes, wo_pspec), + ) + # A `None` means that weight's transition is not a pure all-gather, so the + # whole triple falls back; an empty plan just means that weight is already + # in its target layout, which is fine as long as some other one is not. + gathered_axes = {axis for plan in plans if plan for _, axes in plan for axis in axes} + # `all_gather(to="varying")` makes the shard_map's output vary over every + # gathered axis, so with `check_vma` on the out_specs have to still name + # them. `maybe_replicate_incompatible_batch` can strip exactly those axes + # off the batch dim, and shard_map then rejects the region outright. + out_axes = mesh_axes_in_pspec(output_pspec) + if all(plan is not None for plan in plans) and gathered_axes and gathered_axes <= out_axes: + weight_ag_plans = plans + else: + max_logging.log( + f"moe_sparse_core_offload_targets requests {sparsecore.FSDP_ALL_GATHER!r}, but this sharding does not " + "reduce to an all-gather of the MoE weights inside the sparse_matmul shard_map. Leaving the gather to " + "the partitioner, i.e. on the TensorCore; results are unaffected." + ) + if weight_ag_plans is None: + w0_in_pspec, w1_in_pspec, wo_in_pspec = w0_pspec, w1_pspec, wo_pspec + else: + w0_in_pspec = w1_in_pspec = self._logical_to_mesh_axes(wi_source_axes) + wo_in_pspec = self._logical_to_mesh_axes(wo_source_axes) + + def manual_all_gather_weights(w0, w1, wo): + """Runs `weight_ag_plans` inside the shard_map, tagged for the SparseCore. + + This is the same collective the sharding constraint used to produce, so the + forward value is unchanged. The default `to="varying"` is what keeps the + backward pass correct too: its transpose is the `psum_scatter` that reduces + each shard's partial weight gradient. (`to="invarying"` would transpose to a + bare slice, silently dropping that reduction.) + """ + if weight_ag_plans is None: + return w0, w1, wo + gathered = [] + with self._offload_to_sparse_core(sparsecore.FSDP_ALL_GATHER, collective=sparsecore.ALL_GATHER): + for w, plan in zip((w0, w1, wo), weight_ag_plans): + for dim, axes in plan: + w = jax.lax.all_gather(w, axes, axis=dim, tiled=True) + gathered.append(w) + return tuple(gathered) + def roe_ag_and_route( x, logits, @@ -1870,17 +2007,19 @@ def roe_ag_and_route( # expert shards, and then routes within each shard. # Duplicate inputs to all expert shards. - x, logits, pre_bias_logits = tuple( - jax.lax.all_gather(z, axis_name=self._expert_parallelism_name, tiled=True) for z in (x, logits, pre_bias_logits) - ) - if forced_routed_experts is not None: - # Must follow the same all-gather as logits: routing is done on the - # gathered batch, so a shard-local replay would not line up. - forced_routed_experts = jax.lax.all_gather( - forced_routed_experts, - axis_name=self._expert_parallelism_name, - tiled=True, + with self._offload_to_sparse_core(sparsecore.EP_COLLECTIVES, collective=sparsecore.ALL_GATHER): + x, logits, pre_bias_logits = tuple( + jax.lax.all_gather(z, axis_name=self._expert_parallelism_name, tiled=True) + for z in (x, logits, pre_bias_logits) ) + if forced_routed_experts is not None: + # Must follow the same all-gather as logits: routing is done on the + # gathered batch, so a shard-local replay would not line up. + forced_routed_experts = jax.lax.all_gather( + forced_routed_experts, + axis_name=self._expert_parallelism_name, + tiled=True, + ) # "Route" tokens within each shard. num_experts_per_shard = self.config.num_experts // num_ep @@ -1962,7 +2101,8 @@ def ra2a_and_route( global_group_sizes = group_sizes if is_batch_sharded_by_expert: - all_shards_group_sizes = jax.lax.all_gather(reshaped_group_sizes, axis_name=batch_axis) + with self._offload_to_sparse_core(sparsecore.EP_COLLECTIVES, collective=sparsecore.ALL_GATHER): + all_shards_group_sizes = jax.lax.all_gather(reshaped_group_sizes, axis_name=batch_axis) buffer_size = self.get_ragged_buffer_size( jnp.shape(x)[0], num_ep, @@ -1980,16 +2120,18 @@ def ra2a_and_route( output_shape = jax.lax.empty((buffer_size, self.moe_expert_input_dim), dtype=x.dtype) - x = jax.lax.ragged_all_to_all( - x, - output_shape, - input_offsets, - send_sizes, - output_offsets, - recv_sizes, - axis_name=self._expert_parallelism_name, - ) - global_group_sizes = jax.lax.all_gather(group_sizes, axis_name=self._expert_parallelism_name) + with self._offload_to_sparse_core(sparsecore.EP_COLLECTIVES, collective=sparsecore.RAGGED_ALL_TO_ALL): + x = jax.lax.ragged_all_to_all( + x, + output_shape, + input_offsets, + send_sizes, + output_offsets, + recv_sizes, + axis_name=self._expert_parallelism_name, + ) + with self._offload_to_sparse_core(sparsecore.EP_COLLECTIVES, collective=sparsecore.ALL_GATHER): + global_group_sizes = jax.lax.all_gather(group_sizes, axis_name=self._expert_parallelism_name) x, local_sorted_indices, group_sizes, selected_experts = RoutedMoE.local_permute( x, global_group_sizes, @@ -1999,6 +2141,7 @@ def ra2a_and_route( use_ragged_sort=self.config.use_ragged_sort, ragged_buffer_factor=self.config.ragged_buffer_factor, use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, + offload_index_math=self._offload_ragged_sort_index_math(), ) else: x, local_sorted_indices, group_sizes, selected_experts = RoutedMoE.local_permute( @@ -2012,6 +2155,7 @@ def ra2a_and_route( use_ragged_sort=self.config.use_ragged_sort, ragged_buffer_factor=self.config.ragged_buffer_factor, use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, + offload_index_math=self._offload_ragged_sort_index_math(), ) return ( @@ -2238,6 +2382,7 @@ def unsort_output_and_ra2a( jnp.argsort(route_metadata.local_sorted_indices), # pylint: disable=undefined-variable valid_end, use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore, + offload_index_math=self._offload_ragged_sort_index_math(), ) else: local_output = _sort_activations( @@ -2255,15 +2400,16 @@ def unsort_output_and_ra2a( buffer_size=buffer_size, is_dispatch=False, ) - return jax.lax.ragged_all_to_all( - local_output, - output_shape, - input_offsets, - send_sizes, - output_offsets, - recv_sizes, - axis_name=self._expert_parallelism_name, - ) + with self._offload_to_sparse_core(sparsecore.EP_COLLECTIVES, collective=sparsecore.RAGGED_ALL_TO_ALL): + return jax.lax.ragged_all_to_all( + local_output, + output_shape, + input_offsets, + send_sizes, + output_offsets, + recv_sizes, + axis_name=self._expert_parallelism_name, + ) # If batch is replicated across EP shards then each shard should send # 0..local_shard_size data to the other shards and receive the @@ -2275,15 +2421,16 @@ def unsort_output_and_ra2a( is_batch_sharded=False, is_dispatch=False, ) - return jax.lax.ragged_all_to_all( - intermediate_output, - output_shape, - input_offsets, - send_sizes, - output_offsets, - recv_sizes, - axis_name=self._expert_parallelism_name, - ) + with self._offload_to_sparse_core(sparsecore.EP_COLLECTIVES, collective=sparsecore.RAGGED_ALL_TO_ALL): + return jax.lax.ragged_all_to_all( + intermediate_output, + output_shape, + input_offsets, + send_sizes, + output_offsets, + recv_sizes, + axis_name=self._expert_parallelism_name, + ) def moe_emb_chunking( x, @@ -2474,12 +2621,13 @@ def _moe_body( self.moe_expert_input_dim // self.get_tensor_parallelism_size(), ), ) - output = jax.lax.psum_scatter( - output, - self._expert_parallelism_name, - scatter_dimension=0, - tiled=True, - ) + with self._offload_to_sparse_core(sparsecore.EP_COLLECTIVES, collective=sparsecore.REDUCE_SCATTER): + output = jax.lax.psum_scatter( + output, + self._expert_parallelism_name, + scatter_dimension=0, + tiled=True, + ) return output, routing.lb_loss, routing.bias_updates if self.get_expert_parallelism_size() > 1: @@ -2521,9 +2669,9 @@ def _moe_body( input_partition_pspec, gate_logits_pspec, pre_bias_logits_pspec, - w0_pspec, - w1_pspec, - wo_pspec, + w0_in_pspec, + w1_in_pspec, + wo_in_pspec, w0_bias_pspec, w1_bias_pspec, wo_bias_pspec, @@ -2557,8 +2705,10 @@ def sparse_matmul_route_and_compute( ): # The expert weights (w0/w1/wo) are all-gathered over FSDP once at this # shard_map entry (implicitly, via the `embed_tensor_transpose` pspec which - # drops fsdp -> GSPMD inserts the boundary all-gather) and reused across all - # chunks of the ring-of-experts pipeline below. + # drops fsdp -> GSPMD inserts the boundary all-gather; explicitly here when + # the gather is offloaded to the SparseCore) and reused across all chunks of + # the ring-of-experts pipeline below. + w0, w1, wo = manual_all_gather_weights(w0, w1, wo) n_chunks = self.config.num_moe_token_chunks if n_chunks <= 1 or not self.config.use_ring_of_experts: return _moe_body( @@ -2664,9 +2814,9 @@ def sparse_matmul_route_and_compute( logical_axes=gate_logits_logical_axes, ) - w0_kernel = self._maybe_shard_with_pspec(w0_kernel, w0_pspec) - w1_kernel = self._maybe_shard_with_pspec(w1_kernel, w1_pspec) - wo_kernel = self._maybe_shard_with_pspec(wo_kernel, wo_pspec) + w0_kernel = self._maybe_shard_with_pspec(w0_kernel, w0_in_pspec) + w1_kernel = self._maybe_shard_with_pspec(w1_kernel, w1_in_pspec) + wo_kernel = self._maybe_shard_with_pspec(wo_kernel, wo_in_pspec) if w0_bias is not None: w0_bias = self._maybe_shard_with_pspec(w0_bias, w0_bias_pspec) if w1_bias is not None: diff --git a/src/maxtext/utils/sharding.py b/src/maxtext/utils/sharding.py index e218ca89fb..ad64d5e726 100644 --- a/src/maxtext/utils/sharding.py +++ b/src/maxtext/utils/sharding.py @@ -986,6 +986,59 @@ def get_formatted_sharding_annotations(params, mesh=None): FSDP_MESH_AXES = ("fsdp", "fsdp_transpose") +def _pspec_axes_per_dim(pspec, ndim): + """Normalizes a PartitionSpec into one axis-name tuple per array dim. + + Returns `None` if `pspec` names more dims than `ndim`, since silently ignoring + the extra ones would describe a different sharding than the caller passed. + """ + if len(pspec) > ndim: + return None + per_dim = [] + for i in range(ndim): + axis = pspec[i] if i < len(pspec) else None + if axis is None: + per_dim.append(()) + elif isinstance(axis, str): + per_dim.append((axis,)) + else: + per_dim.append(tuple(axis)) + return per_dim + + +def mesh_axes_in_pspec(pspec): + """Returns the set of mesh axis names a PartitionSpec shards over, on any dim.""" + if pspec is None: + return set() + axes = set() + for entry in pspec: + if entry is None: + continue + axes.update((entry,) if isinstance(entry, str) else entry) + return axes + + +def all_gather_axes_between_pspecs(source_pspec, target_pspec, ndim): + """Returns the `(dim, axes)` all-gathers that take `source_pspec` to `target_pspec`. + + Returns `None` when the transition is not a pure all-gather, i.e. when some dim + gains an axis or drops one from anywhere but the end of its axis tuple. Callers + are expected to fall back to a plain sharding constraint in that case. + """ + source = _pspec_axes_per_dim(source_pspec, ndim) + target = _pspec_axes_per_dim(target_pspec, ndim) + if source is None or target is None: + return None + gathers = [] + for dim, (src_axes, tgt_axes) in enumerate(zip(source, target)): + if src_axes == tgt_axes: + continue + if src_axes[: len(tgt_axes)] != tgt_axes: + return None + gathers.append((dim, src_axes[len(tgt_axes) :])) + return gathers + + def remove_mesh_axes_from_partition_spec(pspec, axes_to_remove, dims=None): """Return `pspec` with `axes_to_remove` stripped from the given dims. diff --git a/src/maxtext/utils/sparsecore.py b/src/maxtext/utils/sparsecore.py new file mode 100644 index 0000000000..64dd0e3095 --- /dev/null +++ b/src/maxtext/utils/sparsecore.py @@ -0,0 +1,394 @@ +# Copyright 2026 Google LLC +# +# 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. + +"""TPU SparseCore capability detection and compute offloading. + +Recent TPUs pair each TensorCore with one or more SparseCores. XLA:TPU can run +some ops there instead of on the TensorCore when they carry the +``tpu_sparsecore`` compute type, which frees TensorCore cycles for the matmuls +and lets the offloaded work overlap with them. + +This module owns three things: + + * :func:`sparse_core_info` / :func:`has_sparse_core` -- whether the *target* + chip has a SparseCore at all. The answer is taken from JAX's own chip table + (``pltpu.get_tpu_info_for_chip``) rather than pattern-matching device + strings, so it stays correct as new chips are added. + * :func:`supports_collective_offload` / :func:`supported_offload_targets` -- + which collectives that SparseCore can actually run. + * :func:`offload` -- the single context manager used to tag ops. + +Offloading is opt-in per target; see ``moe_sparse_core_offload_targets``, which +the config validator refuses to enable unless the *target* chip has a +SparseCore. That check is what lets :func:`offload` annotate unconditionally, +so an ahead-of-time compile on a TPU-less host produces the same HLO the real +run will. + +The annotation never changes results: it picks which core runs an op, not what +it computes. For ordinary compute it is also purely advisory -- XLA offloads the +ops it can lower and silently leaves the rest on the TensorCore. + +*Collectives* are the exception, and the reason +:func:`supports_collective_offload` exists. XLA's SparseCore collective-offload +pass treats the annotation as "force this one", so when it selects an annotated +collective the chip cannot lower it CHECK-fails -- aborting the compiler, not +falling back:: + + F sparse_core_collective_offload.cc:586] Candidate rejected: instruction has + compute type annotation sparseoffload but the operation is currently not + supported on SC. %all-gather-start = ... + +Measured on v5p with libtpu 0.0.46, where the pass runs under its default flags: +an annotated ragged all-to-all and reduce-scatter are offloaded, an annotated +all-reduce is silently ignored, and an annotated all-gather aborts the compile +at every rank. Per ``benchmarks/xla_flags_library.py`` and the original +SparseCore MoE work, all-gather offload arrives with Ironwood, so +:data:`_MIN_GENERATION` only lets an all-gather carry the annotation from chip +generation 7 on -- and, via :data:`_TRANSPOSE_OF`, neither does a reduce-scatter, +whose backward pass *is* an all-gather. +""" + +from __future__ import annotations + +import contextlib +import functools +import inspect + +import jax +from jax.experimental import compute_on + +from maxtext.utils import accelerator_to_spec_map +from maxtext.utils import max_logging + +# XLA compute type that moves an op onto the SparseCore. +SPARSE_CORE_COMPUTE_TYPE = "tpu_sparsecore" + +# Offload targets accepted by ``moe_sparse_core_offload_targets``. +FSDP_ALL_GATHER = "fsdp_all_gather" +EP_COLLECTIVES = "ep_collectives" +RAGGED_SORT = "ragged_sort" +OFFLOAD_TARGETS = (FSDP_ALL_GATHER, EP_COLLECTIVES, RAGGED_SORT) +_ALL_TARGETS = "all" + +# Collective kinds that get annotated, for :func:`supports_collective_offload`. +ALL_GATHER = "all_gather" +RAGGED_ALL_TO_ALL = "ragged_all_to_all" +REDUCE_SCATTER = "reduce_scatter" + +# Lowest chip generation whose SparseCore can run each collective. Anything +# absent is assumed offloadable on every chip that has a SparseCore, which is +# the safe assumption because an annotation XLA does not act on costs nothing -- +# see the module docstring for what happens when it does act on one it cannot +# lower. +_MIN_GENERATION = {ALL_GATHER: 7} + +# What each collective becomes under transposition. The compute type is +# snapshotted onto the jaxpr equation, and AD carries it onto the transposed +# equation, so annotating one end annotates the other: an annotated +# reduce-scatter puts an annotated all-gather in the backward pass, which is +# just as fatal as annotating the all-gather directly. Both ends therefore have +# to be offloadable before either is annotated. +_TRANSPOSE_OF = { + ALL_GATHER: REDUCE_SCATTER, + REDUCE_SCATTER: ALL_GATHER, + RAGGED_ALL_TO_ALL: RAGGED_ALL_TO_ALL, +} + +# The one collective a target exists to annotate. Requesting such a target on a +# chip that cannot offload that collective is pointless, so the target is +# dropped whole rather than left to restructure the graph for nothing. Targets +# that annotate a mix of collectives are absent here and gate per call site. +_TARGET_REQUIRED_COLLECTIVE = {FSDP_ALL_GATHER: ALL_GATHER} + +# ``hardware`` values that can never have a SparseCore. +_NON_TPU_HARDWARE = ("cpu", "gpu", "gpu_multiprocess") + + +def _chip_version(accelerator_name: str): + """Maps a user-facing accelerator name (e.g. ``tpu7x-256``) to a ChipVersion. + + ``accelerator_to_spec_map`` keys are ``-`` and the family + matches ``pltpu.ChipVersion``'s value modulo a ``tpu`` prefix (``tpu7x`` vs + ``7x``), so no per-chip table is needed here. + """ + import jax.experimental.pallas.tpu as pltpu # pylint: disable=import-outside-toplevel + + family = accelerator_name.split("-", maxsplit=1)[0].lower() + for candidate in (family, family.removeprefix("tpu")): + try: + return pltpu.ChipVersion(candidate) + except ValueError: + continue + return None + + +@functools.cache +def _tpu_info(compile_topology: str = "", hardware: str = ""): + """Returns JAX's ``TpuInfo`` for the target chip, or ``None`` if it is not a TPU. + + Args: + compile_topology: AOT target topology (e.g. ``tpu7x-256``). When set, the + answer describes that target rather than the local devices. + hardware: the ``hardware`` config value, used to rule out CPU/GPU runs. + + Returns: + The target chip's ``TpuInfo``, or ``None`` when the target is not a TPU or + cannot be determined. + """ + if hardware in _NON_TPU_HARDWARE: + return None + + import jax.experimental.pallas.tpu as pltpu # pylint: disable=import-outside-toplevel + + if compile_topology: + try: + spec = accelerator_to_spec_map.get_system_characteristics(compile_topology) + except ValueError: + return None + if spec.platform != "tpu": + return None + chip_version = _chip_version(compile_topology) + if chip_version is None: + return None + # Neither SparseCore presence nor the chip generation depends on the + # Megacore split, so 1 core per logical device is a valid probe here. + return pltpu.get_tpu_info_for_chip(chip_version, 1) + + if not is_tpu_runtime(): + return None + try: + return pltpu.get_tpu_info() + except (RuntimeError, ValueError, AttributeError, TypeError, IndexError): + return None + + +def sparse_core_info(compile_topology: str = "", hardware: str = ""): + """Returns the target chip's ``SparseCoreInfo``, or ``None`` if it has none. + + Args: + compile_topology: AOT target topology (e.g. ``tpu7x-256``). When set, the + answer describes that target rather than the local devices. + hardware: the ``hardware`` config value, used to rule out CPU/GPU runs. + + Returns: + JAX's ``SparseCoreInfo`` for the target chip, or ``None`` when the target + has no SparseCore or cannot be determined. + """ + info = _tpu_info(compile_topology, hardware) + return None if info is None else info.sparse_core + + +def has_sparse_core(compile_topology: str = "", hardware: str = "") -> bool: + """Whether the target chip has a SparseCore. See :func:`sparse_core_info`.""" + return sparse_core_info(compile_topology, hardware) is not None + + +def supports_collective_offload(collective: str, compile_topology: str = "", hardware: str = "") -> bool: + """Whether the target's SparseCore can run `collective` when it is annotated. + + Getting this wrong in the permissive direction aborts the compiler rather than + costing performance, so the answer is conservative: a collective is offloadable + only from the chip generation both it and its transpose are known to work on. + See the module docstring. + + Args: + collective: one of :data:`ALL_GATHER`, :data:`RAGGED_ALL_TO_ALL`, + :data:`REDUCE_SCATTER`. + compile_topology: AOT target topology, as in :func:`sparse_core_info`. + hardware: the ``hardware`` config value. + + Returns: + Whether it is safe to annotate that collective for the target chip. + """ + info = _tpu_info(compile_topology, hardware) + if info is None or info.sparse_core is None: + return False + both_ends = {collective, _TRANSPOSE_OF.get(collective, collective)} + return all(info.generation >= _MIN_GENERATION.get(kind, 0) for kind in both_ends) + + +def is_tpu_runtime() -> bool: + """Whether the local devices are TPUs. Mirrors the kernel fallback guards.""" + try: + return jax.devices()[0].platform == "tpu" + except (RuntimeError, IndexError): + return False + + +@functools.cache +def parse_offload_targets(targets: str) -> frozenset[str]: + """Parses ``moe_sparse_core_offload_targets`` into a set of target names. + + Args: + targets: empty, ``"all"``, or a comma-separated subset of + :data:`OFFLOAD_TARGETS`. + + Returns: + The requested targets. + + Raises: + ValueError: if an unrecognized target is requested. + """ + if not targets: + return frozenset() + requested = [t.strip() for t in targets.split(",") if t.strip()] + if _ALL_TARGETS in requested: + return frozenset(OFFLOAD_TARGETS) + unknown = sorted(set(requested) - set(OFFLOAD_TARGETS)) + if unknown: + raise ValueError( + f"Unknown SparseCore offload target(s) {unknown} in " + f"moe_sparse_core_offload_targets={targets!r}. " + f"Supported targets: {list(OFFLOAD_TARGETS)} or '{_ALL_TARGETS}'." + ) + return frozenset(requested) + + +@functools.cache +def _warn_target_unsupported(target: str, collective: str, chip: str) -> None: + """Logs once that a requested offload target cannot run on the target chip.""" + max_logging.log( + f"moe_sparse_core_offload_targets requests {target!r}, but the SparseCore on {chip} cannot run an annotated " + f"{collective} (XLA would abort the compile rather than fall back). Ignoring that target and keeping those ops " + "on the TensorCore; results are unaffected. Other requested targets are unaffected." + ) + + +@functools.cache +def supported_offload_targets(targets: str, compile_topology: str = "", hardware: str = "") -> frozenset[str]: + """Parses ``moe_sparse_core_offload_targets`` and drops what the chip cannot serve. + + A target whose whole purpose is annotating one collective (see + :data:`_TARGET_REQUIRED_COLLECTIVE`) is dropped on a chip that cannot offload + it, with a warning, rather than failing the run: the offload is an + optimization, and the same config should stay usable across chip generations. + Targets that annotate a mix of collectives survive and gate per call site. + + Args: + targets: the ``moe_sparse_core_offload_targets`` config value. + compile_topology: AOT target topology, as in :func:`sparse_core_info`. + hardware: the ``hardware`` config value. + + Returns: + The requested targets this chip can actually serve. + + Raises: + ValueError: if an unrecognized target is requested. + """ + requested = parse_offload_targets(targets) + if not requested: + return requested + supported = set(requested) + for target, collective in _TARGET_REQUIRED_COLLECTIVE.items(): + if target in supported and not supports_collective_offload(collective, compile_topology, hardware): + info = _tpu_info(compile_topology, hardware) + _warn_target_unsupported(target, collective, compile_topology or str(info.chip_version if info else "this chip")) + supported.discard(target) + return frozenset(supported) + + +def _resolve_compute_type_context(): + """Finds the JAX API that stamps a compute type onto every op traced in a block. + + JAX has moved this around. Through 0.11.0, + ``jax.experimental.compute_on.compute_on`` was a context manager taking the + compute type. In 0.11.1 that name became a function transform, + ``compute_on(f, *, compute_type, out_memory_spaces)``, which traces ``f`` into + its own computation and puts the attribute on the call instead. + + The transform cannot express what the MoE needs. It refuses to nest at all -- + ``_compute_on_lowering`` raises "Nesting `compute_on` with different compute + types is not allowed" for *any* nesting, and + ``moe_sparse_core_offload_targets=all`` nests the ragged-sort blocks inside the + expert-parallel ones. It would also turn every annotated block into a separate + computation, changing what XLA is free to fuse and forcing the intermediates + across a call boundary, which is not what per-op offloading is supposed to do. + + The per-op mechanism itself is unchanged in every version: the compute type + lives in a config context that ``jax._src.core.JaxprEqnContext`` snapshots onto + each equation, and ``mlir.wrap_compute_type_in_place`` turns that into the + ``_xla_compute_type`` frontend attribute. So prefer the public context manager + while it still is one, and otherwise fall back to the private helper that both + public spellings are built on. + + Returns: + A callable mapping a compute type to a context manager, or ``None`` if this + JAX exposes neither spelling. + """ + public = getattr(compute_on, "compute_on", None) + if public is not None: + try: + parameters = list(inspect.signature(public).parameters) + except (TypeError, ValueError): + parameters = [] + if parameters == ["compute_type"]: + return public + try: + # pylint: disable-next=import-outside-toplevel + from jax._src.compute_on import extend_compute_type + except ImportError: + return None + return extend_compute_type + + +_COMPUTE_TYPE_CONTEXT = _resolve_compute_type_context() + + +@functools.cache +def _warn_offload_unavailable() -> None: + """Logs once that this JAX offers no way to annotate a block of ops.""" + max_logging.log( + "moe_sparse_core_offload_targets is set, but this JAX exposes no " + "compute-type context manager (looked for a context-manager " + "jax.experimental.compute_on.compute_on and for " + "jax._src.compute_on.extend_compute_type). Running everything on the " + "TensorCore instead; results are unaffected." + ) + + +@contextlib.contextmanager +def offload(enabled: bool, collective: str | None = None, compile_topology: str = "", hardware: str = ""): + """Runs ops traced inside this block on the SparseCore when ``enabled``. + + The annotation picks which core runs an op, not what it computes, so this is + numerically transparent. Non-TPU backends ignore the frontend attribute, and + the config validator already refuses to enable a target whose hardware has no + SparseCore, so this deliberately does *not* re-check the local devices: + ahead-of-time compilation for a SparseCore topology usually runs on a host + that has no TPU at all, and its HLO has to match what the real run compiles. + + Args: + enabled: whether this offload target is turned on. + collective: the collective kind traced inside the block, when it contains + one. The block is left unannotated if the target chip's SparseCore cannot + run it, because XLA aborts the compile on such an annotation instead of + falling back. Blocks of ordinary compute pass ``None``. + compile_topology: AOT target topology, as in :func:`sparse_core_info`. + hardware: the ``hardware`` config value. + + Yields: + None. + """ + if not enabled: + yield + return + if collective is not None and not supports_collective_offload(collective, compile_topology, hardware): + yield + return + if _COMPUTE_TYPE_CONTEXT is None: + _warn_offload_unavailable() + yield + return + with _COMPUTE_TYPE_CONTEXT(SPARSE_CORE_COMPUTE_TYPE): + yield diff --git a/tests/unit/moe_test.py b/tests/unit/moe_test.py index ad3618f8f9..acfe50496f 100644 --- a/tests/unit/moe_test.py +++ b/tests/unit/moe_test.py @@ -34,7 +34,7 @@ from maxtext.layers import nnx_wrappers from maxtext.layers.initializers import NdInitializer, nd_dense_init, variable_to_logically_partitioned from maxtext.layers.quantizations import configure_quantization, Fp8Quantization -from maxtext.utils import max_logging, maxtext_utils +from maxtext.utils import max_logging, maxtext_utils, sparsecore from maxtext.utils.sharding import remove_expert_from_partition_spec from tests.utils.test_helpers import get_test_config_path @@ -466,6 +466,8 @@ def test_sparse_matmul_repairs_batch_specs_only_without_expert_parallelism(exper *(original_batch_partition if axis == "activation_batch" else None for axis in logical_axes) ) fake_moe._maybe_shard_with_pspec = lambda value, _pspec, **_kwargs: value # pylint: disable=protected-access + # No SparseCore offload targets, i.e. the default `moe_sparse_core_offload_targets`. + fake_moe._sparse_core_offload_targets = frozenset # pylint: disable=protected-access inputs = SimpleNamespace(shape=(4, 1024, 2048)) gate_logits = SimpleNamespace(shape=(4, 1024, 256)) @@ -2432,5 +2434,160 @@ def test_prefuse_moe_weights_matches_unfused(self): ) +# Meshes exercised by SparseCoreOffloadTest, spelled out so that enabling one +# axis does not leave `ici_fsdp_parallelism` at its `-1` default. +_FSDP = {"ici_fsdp_parallelism": -1, "ici_expert_parallelism": 1} +_EP = {"ici_fsdp_parallelism": 1, "ici_expert_parallelism": -1} +# The ragged gather/reduce kernels partition the hidden dimension across +# SparseCore lanes, so they need a realistically wide embedding to be legal. +_EP_RAGGED = {**_EP, "use_ragged_sort": True, "base_emb_dim": 4096} + + +@pytest.mark.tpu_only +class SparseCoreOffloadTest(parameterized.TestCase): + """Tests for `moe_sparse_core_offload_targets`. + + Moving an op to the SparseCore is a scheduling hint, so each target must (a) + actually annotate ops in the compiled HLO and (b) leave the loss and every + parameter gradient unchanged. The backward pass is the interesting half: the + offload is applied to real collectives inside the `sparse_matmul` shard_map, + and a collective whose transpose rule drops a reduction would still produce a + correct forward value. + """ + + BASE_CONFIG = { + "enable_checkpointing": False, + "model_name": "mixtral-8x7b", + "override_model_config": True, + "base_emb_dim": 512, + "base_mlp_dim": 256, + "base_moe_mlp_dim": 256, + "dtype": "bfloat16", + "megablox": True, + "sparse_matmul": True, + "per_device_batch_size": 1, + "max_target_length": 64, + "float32_gate_logits": True, + } + + def setUp(self): + super().setUp() + if not sparsecore.has_sparse_core(): + self.skipTest("Requires a TPU with a SparseCore (v5p, v6e, tpu7x or newer).") + + def _loss_and_grads(self, run_name, **overrides): + """Returns `(loss, param_grads, hlo_text)` for one MoE config.""" + cfg = pyconfig.initialize([None, get_test_config_path()], run_name=run_name, **{**self.BASE_CONFIG, **overrides}) + mesh = Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes) + model = moe.get_routed_moe( + name="MoeBlock", + config=cfg, + num_experts=cfg.num_experts, + num_experts_per_tok=cfg.num_experts_per_tok, + mesh=mesh, + kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("embed", "mlp"), + intermediate_dim=cfg.mlp_dim, + dtype=cfg.dtype, + ) + inputs = jax.random.uniform( + jax.random.PRNGKey(1), + (int(cfg.per_device_batch_size) * jax.device_count(), cfg.max_target_length, cfg.base_emb_dim), + dtype=cfg.dtype, + ) + + def loss_fn(params, x): + output, load_balance_loss, _ = model.apply({"params": params}, x) + loss = jnp.mean(output.astype(jnp.float32) ** 2) + if load_balance_loss is not None: + loss += load_balance_loss.astype(jnp.float32) + return loss + + def init(): + return model.init({"params": jax.random.PRNGKey(0), "dropout": jax.random.PRNGKey(0)}, inputs) + + with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg.logical_axis_rules): + var_shardings = nn.logical_to_mesh_sharding( + nn.get_partition_spec(jax.eval_shape(init)), mesh, cfg.logical_axis_rules + ) + variables = jax.jit(init, out_shardings=var_shardings)() + # Constraining the gradients back to the parameter sharding mirrors a real + # train step, which is what makes the weight-gradient collectives appear. + step = jax.jit(jax.value_and_grad(loss_fn), out_shardings=(None, var_shardings["params"])) + hlo_text = step.lower(variables["params"], inputs).compile().as_text() + loss, grads = jax.block_until_ready(step(variables["params"], inputs)) + return float(loss), grads, hlo_text + + @parameterized.named_parameters( + ("fsdp_all_gather", _FSDP, sparsecore.FSDP_ALL_GATHER), + ("ep_collectives", _EP, sparsecore.EP_COLLECTIVES), + ("ragged_sort", _EP_RAGGED, sparsecore.RAGGED_SORT), + ("all_targets", _EP_RAGGED, "all"), + ) + def test_offload_is_numerically_transparent(self, parallelism, targets): + """Enabling a target annotates the HLO without changing loss or gradients.""" + if targets == sparsecore.FSDP_ALL_GATHER and not sparsecore.supports_collective_offload(sparsecore.ALL_GATHER): + # The target is dropped on such a chip, so there would be nothing to + # assert; see `sparsecore.supported_offload_targets`. + self.skipTest("This chip's SparseCore cannot offload an all-gather.") + ref_loss, ref_grads, ref_hlo = self._loss_and_grads( + f"sc_offload_ref_{targets}", moe_sparse_core_offload_targets="", **parallelism + ) + loss, grads, hlo = self._loss_and_grads( + f"sc_offload_{targets}", moe_sparse_core_offload_targets=targets, **parallelism + ) + + annotation = '_xla_compute_type="sparseoffload"' + self.assertEqual(ref_hlo.count(annotation), 0, "The baseline must not offload anything to the SparseCore.") + self.assertGreater(hlo.count(annotation), 0, f"Offload target {targets!r} annotated no ops.") + + self.assertAlmostEqual(loss, ref_loss, places=6) + diff_summary = compare_tree(ref_grads, grads, relative_norm_diff_threshold=1e-5) + max_logging.log("\n" + diff_summary) + + def test_explicit_fsdp_weight_gather_matches_the_implicit_one(self): + """The graph rewrite behind `fsdp_all_gather` is numerically transparent on its own. + + The target replaces the implicit gather at the `sparse_matmul` shard_map + boundary with an explicit `jax.lax.all_gather` inside the manual region, + whose transpose is the reduce-scatter of the weight gradients. Get that + transpose wrong -- `to="invarying"` lowers to a bare slice with no `psum` -- + and the forward value stays correct while the gradients are silently + unreduced, so only a gradient comparison catches it. + + On a chip whose SparseCore cannot offload an all-gather the target is + dropped and the rewrite never happens, which would leave this comparing the + baseline against itself. The capability gate is patched out so the rewrite + runs everywhere; the per-collective gate inside `sparsecore.offload` still + suppresses the annotation, which is what keeps XLA from aborting the + compile, and the count below pins that down. + """ + ref_loss, ref_grads, ref_hlo = self._loss_and_grads("sc_fsdp_ag_ref", moe_sparse_core_offload_targets="", **_FSDP) + + def keep_every_target(targets, *_args, **_kwargs): + return sparsecore.parse_offload_targets(targets) + + with mock.patch.object(sparsecore, "supported_offload_targets", keep_every_target): + loss, grads, hlo = self._loss_and_grads( + "sc_fsdp_ag_forced", moe_sparse_core_offload_targets=sparsecore.FSDP_ALL_GATHER, **_FSDP + ) + + # The explicit gather's transpose is what turns the weight-gradient + # all-reduces into reduce-scatters, so this both proves the rewrite happened + # -- without it the comparison below would be the baseline against itself -- + # and shows the transpose is the collective it is supposed to be. + self.assertEqual(ref_hlo.count("reduce-scatter("), 0) + self.assertGreater(hlo.count("reduce-scatter("), 0) + annotated = hlo.count('_xla_compute_type="sparseoffload"') > 0 + self.assertEqual(annotated, sparsecore.supports_collective_offload(sparsecore.ALL_GATHER)) + self.assertAlmostEqual(loss, ref_loss, places=6) + max_logging.log("\n" + compare_tree(ref_grads, grads, relative_norm_diff_threshold=1e-5)) + + def test_disabled_offload_leaves_the_hlo_untouched(self): + """Backward compatibility: the default config compiles to the same HLO as before.""" + _, _, hlo = self._loss_and_grads("sc_offload_disabled", ici_fsdp_parallelism=-1) + self.assertEqual(hlo.count('_xla_compute_type="sparseoffload"'), 0) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/sparsecore_test.py b/tests/unit/sparsecore_test.py new file mode 100644 index 0000000000..132c9999ec --- /dev/null +++ b/tests/unit/sparsecore_test.py @@ -0,0 +1,318 @@ +# Copyright 2026 Google LLC +# +# 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. +"""Tests for SparseCore capability detection and the MoE offload plumbing. + +Everything here runs on CPU: the detection helpers answer for a *target* chip +named by ``compile_topology``, and :func:`sparsecore.offload` degrades to a +no-op off TPU. The numerical equivalence of the offload itself is covered by +the TPU tests in ``moe_test.py``. +""" + +import contextlib +from unittest import mock + +from absl.testing import absltest +from absl.testing import parameterized + +import jax +import jax.numpy as jnp +from jax.sharding import PartitionSpec + +from maxtext.configs import pyconfig +from maxtext.utils import sharding as sharding_utils +from maxtext.utils import sparsecore + +from tests.utils.test_helpers import get_test_config_path + + +class ParseOffloadTargetsTest(parameterized.TestCase): + """Tests for parsing the ``moe_sparse_core_offload_targets`` config string.""" + + def test_empty_string_disables_every_target(self): + self.assertEqual(sparsecore.parse_offload_targets(""), frozenset()) + + def test_all_expands_to_every_target(self): + self.assertEqual(sparsecore.parse_offload_targets("all"), frozenset(sparsecore.OFFLOAD_TARGETS)) + + @parameterized.parameters(*sparsecore.OFFLOAD_TARGETS) + def test_single_target(self, target): + self.assertEqual(sparsecore.parse_offload_targets(target), frozenset([target])) + + def test_comma_separated_subset_ignores_whitespace(self): + parsed = sparsecore.parse_offload_targets(f" {sparsecore.RAGGED_SORT}, {sparsecore.EP_COLLECTIVES} ,") + self.assertEqual(parsed, frozenset([sparsecore.RAGGED_SORT, sparsecore.EP_COLLECTIVES])) + + def test_unknown_target_raises(self): + with self.assertRaisesRegex(ValueError, "Unknown SparseCore offload target"): + sparsecore.parse_offload_targets("fsdp_all_gather,tensor_core") + + +class SparseCoreDetectionTest(parameterized.TestCase): + """Tests that SparseCore presence is read off JAX's chip table.""" + + @parameterized.named_parameters( + ("v5p", "v5p-8", True), + ("v6e", "v6e-16", True), + ("tpu7x", "tpu7x-256", True), + ("v4", "v4-8", False), + ("v5e", "v5e-16", False), + ) + def test_has_sparse_core_for_compile_topology(self, compile_topology, expected): + self.assertEqual(sparsecore.has_sparse_core(compile_topology=compile_topology), expected) + + def test_gpu_topology_has_no_sparse_core(self): + self.assertFalse(sparsecore.has_sparse_core(compile_topology="a3")) + + @parameterized.parameters(*sparsecore._NON_TPU_HARDWARE) # pylint: disable=protected-access + def test_non_tpu_hardware_has_no_sparse_core(self, hardware): + # The hardware check wins even when the topology names a SparseCore chip. + self.assertFalse(sparsecore.has_sparse_core(compile_topology="tpu7x-256", hardware=hardware)) + + def test_unknown_topology_is_not_an_error(self): + self.assertFalse(sparsecore.has_sparse_core(compile_topology="not-a-real-topology")) + + def test_sparse_core_info_reports_core_count(self): + info = sparsecore.sparse_core_info(compile_topology="v5p-8") + self.assertIsNotNone(info) + self.assertGreater(info.num_cores, 0) + + def test_offload_traces_off_tpu(self): + # CPU/GPU runs must be able to trace a TPU config unchanged. + with sparsecore.offload(True): + pass + with sparsecore.offload(False): + pass + + +class CollectiveCapabilityTest(parameterized.TestCase): + """Tests the gate that keeps MaxText from annotating a collective XLA would abort on.""" + + @parameterized.named_parameters( + ("v5p", "v5p-8", False), + ("v6e", "v6e-16", False), + ("tpu7x", "tpu7x-256", True), + ) + def test_all_gather_offload_needs_ironwood(self, compile_topology, expected): + self.assertEqual( + sparsecore.supports_collective_offload(sparsecore.ALL_GATHER, compile_topology=compile_topology), expected + ) + + def test_ragged_all_to_all_offloads_on_any_sparse_core_chip(self): + self.assertTrue(sparsecore.supports_collective_offload(sparsecore.RAGGED_ALL_TO_ALL, compile_topology="v5p-8")) + + @parameterized.named_parameters(("v5p", "v5p-8", False), ("tpu7x", "tpu7x-256", True)) + def test_reduce_scatter_follows_all_gather(self, compile_topology, expected): + # Its transpose is an all-gather, which would be annotated too. + self.assertEqual( + sparsecore.supports_collective_offload(sparsecore.REDUCE_SCATTER, compile_topology=compile_topology), expected + ) + + def test_no_collective_offloads_without_a_sparse_core(self): + for collective in (sparsecore.ALL_GATHER, sparsecore.RAGGED_ALL_TO_ALL, sparsecore.REDUCE_SCATTER): + self.assertFalse(sparsecore.supports_collective_offload(collective, compile_topology="v4-8")) + self.assertFalse(sparsecore.supports_collective_offload(collective, hardware="cpu")) + + def test_fsdp_target_is_dropped_where_all_gather_cannot_offload(self): + self.assertEqual( + sparsecore.supported_offload_targets("all", compile_topology="v5p-8"), + frozenset([sparsecore.EP_COLLECTIVES, sparsecore.RAGGED_SORT]), + ) + + def test_every_target_survives_on_ironwood(self): + self.assertEqual( + sparsecore.supported_offload_targets("all", compile_topology="tpu7x-256"), + frozenset(sparsecore.OFFLOAD_TARGETS), + ) + + def test_targets_that_need_no_collective_are_kept(self): + self.assertEqual( + sparsecore.supported_offload_targets(sparsecore.RAGGED_SORT, compile_topology="v5p-8"), + frozenset([sparsecore.RAGGED_SORT]), + ) + + def test_unknown_target_still_raises(self): + with self.assertRaisesRegex(ValueError, "Unknown SparseCore offload target"): + sparsecore.supported_offload_targets("tensor_core", compile_topology="tpu7x-256") + + def test_disabled_needs_no_hardware(self): + self.assertEqual(sparsecore.supported_offload_targets("", hardware="cpu"), frozenset()) + + +class ComputeTypeContextTest(parameterized.TestCase): + """Tests that a compute-type context manager is found on the installed JAX. + + ``jax.experimental.compute_on.compute_on`` changed from a context manager into + a function transform in JAX 0.11.1, so the resolution has to be checked + against whatever JAX is actually installed rather than assumed. + """ + + def test_offload_annotates_the_lowered_hlo(self): + """The end-to-end check: ops traced inside `offload` carry the compute type. + + This is the assertion that has to run on CPU. The compute-type API is the + part of this feature most likely to move under us, and every test that could + catch it moving used to be TPU-gated, so a JAX upgrade broke the annotation + with CI green. + """ + + def annotated(x): + with sparsecore.offload(True): + return x + 1 + + def plain(x): + return x + 1 + + x = jnp.zeros((8, 8)) + self.assertIn('_xla_compute_type = "sparseoffload"', jax.jit(annotated).lower(x).as_text()) + self.assertNotIn("sparseoffload", jax.jit(plain).lower(x).as_text()) + + def test_offload_skips_a_collective_the_target_chip_cannot_run(self): + def annotated(x): + with sparsecore.offload(True, collective=sparsecore.ALL_GATHER, compile_topology="v5p-8"): + return x + 1 + + self.assertNotIn("sparseoffload", jax.jit(annotated).lower(jnp.zeros((8, 8))).as_text()) + + def test_a_context_manager_was_resolved(self): + self.assertIsNotNone( + sparsecore._COMPUTE_TYPE_CONTEXT, # pylint: disable=protected-access + "No compute-type context manager found on this JAX; SparseCore offloading would silently do nothing.", + ) + + def test_resolved_context_manager_accepts_the_compute_type(self): + context = sparsecore._COMPUTE_TYPE_CONTEXT # pylint: disable=protected-access + with context(sparsecore.SPARSE_CORE_COMPUTE_TYPE): + pass + + def test_function_transform_spelling_is_not_used_as_a_context_manager(self): + """A 0.11.1-style `compute_on` must fall back, not be called positionally.""" + + def function_transform(f=None, *, compute_type, out_memory_spaces, compiler_options=None): + del f, compute_type, out_memory_spaces, compiler_options + raise AssertionError("the function-transform spelling must not be used as a context manager") + + with mock.patch.object(sparsecore.compute_on, "compute_on", function_transform): + resolved = sparsecore._resolve_compute_type_context() # pylint: disable=protected-access + self.assertIsNotNone(resolved) + self.assertIsNot(resolved, function_transform) + with resolved(sparsecore.SPARSE_CORE_COMPUTE_TYPE): + pass + + def test_context_manager_spelling_is_used_directly(self): + """The pre-0.11.1 public context manager is preferred when present.""" + + @contextlib.contextmanager + def context_manager(compute_type): + del compute_type + yield + + with mock.patch.object(sparsecore.compute_on, "compute_on", context_manager): + self.assertIs(sparsecore._resolve_compute_type_context(), context_manager) # pylint: disable=protected-access + + +class AllGatherAxesBetweenPspecsTest(parameterized.TestCase): + """Tests for deriving the all-gathers that take one PartitionSpec to another.""" + + def test_identical_pspecs_need_no_gather(self): + pspec = PartitionSpec("expert", None, "mlp") + self.assertEqual(sharding_utils.all_gather_axes_between_pspecs(pspec, pspec, 3), []) + + def test_single_dropped_axis(self): + gathers = sharding_utils.all_gather_axes_between_pspecs( + PartitionSpec("fsdp", None, "mlp"), PartitionSpec(None, None, "mlp"), 3 + ) + self.assertEqual(gathers, [(0, ("fsdp",))]) + + def test_dropped_axes_on_several_dims(self): + gathers = sharding_utils.all_gather_axes_between_pspecs( + PartitionSpec(("expert", "fsdp"), "fsdp_transpose", "mlp"), PartitionSpec("expert", None, "mlp"), 3 + ) + self.assertEqual(gathers, [(0, ("fsdp",)), (1, ("fsdp_transpose",))]) + + def test_shorter_pspec_is_padded_with_replicated_dims(self): + gathers = sharding_utils.all_gather_axes_between_pspecs(PartitionSpec("fsdp"), PartitionSpec(), 3) + self.assertEqual(gathers, [(0, ("fsdp",))]) + + def test_gaining_an_axis_is_not_a_pure_all_gather(self): + self.assertIsNone( + sharding_utils.all_gather_axes_between_pspecs( + PartitionSpec(None, None, "mlp"), PartitionSpec("fsdp", None, "mlp"), 3 + ) + ) + + def test_dropping_a_major_axis_is_not_a_pure_all_gather(self): + # `all_gather(tiled=True)` concatenates in device order, so only the minor + # (suffix) axes of a dim can be gathered away. + self.assertIsNone( + sharding_utils.all_gather_axes_between_pspecs( + PartitionSpec(("fsdp", "expert"), None, "mlp"), PartitionSpec("expert", None, "mlp"), 3 + ) + ) + + def test_swapped_axes_are_not_a_pure_all_gather(self): + self.assertIsNone( + sharding_utils.all_gather_axes_between_pspecs(PartitionSpec("fsdp", None), PartitionSpec("expert", None), 2) + ) + + +class SparseCoreConfigValidationTest(parameterized.TestCase): + """Tests that the config rejects offload requests the hardware cannot serve.""" + + def _config(self, **overrides): + return pyconfig.initialize( + [None, get_test_config_path()], + run_name="sparsecore_config_test", + enable_checkpointing=False, + skip_jax_distributed_system=True, + **overrides, + ) + + def test_default_is_disabled(self): + config = self._config() + self.assertEqual(config.moe_sparse_core_offload_targets, "") + + def test_accepted_on_a_sparse_core_topology(self): + config = self._config( + compile_topology="tpu7x-256", compile_topology_num_slices=1, moe_sparse_core_offload_targets="all" + ) + self.assertEqual(config.moe_sparse_core_offload_targets, "all") + + def test_target_the_chip_cannot_serve_is_accepted_and_ignored(self): + # A SparseCore that cannot offload an all-gather must not fail the run: the + # same config has to stay usable across chip generations. + config = self._config( + compile_topology="v5p-8", compile_topology_num_slices=1, moe_sparse_core_offload_targets="fsdp_all_gather" + ) + self.assertEqual(config.moe_sparse_core_offload_targets, "fsdp_all_gather") + self.assertEqual( + sparsecore.supported_offload_targets(config.moe_sparse_core_offload_targets, config.compile_topology), frozenset() + ) + + def test_rejected_on_a_topology_without_a_sparse_core(self): + with self.assertRaisesRegex(ValueError, "requires a TPU with a .*SparseCore"): + self._config(compile_topology="v4-8", compile_topology_num_slices=1, moe_sparse_core_offload_targets="ragged_sort") + + def test_rejected_on_cpu(self): + with self.assertRaisesRegex(ValueError, "requires a TPU with a .*SparseCore"): + self._config(hardware="cpu", moe_sparse_core_offload_targets="ep_collectives") + + def test_unknown_target_is_rejected(self): + with self.assertRaisesRegex(ValueError, "Unknown SparseCore offload target"): + self._config( + compile_topology="tpu7x-256", compile_topology_num_slices=1, moe_sparse_core_offload_targets="everything" + ) + + +if __name__ == "__main__": + absltest.main()