From dfad69f232ffc5dbf33576474672a1132d711bb8 Mon Sep 17 00:00:00 2001 From: Jiahao Chen Zhou Date: Thu, 6 Aug 2026 19:57:49 +0000 Subject: [PATCH] Support Ring Attention with DeepSeek DSA Sparse Indexer - Add dynamic per-ring-step indexer mask slicing and injection to Tokamax Splash Attention forward and backward loops. - Fix backward dK/dV transposition (is_dkv=True) with .swapaxes(0, 1) to match hardware KV-major grid scheduling, resolving TPU network collective deadlocks. - Implement indexer_losses(nnx.Variable) subclass to cleanly bypass Flax NNX layer scan Intermediate filtering with zero blast radius. - Update train.py loss_fn to pop indexer_losses, harvest per-layer auxiliary KL losses, and inject them into the scalar optimization objective for backward gradient flow. --- src/maxtext/common/metric_logger.py | 9 + src/maxtext/configs/types.py | 2 - .../attention/tokamax_ring_attention.py | 41 +++- .../ring_attention_kernel.py | 145 ++++++++++-- src/maxtext/layers/attention_mla.py | 14 +- src/maxtext/layers/attention_op.py | 2 + src/maxtext/trainers/pre_train/train.py | 27 ++- tests/unit/attention_test.py | 219 ++++++++++++++++++ tests/unit/configs_value_test.py | 28 ++- tests/unit/tokamax_ring_attention_test.py | 68 +++++- tests/unit/train_nnx_test.py | 29 +++ 11 files changed, 536 insertions(+), 48 deletions(-) diff --git a/src/maxtext/common/metric_logger.py b/src/maxtext/common/metric_logger.py index d157fcaba5..1cb1ef530e 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: {indexer_l:.3f}") + 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 463a0db10e..7a50bf1bb6 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -3650,8 +3650,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: diff --git a/src/maxtext/kernels/attention/tokamax_ring_attention.py b/src/maxtext/kernels/attention/tokamax_ring_attention.py index 9fd031599a..cb80d11557 100644 --- a/src/maxtext/kernels/attention/tokamax_ring_attention.py +++ b/src/maxtext/kernels/attention/tokamax_ring_attention.py @@ -132,8 +132,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: @@ -304,6 +302,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( @@ -316,11 +315,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): @@ -352,15 +357,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 4e996516e4..f9f4337bf5 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,61 @@ 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] + if local_idx_mask.dtype != jnp.bool_: + local_idx_mask = jnp.isclose(local_idx_mask, 0.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 +171,33 @@ 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 + xs = jnp.arange(0, ring_axis_size) + + def body(carry, i): + 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 + if mask_4d is not None: + # Slice EXACTLY the current KV shard's mask block. + local_idx_mask = jax.lax.dynamic_slice_in_dim(mask_4d, current_kv_shard_idx, 1, axis=-2) + local_idx_mask = jnp.squeeze(local_idx_mask, axis=-2) + else: + local_idx_mask = None + 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 +242,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 +272,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 +303,29 @@ 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 = jax.lax.dynamic_slice_in_dim(mask_4d, step0_mask_idx, 1, axis=-2) + step0_mask = jnp.squeeze(step0_mask, axis=-2) + else: + mask_4d = None + step0_mask = None + + xs = jnp.arange(1, ring_axis_size) + + def compute_step(i: int, 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 +359,17 @@ 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, i): + if mask_4d is not None: + current_kv_shard_idx = (ring_axis_idx - i) % ring_axis_size + local_idx_mask = jax.lax.dynamic_slice_in_dim(mask_4d, current_kv_shard_idx, 1, axis=-2) + local_idx_mask = jnp.squeeze(local_idx_mask, axis=-2) + else: + local_idx_mask = None ( dq_accum, dk_accum, @@ -285,7 +384,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 +412,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 +432,7 @@ def body(carry, i: int): dv.astype(v.dtype), None, dsinks, + None, # indexer_mask ) @@ -344,6 +444,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 +489,7 @@ def _ring_attention_fwd( out, (logsumexp, max_logits) = _ring_attention_forward( fwd_mask_info, + indexer_mask, q, k, v, @@ -405,7 +507,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 +533,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 +571,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 +613,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, @@ -559,6 +664,7 @@ def _ring_attention( v, segment_ids, sinks, + indexer_mask, is_mqa=is_mqa, config=config, mask_value=mask_value, @@ -637,19 +743,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..274c08a14f 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 + """Variable type for storing Indexer loss components -> bypasses nnx.Intermediate scan filters.""" + + class Indexer(nnx.Module): """Indexer for DeepSeek Sparse Attention (DSA). @@ -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 e26d2d903f..efed247519 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -1453,6 +1453,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, ) ) else: @@ -1626,6 +1627,7 @@ def wrap_flash_attention( decoder_segment_ids_q, decoder_segment_ids_kv, splash_kernel, + indexer_mask=indexer_mask, ) return attention_output, None diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index 138bfe1ec7..d8fe7e8c48 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -46,6 +46,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, @@ -196,14 +197,18 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr decoder_target_tokens=data["targets"], decoder_target_mask=data["targets_segmentation"], ) - # mtp_losses and mtp_acceptance subclass nnx.Intermediate, and nnx type filters match - # subclasses. Pop them before the generic Intermediate pop below, which would otherwise - # take them too and leave the MTP loss silently reading as 0. + # Pop dedicated auxiliary variable types (MTP, Indexer) before the generic Intermediate pop below, + # ensuring auxiliary losses are cleanly harvested and stripped from the persistent model state. mtp_losses_state, mtp_acceptance_state = None, None if config.mtp_num_layers > 0: 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() @@ -213,6 +218,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. @@ -275,15 +283,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(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 diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index 6e9c612bd4..ad8207f214 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -2623,6 +2623,225 @@ def ring_loss(lnx): f"context_parallel_load_balance={context_parallel_load_balance}.", ) + @parameterized.named_parameters( + {"testcase_name": "no_load_balance", "context_parallel_load_balance": False}, + {"testcase_name": "load_balance", "context_parallel_load_balance": True}, + ) + @pytest.mark.tpu_only + def test_tpu_flash_attention_ring_context_parallel_with_indexer(self, context_parallel_load_balance): + """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": 256, + "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_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=2, + 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_load_balance", "context_parallel_load_balance": False}, + {"testcase_name": "load_balance", "context_parallel_load_balance": True}, + ) + @pytest.mark.tpu_only + def test_tpu_flash_attention_ring_context_parallel_grad_with_indexer(self, context_parallel_load_balance): + """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": 256, + "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_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=2, + 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 ce5e582390..27cb02574e 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 = [ "", @@ -261,7 +287,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"), @@ -361,6 +386,7 @@ def test_elastic_backup_kind_validation(self): ] with self.assertRaises(pydantic.ValidationError): pyconfig.initialize(argv) + def test_indexer_cutoff_threshold_remat_policy(self): """Tests custom remat policy and validation for indexer_cutoff_threshold.""" # 1. Verify custom remat policy puts indexer_cutoff_threshold on device diff --git a/tests/unit/tokamax_ring_attention_test.py b/tests/unit/tokamax_ring_attention_test.py index 54068ce669..8fdd6193b6 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..2741c1c023 100644 --- a/tests/unit/train_nnx_test.py +++ b/tests/unit/train_nnx_test.py @@ -109,6 +109,16 @@ def __call__(self, decoder_input_tokens, decoder_positions, **kwargs): return out +class _TinyDecoderIndexerLoss(_TinyDecoder): + """`_TinyDecoder` that also sows `indexer_loss` intermediates across its layers.""" + + def __call__(self, decoder_input_tokens, decoder_positions, **kwargs): + out = super().__call__(decoder_input_tokens, decoder_positions, **kwargs) + self.sow(nnx.Intermediate, "indexer_loss", jnp.array([0.25])) + self.sow(nnx.Intermediate, "indexer_loss", jnp.array([0.75])) + return out + + def _make_data(batch=2, seq=4, vocab=8): return { "inputs": jnp.zeros((batch, seq), dtype=jnp.int32), @@ -183,6 +193,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)."""