From 3cd54a90f2f70bd0e27ea330ac0ebd7c32f2b450 Mon Sep 17 00:00:00 2001 From: continuousml Date: Fri, 31 Jul 2026 00:02:25 -0700 Subject: [PATCH] Add TPU Ulysses context parallelism --- docs/guides/optimization/sharding.md | 2 + src/maxtext/configs/base.yml | 2 +- src/maxtext/configs/types.py | 66 ++++- .../attention/context_parallel_utils.py | 50 ++++ .../attention/tokamax_ring_attention.py | 46 +-- .../kernels/attention/ulysses_attention.py | 210 ++++++++++++++ src/maxtext/layers/attention_op.py | 166 ++++++++++- tests/integration/train_tests.py | 15 + tests/unit/attention_test.py | 200 ++++++++++++- tests/unit/configs_value_test.py | 110 ++++++++ tests/unit/hlo_test_utils_test.py | 93 ++++++ tests/unit/ulysses_attention_test.py | 265 ++++++++++++++++++ tests/unit/ulysses_collective_test.py | 155 ++++++++++ tests/utils/hlo_test_utils.py | 80 ++++++ 14 files changed, 1420 insertions(+), 40 deletions(-) create mode 100644 src/maxtext/kernels/attention/context_parallel_utils.py create mode 100644 src/maxtext/kernels/attention/ulysses_attention.py create mode 100644 tests/unit/hlo_test_utils_test.py create mode 100644 tests/unit/ulysses_attention_test.py create mode 100644 tests/unit/ulysses_collective_test.py create mode 100644 tests/utils/hlo_test_utils.py diff --git a/docs/guides/optimization/sharding.md b/docs/guides/optimization/sharding.md index 97abf70070..d0eedc3568 100644 --- a/docs/guides/optimization/sharding.md +++ b/docs/guides/optimization/sharding.md @@ -260,6 +260,8 @@ Note in general there are many flavors of CP such as ring attention, which in th MaxText supports `context_parallel_strategy=all_gather`, and supports `context_parallel_strategy=ring` through GPU Transformer Engine and TPU Tokamax Splash paths; ring performs the computation and communication in chunks and ideally overlaps them in a collective matmul fashion. This strategy requires extending the online softmax trick from only within chip to additionally apply it across chips. +MaxText also supports `context_parallel_strategy=ulysses` ([DeepSpeed Ulysses](https://arxiv.org/abs/2309.14509)) on the TPU Tokamax Splash path for training. Ulysses exchanges sequence ownership for head ownership by communicating the Q, K, V, and output activations through all-to-all collectives: each device computes ordinary full-sequence attention for its head subset, and the inverse all-to-all restores the sequence sharding on the output. It requires explicit positive context parallelism values, `context_sharding=context`, `attention=flash` with Tokamax Splash, global causal attention, query and KV head counts divisible by the context parallel size including after tensor-parallel head sharding, matching Q and KV head-sharding axes, an unsharded head feature dimension, a divisible sequence length, `dq_reduction_steps` of 0 or 3, `context_parallel_load_balance=false` (each device computes full-sequence attention for its head subset, so the work is already balanced and the causal load-balancing reorder must stay off), and ICI-only context parallelism (`dcn_context_parallelism` must equal 1). It does not support MQA, packing, dropout, QK-Clip statistics, ragged attention, attention sinks, sparse indexer masks, chunked prefill, MoBA, or multimodal attention. + ### CP Arithmetic Intensity The main communications are the same as FSDP (all gather weights and synchronize gradients), with an arithmetic intensity of `local_batch` / `sparsity`. diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 6b23e31d40..41d0d5bae7 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -1139,7 +1139,7 @@ cost_estimate_flops_bwd: -1 # -1 means using splash default cost estmiation, any dq_reduction_steps: 0 #the number of reduction steps. For now, only 3 or all the kv steps are supported. ### Determine if we want to use load balance for context parallelism context_parallel_load_balance: true -context_parallel_strategy: "all_gather" # "all_gather" or "ring" +context_parallel_strategy: "all_gather" # "all_gather", "ring", or "ulysses" context_parallel_reorder_strategy: "auto" # "auto", "dual_chunk_swap", or "striped" diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 72cc3de96a..9828637958 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1082,7 +1082,7 @@ class HardwareAndMesh(BaseModel): context_parallel_load_balance: bool = Field(True, description="Whether to use load balancing for context parallelism.") context_parallel_strategy: str = Field( "all_gather", - description="Strategy for context parallelism ('all_gather' or 'ring').", + description="Strategy for context parallelism ('all_gather', 'ring', or 'ulysses').", ) context_parallel_reorder_strategy: ReorderStrategy = Field( ReorderStrategy.AUTO, @@ -3573,6 +3573,9 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de self, f"dcn_{self.context_sharding}_parallelism", 1 ) context_parallel_strategy = self.context_parallel_strategy.lower() + if context_parallel_strategy not in ("all_gather", "ring", "ulysses"): + raise ValueError("context_parallel_strategy must be one of 'all_gather', 'ring', or 'ulysses'.") + self.context_parallel_strategy = context_parallel_strategy if ( context_parallel_strategy == "ring" and "gpu" not in self.hardware @@ -3628,6 +3631,67 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de f"ring_scan_unroll={self.ring_scan_unroll} was specified, but is only supported when " "context_parallel_strategy='ring'." ) + if context_parallel_strategy == "ulysses": + if self.hardware != "tpu": + raise ValueError("Ulysses context parallelism (context_parallel_strategy='ulysses') is only supported on TPU.") + if self.context_sharding != "context": + raise ValueError("TPU Ulysses attention requires context_sharding='context'.") + ici_context_parallel_size = self.ici_context_parallelism + dcn_context_parallel_size = self.dcn_context_parallelism + if ici_context_parallel_size <= 0 or dcn_context_parallel_size <= 0: + raise ValueError( + "TPU Ulysses attention requires explicit positive ici/dcn context parallelism values; " + "inferred (-1) sizes are not supported." + ) + if context_parallel_size <= 1: + raise ValueError("TPU Ulysses attention requires context_parallel_size > 1.") + if dcn_context_parallel_size != 1: + raise ValueError("TPU Ulysses attention does not support dcn context parallelism yet.") + if self.attention != "flash": + raise ValueError("TPU Ulysses attention requires attention=flash.") + if not self.use_tokamax_splash: + raise ValueError("TPU Ulysses attention requires use_tokamax_splash=True.") + if self.use_jax_splash: + raise ValueError("TPU Ulysses attention requires use_jax_splash=False.") + if self.attention_type != "global": + raise ValueError("TPU Ulysses attention is initially supported only for global causal attention.") + if self.packing: + raise ValueError("TPU Ulysses attention does not support packing yet.") + if self.context_parallel_load_balance: + raise ValueError("TPU Ulysses attention does not support context_parallel_load_balance=True.") + if self.use_ragged_attention: + raise ValueError("TPU Ulysses attention does not support ragged attention.") + if self.attention_sink: + raise ValueError("TPU Ulysses attention does not support attention sinks.") + if self.use_indexer: + raise ValueError("TPU Ulysses attention does not support sparse indexer masks.") + if self.use_chunked_prefill: + raise ValueError("TPU Ulysses attention does not support chunked prefill yet.") + if self.use_multimodal: + raise ValueError("TPU Ulysses attention does not support multimodal attention.") + if self.enable_dropout and self.dropout_rate > 0.0: + raise ValueError("TPU Ulysses attention does not support dropout yet.") + if self.dq_reduction_steps not in (0, 3): + raise ValueError("TPU Ulysses attention requires dq_reduction_steps to be 0 or 3.") + if self.use_qk_clip: + raise ValueError("TPU Ulysses attention does not support QK-Clip statistics yet.") + if self.max_target_length % context_parallel_size != 0: + raise ValueError( + "TPU Ulysses attention requires max_target_length " + f"({self.max_target_length}) to be divisible by context_parallel_size ({context_parallel_size})." + ) + if self.num_query_heads % context_parallel_size != 0: + raise ValueError( + "TPU Ulysses attention requires num_query_heads " + f"({self.num_query_heads}) to be divisible by context_parallel_size ({context_parallel_size})." + ) + if self.num_kv_heads == 1: + raise ValueError("TPU Ulysses attention does not support MQA with context_parallel_size > 1.") + if self.num_kv_heads % context_parallel_size != 0: + raise ValueError( + "TPU Ulysses attention requires num_kv_heads " + f"({self.num_kv_heads}) to be divisible by context_parallel_size ({context_parallel_size})." + ) # STRIPED reorder strategy is a Transformer Engine feature and is GPU-only. # AUTO is resolved in training because test code paths may load the same # config but use a different reorder path. diff --git a/src/maxtext/kernels/attention/context_parallel_utils.py b/src/maxtext/kernels/attention/context_parallel_utils.py new file mode 100644 index 0000000000..a732fbc9f4 --- /dev/null +++ b/src/maxtext/kernels/attention/context_parallel_utils.py @@ -0,0 +1,50 @@ +# 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. +"""Shared helpers for attention context-parallel sharding metadata.""" + +from __future__ import annotations + +from typing import Any + +import jax + + +def mesh_axes_for_dim(axis_names: Any) -> tuple[Any, ...]: + """Returns the mesh axes attached to one tensor dimension.""" + if axis_names is None: + return () + if isinstance(axis_names, str): + return (axis_names,) + return tuple(axis for axis in axis_names if axis is not None) + + +def mesh_axes_size(mesh: Any, axes: tuple[Any, ...], *, label: str) -> int: + """Returns the product of mesh sizes for a set of axes.""" + size = 1 + for axis in axes: + if axis not in mesh.shape: + raise ValueError(f"{label} requires mesh axis {axis!r} to exist.") + size *= mesh.shape[axis] + return size + + +def with_axis_on_dim(axis_names: Any, axis: Any, dim: int) -> Any: + """Returns sharding axis names with one dimension replaced.""" + axes = list(axis_names) + axes[dim] = axis + if isinstance(axis_names, jax.sharding.PartitionSpec): + return jax.sharding.PartitionSpec(*axes, unreduced=axis_names.unreduced, reduced=axis_names.reduced) + if isinstance(axis_names, tuple): + return tuple(axes) + return axes diff --git a/src/maxtext/kernels/attention/tokamax_ring_attention.py b/src/maxtext/kernels/attention/tokamax_ring_attention.py index 9fd031599a..df4578702c 100644 --- a/src/maxtext/kernels/attention/tokamax_ring_attention.py +++ b/src/maxtext/kernels/attention/tokamax_ring_attention.py @@ -28,6 +28,7 @@ import numpy as np from maxtext.common.common_types import MODEL_MODE_TRAIN +from maxtext.kernels.attention import context_parallel_utils from maxtext.kernels.tokamax_splash_attention import ring_attention_kernel from maxtext.kernels.tokamax_splash_attention import splash_attention_kernel as tokamax_splash_kernel from maxtext.kernels.tokamax_splash_attention import splash_attention_mask as tokamax_splash_mask @@ -39,42 +40,19 @@ def is_context_parallel_ring_requested(config: Any) -> bool: return config.context_parallel_strategy.lower() == "ring" -def _mesh_axes_for_dim(axis_names: Any) -> tuple[Any, ...]: - if axis_names is None: - return () - if isinstance(axis_names, str): - return (axis_names,) - return tuple(axis for axis in axis_names if axis is not None) - - -def _mesh_axes_size(mesh: Any, axes: tuple[Any, ...]) -> int: - size = 1 - for axis in axes: - if axis not in mesh.shape: - raise ValueError(f"TPU Tokamax ring attention requires mesh axis {axis!r} to exist.") - size *= mesh.shape[axis] - return size - - def with_sequence_axis(axis_names: Any, ring_axis: str, sequence_dim: int) -> Any: """Returns axis names with the sequence dimension set to the ring axis.""" if axis_names is None: return None if len(axis_names) <= sequence_dim: raise ValueError("TPU Tokamax ring attention expects a sequence sharding dimension.") - axes = list(axis_names) - existing_sequence_axes = _mesh_axes_for_dim(axes[sequence_dim]) + existing_sequence_axes = context_parallel_utils.mesh_axes_for_dim(axis_names[sequence_dim]) if existing_sequence_axes and existing_sequence_axes != (ring_axis,): raise ValueError( "TPU Tokamax ring attention expects the existing sequence sharding to be " f"unsharded or exactly {(ring_axis,)}, got {existing_sequence_axes}." ) - axes[sequence_dim] = ring_axis - if isinstance(axis_names, jax.sharding.PartitionSpec): - return jax.sharding.PartitionSpec(*axes, unreduced=axis_names.unreduced, reduced=axis_names.reduced) - if isinstance(axis_names, tuple): - return tuple(axes) - return axes + return context_parallel_utils.with_axis_on_dim(axis_names, ring_axis, sequence_dim) def _validate_ring_axis_only_on_sequence( @@ -88,7 +66,7 @@ def _validate_ring_axis_only_on_sequence( for dim, axis_name in enumerate(axis_names): if dim == sequence_dim: continue - dim_axes = _mesh_axes_for_dim(axis_name) + dim_axes = context_parallel_utils.mesh_axes_for_dim(axis_name) if ring_axis in dim_axes: raise ValueError( "TPU Tokamax ring attention requires the context axis to appear only " @@ -104,8 +82,8 @@ def validate_dkv_sharding( dkv_dim_kv: int, ) -> None: """Validates that the head-dim/D_KV dimension stays local for ring attention.""" - q_dkv_axes = _mesh_axes_for_dim(axis_names_q[dkv_dim_q]) - kv_dkv_axes = _mesh_axes_for_dim(axis_names_kv[dkv_dim_kv]) + q_dkv_axes = context_parallel_utils.mesh_axes_for_dim(axis_names_q[dkv_dim_q]) + kv_dkv_axes = context_parallel_utils.mesh_axes_for_dim(axis_names_kv[dkv_dim_kv]) if q_dkv_axes or kv_dkv_axes: raise ValueError( "TPU Tokamax ring attention does not support sharding the D_KV/head-dim " @@ -168,8 +146,8 @@ def validate_ring_mesh_axis( ring_axis=ring_axis, ) expected_axes = (ring_axis,) - q_sequence_axes = _mesh_axes_for_dim(axis_names_q[sequence_dim_q]) - key_value_sequence_axes = _mesh_axes_for_dim(axis_names_kv[sequence_dim_kv]) + q_sequence_axes = context_parallel_utils.mesh_axes_for_dim(axis_names_q[sequence_dim_q]) + key_value_sequence_axes = context_parallel_utils.mesh_axes_for_dim(axis_names_kv[sequence_dim_kv]) if q_sequence_axes != expected_axes: raise ValueError( "TPU Tokamax ring attention requires Q sequence sharding to be exactly " @@ -193,10 +171,10 @@ def validate_head_sharding( head_dim_kv: int, ) -> None: """Validates that local head layout preserves GQA/MQA head mapping.""" - q_head_axes = _mesh_axes_for_dim(axis_names_q[head_dim_q]) - kv_head_axes = _mesh_axes_for_dim(axis_names_kv[head_dim_kv]) - q_head_shards = _mesh_axes_size(mesh, q_head_axes) - kv_head_shards = _mesh_axes_size(mesh, kv_head_axes) + q_head_axes = context_parallel_utils.mesh_axes_for_dim(axis_names_q[head_dim_q]) + kv_head_axes = context_parallel_utils.mesh_axes_for_dim(axis_names_kv[head_dim_kv]) + q_head_shards = context_parallel_utils.mesh_axes_size(mesh, q_head_axes, label="TPU Tokamax ring attention") + kv_head_shards = context_parallel_utils.mesh_axes_size(mesh, kv_head_axes, label="TPU Tokamax ring attention") if num_query_heads % q_head_shards != 0: raise ValueError( "TPU Tokamax ring attention requires num_query_heads " diff --git a/src/maxtext/kernels/attention/ulysses_attention.py b/src/maxtext/kernels/attention/ulysses_attention.py new file mode 100644 index 0000000000..68d3c6fac6 --- /dev/null +++ b/src/maxtext/kernels/attention/ulysses_attention.py @@ -0,0 +1,210 @@ +# 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. +"""Ulysses attention layout helpers.""" + +from __future__ import annotations + +from typing import Any + +import jax + +from maxtext.common.common_types import MODEL_MODE_TRAIN +from maxtext.kernels.attention import context_parallel_utils + + +def is_context_parallel_ulysses_requested(config: Any) -> bool: + """Returns True when the config requests Ulysses context parallelism.""" + return config.context_parallel_strategy == "ulysses" + + +def validate_ulysses_runtime( + *, + model_mode: str, + use_ragged_attention: bool = False, + previous_chunk: Any = None, + sinks: Any = None, + indexer_mask: Any = None, + bidirectional_mask: Any = None, + record_max_logits: bool = False, +) -> None: + """Validates runtime-only constraints for the Ulysses path.""" + if model_mode != MODEL_MODE_TRAIN: + raise ValueError("TPU Ulysses attention is supported only for train mode.") + if use_ragged_attention: + raise ValueError("TPU Ulysses attention does not support ragged attention.") + if previous_chunk is not None: + raise ValueError("TPU Ulysses attention does not support chunked prefill yet.") + if sinks is not None: + raise ValueError("TPU Ulysses attention does not support attention sinks.") + if indexer_mask is not None: + raise ValueError("TPU Ulysses attention does not support indexer masks.") + if bidirectional_mask is not None: + raise ValueError("TPU Ulysses attention does not support bidirectional masks.") + if record_max_logits: + raise NotImplementedError("TPU Ulysses attention does not support record_max_logits yet.") + + +def with_sequence_axis(axis_names: Any, sequence_axis: str, sequence_dim: int) -> Any: + """Returns axis names with the sequence dimension set to Ulysses.""" + if axis_names is None: + return None + if len(axis_names) <= sequence_dim: + raise ValueError("TPU Ulysses attention expects a sequence sharding dimension.") + existing_sequence_axes = context_parallel_utils.mesh_axes_for_dim(axis_names[sequence_dim]) + if existing_sequence_axes and existing_sequence_axes != (sequence_axis,): + raise ValueError( + "TPU Ulysses attention expects the existing sequence sharding to be " + f"unsharded or exactly {(sequence_axis,)}, got {existing_sequence_axes}." + ) + return context_parallel_utils.with_axis_on_dim(axis_names, sequence_axis, sequence_dim) + + +def _validate_ulysses_axis_only_on_sequence( + axis_names: Any, + *, + tensor_name: str, + sequence_dim: int, + ulysses_axis: str, +) -> None: + """Raises if the Ulysses mesh axis appears outside the sequence dimension.""" + for dim, axis_name in enumerate(axis_names): + if dim == sequence_dim: + continue + dim_axes = context_parallel_utils.mesh_axes_for_dim(axis_name) + if ulysses_axis in dim_axes: + raise ValueError( + "TPU Ulysses attention requires the context axis to appear only " + f"on the sequence dimension; got {ulysses_axis!r} on {tensor_name} dim {dim}." + ) + + +def validate_ulysses_mesh_axis( + *, + axis_names_q: Any, + axis_names_kv: Any, + sequence_dim_q: int, + sequence_dim_kv: int, + mesh: Any, + ulysses_axis: str, +) -> None: + """Validates sequence sharding before the Ulysses all-to-all.""" + if not ulysses_axis: + raise ValueError("TPU Ulysses attention requires a non-empty context_sharding axis.") + if ulysses_axis not in mesh.shape: + raise ValueError(f"TPU Ulysses attention requires mesh axis {ulysses_axis!r} to exist.") + _validate_ulysses_axis_only_on_sequence( + axis_names_q, + tensor_name="Q", + sequence_dim=sequence_dim_q, + ulysses_axis=ulysses_axis, + ) + _validate_ulysses_axis_only_on_sequence( + axis_names_kv, + tensor_name="K/V", + sequence_dim=sequence_dim_kv, + ulysses_axis=ulysses_axis, + ) + + expected_axes = (ulysses_axis,) + q_sequence_axes = context_parallel_utils.mesh_axes_for_dim(axis_names_q[sequence_dim_q]) + kv_sequence_axes = context_parallel_utils.mesh_axes_for_dim(axis_names_kv[sequence_dim_kv]) + if q_sequence_axes != expected_axes: + raise ValueError( + f"TPU Ulysses attention requires Q sequence sharding to be exactly {expected_axes}, got {q_sequence_axes}." + ) + if kv_sequence_axes != expected_axes: + raise ValueError( + f"TPU Ulysses attention requires K/V sequence sharding to be exactly {expected_axes}, got {kv_sequence_axes}." + ) + + +def validate_dkv_sharding( + *, + axis_names_q: Any, + axis_names_kv: Any, + dkv_dim_q: int, + dkv_dim_kv: int, +) -> None: + """Validates that the head-dim/D_KV dimension stays local for Ulysses attention.""" + q_dkv_axes = context_parallel_utils.mesh_axes_for_dim(axis_names_q[dkv_dim_q]) + kv_dkv_axes = context_parallel_utils.mesh_axes_for_dim(axis_names_kv[dkv_dim_kv]) + if q_dkv_axes or kv_dkv_axes: + raise ValueError( + "TPU Ulysses attention does not support sharding the D_KV/head-dim " + f"dimension; got Q axes {q_dkv_axes} and K/V axes {kv_dkv_axes}." + ) + + +def validate_head_sharding( + *, + axis_names_q: Any, + axis_names_kv: Any, + mesh: Any, + num_query_heads: int, + num_kv_heads: int, + head_dim_q: int, + head_dim_kv: int, + ulysses_size: int, +) -> None: + """Validates local head counts before the Ulysses head/sequence exchange.""" + q_head_axes = context_parallel_utils.mesh_axes_for_dim(axis_names_q[head_dim_q]) + kv_head_axes = context_parallel_utils.mesh_axes_for_dim(axis_names_kv[head_dim_kv]) + q_head_shards = context_parallel_utils.mesh_axes_size(mesh, q_head_axes, label="TPU Ulysses attention") + kv_head_shards = context_parallel_utils.mesh_axes_size(mesh, kv_head_axes, label="TPU Ulysses attention") + if num_query_heads % q_head_shards != 0: + raise ValueError( + "TPU Ulysses attention requires num_query_heads " + f"({num_query_heads}) to be divisible by Q head shards ({q_head_shards})." + ) + if num_kv_heads % kv_head_shards != 0: + raise ValueError( + "TPU Ulysses attention requires num_kv_heads " + f"({num_kv_heads}) to be divisible by KV head shards ({kv_head_shards})." + ) + + if num_kv_heads == 1: + raise ValueError("TPU Ulysses attention does not support MQA with context_parallel_size > 1.") + if q_head_axes != kv_head_axes: + raise ValueError( + "TPU Ulysses attention requires Q and KV head sharding to match for MHA/GQA, " + f"got Q head axes {q_head_axes} and KV head axes {kv_head_axes}." + ) + + local_query_heads = num_query_heads // q_head_shards + local_kv_heads = num_kv_heads // kv_head_shards + if local_query_heads % local_kv_heads != 0: + raise ValueError( + "TPU Ulysses attention requires local query heads " + f"({local_query_heads}) to be divisible by local KV heads ({local_kv_heads})." + ) + if local_query_heads % ulysses_size != 0: + raise ValueError( + "TPU Ulysses attention requires local query heads " + f"({local_query_heads}) to be divisible by context_parallel_size ({ulysses_size})." + ) + if local_kv_heads % ulysses_size != 0: + raise ValueError( + "TPU Ulysses attention requires local KV heads " + f"({local_kv_heads}) to be divisible by context_parallel_size ({ulysses_size})." + ) + + +def ulysses_all_to_all(tensor: Any, ulysses_axis: str): + """Moves `[B, H, S/U, D]` to `[B, H/U, S, D]`.""" + return jax.lax.all_to_all(tensor, ulysses_axis, split_axis=1, concat_axis=2, tiled=True) + + +def inverse_ulysses_all_to_all(tensor: Any, ulysses_axis: str): + """Moves `[B, H/U, S, D]` back to `[B, H, S/U, D]`.""" + return jax.lax.all_to_all(tensor, ulysses_axis, split_axis=2, concat_axis=1, tiled=True) diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index 1dd29c1a5e..03651238fe 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -65,6 +65,7 @@ from maxtext.inference.kvcache import KVQuant, KVTensor from maxtext.kernels.attention import jax_flash_attention from maxtext.kernels.attention import tokamax_ring_attention +from maxtext.kernels.attention import ulysses_attention from maxtext.kernels.attention.ragged_attention import ragged_gqa from maxtext.kernels.attention.ragged_attention import ragged_mha from maxtext.layers import nnx_wrappers @@ -568,6 +569,55 @@ def __init__( dkv_dim_q=3, dkv_dim_kv=3, ) + if self.attention_kernel == "flash" and ulysses_attention.is_context_parallel_ulysses_requested(self.config): + target_hardware = self.mesh.devices[(0,) * self.mesh.devices.ndim].platform + if target_hardware != "tpu": + raise ValueError("Ulysses context parallelism (context_parallel_strategy='ulysses') is only supported on TPU.") + if not self.config.use_tokamax_splash: + raise ValueError("TPU Ulysses attention requires use_tokamax_splash=True.") + if self.config.use_jax_splash: + raise ValueError("TPU Ulysses attention requires use_jax_splash=False.") + if self.attention_type != AttentionType.GLOBAL: + raise ValueError("TPU Ulysses attention is initially supported only for global causal attention.") + if self.config.enable_dropout and self.dropout_rate > 0.0: + raise ValueError("TPU Ulysses attention does not support dropout yet.") + if self.use_ragged_attention: + raise ValueError("TPU Ulysses attention does not support ragged attention.") + + context_axis = self.config.context_sharding + axis_names_q = self._logical_to_mesh_axes(self.flash_axis_names_q) + axis_names_kv = self._logical_to_mesh_axes(self.flash_axis_names_kv) + axis_names_kv = ulysses_attention.with_sequence_axis( + axis_names_kv, + context_axis, + sequence_dim=2, + ) + ulysses_attention.validate_ulysses_mesh_axis( + axis_names_q=axis_names_q, + axis_names_kv=axis_names_kv, + sequence_dim_q=2, + sequence_dim_kv=2, + mesh=self.mesh, + ulysses_axis=context_axis, + ) + if self.mesh.shape[context_axis] <= 1: + raise ValueError("TPU Ulysses attention requires a context parallel mesh axis larger than one.") + ulysses_attention.validate_head_sharding( + axis_names_q=axis_names_q, + axis_names_kv=axis_names_kv, + mesh=self.mesh, + num_query_heads=self.num_query_heads, + num_kv_heads=self.num_kv_heads, + head_dim_q=1, + head_dim_kv=1, + ulysses_size=self.mesh.shape[context_axis], + ) + ulysses_attention.validate_dkv_sharding( + axis_names_q=axis_names_q, + axis_names_kv=axis_names_kv, + dkv_dim_q=3, + dkv_dim_kv=3, + ) def maybe_create_nnx(einsum, *args): if isinstance(einsum, nn.Module): @@ -1018,6 +1068,28 @@ def _validate_tpu_tokamax_ring_runtime( record_max_logits=record_max_logits, ) + def _validate_tpu_ulysses_runtime( + self, + *, + model_mode: str, + previous_chunk: Any = None, + bidirectional_mask: Any = None, + sinks: Array | None = None, + indexer_mask: Array | None = None, + use_ragged_attention: bool = False, + record_max_logits: bool = False, + ) -> None: + """Validates runtime constraints for the TPU Ulysses path.""" + ulysses_attention.validate_ulysses_runtime( + model_mode=model_mode, + previous_chunk=previous_chunk, + sinks=sinks, + indexer_mask=indexer_mask, + use_ragged_attention=use_ragged_attention, + bidirectional_mask=bidirectional_mask, + record_max_logits=record_max_logits, + ) + def apply_attention( self, query: Array, @@ -1048,6 +1120,11 @@ def apply_attention( and self.attention_kernel != "flash" ): raise ValueError("TPU Tokamax ring attention requires attention_kernel='flash'.") + if ulysses_attention.is_context_parallel_ulysses_requested(self.config): + if target_hardware != "tpu": + raise ValueError("Ulysses context parallelism (context_parallel_strategy='ulysses') is only supported on TPU.") + if self.attention_kernel != "flash": + raise ValueError("TPU Ulysses attention requires attention_kernel='flash'.") if use_ragged_attention and model_mode == MODEL_MODE_AUTOREGRESSIVE: if lengths is None: @@ -1288,6 +1365,7 @@ def tpu_flash_attention( """TPU Flash Attention.""" use_tokamax_ring = tokamax_ring_attention.is_context_parallel_ring_requested(self.config) + use_ulysses = ulysses_attention.is_context_parallel_ulysses_requested(self.config) cp_size = self.mesh.shape.get(self.config.context_sharding, 1) load_balanced_context_parallel = self.config.context_parallel_load_balance if use_tokamax_ring: @@ -1300,6 +1378,16 @@ def tpu_flash_attention( use_ragged_attention=use_ragged_attention, record_max_logits=record_max_logits, ) + elif use_ulysses: + self._validate_tpu_ulysses_runtime( + model_mode=model_mode, + previous_chunk=previous_chunk, + bidirectional_mask=bidirectional_mask, + sinks=sinks, + indexer_mask=indexer_mask, + use_ragged_attention=use_ragged_attention, + record_max_logits=record_max_logits, + ) # Transpose to ('batch', 'heads', 'length', 'kv') query = jnp.transpose(query, axes=(0, 2, 1, 3)) @@ -1333,6 +1421,23 @@ def tpu_flash_attention( context_axis, sequence_dim=1, ) + elif use_ulysses: + context_axis = self.config.context_sharding + segment_axis_names_q = ulysses_attention.with_sequence_axis( + segment_axis_names_q, + context_axis, + sequence_dim=1, + ) + axis_names_kv = ulysses_attention.with_sequence_axis( + axis_names_kv, + context_axis, + sequence_dim=2, + ) + segment_axis_names_kv = ulysses_attention.with_sequence_axis( + segment_axis_names_kv, + context_axis, + sequence_dim=1, + ) devices_in_data_fsdp = self.mesh.shape.get("data", 1) * self.mesh.shape.get("fsdp", 1) assert (query.shape[0] / devices_in_data_fsdp).is_integer(), ( @@ -1405,6 +1510,33 @@ def create_sa_config(config, query, key, attn_logits_soft_cap): maybe_shard_with_pspec=self._maybe_shard_with_pspec, ) ) + elif use_ulysses: + sa_config = create_sa_config(self.config, query, key, attn_logits_soft_cap) + if self.config.use_max_logit_estimate > 0: + sa_config = dataclasses.replace(sa_config, max_logit_const=self.config.use_max_logit_estimate) + mask_shape = (query.shape[2], key.shape[2]) # (q_seq_len, kv_seq_len) + mask = tokamax_splash_mask.CausalMask(shape=mask_shape) + + @partial( + jax.jit, + static_argnames=[ + "single_head_mask", + ], + ) + def wrap_ulysses_splash_kernel(single_head_mask): + splash_kernel = tokamax_splash_kernel.make_splash_mha( + mask=single_head_mask, + config=sa_config, + q_seq_shards=1, + ) + return splash_kernel + + splash_kernel = wrap_ulysses_splash_kernel(mask) + # After the all-to-all every device runs the kernel over the full + # sequence, so its mask metadata is replicated (q_seq_shards=1) instead + # of sequence-sharded as in the all-gather path. + segment_axis_names_splash_kernel = jax.sharding.PartitionSpec(None) + splash_kernel = self._maybe_shard_with_pspec(splash_kernel, segment_axis_names_splash_kernel) else: sa_config = create_sa_config(self.config, query, key, attn_logits_soft_cap) mask_shape = (query.shape[2], key.shape[2]) # (q_seq_len, kv_seq_len) @@ -1453,7 +1585,7 @@ def create_sa_config(config, query, key, attn_logits_soft_cap): ) max_logit_value = None - if not use_tokamax_ring and self.config.use_tokamax_splash: + if not use_tokamax_ring and not use_ulysses and self.config.use_tokamax_splash: # Create mask single_head_mask = mask # tokamax now just uses a single mask and assumes broadcast to all heads if self.config.use_max_logit_estimate > 0: @@ -1477,11 +1609,11 @@ def wrap_tokamax_splash_kernel(single_head_mask): splash_kernel = wrap_tokamax_splash_kernel(single_head_mask) segment_axis_names_splash_kernel = self._logical_to_mesh_axes((Q_LENGTH,)) splash_kernel = self._maybe_shard_with_pspec(splash_kernel, segment_axis_names_splash_kernel) - elif not use_tokamax_ring and self.config.use_jax_splash: + elif not use_tokamax_ring and not use_ulysses and self.config.use_jax_splash: if self.config.use_max_logit_estimate > 0: sa_config = dataclasses.replace(sa_config, max_logit_const=self.config.use_max_logit_estimate) segment_axis_names_splash_kernel = nn.logical_to_mesh_axes((Q_LENGTH,)) - elif not use_tokamax_ring: + elif not use_tokamax_ring and not use_ulysses: # Create multi-head mask multi_head_mask = splash_attention_mask.MultiHeadMask(masks=(mask,) * query.shape[1]) @@ -1521,6 +1653,8 @@ def wrap_jax_splash_kernel(multi_head_mask, shard_head_size=1): # specified in the shard_map in_specs below. For the all-gather path Q is # sequence-sharded and K/V are replicated. For the Tokamax ring path Q, K, # V, and segment IDs are all sequence-sharded over the context axis. + # For Ulysses Q/K/V are sequence-sharded at the boundary and head-sharded + # inside the local Splash call. if record_max_logits: # max_logits will share similar sharding as query but last dim is unrelated to model @@ -1579,6 +1713,32 @@ def wrap_flash_attention( ) return attention_output, None + if use_ulysses: + query = ulysses_attention.ulysses_all_to_all(query, context_axis) + key = ulysses_attention.ulysses_all_to_all(key, context_axis) + value = ulysses_attention.ulysses_all_to_all(value, context_axis) + if decoder_segment_ids_q is not None: + # Q and KV segment IDs are the same tensor in this train-only + # self-attention path, so one gather serves both kernel operands. + full_segment_ids = jax.lax.all_gather( + decoder_segment_ids_q, + context_axis, + axis=1, + tiled=True, + ) + decoder_segment_ids_tuple = tokamax_splash_kernel.SegmentIds( + full_segment_ids, + full_segment_ids, + ) + else: + decoder_segment_ids_tuple = None + kernel = partial(splash_kernel, max_logit_value=max_logit_value) + attention_output = jax.vmap(lambda q, k, v, d, s: kernel(q, k, v, d, sinks=s), in_axes=(0, 0, 0, 0, None))( + query, key, value, decoder_segment_ids_tuple, sinks + ) + attention_output = ulysses_attention.inverse_ulysses_all_to_all(attention_output, context_axis) + return attention_output, None + # The load-balanced all-gather path restores K/V to contiguous order # before calling Splash attention. if cp_size > 1 and load_balanced_context_parallel: diff --git a/tests/integration/train_tests.py b/tests/integration/train_tests.py index e984b58667..fa629af474 100644 --- a/tests/integration/train_tests.py +++ b/tests/integration/train_tests.py @@ -215,6 +215,21 @@ def test_tpu_base(self): def test_tpu_tokamax(self): train_main(TrainTests.CONFIGS["synthetic"] + ["use_tokamax_splash=true"]) + @pytest.mark.integration_test + @pytest.mark.tpu_only + def test_tpu_ulysses_context_parallelism(self): + train_main( + TrainTests.CONFIGS["synthetic"] + + [ + "attention=flash", + "use_tokamax_splash=true", + "ici_context_parallelism=4", + "context_parallel_strategy=ulysses", + "context_parallel_load_balance=false", + "packing=false", + ] + ) + @pytest.mark.integration_test @pytest.mark.gpu_only def test_gpu_base(self): diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index 03b4571b73..e03652eab4 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -28,7 +28,7 @@ import jax import jax.numpy as jnp from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask -from jax.sharding import AxisType, Mesh +from jax.sharding import AxisType, Mesh, NamedSharding from maxtext.utils import max_utils from maxtext.utils import maxtext_utils from maxtext.common.gcloud_stub import is_decoupled @@ -58,6 +58,7 @@ import pytest from tests.utils import attention_test_util +from tests.utils import hlo_test_utils from tests.utils.test_helpers import get_test_config_path @@ -1398,6 +1399,203 @@ def ring_loss(lnx): f"dq_reduction_steps={dq_reduction_steps}, ring_scan_unroll={ring_scan_unroll}, packing={packing}.", ) + def _ulysses_test_config(self, ici_context_parallelism): + return pyconfig.initialize( + [sys.argv[0], get_test_config_path()], + **self.config_arguments, + attention="flash", + context_parallel_strategy="ulysses", + context_parallel_load_balance=False, + ici_context_parallelism=ici_context_parallelism, + use_tokamax_splash=True, + use_jax_splash=False, + packing=False, + dtype="float32", + ) + + def _ulysses_test_modules(self, cfg_cp, mesh_cp, lnx): + """Builds the dot-product reference and the Ulysses flash attention modules.""" + attention_as_mha_generic = Attention( + config=self.cfg, + num_query_heads=cfg_cp.num_query_heads, + num_kv_heads=cfg_cp.num_kv_heads, + head_dim=cfg_cp.head_dim, + max_target_length=cfg_cp.max_target_length, + max_prefill_predict_length=cfg_cp.max_prefill_predict_length, + inputs_q_shape=lnx.shape, + inputs_kv_shape=lnx.shape, + mesh=self.mesh, + attention_kernel="dot_product", + dtype=cfg_cp.dtype, + dropout_rate=cfg_cp.dropout_rate, + rngs=self.nnx_rng, + ) + with nn_partitioning.axis_rules(cfg_cp.logical_axis_rules): + attention_as_mha_flash_cp = Attention( + config=cfg_cp, + num_query_heads=cfg_cp.num_query_heads, + num_kv_heads=cfg_cp.num_kv_heads, + head_dim=cfg_cp.head_dim, + max_target_length=cfg_cp.max_target_length, + max_prefill_predict_length=cfg_cp.max_prefill_predict_length, + inputs_q_shape=lnx.shape, + inputs_kv_shape=lnx.shape, + mesh=mesh_cp, + attention_kernel="flash", + dtype=cfg_cp.dtype, + dropout_rate=cfg_cp.dropout_rate, + model_mode=MODEL_MODE_PREFILL, + rngs=self.nnx_rng, + ) + return attention_as_mha_generic, attention_as_mha_flash_cp + + @parameterized.named_parameters( + {"testcase_name": "ulysses_size_2", "ici_context_parallelism": 2}, + {"testcase_name": "ulysses_size_4", "ici_context_parallelism": 4}, + ) + @pytest.mark.tpu_only + def test_tpu_flash_attention_ulysses_context_parallel(self, ici_context_parallelism): + """Test equivalence between dot_product and flash attention + Ulysses context parallelism""" + + cfg_cp = self._ulysses_test_config(ici_context_parallelism) + devices_array_cp = maxtext_utils.create_device_mesh(cfg_cp) + mesh_cp = Mesh(devices_array_cp, cfg_cp.mesh_axes) + lnx, decoder_segment_ids, decoder_positions = self.get_data(cfg_cp.dtype) + attention_as_mha_generic, attention_as_mha_flash_cp = self._ulysses_test_modules(cfg_cp, mesh_cp, lnx) + mha_generic_output, _ = attention_as_mha_generic( + lnx, + lnx, + decoder_segment_ids=decoder_segment_ids, + inputs_positions=decoder_positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + nnx.update(attention_as_mha_flash_cp, nnx.state(attention_as_mha_generic)) + + mha_generic_flash_cp_output = attention_test_util.forward_with_context_expert_parallelism( + cfg_cp, + mesh_cp, + attention_as_mha_flash_cp, + lnx, + decoder_segment_ids, + decoder_positions, + ) + + mha_generic_output = jax.device_get(mha_generic_output) + mha_generic_flash_cp_output = jax.device_get(mha_generic_flash_cp_output) + + self.assertTrue( + jax.numpy.allclose(mha_generic_output, mha_generic_flash_cp_output, rtol=1e-02, atol=1e-02, equal_nan=False), + msg="Logits from generic dot product and flash attention + Ulysses context parallelism are not close. " + f"ici_context_parallelism={ici_context_parallelism}.", + ) + + @parameterized.named_parameters( + {"testcase_name": "ulysses_size_2", "ici_context_parallelism": 2}, + {"testcase_name": "ulysses_size_4", "ici_context_parallelism": 4}, + ) + @pytest.mark.tpu_only + def test_tpu_flash_attention_ulysses_context_parallel_grad(self, ici_context_parallelism): + """Test input-gradient equivalence between dot_product and flash attention + Ulysses context parallelism""" + + cfg_cp = self._ulysses_test_config(ici_context_parallelism) + devices_array_cp = maxtext_utils.create_device_mesh(cfg_cp) + mesh_cp = Mesh(devices_array_cp, cfg_cp.mesh_axes) + lnx, decoder_segment_ids, decoder_positions = self.get_data(cfg_cp.dtype) + attention_as_mha_generic, attention_as_mha_flash_cp = self._ulysses_test_modules(cfg_cp, mesh_cp, lnx) + nnx.update(attention_as_mha_flash_cp, nnx.state(attention_as_mha_generic)) + + def generic_loss(lnx): + output, _ = attention_as_mha_generic( + lnx, + lnx, + decoder_segment_ids=decoder_segment_ids, + inputs_positions=decoder_positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + return jnp.mean(output.astype(jnp.float32) ** 2) + + def ulysses_loss(lnx): + output, _ = attention_as_mha_flash_cp( + lnx, + lnx, + decoder_segment_ids=decoder_segment_ids, + inputs_positions=decoder_positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + return jnp.mean(output.astype(jnp.float32) ** 2) + + generic_grad = jax.grad(generic_loss)(lnx) + with jax.set_mesh(mesh_cp), nn_partitioning.axis_rules(cfg_cp.logical_axis_rules): + ulysses_grad = jax.grad(ulysses_loss)(lnx) + generic_grad = jax.device_get(generic_grad) + ulysses_grad = jax.device_get(ulysses_grad) + + self.assertTrue( + jax.numpy.allclose(generic_grad, ulysses_grad, rtol=1e-02, atol=1e-07, equal_nan=False), + msg="Input gradients from generic dot product and flash attention + Ulysses context parallelism are not " + f"close. ici_context_parallelism={ici_context_parallelism}.", + ) + + @pytest.mark.tpu_only + def test_tpu_flash_attention_ulysses_hlo_uses_all_to_all(self): + """Checks compiled TPU Ulysses attention HLO uses all-to-all collectives.""" + + cfg_cp = self._ulysses_test_config(4) + devices_array_cp = maxtext_utils.create_device_mesh(cfg_cp) + mesh_cp = Mesh(devices_array_cp, cfg_cp.mesh_axes) + lnx, decoder_segment_ids, decoder_positions = self.get_data(cfg_cp.dtype) + _, attention_as_mha_flash_cp = self._ulysses_test_modules(cfg_cp, mesh_cp, lnx) + + def attention_forward(x, pos, seg): + output, _ = attention_as_mha_flash_cp( + x, + x, + decoder_segment_ids=seg, + inputs_positions=pos, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + return output + + def attention_loss(x, pos, seg): + return jnp.sum(attention_forward(x, pos, seg).astype(jnp.float32)) + + hlo_texts = [] + for lowered_fn in (attention_forward, jax.grad(attention_loss)): + # The mesh and axis-rules contexts wrap the jit from outside because + # jax.set_mesh raises inside a traced function, and the output keeps its + # natural sequence sharding so the only full-sequence gathers in the + # program are the ones the attention path itself emits. + with jax.set_mesh(mesh_cp), nn_partitioning.axis_rules(cfg_cp.logical_axis_rules): + input_sharding = NamedSharding( + mesh_cp, + nn_partitioning.logical_to_mesh_axes( + ("activation_batch", "activation_length", "activation_embed"), nn_partitioning.get_axis_rules() + ), + ) + metadata_sharding = NamedSharding( + mesh_cp, nn_partitioning.logical_to_mesh_axes((None, "activation_length"), nn_partitioning.get_axis_rules()) + ) + lowered = jax.jit(lowered_fn).lower( + jax.device_put(lnx, input_sharding), + jax.device_put(decoder_positions, metadata_sharding), + jax.device_put(decoder_segment_ids, metadata_sharding), + ) + hlo_texts.append(lowered.compile().as_text()) + + sequence_lengths = (cfg_cp.max_target_length,) + for hlo_text in hlo_texts: + self.assertGreater(len(hlo_test_utils.collective_lines(hlo_text, "all-to-all")), 0) + self.assertLen(hlo_test_utils.attention_sequence_all_gather_lines(hlo_text, sequence_lengths), 0) + # The int32 segment-ID gathers are the only intended full-sequence gathers. + self.assertGreater( + len(hlo_test_utils.attention_sequence_all_gather_lines(hlo_text, sequence_lengths, dtypes=("s32",))), 0 + ) + self.assertLen(hlo_test_utils.collective_lines(hlo_text, "collective-permute"), 0) + @pytest.mark.tpu_only def test_dot_product_cache_axis_order(self): all_axis_orders = tuple(itertools.permutations(range(4))) diff --git a/tests/unit/configs_value_test.py b/tests/unit/configs_value_test.py index fbd1c6685e..dd49bf58ce 100644 --- a/tests/unit/configs_value_test.py +++ b/tests/unit/configs_value_test.py @@ -249,6 +249,116 @@ def test_tpu_tokamax_ring_config_validation_rejects_unsupported_configs(self): with self.assertRaisesRegex((ValueError, pydantic.ValidationError), expected_regex): pyconfig.initialize(argv) + def test_tpu_ulysses_config_validation_accepts_initial_config(self): + argv = [ + "", + _BASE_CONFIG_PATH, + "run_name=test", + "attention=flash", + "use_tokamax_splash=True", + "use_jax_splash=False", + "context_parallel_strategy=ulysses", + "context_parallel_load_balance=False", + "ici_context_parallelism=4", + "hardware=tpu", + "packing=False", + "dataset_type=synthetic", + "skip_jax_distributed_system=True", + ] + mock_devices = [unittest.mock.MagicMock(slice_index=0) for _ in range(8)] + with unittest.mock.patch("jax.devices", return_value=mock_devices): + config = pyconfig.initialize(argv) + + self.assertEqual(config.context_parallel_strategy, "ulysses") + self.assertEqual(config.ici_context_parallelism, 4) + self.assertFalse(config.context_parallel_load_balance) + + def test_context_parallel_strategy_is_normalized(self): + argv = [ + "", + _BASE_CONFIG_PATH, + "run_name=test", + "attention=flash", + "use_tokamax_splash=True", + "use_jax_splash=False", + "context_parallel_strategy=Ulysses", + "context_parallel_load_balance=False", + "ici_context_parallelism=4", + "hardware=tpu", + "packing=False", + "dataset_type=synthetic", + "skip_jax_distributed_system=True", + ] + mock_devices = [unittest.mock.MagicMock(slice_index=0) for _ in range(8)] + with unittest.mock.patch("jax.devices", return_value=mock_devices): + config = pyconfig.initialize(argv) + + self.assertEqual(config.context_parallel_strategy, "ulysses") + + def test_tpu_ulysses_config_validation_rejects_unsupported_configs(self): + base_args = [ + "", + _BASE_CONFIG_PATH, + "run_name=test", + "attention=flash", + "use_tokamax_splash=True", + "use_jax_splash=False", + "context_parallel_strategy=ulysses", + "context_parallel_load_balance=False", + "ici_context_parallelism=4", + "hardware=tpu", + "packing=False", + "dataset_type=synthetic", + "skip_jax_distributed_system=True", + ] + cases = [ + (["context_parallel_load_balance=True"], ["context_parallel_load_balance=False"], "load_balance"), + (["base_num_kv_heads=1"], [], "MQA"), + (["base_num_query_heads=18"], [], "requires num_query_heads"), + (["base_num_kv_heads=10"], [], "requires num_kv_heads"), + (["attention_type=mla"], [], "global causal attention"), + (["attention_type=local_sliding", "sliding_window_size=128"], [], "global causal attention"), + (["attention_type=chunk", "chunk_attn_window_size=128"], [], "global causal attention"), + (["attention_type=full"], [], "global causal attention"), + (["attention_type=compressed"], [], "global causal attention"), + (["use_qk_clip=True"], [], "QK-Clip"), + (["dq_reduction_steps=2"], [], "dq_reduction_steps"), + (["attention=dot_product"], ["attention=flash"], "attention=flash"), + (["use_tokamax_splash=False"], ["use_tokamax_splash=True"], "use_tokamax_splash"), + (["use_jax_splash=True"], ["use_jax_splash=False"], "use_jax_splash"), + (["max_target_length=2050"], [], "divisible by context_parallel_size"), + (["ici_context_parallelism=-1"], ["ici_context_parallelism=4"], "explicit positive"), + (["dcn_context_parallelism=-1"], [], "explicit positive"), + (["dcn_context_parallelism=2"], [], "dcn context parallelism"), + ( + ["ici_context_parallelism=-1", "dcn_context_parallelism=-1"], + ["ici_context_parallelism=4"], + "explicit positive", + ), + (["ici_context_parallelism=1"], ["ici_context_parallelism=4"], "context_parallel_size > 1"), + (["context_sharding=expert"], [], "context_sharding"), + (["packing=True", "dataset_type=tfds"], ["packing=False", "dataset_type=synthetic"], "packing"), + (["use_ragged_attention=True"], [], "ragged attention"), + (["attention_sink=True"], [], "attention sinks"), + (["use_indexer=True", "q_lora_rank=1"], [], "sparse indexer"), + (["use_chunked_prefill=True"], [], "chunked prefill"), + (["moba=True"], [], "MoBA"), + (["use_multimodal=True"], [], "multimodal"), + (["dropout_rate=0.1"], [], "dropout"), + (["context_parallel_strategy=ulysess"], ["context_parallel_strategy=ulysses"], "context_parallel_strategy"), + (["hardware=gpu"], ["hardware=tpu"], "only supported on TPU"), + (["hardware=gpu_multiprocess"], ["hardware=tpu"], "only supported on TPU"), + (["hardware=cpu"], ["hardware=tpu"], "only supported on TPU"), + ] + mock_devices = [unittest.mock.MagicMock(slice_index=0) for _ in range(8)] + for bad_args, args_to_remove, expected_regex in cases: + with self.subTest(bad_args=bad_args): + argv = [arg for arg in base_args if arg not in args_to_remove] + argv.extend(bad_args) + with unittest.mock.patch("jax.devices", return_value=mock_devices): + with self.assertRaisesRegex((ValueError, pydantic.ValidationError), expected_regex): + pyconfig.initialize(argv) + def test_load_balanced_chunk_context_parallel_config(self): argv = [ "", diff --git a/tests/unit/hlo_test_utils_test.py b/tests/unit/hlo_test_utils_test.py new file mode 100644 index 0000000000..5f2ed1f5f0 --- /dev/null +++ b/tests/unit/hlo_test_utils_test.py @@ -0,0 +1,93 @@ +# 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 the HLO collective text helpers. + +The fixture lines are captured from real jax.jit lowerings of shard_map +collectives, plus the asynchronous -start/-done forms compiled HLO uses. +""" + +from absl.testing import absltest + +from tests.utils import hlo_test_utils + +_PPERMUTE = ( + " ppermute.1 = f32[2,16]{1,0} collective-permute(shard_map.2), channel_id=1," + " source_target_pairs={{0,1},{1,2},{2,3},{3,0}}" +) +_PPERMUTE_START = ( + " collective-permute-start.1 = (f32[2,16]{1,0}, f32[2,16]{1,0}) collective-permute-start(shard_map.2)," + " channel_id=1, source_target_pairs={{0,1},{1,2},{2,3},{3,0}}" +) +_ALL_TO_ALL = ( + " all_to_all.5 = bf16[8,8,4]{2,1,0} all-to-all(shard_map.2), channel_id=1," + " replica_groups={{0,1,2,3,4,5,6,7}}, dimensions={0}" +) +_SEGMENT_ID_ALL_GATHER = ( + " all_gather.1 = s32[4,512]{1,0} all-gather(broadcast.1), channel_id=1," + " replica_groups={{0,1,2,3}}, dimensions={1}, use_global_device_ids=true" +) +_KV_ALL_GATHER = ( + " all-gather.3 = bf16[4,512,8,128]{3,2,1,0} all-gather(param.1), channel_id=2," + " replica_groups={{0,1,2,3}}, dimensions={1}, use_global_device_ids=true" +) +_KV_ALL_GATHER_START = ( + " all-gather-start.1 = (bf16[4,128,8,128]{3,2,1,0}, bf16[4,512,8,128]{3,2,1,0})" + " all-gather-start(param.1), channel_id=3, replica_groups={{0,1,2,3}}, dimensions={1}" +) +_KV_ALL_GATHER_DONE = " all-gather-done.1 = bf16[4,512,8,128]{3,2,1,0} all-gather-done(all-gather-start.1)" +_FUSION_WITH_COLLECTIVE_OPERAND = " fusion.1 = bf16[4,512]{1,0} fusion(all-gather-done.1), kind=kLoop" +_NON_SEQUENCE_DIM_ALL_GATHER = ( + " all-gather.9 = bf16[8,512,16]{2,1,0} all-gather(param.2), channel_id=4," + " replica_groups={{0,1,2,3,4,5,6,7}}, dimensions={0}, use_global_device_ids=true" +) + + +class HloTestUtilsTest(absltest.TestCase): + """Tests for collective_lines and attention_sequence_all_gather_lines.""" + + def test_collective_lines_matches_sync_and_async_forms(self): + hlo_text = "\n".join([_PPERMUTE, _PPERMUTE_START, _ALL_TO_ALL]) + self.assertLen(hlo_test_utils.collective_lines(hlo_text, "collective-permute"), 2) + self.assertLen(hlo_test_utils.collective_lines(hlo_text, "all-to-all"), 1) + self.assertLen(hlo_test_utils.collective_lines(hlo_text, "all-gather"), 0) + + def test_collective_lines_counts_each_async_collective_once(self): + hlo_text = "\n".join([_KV_ALL_GATHER_START, _KV_ALL_GATHER_DONE, _FUSION_WITH_COLLECTIVE_OPERAND]) + self.assertLen(hlo_test_utils.collective_lines(hlo_text, "all-gather"), 1) + + def test_sequence_all_gather_lines_excludes_segment_id_gathers(self): + hlo_text = "\n".join([_SEGMENT_ID_ALL_GATHER, _KV_ALL_GATHER]) + lines = hlo_test_utils.attention_sequence_all_gather_lines(hlo_text, (512,)) + self.assertLen(lines, 1) + self.assertIn("bf16", lines[0]) + + def test_sequence_all_gather_lines_counts_segment_id_gathers_for_s32(self): + hlo_text = "\n".join([_SEGMENT_ID_ALL_GATHER, _KV_ALL_GATHER]) + lines = hlo_test_utils.attention_sequence_all_gather_lines(hlo_text, (512,), dtypes=("s32",)) + self.assertLen(lines, 1) + self.assertIn("s32", lines[0]) + + def test_sequence_all_gather_lines_detects_full_shape_in_async_tuple(self): + lines = hlo_test_utils.attention_sequence_all_gather_lines(_KV_ALL_GATHER_START, (512,)) + self.assertLen(lines, 1) + + def test_sequence_all_gather_lines_ignores_other_sequence_lengths(self): + self.assertLen(hlo_test_utils.attention_sequence_all_gather_lines(_KV_ALL_GATHER, (1024,)), 0) + + def test_sequence_all_gather_lines_ignores_non_sequence_gather_dimensions(self): + self.assertLen(hlo_test_utils.attention_sequence_all_gather_lines(_NON_SEQUENCE_DIM_ALL_GATHER, (512,)), 0) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/unit/ulysses_attention_test.py b/tests/unit/ulysses_attention_test.py new file mode 100644 index 0000000000..be5995adbc --- /dev/null +++ b/tests/unit/ulysses_attention_test.py @@ -0,0 +1,265 @@ +# 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 Ulysses attention layout helpers.""" + +from __future__ import annotations + +import types +from unittest import mock + +from absl.testing import absltest +import jax +import jax.numpy as jnp + +from maxtext.common.common_types import MODEL_MODE_PREFILL +from maxtext.common.common_types import MODEL_MODE_TRAIN +from maxtext.kernels.attention import ulysses_attention + + +class UlyssesAttentionTest(absltest.TestCase): + + def test_context_parallel_strategy_helper_identifies_ulysses(self): + self.assertTrue( + ulysses_attention.is_context_parallel_ulysses_requested( + types.SimpleNamespace(context_parallel_strategy="ulysses") + ) + ) + self.assertFalse( + ulysses_attention.is_context_parallel_ulysses_requested(types.SimpleNamespace(context_parallel_strategy="ring")) + ) + + def test_validate_ulysses_runtime_allows_train_mode(self): + ulysses_attention.validate_ulysses_runtime(model_mode=MODEL_MODE_TRAIN) + + def test_validate_ulysses_runtime_rejects_unsupported_runtime_features(self): + with self.assertRaisesRegex(ValueError, "train mode"): + ulysses_attention.validate_ulysses_runtime(model_mode=MODEL_MODE_PREFILL) + with self.assertRaisesRegex(ValueError, "ragged attention"): + ulysses_attention.validate_ulysses_runtime(model_mode=MODEL_MODE_TRAIN, use_ragged_attention=True) + with self.assertRaisesRegex(ValueError, "chunked prefill"): + ulysses_attention.validate_ulysses_runtime(model_mode=MODEL_MODE_TRAIN, previous_chunk=object()) + with self.assertRaisesRegex(ValueError, "attention sinks"): + ulysses_attention.validate_ulysses_runtime(model_mode=MODEL_MODE_TRAIN, sinks=object()) + with self.assertRaisesRegex(ValueError, "indexer"): + ulysses_attention.validate_ulysses_runtime(model_mode=MODEL_MODE_TRAIN, indexer_mask=object()) + with self.assertRaisesRegex(ValueError, "bidirectional"): + ulysses_attention.validate_ulysses_runtime(model_mode=MODEL_MODE_TRAIN, bidirectional_mask=object()) + with self.assertRaisesRegex(NotImplementedError, "record_max_logits"): + ulysses_attention.validate_ulysses_runtime(model_mode=MODEL_MODE_TRAIN, record_max_logits=True) + + def test_with_sequence_axis_preserves_partition_spec_type(self): + spec = jax.sharding.PartitionSpec("data", None, None, "tensor") + + out = ulysses_attention.with_sequence_axis(spec, "context", sequence_dim=2) + + self.assertIsInstance(out, jax.sharding.PartitionSpec) + self.assertEqual(tuple(out), ("data", None, "context", "tensor")) + + def test_validate_ulysses_mesh_axis_requires_sequence_sharding(self): + mesh = types.SimpleNamespace(shape={"context": 4}) + + with self.assertRaisesRegex(ValueError, "K/V sequence"): + ulysses_attention.validate_ulysses_mesh_axis( + axis_names_q=(None, None, "context", None), + axis_names_kv=(None, None, None, None), + sequence_dim_q=2, + sequence_dim_kv=2, + mesh=mesh, + ulysses_axis="context", + ) + + def test_layout_validators_reject_invalid_shardings(self): + mesh = types.SimpleNamespace(shape={"context": 4, "tensor": 2}) + cases = [ + ( + "unsharded or exactly", + lambda: ulysses_attention.with_sequence_axis((None, None, "tensor", None), "context", sequence_dim=2), + ), + ( + "mesh axis 'context' to exist", + lambda: ulysses_attention.validate_ulysses_mesh_axis( + axis_names_q=(None, None, "context", None), + axis_names_kv=(None, None, "context", None), + sequence_dim_q=2, + sequence_dim_kv=2, + mesh=types.SimpleNamespace(shape={"tensor": 2}), + ulysses_axis="context", + ), + ), + ( + "only on the sequence dimension", + lambda: ulysses_attention.validate_ulysses_mesh_axis( + axis_names_q=(None, "context", "context", None), + axis_names_kv=(None, None, "context", None), + sequence_dim_q=2, + sequence_dim_kv=2, + mesh=mesh, + ulysses_axis="context", + ), + ), + ( + "Q sequence sharding to be exactly", + lambda: ulysses_attention.validate_ulysses_mesh_axis( + axis_names_q=(None, None, None, None), + axis_names_kv=(None, None, "context", None), + sequence_dim_q=2, + sequence_dim_kv=2, + mesh=mesh, + ulysses_axis="context", + ), + ), + ( + "D_KV/head-dim", + lambda: ulysses_attention.validate_dkv_sharding( + axis_names_q=(None, None, "context", "tensor"), + axis_names_kv=(None, None, "context", None), + dkv_dim_q=3, + dkv_dim_kv=3, + ), + ), + ( + "divisible by Q head shards", + lambda: ulysses_attention.validate_head_sharding( + axis_names_q=(None, "tensor", "context", None), + axis_names_kv=(None, "tensor", "context", None), + mesh=mesh, + num_query_heads=9, + num_kv_heads=4, + head_dim_q=1, + head_dim_kv=1, + ulysses_size=4, + ), + ), + ( + "divisible by KV head shards", + lambda: ulysses_attention.validate_head_sharding( + axis_names_q=(None, "tensor", "context", None), + axis_names_kv=(None, "tensor", "context", None), + mesh=mesh, + num_query_heads=64, + num_kv_heads=9, + head_dim_q=1, + head_dim_kv=1, + ulysses_size=4, + ), + ), + ( + "divisible by local KV heads", + lambda: ulysses_attention.validate_head_sharding( + axis_names_q=(None, "tensor", "context", None), + axis_names_kv=(None, "tensor", "context", None), + mesh=mesh, + num_query_heads=64, + num_kv_heads=24, + head_dim_q=1, + head_dim_kv=1, + ulysses_size=4, + ), + ), + ( + r"local query heads \(8\) to be divisible by context_parallel_size", + lambda: ulysses_attention.validate_head_sharding( + axis_names_q=(None, None, "context", None), + axis_names_kv=(None, None, "context", None), + mesh=mesh, + num_query_heads=8, + num_kv_heads=4, + head_dim_q=1, + head_dim_kv=1, + ulysses_size=16, + ), + ), + ] + for expected_regex, invoke in cases: + with self.subTest(expected_regex=expected_regex): + with self.assertRaisesRegex(ValueError, expected_regex): + invoke() + + def test_validate_head_sharding_uses_local_heads_after_tensor_sharding(self): + mesh = types.SimpleNamespace(shape={"context": 4, "tensor": 2}) + + ulysses_attention.validate_head_sharding( + axis_names_q=(None, "tensor", "context", None), + axis_names_kv=(None, "tensor", "context", None), + mesh=mesh, + num_query_heads=64, + num_kv_heads=16, + head_dim_q=1, + head_dim_kv=1, + ulysses_size=4, + ) + + with self.assertRaisesRegex(ValueError, "local KV heads"): + ulysses_attention.validate_head_sharding( + axis_names_q=(None, "tensor", "context", None), + axis_names_kv=(None, "tensor", "context", None), + mesh=mesh, + num_query_heads=64, + num_kv_heads=4, + head_dim_q=1, + head_dim_kv=1, + ulysses_size=4, + ) + + def test_validate_head_sharding_rejects_mqa(self): + mesh = types.SimpleNamespace(shape={"context": 4}) + + with self.assertRaisesRegex(ValueError, "MQA"): + ulysses_attention.validate_head_sharding( + axis_names_q=(None, None, "context", None), + axis_names_kv=(None, None, "context", None), + mesh=mesh, + num_query_heads=16, + num_kv_heads=1, + head_dim_q=1, + head_dim_kv=1, + ulysses_size=4, + ) + + def test_validate_head_sharding_requires_q_and_kv_head_axes_to_match(self): + mesh = types.SimpleNamespace(shape={"context": 4, "tensor": 2}) + + with self.assertRaisesRegex(ValueError, "head sharding to match"): + ulysses_attention.validate_head_sharding( + axis_names_q=(None, "tensor", "context", None), + axis_names_kv=(None, None, "context", None), + mesh=mesh, + num_query_heads=64, + num_kv_heads=16, + head_dim_q=1, + head_dim_kv=1, + ulysses_size=4, + ) + + def test_ulysses_all_to_all_moves_heads_to_sequence(self): + tensor = jnp.ones((1, 16, 8, 2)) + + with mock.patch.object(ulysses_attention.jax.lax, "all_to_all", return_value="out") as all_to_all: + out = ulysses_attention.ulysses_all_to_all(tensor, "context") + + self.assertEqual(out, "out") + all_to_all.assert_called_once_with(tensor, "context", split_axis=1, concat_axis=2, tiled=True) + + def test_inverse_ulysses_all_to_all_moves_sequence_to_heads(self): + tensor = jnp.ones((1, 4, 32, 2)) + + with mock.patch.object(ulysses_attention.jax.lax, "all_to_all", return_value="out") as all_to_all: + out = ulysses_attention.inverse_ulysses_all_to_all(tensor, "context") + + self.assertEqual(out, "out") + all_to_all.assert_called_once_with(tensor, "context", split_axis=2, concat_axis=1, tiled=True) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/unit/ulysses_collective_test.py b/tests/unit/ulysses_collective_test.py new file mode 100644 index 0000000000..1397e8a891 --- /dev/null +++ b/tests/unit/ulysses_collective_test.py @@ -0,0 +1,155 @@ +# 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. +"""Executes the Ulysses collectives on a forced multi-device CPU mesh. + +Runs as a subprocess so the forced device count takes effect before JAX +initializes; the parent pytest process has already initialized JAX with the +default device count. The child checks, against a dense single-device +reference, the sequence-to-head exchange layout, the inverse round trip, GQA +grouping, segment-ID gather ordering across shard boundaries with padding, +and independent Q, K, and V gradients, on a 1-D context mesh and on a 2-D +fsdp x context mesh with the batch sharded over fsdp. +""" + +import os +import subprocess +import sys +from functools import partial + +import jax +import jax.numpy as jnp +from jax.sharding import Mesh, PartitionSpec as P +import numpy as np +import pytest + +from maxtext.kernels.attention import ulysses_attention + + +@pytest.mark.cpu_only +def test_ulysses_collectives_match_dense_reference_on_cpu_mesh(): + env = os.environ.copy() + env["XLA_FLAGS"] = env.get("XLA_FLAGS", "") + " --xla_force_host_platform_device_count=4" + env["JAX_PLATFORMS"] = "cpu" + result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False) + assert result.returncode == 0, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + assert "ULYSSES_COLLECTIVE_CHECKS_PASSED" in result.stdout + + +def _dense_reference_attention(query, key, value, segment_ids): + """Causal segment-masked GQA attention computed on one device.""" + _, num_query_heads, seq_len, _ = query.shape + num_kv_heads = key.shape[1] + group_size = num_query_heads // num_kv_heads + key = jnp.repeat(key, group_size, axis=1) + value = jnp.repeat(value, group_size, axis=1) + + logits = jnp.einsum("bhqd,bhkd->bhqk", query, key) + causal = jnp.tril(jnp.ones((seq_len, seq_len), dtype=bool)) + same_segment = segment_ids[:, :, None] == segment_ids[:, None, :] + not_padding = segment_ids != 0 + mask = causal[None, None, :, :] & same_segment[:, None, :, :] & not_padding[:, None, None, :] + logits = jnp.where(mask, logits, -1e30) + weights = jnp.exp(logits - jnp.max(logits, axis=-1, keepdims=True)) + weights = weights * mask + weights = weights / jnp.maximum(jnp.sum(weights, axis=-1, keepdims=True), 1e-30) + return jnp.einsum("bhqk,bhkd->bhqd", weights, value) + + +def _run_collective_checks(mesh, batch_axis): + """Runs the exchange layout, round-trip, attention, and gradient checks on one mesh.""" + batch, num_query_heads, num_kv_heads, seq_len, head_dim = 2, 8, 4, 32, 4 + assert mesh.size == len(jax.devices()), jax.devices() + ulysses_axis = "context" + data_spec = P(batch_axis, None, "context", None) + exchanged_spec = P(batch_axis, "context", None, None) + segment_spec = P(batch_axis, "context") + + # Rank-coded values make any head or sequence misordering visible exactly. + def coded(num_heads, offset): + values = np.arange(batch * num_heads * seq_len * head_dim, dtype=np.float32) + return jnp.asarray(values.reshape(batch, num_heads, seq_len, head_dim) / 100.0 + offset) + + query = coded(num_query_heads, 1.0) + key = coded(num_kv_heads, 2.0) + value = coded(num_kv_heads, 3.0) + # Segments begin and end inside different context shards, with trailing + # padding zeros. + segment_ids = jnp.broadcast_to(jnp.asarray([1] * 10 + [2] * 12 + [0] * 10, dtype=jnp.int32)[None, :], (batch, seq_len)) + + # Round trip through the real helpers is exact. + @partial( + jax.shard_map, + mesh=mesh, + in_specs=data_spec, + out_specs=data_spec, + check_vma=False, + ) + def round_trip(tensor): + return ulysses_attention.inverse_ulysses_all_to_all( + ulysses_attention.ulysses_all_to_all(tensor, ulysses_axis), ulysses_axis + ) + + np.testing.assert_array_equal(jax.device_get(round_trip(query)), jax.device_get(query)) + + # The forward exchange produces each rank's head subset over the full sequence. + @partial( + jax.shard_map, + mesh=mesh, + in_specs=data_spec, + out_specs=exchanged_spec, + check_vma=False, + ) + def exchange(tensor): + return ulysses_attention.ulysses_all_to_all(tensor, ulysses_axis) + + np.testing.assert_array_equal(jax.device_get(exchange(query)), jax.device_get(query)) + + @partial( + jax.shard_map, + mesh=mesh, + in_specs=(data_spec, data_spec, data_spec, segment_spec), + out_specs=data_spec, + check_vma=False, + ) + def ulysses_attention_fn(query, key, value, segment_ids): + query = ulysses_attention.ulysses_all_to_all(query, ulysses_axis) + key = ulysses_attention.ulysses_all_to_all(key, ulysses_axis) + value = ulysses_attention.ulysses_all_to_all(value, ulysses_axis) + segment_ids = jax.lax.all_gather(segment_ids, ulysses_axis, axis=1, tiled=True) + output = _dense_reference_attention(query, key, value, segment_ids) + return ulysses_attention.inverse_ulysses_all_to_all(output, ulysses_axis) + + def dense_loss(query, key, value): + output = _dense_reference_attention(query, key, value, segment_ids) + return jnp.sum(output * jnp.cos(output)) + + def ulysses_loss(query, key, value): + output = ulysses_attention_fn(query, key, value, segment_ids) + return jnp.sum(output * jnp.cos(output)) + + dense_output = _dense_reference_attention(query, key, value, segment_ids) + ulysses_output = ulysses_attention_fn(query, key, value, segment_ids) + np.testing.assert_allclose(jax.device_get(ulysses_output), jax.device_get(dense_output), atol=1e-5) + + dense_grads = jax.grad(dense_loss, argnums=(0, 1, 2))(query, key, value) + ulysses_grads = jax.grad(ulysses_loss, argnums=(0, 1, 2))(query, key, value) + for name, dense_grad, ulysses_grad in zip(("dQ", "dK", "dV"), dense_grads, ulysses_grads): + np.testing.assert_allclose(jax.device_get(ulysses_grad), jax.device_get(dense_grad), atol=1e-5, err_msg=name) + + +if __name__ == "__main__": + _devices = np.array(jax.devices()) + _run_collective_checks(Mesh(_devices, ("context",)), batch_axis=None) + _run_collective_checks(Mesh(_devices.reshape(2, 2), ("fsdp", "context")), batch_axis="fsdp") + print("ULYSSES_COLLECTIVE_CHECKS_PASSED") diff --git a/tests/utils/hlo_test_utils.py b/tests/utils/hlo_test_utils.py new file mode 100644 index 0000000000..9fef4e3927 --- /dev/null +++ b/tests/utils/hlo_test_utils.py @@ -0,0 +1,80 @@ +# 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. +"""Helpers for asserting on the collectives in attention HLO text. + +The helpers work on HLO text lines such as: + + ppermute.1 = f32[2,16]{1,0} collective-permute(shard_map.2), ... + all_gather.1 = s32[64]{0} all-gather(broadcast.1), ... + +Compiled HLO may rewrite a collective into its asynchronous -start/-done +pair; the -start instruction is counted and the -done instruction is not, so +each collective is counted once. +""" + +import re + +_FLOAT_TYPES = ("bf16", "f16", "f32", "f64") + + +def collective_lines(hlo_text, collective): + """Returns the HLO instruction lines that call the given collective op.""" + pattern = re.compile(rf"\b{re.escape(collective)}(-start)?\(") + return [line for line in hlo_text.splitlines() if "=" in line and pattern.search(line)] + + +def _result_shapes(line, collective): + """Parses the result shapes of a collective instruction line. + + Everything bracketed before the op call belongs to the result; an + asynchronous -start instruction has a tuple result whose elements include + both the sharded operand buffer and the full output buffer. + """ + call = re.search(rf"\b{re.escape(collective)}(-start)?\(", line) + result_part = line[: call.start()] if call else line + shapes = [] + for match in re.finditer(r"([a-z][a-z0-9]*)\[([0-9,]*)\]", result_part): + dims = tuple(int(dim) for dim in match.group(2).split(",") if dim) + shapes.append((match.group(1), dims)) + return shapes + + +def _collective_dimensions(line): + """Parses the dimensions={...} attribute of a collective instruction line.""" + match = re.search(r"dimensions=\{([0-9,]*)\}", line) + if not match: + return () + return tuple(int(dim) for dim in match.group(1).split(",") if dim) + + +def attention_sequence_all_gather_lines(hlo_text, sequence_lengths, dtypes=_FLOAT_TYPES): + """Returns all-gather lines whose gathered dimension is a full-sequence dimension. + + A line matches only when the dimension named in its dimensions={...} + attribute has a full-sequence result size, so a sequence-sized size on a + non-gathered dimension does not match. dtypes restricts matches by result + element type. The float default excludes the intended int32 segment-ID + gathers, so a match means a full-sequence gather of activations or + gradients; pass ("s32",) to count the segment-ID gathers instead. + """ + lines = [] + for line in collective_lines(hlo_text, "all-gather"): + gather_dims = _collective_dimensions(line) + for result_type, dims in _result_shapes(line, "all-gather"): + if result_type in dtypes and any( + gather_dim < len(dims) and dims[gather_dim] in sequence_lengths for gather_dim in gather_dims + ): + lines.append(line) + break + return lines