Support packing for TPU Ulysses attention

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
 

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 "

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")

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

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"),