diff --git a/src/maxtext/common/metric_logger.py b/src/maxtext/common/metric_logger.py index a976af1698..687c35963c 100644 --- a/src/maxtext/common/metric_logger.py +++ b/src/maxtext/common/metric_logger.py @@ -223,6 +223,10 @@ def _log_training_metrics(self, metrics, step): log_parts.append(f"main_model_loss: {loss - mtp_loss:.3f}") log_parts.append(f"mtp_loss: {mtp_loss:.3f}") + if getattr(self.config, "use_indexer", False): + indexer_l = scalars.get("learning/indexer_loss", 0.0) + log_parts.append(f"indexer_loss: {float(indexer_l)}") + max_logging.log(", ".join(log_parts)) def _log_eval_metrics(self, metrics, step): @@ -246,6 +250,11 @@ def _log_eval_metrics(self, metrics, step): ) if "eval/avg_dpo_reward_accuracy" in scalars: log_parts.append(f"dpo_reward_accuracy={scalars['eval/avg_dpo_reward_accuracy']:.3f}") + + if getattr(self.config, "use_indexer", False): + indexer_l = scalars.get("eval/avg_indexer_loss", 0.0) + log_parts.append(f"avg_indexer_loss={indexer_l:.3f}") + max_logging.log(", ".join(log_parts)) def _log_running_eval_metrics(self, metrics, step): diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index acaf519a3d..6d5d188060 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -3497,6 +3497,11 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de f"`indexer_n_heads` ({self.indexer_n_heads})." ) if self.use_indexer: + if self.attention_type != AttentionType.MLA.value: + raise ValueError( + f"`use_indexer=True` requires `attention_type='{AttentionType.MLA.value}'`, since only the " + "MLA indexer produces this mask." + ) if self.q_lora_rank == 0: raise NotImplementedError("Sparse indexer has not implemented for q_lora_rank = 0.") supports_dot_product = self.attention == "dot_product" @@ -3699,8 +3704,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de raise ValueError("TPU Tokamax ring attention does not support ragged attention.") if self.attention_sink: raise ValueError("TPU Tokamax ring attention does not support attention sinks.") - if self.use_indexer: - raise ValueError("TPU Tokamax ring attention does not support sparse indexer masks.") if self.use_chunked_prefill: raise ValueError("TPU Tokamax ring attention does not support chunked prefill yet.") if self.moba: @@ -3738,6 +3741,8 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de 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.use_indexer: + raise ValueError("TPU Ulysses attention does not support sparse indexer masks.") if self.attention_type != "global": raise ValueError("TPU Ulysses attention is initially supported only for global causal attention.") if self.context_parallel_load_balance: @@ -3750,8 +3755,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de 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: diff --git a/src/maxtext/kernels/attention/tokamax_ring_attention.py b/src/maxtext/kernels/attention/tokamax_ring_attention.py index 05a23fffec..0aaf82cdd0 100644 --- a/src/maxtext/kernels/attention/tokamax_ring_attention.py +++ b/src/maxtext/kernels/attention/tokamax_ring_attention.py @@ -110,8 +110,6 @@ def validate_tokamax_ring_runtime( raise ValueError("TPU Tokamax ring attention does not support chunked prefill yet.") if sinks is not None: raise ValueError("TPU Tokamax ring attention does not support attention sinks.") - if indexer_mask is not None: - raise ValueError("TPU Tokamax ring attention does not support indexer masks.") if bidirectional_mask is not None: raise ValueError("TPU Tokamax ring attention does not support bidirectional masks.") if record_max_logits: @@ -283,6 +281,7 @@ def make_sharded_ring_attention_kernel( ring_axis: str, attn_logits_soft_cap: float | None, maybe_shard_with_pspec: Any, + mask: Any = None, ): """Builds and shards the Tokamax ring attention kernel for MaxText.""" splash_config = build_splash_config( @@ -295,11 +294,17 @@ def make_sharded_ring_attention_kernel( if config.use_max_logit_estimate > 0: splash_config = dataclasses.replace(splash_config, max_logit_const=config.use_max_logit_estimate) - mask = _make_causal_mask( - (query.shape[2], key.shape[2]), - context_parallel_size, - load_balanced=config.context_parallel_load_balance, - ) + if mask is None: + # When using the indexer, causal masking is unified into the dynamic indexer_mask + # and applied dynamically per block; use FullMask to avoid duplicate static masks. + if getattr(config, "use_indexer", False): + mask = tokamax_splash_mask.FullMask((query.shape[2], key.shape[2])) + else: + mask = _make_causal_mask( + (query.shape[2], key.shape[2]), + context_parallel_size, + load_balanced=config.context_parallel_load_balance, + ) @functools.partial(jax.jit, static_argnames=["single_head_mask"]) def wrap_ring_kernel(single_head_mask): @@ -331,15 +336,27 @@ def call_ring_attention( decoder_segment_ids_q: Any, decoder_segment_ids_kv: Any, ring_kernel: Any, + indexer_mask: Any = None, ): """Calls a Tokamax ring attention kernel over the MaxText batch dimension.""" if (decoder_segment_ids_q is None) != (decoder_segment_ids_kv is None): raise ValueError("decoder_segment_ids_q and decoder_segment_ids_kv must both be set or both be None.") + # Vectorize execution across batch dimension, threading indexer_mask when present. + # Note: ring_kernel expects positional arguments (q, k, v, segment_ids, sinks, indexer_mask). if decoder_segment_ids_q is None: - return jax.vmap(lambda q, k, v: ring_kernel(q, k, v, None), in_axes=(0, 0, 0))(query, key, value) - - def call_one(q, k, v, q_segment_ids, kv_segment_ids): + if indexer_mask is None: + return jax.vmap(lambda q, k, v: ring_kernel(q, k, v, None, None, None), in_axes=(0, 0, 0))(query, key, value) + return jax.vmap( + lambda q, k, v, im: ring_kernel(q, k, v, None, None, im), + in_axes=(0, 0, 0, 0), + )(query, key, value, indexer_mask) + + def call_one(q, k, v, q_segment_ids, kv_segment_ids, im=None): segment_ids = ring_attention_kernel.SegmentIds(q_segment_ids, kv_segment_ids) - return ring_kernel(q, k, v, segment_ids) + return ring_kernel(q, k, v, segment_ids, None, im) - return jax.vmap(call_one, in_axes=(0, 0, 0, 0, 0))(query, key, value, decoder_segment_ids_q, decoder_segment_ids_kv) + if indexer_mask is None: + return jax.vmap(call_one, in_axes=(0, 0, 0, 0, 0))(query, key, value, decoder_segment_ids_q, decoder_segment_ids_kv) + return jax.vmap(call_one, in_axes=(0, 0, 0, 0, 0, 0))( + query, key, value, decoder_segment_ids_q, decoder_segment_ids_kv, indexer_mask + ) diff --git a/src/maxtext/kernels/tokamax_splash_attention/ring_attention_kernel.py b/src/maxtext/kernels/tokamax_splash_attention/ring_attention_kernel.py index de0ce2b042..3daae6405c 100644 --- a/src/maxtext/kernels/tokamax_splash_attention/ring_attention_kernel.py +++ b/src/maxtext/kernels/tokamax_splash_attention/ring_attention_kernel.py @@ -60,8 +60,59 @@ def _validate_ring_axis_size(ring_axis: str, ring_axis_size: int, expected_ring_ ) +def _inject_local_indexer_mask( + local_mask_info: MaskInfo, + local_idx_mask: jax.Array | None, + block_shape: tuple[int, int] = (128, 128), + is_dkv: bool = False, +) -> MaskInfo: + """Injects a pre-sliced dynamic Indexer mask shard into local MaskInfo for the current ring step.""" + if local_idx_mask is None: + return local_mask_info + + bq, bkv = block_shape + if local_idx_mask.ndim == 3: + local_idx_mask = local_idx_mask[0] + + q_len, kv_len = local_idx_mask.shape + # Since causal and padding masks are already fully integrated into the global indexer_mask + # tensor (via indexer_mask += attention_mask in attention_mla.py) before slicing, + # local_idx_mask already contains the complete causally-masked top-k selection for this ring hop. + combined_mask = local_idx_mask + + # Tile 2D mask into hardware block chunks [bq, bkv] + q_blocks = q_len // bq + kv_blocks = kv_len // bkv + num_blocks = q_blocks * kv_blocks + + blocks = combined_mask.reshape(q_blocks, bq, kv_blocks, bkv) + blocks = blocks.swapaxes(1, 2) # [q_blocks, kv_blocks, bq, bkv] + + if is_dkv: + # SplashAttention dkv grids are scheduled as KV-major (kv_blocks, q_blocks). + # We transpose both block grid and intra-block axes to match Pallas grid_idx order. + blocks = blocks.swapaxes(0, 1) # [kv_blocks, q_blocks, bq, bkv] + blocks = blocks.swapaxes(-1, -2) # [kv_blocks, q_blocks, bkv, bq] + + blocks = blocks.reshape(num_blocks, blocks.shape[-2], blocks.shape[-1]) + blocks = blocks.astype(jnp.int8) + + mask_next = jnp.arange(num_blocks, dtype=jnp.int32) + return local_mask_info._replace( + mask_next=mask_next, + active_rows=None, + active_cols=None, + block_mask=None, + num_active_blocks=None, + partial_mask_blocks=blocks, + q_sequence=None, + kv_sequence=None, + ) + + def _ring_attention_forward( fwd_mask_info: MaskInfo, + indexer_mask: jax.Array | None, q: jax.Array, k: jax.Array, v: jax.Array, @@ -118,12 +169,40 @@ def _ring_attention_forward( l_init = jnp.zeros((o_shape[0], o_shape[1]), jnp.float32) m_init = jnp.full_like(l_init, mask_value, dtype=jnp.float32) - def body(carry, i: int): - m_prev, l_prev, o_prev, k_current, v_current, segment_ids_current = carry + if indexer_mask is not None: + # Reshape global indexer mask to [..., ring_axis_size, kv_shard_len] for dynamic step slicing. + kv_shard_len = k.shape[-2] + mask_4d = indexer_mask.reshape(*indexer_mask.shape[:-1], ring_axis_size, kv_shard_len) + else: + mask_4d = None + + if mask_4d is not None: + i_vals = jnp.arange(0, ring_axis_size) + target_indices = (ring_axis_idx - i_vals) % ring_axis_size + mask_ring_steps = mask_4d[..., target_indices, :] + mask_ring_steps = jnp.moveaxis(mask_ring_steps, -2, 0) + xs = (i_vals, mask_ring_steps) + else: + xs = jnp.arange(0, ring_axis_size) + + def body(carry, xs_arg): + if mask_4d is not None: + i, local_idx_mask = xs_arg + else: + i = xs_arg + local_idx_mask = None + m_prev, l_prev, o_prev, k_current, v_current, segment_ids_current = carry current_kv_shard_idx = (ring_axis_idx - i) % ring_axis_size + local_fwd_mask_info = _dynamic_slice_mask_info(fwd_mask_info, current_kv_shard_idx, ring_axis_size) local_fwd_mask_info = _offset_q_sequence_for_kv_shard(local_fwd_mask_info, current_kv_shard_idx, k_current.shape[-2]) + local_fwd_mask_info = _inject_local_indexer_mask( + local_fwd_mask_info, + local_idx_mask, + block_shape=(config.block_q, config.block_kv), + is_dkv=False, + ) k_next = shift(k_current) v_next = shift(v_current) @@ -168,7 +247,7 @@ def body(carry, i: int): (m_final, l_final, o_final, _, _, _), _ = lax.scan( body, initial_carry, - xs=jnp.arange(0, ring_axis_size), + xs=xs, length=ring_axis_size, unroll=config.ring_scan_unroll, ) # type: ignore[arg-type] @@ -198,7 +277,7 @@ def _ring_attention_bwd( do: jax.Array, ): del save_residuals - (q, k, v, segment_ids, sinks, out, logsumexp, dkv_mask_info) = res + (q, k, v, segment_ids, sinks, out, logsumexp, dkv_mask_info, indexer_mask) = res do = do.astype(jnp.float32) if dkv_mask_info is None: raise ValueError("Need to specify backward blocks.") @@ -229,10 +308,34 @@ def rotate_kv(k_current, v_current, segment_ids_current): segment_ids_next = None return k_next, v_next, segment_ids_next - def compute_step(i: int, k_current, v_current, segment_ids_current, dq_accum): + if indexer_mask is not None: + # Reshape global indexer mask to [..., ring_axis_size, kv_shard_len] for backward rotation slicing. + kv_shard_len = k.shape[-2] + mask_4d = indexer_mask.reshape(*indexer_mask.shape[:-1], ring_axis_size, kv_shard_len) + + step0_mask_idx = (ring_axis_idx - 0) % ring_axis_size + step0_mask = mask_4d[..., step0_mask_idx, :] + + i_vals = jnp.arange(1, ring_axis_size) + target_indices = (ring_axis_idx - i_vals) % ring_axis_size + mask_ring_steps = mask_4d[..., target_indices, :] + mask_ring_steps = jnp.moveaxis(mask_ring_steps, -2, 0) + xs = (i_vals, mask_ring_steps) + else: + mask_4d = None + step0_mask = None + xs = jnp.arange(1, ring_axis_size) + + def compute_step(i, local_idx_mask, k_current, v_current, segment_ids_current, dq_accum): current_kv_shard_idx = (ring_axis_idx - i) % ring_axis_size local_dkv_mask_info = _dynamic_slice_mask_info(dkv_mask_info, current_kv_shard_idx, ring_axis_size) local_dkv_mask_info = _offset_q_sequence_for_kv_shard(local_dkv_mask_info, current_kv_shard_idx, k_current.shape[-2]) + local_dkv_mask_info = _inject_local_indexer_mask( + local_dkv_mask_info, + local_idx_mask, + block_shape=(config.block_q_dkv, config.block_kv_dkv), + is_dkv=True, + ) residuals_for_chunk = ( q, @@ -266,11 +369,16 @@ def compute_step(i: int, k_current, v_current, segment_ids_current, dq_accum): dq_i = dq_accum + dq_i.astype(jnp.float32) return dq_i, dk_i, dv_i, dsinks - dq_i, dk_pending, dv_pending, dsinks = compute_step(0, k, v, segment_ids, dq_accum) + dq_i, dk_pending, dv_pending, dsinks = compute_step(0, step0_mask, k, v, segment_ids, dq_accum) dq_accum = dq_i k_current, v_current, segment_ids_current = rotate_kv(k, v, segment_ids) - def body(carry, i: int): + def body(carry, xs_arg): + if mask_4d is not None: + i, local_idx_mask = xs_arg + else: + i = xs_arg + local_idx_mask = None ( dq_accum, dk_accum, @@ -285,7 +393,7 @@ def body(carry, i: int): dk_next = shift(dk_accum + dk_pending.astype(jnp.float32)) dv_next = shift(dv_accum + dv_pending.astype(jnp.float32)) k_next, v_next, segment_ids_next = rotate_kv(k_current, v_current, segment_ids_current) - dq_i, dk_i, dv_i, dsinks = compute_step(i, k_current, v_current, segment_ids_current, dq_accum) + dq_i, dk_i, dv_i, dsinks = compute_step(i, local_idx_mask, k_current, v_current, segment_ids_current, dq_accum) dq_accum = dq_i return ( dq_accum, @@ -313,7 +421,7 @@ def body(carry, i: int): (dq, dk, dv, dk_pending, dv_pending, _, _, _, dsinks), _ = lax.scan( body, initial_carry, - xs=jnp.arange(1, ring_axis_size), + xs=xs, length=ring_axis_size - 1, unroll=config.ring_scan_unroll, ) @@ -333,6 +441,7 @@ def body(carry, i: int): dv.astype(v.dtype), None, dsinks, + None, # indexer_mask ) @@ -344,6 +453,7 @@ def _ring_attention_fwd( v: jax.Array, segment_ids: SegmentIds | None, sinks: jax.Array | None, + indexer_mask: jax.Array | None, # nondiff_args mask_value: float, # 1 is_mqa: bool, # 2 @@ -388,6 +498,7 @@ def _ring_attention_fwd( out, (logsumexp, max_logits) = _ring_attention_forward( fwd_mask_info, + indexer_mask, q, k, v, @@ -405,7 +516,7 @@ def _ring_attention_fwd( if config.residual_checkpoint_name is not None: out = ad_checkpoint.checkpoint_name(out, name=config.residual_checkpoint_name) logsumexp = ad_checkpoint.checkpoint_name(logsumexp, name=config.residual_checkpoint_name) - residuals = (q, k, v, segment_ids, sinks, out, logsumexp, dkv_mask_info) + residuals = (q, k, v, segment_ids, sinks, out, logsumexp, dkv_mask_info, indexer_mask) return out, residuals @@ -431,6 +542,7 @@ def _ring_attention_custom( v: jax.Array, segment_ids: SegmentIds | None, sinks: jax.Array | None, + indexer_mask: jax.Array | None, mask_value: float, is_mqa: bool, config: SplashConfig, @@ -468,6 +580,7 @@ def _ring_attention_custom( del dkv_mask_info, dkv_mask_sparsity out, _ = _ring_attention_forward( fwd_mask_info, + indexer_mask, q, k, v, @@ -509,6 +622,7 @@ def _ring_attention( v: jax.Array, segment_ids: SegmentIds | None = None, sinks: jax.Array | None = None, + indexer_mask: jax.Array | None = None, *, is_mqa: bool, config: SplashConfig, @@ -548,6 +662,9 @@ def _ring_attention( Raises: ValueError: If the specified `ring_axis` does not exist. """ + if indexer_mask is not None and indexer_mask.dtype != jnp.bool_: + indexer_mask = jnp.isclose(indexer_mask, 0.0) + if not _has_axis(ring_axis): raise ValueError(f"Ring axis {ring_axis} does not exist") @@ -559,6 +676,7 @@ def _ring_attention( v, segment_ids, sinks, + indexer_mask, is_mqa=is_mqa, config=config, mask_value=mask_value, @@ -637,19 +755,16 @@ def mask_info_spec(mask_info): if mask_info is None: return None return MaskInfo( # pytype: disable=wrong-arg-types - mask_next=_resolve_spec(mask_info.mask_next), # pyrefly: ignore[bad-argument-type] - active_rows=_resolve_spec(mask_info.active_rows), # pyrefly: ignore[bad-argument-type] - active_cols=_resolve_spec(mask_info.active_cols), # pyrefly: ignore[bad-argument-type] - num_active_blocks=_resolve_spec(mask_info.num_active_blocks), # pyrefly: ignore[bad-argument-type] - block_mask=_resolve_spec(mask_info.block_mask), # pyrefly: ignore[bad-argument-type] - partial_mask_blocks=jax.sharding.PartitionSpec() # replicated # pyrefly: ignore[bad-argument-type] + mask_next=_resolve_spec(mask_info.mask_next), + active_rows=_resolve_spec(mask_info.active_rows), + active_cols=_resolve_spec(mask_info.active_cols), + num_active_blocks=_resolve_spec(mask_info.num_active_blocks), + block_mask=_resolve_spec(mask_info.block_mask), + partial_mask_blocks=jax.sharding.PartitionSpec() # replicated if mask_info.partial_mask_blocks is not None else None, - q_sequence=_resolve_spec(mask_info.q_sequence), # pyrefly: ignore[bad-argument-type] - # pyrefly: ignore[bad-argument-type] - kv_sequence=jax.sharding.PartitionSpec() - if mask_info.kv_sequence is not None - else None, # pyrefly: ignore[bad-argument-type] + q_sequence=_resolve_spec(mask_info.q_sequence), + kv_sequence=jax.sharding.PartitionSpec() if mask_info.kv_sequence is not None else None, ) return RingSplashAttentionKernel( diff --git a/src/maxtext/layers/attention_mla.py b/src/maxtext/layers/attention_mla.py index 9d3c64f05b..3fc1a3c69d 100644 --- a/src/maxtext/layers/attention_mla.py +++ b/src/maxtext/layers/attention_mla.py @@ -78,6 +78,10 @@ PLACEHOLDER_SEQ_LEN = 1 +class indexer_losses(nnx.Variable): # pylint: disable=invalid-name,abstract-method + """Variable type for storing Indexer loss components -> bypasses nnx.Intermediate scan filters.""" + + class Indexer(nnx.Module): """Indexer for DeepSeek Sparse Attention (DSA). @@ -361,7 +365,7 @@ def __call__( # NOTE: If the total available sequence length <= topk, indexer always selects all tokens. if k.shape[1] <= self.indexer_topk: - return None, None, None + return attention_mask, cached_s, attention_mask # Compute head weights: project from input, [b, t, embed_dim] -> [b, t, h] weights = self.weights_proj(inputs_q) @@ -1287,7 +1291,13 @@ def __call__( if self.use_indexer: # generate mask: with 0 and large negative, [b, 1, 1, q_len, kv_len] -> [b, q_len, kv_len] attention_mask = self.attention_op.generate_attention_mask( - query, key, decoder_segment_ids, model_mode, previous_chunk, bidirectional_mask + query, + key, + decoder_segment_ids, + model_mode, + previous_chunk, + bidirectional_mask, + segment_positions=inputs_positions, ) if attention_mask is not None: attention_mask = attention_mask.squeeze(axis=(1, 2)) @@ -1314,7 +1324,7 @@ def __call__( sparse_loss=self.config.indexer_sparse_training, scaling_factor=self.config.indexer_loss_scaling_factor, ) - self.indexer_loss = nnx.Intermediate(indexer_loss) + self.indexer_loss = indexer_losses(indexer_loss) # Check if we need QK Clip stats use_qk_clip = self.model_mode == MODEL_MODE_TRAIN and self.config.use_qk_clip diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index 25dc782ab7..3b15f623cc 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -1211,6 +1211,12 @@ def _validate_tpu_tokamax_ring_runtime( record_max_logits: bool = False, ) -> None: """Validates runtime constraints for the TPU Tokamax ring path.""" + if getattr(self.config, "use_indexer", False) and indexer_mask is None: + raise ValueError( + "`indexer_mask` cannot be None when `use_indexer` is True. " + "Ring Attention drops its static causal mask when use_indexer is True, " + "so passing None would result in no causal masking." + ) tokamax_ring_attention.validate_tokamax_ring_runtime( model_mode=model_mode, previous_chunk=previous_chunk, @@ -1666,6 +1672,7 @@ def create_sa_config(config, query, key, attn_logits_soft_cap): ring_axis=self.config.context_sharding, attn_logits_soft_cap=attn_logits_soft_cap, maybe_shard_with_pspec=self._maybe_shard_with_pspec, + mask=None, ) ) elif use_ulysses: @@ -1879,6 +1886,7 @@ def wrap_flash_attention( decoder_segment_ids_q, decoder_segment_ids_kv, splash_kernel, + indexer_mask=indexer_mask, ) return attention_output, None @@ -1913,6 +1921,7 @@ def wrap_flash_attention( if cp_size > 1 and load_balanced_context_parallel: key = max_utils.reorder_sequence(tensor=key, cp_size=cp_size, seq_dim=2, to_contiguous=True) value = max_utils.reorder_sequence(tensor=value, cp_size=cp_size, seq_dim=2, to_contiguous=True) + decoder_segment_ids_unpermuted = max_utils.reorder_sequence( tensor=decoder_segment_ids_kv, cp_size=cp_size, @@ -1920,6 +1929,14 @@ def wrap_flash_attention( to_contiguous=True, ) + if indexer_mask is not None: + indexer_mask = max_utils.reorder_sequence( + tensor=indexer_mask, + cp_size=cp_size, + seq_dim=2, + to_contiguous=True, + ) + if decoder_segment_ids_q is not None: if cp_size > 1 and load_balanced_context_parallel: decoder_segment_ids_tuple = splash_attention_kernel.SegmentIds( diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index a992ca6fab..6a6cb920d7 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -51,6 +51,7 @@ # pylint: disable=too-many-positional-arguments from maxtext.layers.multi_token_prediction import calculate_mtp_acceptance_rate, calculate_mtp_loss, mtp_acceptance, mtp_losses +from maxtext.layers.attention_mla import indexer_losses from maxtext.common import checkpointing, profiler from maxtext.common.goodput import ( GoodputEvent, @@ -124,6 +125,9 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr # make its specific collection mutable so the MTPBlock can sow into it. if config.mtp_eval_target_module > 0 and not is_train: mutable_collections.append("mtp_acceptance") + if config.use_indexer and is_train: + mutable_collections.append("indexer_losses") + sparsity_enabled = is_train and config.weight_sparsity_n and config.weight_sparsity_m if sparsity_enabled: mutable_collections.append("batch_stats") @@ -209,6 +213,11 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr mtp_losses_state = nnx.pop(model, mtp_losses) mtp_acceptance_state = nnx.pop(model, mtp_acceptance) + indexer_losses_state = None + if config.use_indexer: + # Pop dedicated indexer_losses to harvest auxiliary KL loss and prevent model state PyTree mismatches. + indexer_losses_state = nnx.pop(model, indexer_losses) + intermediates = nnx.pop(model, nnx.Intermediate) intermediate_outputs = intermediates.to_pure_dict() @@ -218,6 +227,9 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr intermediate_outputs["mtp_losses"] = mtp_losses_state.to_pure_dict() intermediate_outputs["mtp_acceptance"] = mtp_acceptance_state.to_pure_dict() + if indexer_losses_state is not None: + intermediate_outputs["indexer_losses"] = indexer_losses_state.to_pure_dict() + if (config.use_indexer and not config.indexer_sparse_training) and is_train: # In Dense Warm-up stage, we skip main model loss calculation for efficiency. # The main model parameters are frozen and only the indexer is trained via KL divergence. @@ -280,15 +292,16 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr mtp_loss = calculate_mtp_loss(intermediate_outputs, config) loss += mtp_loss - # get indexer loss + # Calculate and add auxiliary Indexer loss indexer_loss = 0.0 if config.use_indexer and config.indexer_loss_scaling_factor > 0.0: - indexer_losses = maxtext_utils.collect_intermediates_by_suffix(intermediate_outputs, "self_attention", "indexer_loss") - if indexer_losses: - indexer_loss = jnp.mean(jnp.concatenate(indexer_losses)) - loss += indexer_loss + # Recursively collect per-layer indexer losses across all scanned transformer layers. + indexer_losses_list = maxtext_utils.collect_intermediates_by_suffix(intermediate_outputs, "indexer_loss") + if indexer_losses_list: + indexer_loss = jnp.mean(jnp.concatenate([jnp.atleast_1d(x) for x in indexer_losses_list])) + loss += indexer_loss # Injects loss into scalar objective to drive backward gradients for indexer weights. else: - max_logging.debug("No indexer loss found.") + max_logging.debug("No Indexer loss found. Defaulting to 0.0.") # get MoE load balance loss moe_lb_loss = 0.0 @@ -562,6 +575,7 @@ def move(path, value): new_state = state # Apply updates for Auxiliary-Loss-Free load balancing for DeepSeek family + # pylint: disable=too-many-nested-blocks if config.routed_bias and config.routed_bias_update_rate > 0.0: if config.model_name.startswith("deepseek4"): max_logging.log("DeepSeek V4: Applying auxiliary-loss-free routing bias via pure NNX MoEBiasVar.") diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index dc7ea03f43..0d7bfea95a 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -3167,6 +3167,466 @@ def ring_loss(lnx): f"context_parallel_load_balance={context_parallel_load_balance}.", ) + @parameterized.named_parameters( + { + "testcase_name": "no_lb_cp2", + "context_parallel_load_balance": False, + "ici_context_parallelism": 2, + "indexer_topk": 256, + }, + { + "testcase_name": "lb_cp4_smallk", + "context_parallel_load_balance": True, + "ici_context_parallelism": 4, + "indexer_topk": 32, + }, + ) + @pytest.mark.tpu_only + def test_tpu_dot_product_context_parallel_with_indexer( + self, context_parallel_load_balance, ici_context_parallelism=2, indexer_topk=256 + ): + """Test equivalence between single-device dot_product MLA + Indexer and multi-device dot_product + CP + Indexer""" + config_arguments = { + "per_device_batch_size": 1.0, + "run_name": "test", + "enable_checkpointing": False, + "max_target_length": 512, + "attention_type": AttentionType.MLA.value, + "use_indexer": True, + "indexer_loss_scaling_factor": 0.0, + "indexer_topk": indexer_topk, + "q_lora_rank": 4, + "kv_lora_rank": 4, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 64, + "v_head_dim": 128, + "dtype": "float32", + } + + cfg, mla = self.init_mla({**config_arguments, "attention": "dot_product"}, rope_type="default") + lnx, decoder_segment_ids, decoder_positions = self.get_structured_data(cfg, cfg.dtype) + mla_generic_output, _ = mla( + lnx, + lnx, + decoder_segment_ids=decoder_segment_ids, + inputs_positions=decoder_positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + generic_state = nnx.state(mla) + + cfg_cp = pyconfig.initialize( + [sys.argv[0], get_test_config_path()], + **config_arguments, + attention="dot_product", + rope_type=cfg.rope_type, + context_parallel_strategy="all_gather", + context_parallel_load_balance=context_parallel_load_balance, + ici_context_parallelism=ici_context_parallelism, + ) + devices_array_cp = maxtext_utils.create_device_mesh(cfg_cp) + mesh_cp = Mesh(devices_array_cp, cfg_cp.mesh_axes) + with nn_partitioning.axis_rules(cfg_cp.logical_axis_rules): + attention_as_mla_cp = MLA( + 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, + inputs_q_shape=lnx.shape, + inputs_kv_shape=lnx.shape, + max_target_length=cfg_cp.max_target_length, + max_prefill_predict_length=cfg_cp.max_prefill_predict_length, + mesh=mesh_cp, + attention_kernel="dot_product", + dtype=cfg_cp.dtype, + dropout_rate=cfg_cp.dropout_rate, + attention_type=AttentionType(cfg_cp.attention_type), + q_lora_rank=cfg_cp.q_lora_rank, + kv_lora_rank=cfg_cp.kv_lora_rank, + qk_nope_head_dim=cfg_cp.qk_nope_head_dim, + qk_rope_head_dim=cfg_cp.qk_rope_head_dim, + v_head_dim=cfg_cp.v_head_dim, + model_mode=MODEL_MODE_PREFILL, + rngs=self.nnx_rng, + ) + nnx.update(attention_as_mla_cp, generic_state) + + mla_cp_output = attention_test_util.forward_with_context_expert_parallelism( + cfg_cp, + mesh_cp, + attention_as_mla_cp, + lnx, + decoder_segment_ids, + decoder_positions, + ) + + mla_generic_output = jax.device_get(mla_generic_output) + mla_cp_output = jax.device_get(mla_cp_output) + + self.assertTrue( + jax.numpy.allclose(mla_generic_output, mla_cp_output, rtol=1e-02, atol=1e-02, equal_nan=False), + msg=( + "MLA+Indexer logits from single-device dot product and multi-device dot product context parallelism are" + f" not close. context_parallel_load_balance={context_parallel_load_balance}." + ), + ) + + @parameterized.named_parameters( + { + "testcase_name": "no_lb_cp2", + "context_parallel_load_balance": False, + "ici_context_parallelism": 2, + "indexer_topk": 256, + }, + { + "testcase_name": "lb_cp4_smallk", + "context_parallel_load_balance": True, + "ici_context_parallelism": 4, + "indexer_topk": 32, + }, + ) + @pytest.mark.tpu_only + def test_tpu_flash_attention_context_parallel_with_indexer( + self, context_parallel_load_balance, ici_context_parallelism=2, indexer_topk=256 + ): + """Test equivalence between dot_product MLA + Indexer and all-gather flash attention + context parallelism + Indexer""" + config_arguments = { + "per_device_batch_size": 1.0, + "run_name": "test", + "enable_checkpointing": False, + "max_target_length": 512, + "sa_block_q": 128, + "sa_block_kv": 128, + "sa_block_kv_compute": 128, + "sa_block_q_dkv": 128, + "sa_block_kv_dkv": 128, + "sa_block_kv_dkv_compute": 128, + "attention_type": AttentionType.MLA.value, + "use_indexer": True, + "indexer_loss_scaling_factor": 0.1, + "indexer_topk": indexer_topk, + "q_lora_rank": 4, + "kv_lora_rank": 4, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 64, + "v_head_dim": 128, + "dtype": "float32", + } + + cfg, mla = self.init_mla({**config_arguments, "attention": "dot_product"}, rope_type="default") + lnx, decoder_segment_ids, decoder_positions = self.get_structured_data(cfg, cfg.dtype) + mla_generic_output, _ = mla( + lnx, + lnx, + decoder_segment_ids=decoder_segment_ids, + inputs_positions=decoder_positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + generic_state = nnx.state(mla) + + cfg_cp = pyconfig.initialize( + [sys.argv[0], get_test_config_path()], + **config_arguments, + attention="flash", + rope_type=cfg.rope_type, + context_parallel_strategy="all_gather", + context_parallel_load_balance=context_parallel_load_balance, + ici_context_parallelism=ici_context_parallelism, + use_tokamax_splash=True, + use_jax_splash=False, + packing=False, + ) + devices_array_cp = maxtext_utils.create_device_mesh(cfg_cp) + mesh_cp = Mesh(devices_array_cp, cfg_cp.mesh_axes) + with nn_partitioning.axis_rules(cfg_cp.logical_axis_rules): + attention_as_mla_flash_cp = MLA( + 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, + inputs_q_shape=lnx.shape, + inputs_kv_shape=lnx.shape, + max_target_length=cfg_cp.max_target_length, + max_prefill_predict_length=cfg_cp.max_prefill_predict_length, + mesh=mesh_cp, + attention_kernel="flash", + dtype=cfg_cp.dtype, + dropout_rate=cfg_cp.dropout_rate, + attention_type=AttentionType(cfg_cp.attention_type), + q_lora_rank=cfg_cp.q_lora_rank, + kv_lora_rank=cfg_cp.kv_lora_rank, + qk_nope_head_dim=cfg_cp.qk_nope_head_dim, + qk_rope_head_dim=cfg_cp.qk_rope_head_dim, + v_head_dim=cfg_cp.v_head_dim, + model_mode=MODEL_MODE_PREFILL, + rngs=self.nnx_rng, + ) + nnx.update(attention_as_mla_flash_cp, generic_state) + + mla_generic_flash_cp_output = attention_test_util.forward_with_context_expert_parallelism( + cfg_cp, + mesh_cp, + attention_as_mla_flash_cp, + lnx, + decoder_segment_ids, + decoder_positions, + ) + + mla_generic_output = jax.device_get(mla_generic_output) + mla_generic_flash_cp_output = jax.device_get(mla_generic_flash_cp_output) + + self.assertTrue( + jax.numpy.allclose(mla_generic_output, mla_generic_flash_cp_output, rtol=1e-02, atol=1e-02, equal_nan=False), + msg=( + "MLA+Indexer logits from generic dot product and flash attention + all-gather context parallelism are not" + f" close. context_parallel_load_balance={context_parallel_load_balance}." + ), + ) + + @parameterized.named_parameters( + { + "testcase_name": "no_lb_cp2", + "context_parallel_load_balance": False, + "ici_context_parallelism": 2, + "indexer_topk": 256, + }, + { + "testcase_name": "lb_cp4_smallk", + "context_parallel_load_balance": True, + "ici_context_parallelism": 4, + "indexer_topk": 32, + }, + ) + @pytest.mark.tpu_only + def test_tpu_flash_attention_ring_context_parallel_with_indexer( + self, context_parallel_load_balance, ici_context_parallelism=2, indexer_topk=256 + ): + """Test equivalence between dot_product MLA + Indexer and flash attention + ring context parallelism + Indexer""" + config_arguments = { + "per_device_batch_size": 1.0, + "run_name": "test", + "enable_checkpointing": False, + "max_target_length": 512, + "sa_block_q": 128, + "sa_block_kv": 128, + "sa_block_kv_compute": 128, + "sa_block_q_dkv": 128, + "sa_block_kv_dkv": 128, + "sa_block_kv_dkv_compute": 128, + "attention_type": AttentionType.MLA.value, + "use_indexer": True, + "indexer_loss_scaling_factor": 0.1, + "indexer_topk": indexer_topk, + "q_lora_rank": 4, + "kv_lora_rank": 4, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 64, + "v_head_dim": 128, + "dtype": "float32", + } + + cfg, mla = self.init_mla({**config_arguments, "attention": "dot_product"}, rope_type="default") + lnx, decoder_segment_ids, decoder_positions = self.get_structured_data(cfg, cfg.dtype) + mla_generic_output, _ = mla( + lnx, + lnx, + decoder_segment_ids=decoder_segment_ids, + inputs_positions=decoder_positions, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + generic_state = nnx.state(mla) + + cfg_cp = pyconfig.initialize( + [sys.argv[0], get_test_config_path()], + **config_arguments, + attention="flash", + rope_type=cfg.rope_type, + context_parallel_strategy="ring", + context_parallel_load_balance=context_parallel_load_balance, + ici_context_parallelism=ici_context_parallelism, + use_tokamax_splash=True, + use_jax_splash=False, + packing=False, + ) + devices_array_cp = maxtext_utils.create_device_mesh(cfg_cp) + mesh_cp = Mesh(devices_array_cp, cfg_cp.mesh_axes) + with nn_partitioning.axis_rules(cfg_cp.logical_axis_rules): + attention_as_mla_flash_cp = MLA( + 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, + inputs_q_shape=lnx.shape, + inputs_kv_shape=lnx.shape, + max_target_length=cfg_cp.max_target_length, + max_prefill_predict_length=cfg_cp.max_prefill_predict_length, + mesh=mesh_cp, + attention_kernel="flash", + dtype=cfg_cp.dtype, + dropout_rate=cfg_cp.dropout_rate, + attention_type=AttentionType(cfg_cp.attention_type), + q_lora_rank=cfg_cp.q_lora_rank, + kv_lora_rank=cfg_cp.kv_lora_rank, + qk_nope_head_dim=cfg_cp.qk_nope_head_dim, + qk_rope_head_dim=cfg_cp.qk_rope_head_dim, + v_head_dim=cfg_cp.v_head_dim, + model_mode=MODEL_MODE_PREFILL, + rngs=self.nnx_rng, + ) + nnx.update(attention_as_mla_flash_cp, generic_state) + + mla_generic_flash_cp_output = attention_test_util.forward_with_context_expert_parallelism( + cfg_cp, + mesh_cp, + attention_as_mla_flash_cp, + lnx, + decoder_segment_ids, + decoder_positions, + ) + + mla_generic_output = jax.device_get(mla_generic_output) + mla_generic_flash_cp_output = jax.device_get(mla_generic_flash_cp_output) + + self.assertTrue( + jax.numpy.allclose(mla_generic_output, mla_generic_flash_cp_output, rtol=1e-02, atol=1e-02, equal_nan=False), + msg="MLA+Indexer logits from generic dot product and flash attention + ring context parallelism are not close. " + f"context_parallel_load_balance={context_parallel_load_balance}.", + ) + + @parameterized.named_parameters( + { + "testcase_name": "no_lb_cp2", + "context_parallel_load_balance": False, + "ici_context_parallelism": 2, + "indexer_topk": 256, + }, + { + "testcase_name": "lb_cp4_smallk", + "context_parallel_load_balance": True, + "ici_context_parallelism": 4, + "indexer_topk": 32, + }, + ) + @pytest.mark.tpu_only + def test_tpu_flash_attention_ring_context_parallel_grad_with_indexer( + self, context_parallel_load_balance, ici_context_parallelism=2, indexer_topk=256 + ): + """Test gradient equivalence between dot_product and flash attention + ring context parallelism with Indexer""" + config_arguments = { + "per_device_batch_size": 1.0, + "run_name": "test", + "enable_checkpointing": False, + "max_target_length": 512, + "sa_block_q": 128, + "sa_block_kv": 128, + "sa_block_kv_compute": 128, + "sa_block_q_dkv": 128, + "sa_block_kv_dkv": 128, + "sa_block_kv_dkv_compute": 128, + "attention_type": AttentionType.MLA.value, + "use_indexer": True, + "indexer_loss_scaling_factor": 0.1, + "indexer_topk": indexer_topk, + "q_lora_rank": 4, + "kv_lora_rank": 4, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 64, + "v_head_dim": 128, + "dtype": "float32", + } + + cfg, mla = self.init_mla({**config_arguments, "attention": "dot_product"}, rope_type="default") + lnx, decoder_segment_ids, decoder_positions = self.get_structured_data(cfg, cfg.dtype) + + cfg_cp = pyconfig.initialize( + [sys.argv[0], get_test_config_path()], + **config_arguments, + attention="flash", + rope_type=cfg.rope_type, + context_parallel_strategy="ring", + context_parallel_load_balance=context_parallel_load_balance, + ici_context_parallelism=ici_context_parallelism, + use_tokamax_splash=True, + use_jax_splash=False, + packing=False, + ) + devices_array_cp = maxtext_utils.create_device_mesh(cfg_cp) + mesh_cp = Mesh(devices_array_cp, cfg_cp.mesh_axes) + with nn_partitioning.axis_rules(cfg_cp.logical_axis_rules): + attention_as_mla_flash_cp = MLA( + 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, + inputs_q_shape=lnx.shape, + inputs_kv_shape=lnx.shape, + max_target_length=cfg_cp.max_target_length, + max_prefill_predict_length=cfg_cp.max_prefill_predict_length, + mesh=mesh_cp, + attention_kernel="flash", + dtype=cfg_cp.dtype, + dropout_rate=cfg_cp.dropout_rate, + attention_type=AttentionType(cfg_cp.attention_type), + q_lora_rank=cfg_cp.q_lora_rank, + kv_lora_rank=cfg_cp.kv_lora_rank, + qk_nope_head_dim=cfg_cp.qk_nope_head_dim, + qk_rope_head_dim=cfg_cp.qk_rope_head_dim, + v_head_dim=cfg_cp.v_head_dim, + model_mode=MODEL_MODE_PREFILL, + rngs=self.nnx_rng, + ) + nnx.update(attention_as_mla_flash_cp, nnx.state(mla)) + generic_graphdef, generic_state = nnx.split(mla) + ring_graphdef, ring_state = nnx.split(attention_as_mla_flash_cp) + + def generic_loss(lnx): + mla_merged = nnx.merge(generic_graphdef, generic_state) + output, _ = mla_merged( + 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 ring_loss(lnx): + if context_parallel_load_balance: + context_parallel_size = cfg_cp.ici_context_parallelism + lnx = max_utils.reorder_sequence(lnx, cp_size=context_parallel_size) + ring_decoder_segment_ids = max_utils.reorder_sequence(decoder_segment_ids, cp_size=context_parallel_size) + ring_decoder_positions = max_utils.reorder_sequence(decoder_positions, cp_size=context_parallel_size) + else: + ring_decoder_segment_ids = decoder_segment_ids + ring_decoder_positions = decoder_positions + ring_merged = nnx.merge(ring_graphdef, ring_state) + output, _ = ring_merged( + lnx, + lnx, + decoder_segment_ids=ring_decoder_segment_ids, + inputs_positions=ring_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): + ring_grad = jax.grad(ring_loss)(lnx) + generic_grad = jax.device_get(generic_grad) + ring_grad = jax.device_get(ring_grad) + + self.assertTrue( + jax.numpy.allclose(generic_grad, ring_grad, rtol=1e-02, atol=1e-06, equal_nan=False), + msg=( + "MLA+Indexer input gradients from generic dot product and flash attention + ring context parallelism are" + f" not close. context_parallel_load_balance={context_parallel_load_balance}." + ), + ) + def get_indexer_test_data(self, batch_size, q_len, kv_len, num_heads, head_dim): """Helper to generate random data for indexer tests.""" key_q, key_k, key_is = jax.random.split(self.rng, 3) diff --git a/tests/unit/configs_value_test.py b/tests/unit/configs_value_test.py index 89db8cd7b0..8c0fea2f3d 100644 --- a/tests/unit/configs_value_test.py +++ b/tests/unit/configs_value_test.py @@ -211,6 +211,32 @@ def test_tpu_tokamax_ring_config_validation_accepts_packed_load_balance(self): self.assertTrue(config.context_parallel_load_balance) self.assertTrue(config.packing) + def test_tpu_tokamax_ring_config_validation_accepts_indexer(self): + argv = [ + "", + _BASE_CONFIG_PATH, + "run_name=test", + "attention=flash", + "attention_type=mla", + "use_indexer=True", + "q_lora_rank=1", + "use_tokamax_splash=True", + "use_jax_splash=False", + "context_parallel_strategy=ring", + "context_parallel_load_balance=False", + "ici_context_parallelism=2", + "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.assertTrue(config.use_indexer) + self.assertEqual(config.attention_type, "mla") + def test_tpu_tokamax_ring_config_validation_rejects_unsupported_configs(self): base_args = [ "", @@ -262,7 +288,6 @@ def test_tpu_tokamax_ring_config_validation_rejects_unsupported_configs(self): ), (["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"), @@ -389,7 +414,7 @@ def test_tpu_ulysses_config_validation_rejects_unsupported_configs(self): (["context_sharding=expert"], [], "context_sharding"), (["use_ragged_attention=True"], [], "ragged attention"), (["attention_sink=True"], [], "attention sinks"), - (["use_indexer=True", "q_lora_rank=1"], [], "sparse indexer"), + (["use_indexer=True", "attention_type=mla", "q_lora_rank=1"], [], "sparse indexer"), (["use_chunked_prefill=True"], [], "chunked prefill"), (["moba=True"], [], "MoBA"), (["use_multimodal=True"], [], "multimodal"), @@ -554,6 +579,7 @@ def test_indexer_cutoff_threshold_remat_policy(self): _BASE_CONFIG_PATH, "run_name=test", "use_indexer=true", + "attention_type=mla", "q_lora_rank=1536", "attention=dot_product", "remat_policy=custom", @@ -568,6 +594,7 @@ def test_indexer_cutoff_threshold_remat_policy(self): _BASE_CONFIG_PATH, "run_name=test", "use_indexer=true", + "attention_type=mla", "q_lora_rank=1536", "attention=dot_product", "remat_policy=custom", diff --git a/tests/unit/tokamax_ring_attention_test.py b/tests/unit/tokamax_ring_attention_test.py index 8729d56928..cb5c301619 100644 --- a/tests/unit/tokamax_ring_attention_test.py +++ b/tests/unit/tokamax_ring_attention_test.py @@ -101,10 +101,12 @@ def __init__(self, q, kv): self.q = q self.kv = kv - def kernel(q, k, v, segment_ids): + def kernel(q, k, v, segment_ids, sinks, indexer_mask): captured["segment_ids_type"] = type(segment_ids) captured["q_segment_shape"] = segment_ids.q.shape captured["kv_segment_shape"] = segment_ids.kv.shape + captured["sinks"] = sinks + captured["indexer_mask"] = indexer_mask return q + k + v query = jnp.ones((1, 2, 4, 2)) @@ -126,6 +128,70 @@ def kernel(q, k, v, segment_ids): self.assertIs(captured["segment_ids_type"], RingSegmentIds) self.assertEqual(captured["q_segment_shape"], (4,)) self.assertEqual(captured["kv_segment_shape"], (4,)) + self.assertIsNone(captured["sinks"]) + self.assertIsNone(captured["indexer_mask"]) + + def test_call_ring_attention_threads_indexer_mask_without_segment_ids(self): + captured = {} + + def kernel(q, k, v, segment_ids, sinks, indexer_mask): + captured["has_indexer_mask"] = indexer_mask is not None + captured["indexer_mask_shape"] = indexer_mask.shape + return q + k + v + + query = jnp.ones((2, 2, 4, 2)) + key = jnp.ones((2, 2, 4, 2)) + value = jnp.ones((2, 2, 4, 2)) + indexer_mask = jnp.ones((2, 4, 4), dtype=jnp.bool_) + + out = tokamax_ring_attention.call_ring_attention( + query, + key, + value, + None, + None, + kernel, + indexer_mask=indexer_mask, + ) + + self.assertEqual(out.shape, query.shape) + self.assertTrue(captured["has_indexer_mask"]) + self.assertEqual(captured["indexer_mask_shape"], (4, 4)) + + def test_call_ring_attention_threads_indexer_mask_with_segment_ids(self): + captured = {} + + class RingSegmentIds: + + def __init__(self, q, kv): + self.q = q + self.kv = kv + + def kernel(q, k, v, segment_ids, sinks, indexer_mask): + captured["segment_ids_type"] = type(segment_ids) + captured["indexer_mask_shape"] = indexer_mask.shape + return q + k + v + + query = jnp.ones((2, 2, 4, 2)) + key = jnp.ones((2, 2, 4, 2)) + value = jnp.ones((2, 2, 4, 2)) + segment_ids = jnp.ones((2, 4), dtype=jnp.int32) + indexer_mask = jnp.ones((2, 4, 4), dtype=jnp.bool_) + + with mock.patch.object(tokamax_ring_attention.ring_attention_kernel, "SegmentIds", RingSegmentIds): + out = tokamax_ring_attention.call_ring_attention( + query, + key, + value, + segment_ids, + segment_ids, + kernel, + indexer_mask=indexer_mask, + ) + + self.assertEqual(out.shape, query.shape) + self.assertIs(captured["segment_ids_type"], RingSegmentIds) + self.assertEqual(captured["indexer_mask_shape"], (4, 4)) def test_with_sequence_axis_preserves_partition_spec_type(self): spec = jax.sharding.PartitionSpec("data", None, None, "model") diff --git a/tests/unit/train_nnx_test.py b/tests/unit/train_nnx_test.py index c2d642fba7..1f0df8cdd3 100644 --- a/tests/unit/train_nnx_test.py +++ b/tests/unit/train_nnx_test.py @@ -26,6 +26,7 @@ from flax import nnx import jax import jax.numpy as jnp +from maxtext.layers import nnx_scan import numpy as np from maxtext.common import train_state_nnx from maxtext.common.metric_logger import record_activation_metrics @@ -109,6 +110,48 @@ def __call__(self, decoder_input_tokens, decoder_positions, **kwargs): return out +from maxtext.layers.attention_mla import indexer_losses + + +class _MockIndexerLayer(nnx.Module): + + def __init__(self, rngs): + self.mock_val = nnx.Param(jnp.zeros(())) + + def __call__(self, carry): + self.sow(indexer_losses, "indexer_loss", self.mock_val.get_value()) + return carry + + +class _TinyDecoderIndexerLoss(_TinyDecoder): + """_TinyDecoder that also sows indexer_loss via a scanned layer.""" + + def __init__(self, vocab_size: int, hidden: int, rngs: nnx.Rngs): + super().__init__(vocab_size, hidden, rngs) + + self.layers = nnx_scan.create_scanned_layers( + _MockIndexerLayer, + length=2, + param_scan_axis=0, + metadata_axis_name="layer", + rngs=rngs, + ) + + # Overwrite the empty parameters generated with our mock test metrics! + _, params, other = nnx.split(self.layers, nnx.Param, ...) + params.mock_val.value = jnp.array([0.25, 0.75]) + nnx.update(self.layers, params, other) + + def __call__(self, decoder_input_tokens, decoder_positions, **kwargs): + out = super().__call__(decoder_input_tokens, decoder_positions, **kwargs) + + def apply_fn(module, carry): + return module(carry) + + nnx_scan.apply_scanned_layers(self.layers, carry=None, length=2, param_scan_axis=0, apply_fn=apply_fn) + return out + + def _make_data(batch=2, seq=4, vocab=8): return { "inputs": jnp.zeros((batch, seq), dtype=jnp.int32), @@ -183,6 +226,25 @@ def test_indexer_warmup_precedes_vocab_tiling(self): self.assertEqual(float(aux["xent_sum"]), 0.0) self.assertEqual(float(loss), 0.0) + def test_indexer_losses_harvested_and_injected_into_loss(self): + cfg = _Cfg() + cfg.use_indexer = True + cfg.indexer_sparse_training = True + cfg.indexer_loss_scaling_factor = 0.1 + model = _TinyDecoderIndexerLoss(cfg.vocab_size, hidden=4, rngs=nnx.Rngs(0)) + data = _make_data(batch=cfg.micro_batch_size_to_train_on, vocab=cfg.vocab_size) + + loss_without_indexer, _ = pre_train.loss_fn( + _TinyDecoder(cfg.vocab_size, hidden=4, rngs=nnx.Rngs(0)), cfg, data, None, None, is_train=True + ) + + loss, aux = pre_train.loss_fn(model, cfg, data, None, None, is_train=True) + expected_indexer_loss = 0.5 # mean of 0.25 and 0.75 + + self.assertTrue(jnp.isfinite(loss)) + self.assertAlmostEqual(float(aux["indexer_loss"]), expected_indexer_loss, places=5) + self.assertAlmostEqual(float(loss), float(loss_without_indexer) + expected_indexer_loss, places=5) + class TestTrainStepNNX(unittest.TestCase): """Cover the NNX branch of train_step (the diff_wrapper / nnx.update path)."""