From 3afe7b3bd456616a1e1a0a2d4c163753452189a8 Mon Sep 17 00:00:00 2001 From: continuousml Date: Mon, 27 Jul 2026 01:47:30 -0700 Subject: [PATCH] Support packed TPU Tokamax ring attention --- src/maxtext/configs/types.py | 7 ++--- src/maxtext/layers/attention_op.py | 2 -- src/maxtext/utils/train_utils.py | 8 +++-- tests/unit/attention_test.py | 48 ++++++++++++++++++++++-------- tests/unit/configs_value_test.py | 44 ++++++++++++++++++++++++++- 5 files changed, 87 insertions(+), 22 deletions(-) diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index f8f1c7e021..595b6613d2 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -3450,8 +3450,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de 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.packing: - raise ValueError("TPU Tokamax ring attention does not support packing yet.") 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.") @@ -3474,9 +3472,8 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de if self.enable_dropout and self.dropout_rate > 0.0: raise ValueError("TPU Tokamax ring attention does not support dropout yet.") # STRIPED reorder strategy is a Transformer Engine feature and is GPU-only. - # The AUTO + packing case, which training resolves to STRIPED, is not - # validated here because test code paths may load the same config but use a - # different reorder path. Training's runtime path enforces this. + # AUTO is resolved in training because test code paths may load the same + # config but use a different reorder path. if ( context_parallel_size > 1 and "gpu" not in self.hardware diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index 8c40613266..12122871fa 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -534,8 +534,6 @@ def __init__( raise ValueError("TPU Tokamax ring attention requires use_tokamax_splash=True.") if self.config.use_jax_splash: 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.") diff --git a/src/maxtext/utils/train_utils.py b/src/maxtext/utils/train_utils.py index 2cbea04aa3..b602488ace 100644 --- a/src/maxtext/utils/train_utils.py +++ b/src/maxtext/utils/train_utils.py @@ -281,11 +281,15 @@ def create_train_state_fn(): with jax.set_mesh(mesh): if context_parallel_size > 1 and config.context_parallel_load_balance: - # Determine load balancing reorder strategy based on whether packing is enabled + # Determine load balancing reorder strategy. if config.context_parallel_reorder_strategy == ReorderStrategy.AUTO: reorder_strategy = ( ReorderStrategy.STRIPED - if config.packing and context_parallel_strategy == "ring" + if ( + config.packing + and context_parallel_strategy == "ring" + and config.hardware in ("gpu", "gpu_multiprocess") + ) else ReorderStrategy.DUAL_CHUNK_SWAP ) else: diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index 6c0d55c558..3655f717aa 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -656,6 +656,20 @@ def get_data(self, dtype): return lnx, decoder_segment_ids, decoder_positions + def get_packed_data(self, dtype): + """get packed data""" + lnx, _, _ = self.get_data(dtype) + # Uneven segment lengths so boundaries don't line up with splash blocks, reorder chunks, or the CP shard split. + segment_lengths = (80, 240, 112, 80) + segment_ids = jnp.concatenate( + [jnp.full((length,), segment, dtype=jnp.int32) for segment, length in enumerate(segment_lengths, start=1)] + ) + positions = jnp.concatenate([jnp.arange(length, dtype=jnp.int32) for length in segment_lengths]) + decoder_segment_ids = jnp.broadcast_to(segment_ids, (self.global_batch_size, self.max_target_length)) + decoder_positions = jnp.broadcast_to(positions, (self.global_batch_size, self.max_target_length)) + + return lnx, decoder_segment_ids, decoder_positions + def get_structured_data(self, dtype): """get structured data""" lnx = jax.random.normal( @@ -1143,11 +1157,13 @@ def test_tpu_flash_attention_packed_all_gather_context_parallel(self, context_pa ) @parameterized.named_parameters( - {"testcase_name": "no_load_balance", "context_parallel_load_balance": False}, - {"testcase_name": "load_balance", "context_parallel_load_balance": True}, + {"testcase_name": "no_load_balance", "context_parallel_load_balance": False, "packing": False}, + {"testcase_name": "load_balance", "context_parallel_load_balance": True, "packing": False}, + {"testcase_name": "packed", "context_parallel_load_balance": False, "packing": True}, + {"testcase_name": "packed_load_balance", "context_parallel_load_balance": True, "packing": True}, ) @pytest.mark.tpu_only - def test_tpu_flash_attention_ring_context_parallel(self, context_parallel_load_balance): + def test_tpu_flash_attention_ring_context_parallel(self, context_parallel_load_balance, packing): """Test equivalence between dot_product and flash attention + ring context parallelism""" cfg_cp = pyconfig.initialize( @@ -1159,12 +1175,15 @@ def test_tpu_flash_attention_ring_context_parallel(self, context_parallel_load_b ici_context_parallelism=2, use_tokamax_splash=True, use_jax_splash=False, - packing=False, + packing=packing, dtype="float32", ) devices_array_cp = maxtext_utils.create_device_mesh(cfg_cp) mesh_cp = Mesh(devices_array_cp, cfg_cp.mesh_axes) - lnx, decoder_segment_ids, decoder_positions = self.get_data(cfg_cp.dtype) + if packing: + lnx, decoder_segment_ids, decoder_positions = self.get_packed_data(cfg_cp.dtype) + else: + lnx, decoder_segment_ids, decoder_positions = self.get_data(cfg_cp.dtype) attention_as_mha_generic = Attention( config=self.cfg, num_query_heads=cfg_cp.num_query_heads, @@ -1224,15 +1243,17 @@ def test_tpu_flash_attention_ring_context_parallel(self, context_parallel_load_b self.assertTrue( jax.numpy.allclose(mha_generic_output, mha_generic_flash_cp_output, rtol=1e-02, atol=1e-02, equal_nan=False), msg="Logits from generic dot product and flash attention + ring context parallelism are not close. " - f"context_parallel_load_balance={context_parallel_load_balance}.", + f"context_parallel_load_balance={context_parallel_load_balance}, packing={packing}.", ) @parameterized.named_parameters( - {"testcase_name": "no_load_balance", "context_parallel_load_balance": False}, - {"testcase_name": "load_balance", "context_parallel_load_balance": True}, + {"testcase_name": "no_load_balance", "context_parallel_load_balance": False, "packing": False}, + {"testcase_name": "load_balance", "context_parallel_load_balance": True, "packing": False}, + {"testcase_name": "packed", "context_parallel_load_balance": False, "packing": True}, + {"testcase_name": "packed_load_balance", "context_parallel_load_balance": True, "packing": True}, ) @pytest.mark.tpu_only - def test_tpu_flash_attention_ring_context_parallel_grad(self, context_parallel_load_balance): + def test_tpu_flash_attention_ring_context_parallel_grad(self, context_parallel_load_balance, packing): """Test gradient equivalence between dot_product and flash attention + ring context parallelism""" cfg_cp = pyconfig.initialize( @@ -1244,12 +1265,15 @@ def test_tpu_flash_attention_ring_context_parallel_grad(self, context_parallel_l ici_context_parallelism=2, use_tokamax_splash=True, use_jax_splash=False, - packing=False, + packing=packing, dtype="float32", ) devices_array_cp = maxtext_utils.create_device_mesh(cfg_cp) mesh_cp = Mesh(devices_array_cp, cfg_cp.mesh_axes) - lnx, decoder_segment_ids, decoder_positions = self.get_data(cfg_cp.dtype) + if packing: + lnx, decoder_segment_ids, decoder_positions = self.get_packed_data(cfg_cp.dtype) + else: + lnx, decoder_segment_ids, decoder_positions = self.get_data(cfg_cp.dtype) attention_as_mha_generic = Attention( config=self.cfg, num_query_heads=cfg_cp.num_query_heads, @@ -1325,7 +1349,7 @@ def ring_loss(lnx): self.assertTrue( jax.numpy.allclose(generic_grad, ring_grad, rtol=1e-02, atol=1e-07, equal_nan=False), msg="Input gradients from generic dot product and flash attention + ring context parallelism are not close. " - f"context_parallel_load_balance={context_parallel_load_balance}.", + f"context_parallel_load_balance={context_parallel_load_balance}, packing={packing}.", ) @pytest.mark.tpu_only diff --git a/tests/unit/configs_value_test.py b/tests/unit/configs_value_test.py index 465adbc307..7367acfabe 100644 --- a/tests/unit/configs_value_test.py +++ b/tests/unit/configs_value_test.py @@ -143,6 +143,49 @@ 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_packing(self): + argv = [ + "", + _BASE_CONFIG_PATH, + "run_name=test", + "attention=flash", + "use_tokamax_splash=True", + "use_jax_splash=False", + "context_parallel_strategy=ring", + "context_parallel_load_balance=False", + "ici_context_parallelism=2", + "hardware=tpu", + "packing=True", + "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.packing) + + def test_tpu_tokamax_ring_config_validation_accepts_packed_load_balance(self): + argv = [ + "", + _BASE_CONFIG_PATH, + "run_name=test", + "attention=flash", + "use_tokamax_splash=True", + "use_jax_splash=False", + "context_parallel_strategy=ring", + "context_parallel_load_balance=True", + "ici_context_parallelism=2", + "hardware=tpu", + "packing=True", + "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.context_parallel_load_balance) + self.assertTrue(config.packing) + def test_tpu_tokamax_ring_config_validation_rejects_unsupported_configs(self): base_args = [ "", @@ -168,7 +211,6 @@ def test_tpu_tokamax_ring_config_validation_rejects_unsupported_configs(self): (["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"), - (["packing=True"], ["packing=False"], "packing"), ( [ "context_parallel_load_balance=True",