From 11f80dfe61cbf9d5e6ef3f398996f51536b8bfd4 Mon Sep 17 00:00:00 2001 From: continuousml Date: Mon, 10 Aug 2026 17:18:09 -0700 Subject: [PATCH] Support packing for TPU Ulysses attention --- docs/guides/optimization/sharding.md | 2 +- src/maxtext/configs/types.py | 2 -- src/maxtext/utils/train_utils.py | 5 ++-- tests/unit/attention_test.py | 36 +++++++++++++++++----------- tests/unit/configs_value_test.py | 22 ++++++++++++++++- 5 files changed, 47 insertions(+), 20 deletions(-) diff --git a/docs/guides/optimization/sharding.md b/docs/guides/optimization/sharding.md index d0eedc3568..5927914995 100644 --- a/docs/guides/optimization/sharding.md +++ b/docs/guides/optimization/sharding.md @@ -260,7 +260,7 @@ Note in general there are many flavors of CP such as ring attention, which in th MaxText supports `context_parallel_strategy=all_gather`, and supports `context_parallel_strategy=ring` through GPU Transformer Engine and TPU Tokamax Splash paths; ring performs the computation and communication in chunks and ideally overlaps them in a collective matmul fashion. This strategy requires extending the online softmax trick from only within chip to additionally apply it across chips. -MaxText also supports `context_parallel_strategy=ulysses` ([DeepSpeed Ulysses](https://arxiv.org/abs/2309.14509)) on the TPU Tokamax Splash path for training. Ulysses exchanges sequence ownership for head ownership by communicating the Q, K, V, and output activations through all-to-all collectives: each device computes ordinary full-sequence attention for its head subset, and the inverse all-to-all restores the sequence sharding on the output. It requires explicit positive context parallelism values, `context_sharding=context`, `attention=flash` with Tokamax Splash, global causal attention, query and KV head counts divisible by the context parallel size including after tensor-parallel head sharding, matching Q and KV head-sharding axes, an unsharded head feature dimension, a divisible sequence length, `dq_reduction_steps` of 0 or 3, `context_parallel_load_balance=false` (each device computes full-sequence attention for its head subset, so the work is already balanced and the causal load-balancing reorder must stay off), and ICI-only context parallelism (`dcn_context_parallelism` must equal 1). It does not support MQA, packing, dropout, QK-Clip statistics, ragged attention, attention sinks, sparse indexer masks, chunked prefill, MoBA, or multimodal attention. +MaxText also supports `context_parallel_strategy=ulysses` ([DeepSpeed Ulysses](https://arxiv.org/abs/2309.14509)) on the TPU Tokamax Splash path for training. Ulysses exchanges sequence ownership for head ownership by communicating the Q, K, V, and output activations through all-to-all collectives: each device computes ordinary full-sequence attention for its head subset, and the inverse all-to-all restores the sequence sharding on the output. It requires explicit positive context parallelism values, `context_sharding=context`, `attention=flash` with Tokamax Splash, global causal attention, query and KV head counts divisible by the context parallel size including after tensor-parallel head sharding, matching Q and KV head-sharding axes, an unsharded head feature dimension, a divisible sequence length, `dq_reduction_steps` of 0 or 3, `context_parallel_load_balance=false` (each device computes full-sequence attention for its head subset, so the work is already balanced and the causal load-balancing reorder must stay off), and ICI-only context parallelism (`dcn_context_parallelism` must equal 1). It does not support MQA, dropout, QK-Clip statistics, ragged attention, attention sinks, sparse indexer masks, chunked prefill, MoBA, or multimodal attention. ### CP Arithmetic Intensity diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 68684cb38a..acaf519a3d 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -3740,8 +3740,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de raise ValueError("TPU Ulysses attention requires use_jax_splash=False.") if self.attention_type != "global": raise ValueError("TPU Ulysses attention is initially supported only for global causal attention.") - if self.packing: - raise ValueError("TPU Ulysses attention does not support packing yet.") if self.context_parallel_load_balance: raise ValueError( "TPU Ulysses attention requires context_parallel_load_balance=False: after the all-to-all every device " diff --git a/src/maxtext/utils/train_utils.py b/src/maxtext/utils/train_utils.py index b063ebdb0c..c1e43afef9 100644 --- a/src/maxtext/utils/train_utils.py +++ b/src/maxtext/utils/train_utils.py @@ -264,9 +264,10 @@ def create_train_state_fn(): # Validate context parallelism with packing configuration context_parallel_strategy = config.context_parallel_strategy.lower() if context_parallel_size > 1 and config.packing: - if context_parallel_strategy not in ("all_gather", "ring"): + if context_parallel_strategy not in ("all_gather", "ring", "ulysses"): raise ValueError( - "Context parallelism with sequence packing supports context_parallel_strategy='all_gather' or 'ring'." + "Context parallelism with sequence packing supports context_parallel_strategy='all_gather', 'ring', " + "or 'ulysses'." ) if ( config.hardware in ("gpu", "gpu_multiprocess") diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index 0395feb950..dc7ea03f43 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -1737,7 +1737,7 @@ def ring_loss(lnx): f"dq_reduction_steps={dq_reduction_steps}, ring_scan_unroll={ring_scan_unroll}, packing={packing}.", ) - def _ulysses_test_config(self, ici_context_parallelism): + def _ulysses_test_config(self, ici_context_parallelism, packing=False): return pyconfig.initialize( [sys.argv[0], get_test_config_path()], **self.config_arguments, @@ -1747,7 +1747,7 @@ def _ulysses_test_config(self, ici_context_parallelism): ici_context_parallelism=ici_context_parallelism, use_tokamax_splash=True, use_jax_splash=False, - packing=False, + packing=packing, dtype="float32", ) @@ -1788,17 +1788,21 @@ def _ulysses_test_modules(self, cfg_cp, mesh_cp, lnx): return attention_as_mha_generic, attention_as_mha_flash_cp @parameterized.named_parameters( - {"testcase_name": "ulysses_size_2", "ici_context_parallelism": 2}, - {"testcase_name": "ulysses_size_4", "ici_context_parallelism": 4}, + {"testcase_name": "ulysses_size_2", "ici_context_parallelism": 2, "packing": False}, + {"testcase_name": "ulysses_size_4", "ici_context_parallelism": 4, "packing": False}, + {"testcase_name": "ulysses_size_4_packed", "ici_context_parallelism": 4, "packing": True}, ) @pytest.mark.tpu_only - def test_tpu_flash_attention_ulysses_context_parallel(self, ici_context_parallelism): + def test_tpu_flash_attention_ulysses_context_parallel(self, ici_context_parallelism, packing): """Test equivalence between dot_product and flash attention + Ulysses context parallelism""" - cfg_cp = self._ulysses_test_config(ici_context_parallelism) + cfg_cp = self._ulysses_test_config(ici_context_parallelism, packing=packing) 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_as_mha_flash_cp = self._ulysses_test_modules(cfg_cp, mesh_cp, lnx) mha_generic_output, _ = attention_as_mha_generic( lnx, @@ -1825,21 +1829,25 @@ def test_tpu_flash_attention_ulysses_context_parallel(self, ici_context_parallel self.assertTrue( jax.numpy.allclose(mha_generic_output, mha_generic_flash_cp_output, rtol=1e-02, atol=1e-02, equal_nan=False), msg="Logits from generic dot product and flash attention + Ulysses context parallelism are not close. " - f"ici_context_parallelism={ici_context_parallelism}.", + f"ici_context_parallelism={ici_context_parallelism}, packing={packing}.", ) @parameterized.named_parameters( - {"testcase_name": "ulysses_size_2", "ici_context_parallelism": 2}, - {"testcase_name": "ulysses_size_4", "ici_context_parallelism": 4}, + {"testcase_name": "ulysses_size_2", "ici_context_parallelism": 2, "packing": False}, + {"testcase_name": "ulysses_size_4", "ici_context_parallelism": 4, "packing": False}, + {"testcase_name": "ulysses_size_4_packed", "ici_context_parallelism": 4, "packing": True}, ) @pytest.mark.tpu_only - def test_tpu_flash_attention_ulysses_context_parallel_grad(self, ici_context_parallelism): + def test_tpu_flash_attention_ulysses_context_parallel_grad(self, ici_context_parallelism, packing): """Test input-gradient equivalence between dot_product and flash attention + Ulysses context parallelism""" - cfg_cp = self._ulysses_test_config(ici_context_parallelism) + cfg_cp = self._ulysses_test_config(ici_context_parallelism, packing=packing) 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_as_mha_flash_cp = self._ulysses_test_modules(cfg_cp, mesh_cp, lnx) nnx.update(attention_as_mha_flash_cp, nnx.state(attention_as_mha_generic)) @@ -1874,7 +1882,7 @@ def ulysses_loss(lnx): self.assertTrue( jax.numpy.allclose(generic_grad, ulysses_grad, rtol=1e-02, atol=1e-07, equal_nan=False), msg="Input gradients from generic dot product and flash attention + Ulysses context parallelism are not " - f"close. ici_context_parallelism={ici_context_parallelism}.", + f"close. ici_context_parallelism={ici_context_parallelism}, packing={packing}.", ) @pytest.mark.tpu_only diff --git a/tests/unit/configs_value_test.py b/tests/unit/configs_value_test.py index bd1508ea59..89db8cd7b0 100644 --- a/tests/unit/configs_value_test.py +++ b/tests/unit/configs_value_test.py @@ -302,6 +302,27 @@ def test_tpu_ulysses_config_validation_accepts_initial_config(self): self.assertEqual(config.ici_context_parallelism, 4) self.assertFalse(config.context_parallel_load_balance) + def test_tpu_ulysses_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=ulysses", + "context_parallel_load_balance=False", + "ici_context_parallelism=4", + "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_context_parallel_strategy_is_normalized(self): argv = [ "", @@ -366,7 +387,6 @@ def test_tpu_ulysses_config_validation_rejects_unsupported_configs(self): ), (["ici_context_parallelism=1"], ["ici_context_parallelism=4"], "context_parallel_size > 1"), (["context_sharding=expert"], [], "context_sharding"), - (["packing=True", "dataset_type=tfds"], ["packing=False", "dataset_type=synthetic"], "packing"), (["use_ragged_attention=True"], [], "ragged attention"), (["attention_sink=True"], [], "attention sinks"), (["use_indexer=True", "q_lora_rank=1"], [], "sparse indexer"),