diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 83aeea7e90..d9868c445a 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -3448,10 +3448,12 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de raise ValueError("TPU ring context parallelism requires use_tokamax_splash=True.") if self.use_jax_splash: raise ValueError("TPU ring context parallelism requires use_jax_splash=False.") - if self.attention_type != "global": - raise ValueError("TPU Tokamax ring attention is initially supported only for global causal attention.") + if self.attention_type not in ("global", "mla"): + raise ValueError("TPU Tokamax ring attention supports only attention_type='global' or 'mla'.") if self.packing: raise ValueError("TPU Tokamax ring attention does not support packing yet.") + if self.use_batch_split_schedule: + raise ValueError("TPU Tokamax ring attention does not support the DeepSeek batch-split schedule.") if self.context_parallel_load_balance: if context_parallel_size % 2 != 0: raise ValueError("TPU Tokamax ring load balancing requires an even context_parallel_size.") diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index cf076b4794..e1799ecd1c 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -536,8 +536,10 @@ def __init__( raise ValueError("TPU Tokamax ring attention requires use_jax_splash=False.") if self.config.packing: raise ValueError("TPU Tokamax ring attention does not support packing yet.") - if self.attention_type != AttentionType.GLOBAL: - raise ValueError("TPU Tokamax ring attention is initially supported only for global causal attention.") + if self.config.use_batch_split_schedule: + raise ValueError("TPU Tokamax ring attention does not support the DeepSeek batch-split schedule.") + if self.attention_type not in (AttentionType.GLOBAL, AttentionType.MLA): + raise ValueError("TPU Tokamax ring attention supports only attention_type='global' or 'mla'.") context_axis = self.config.context_sharding axis_names_q = self._logical_to_mesh_axes(self.flash_axis_names_q) diff --git a/src/maxtext/models/deepseek.py b/src/maxtext/models/deepseek.py index 53166f00ee..de20788e77 100644 --- a/src/maxtext/models/deepseek.py +++ b/src/maxtext/models/deepseek.py @@ -24,7 +24,7 @@ from jax.ad_checkpoint import checkpoint_name import jax.numpy as jnp from jax.sharding import Mesh -from maxtext.common.common_types import Config +from maxtext.common.common_types import Config, AttentionType from maxtext.common.common_types import HyperConnectionType, MODEL_MODE_PREFILL, DecoderBlockType from maxtext.layers import attention_mla from maxtext.layers import initializers @@ -149,7 +149,7 @@ def __init__( max_target_length=self.config.max_target_length, max_prefill_predict_length=self.config.max_prefill_predict_length, attention_kernel=self.config.attention, - attention_type=self.config.attention_type, + attention_type=AttentionType(self.config.attention_type), inputs_q_shape=self.dummy_inputs_shape, inputs_kv_shape=self.dummy_inputs_shape, mesh=mesh, diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index 6c0d55c558..6b53946901 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -2238,6 +2238,219 @@ def test_tpu_flash_attention_context_parallel( f" ici_expert_parallelism={ici_expert_parallelism}.", ) + @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(self, context_parallel_load_balance): + """Test equivalence between dot_product and flash attention + ring context parallelism""" + + 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, + "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, 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 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(self, context_parallel_load_balance): + """Test gradient equivalence between dot_product and flash attention + ring context parallelism""" + + 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, + "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, 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 input gradients from generic dot product and flash attention + ring context parallelism are not close. " + f"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 465adbc307..5b93d37312 100644 --- a/tests/unit/configs_value_test.py +++ b/tests/unit/configs_value_test.py @@ -143,6 +143,29 @@ def test_tpu_tokamax_ring_config_validation_accepts_load_balance(self): self.assertTrue(config.context_parallel_load_balance) + def test_tpu_tokamax_ring_config_validation_accepts_mla(self): + argv = [ + "", + _BASE_CONFIG_PATH, + "run_name=test", + "attention=flash", + "attention_type=mla", + "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.assertEqual(config.attention_type, "mla") + def test_tpu_tokamax_ring_config_validation_rejects_unsupported_configs(self): base_args = [ "", @@ -167,8 +190,12 @@ def test_tpu_tokamax_ring_config_validation_rejects_unsupported_configs(self): (["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"), - (["attention_type=full"], [], "global causal"), + (["attention_type=full"], [], "attention_type"), + (["attention_type=local_sliding", "sliding_window_size=128"], [], "attention_type"), + (["attention_type=chunk", "chunk_attn_window_size=128"], [], "attention_type"), + (["attention_type=compressed"], [], "attention_type"), (["packing=True"], ["packing=False"], "packing"), + (["use_batch_split_schedule=True"], [], "batch-split"), ( [ "context_parallel_load_balance=True",