Support load balancing for TPU USP attention

Author: huytransformer

docs/guides/optimization/sharding.md

@@ -262,7 +262,7 @@ MaxText supports `context_parallel_strategy=all_gather`, and supports `context_p
 
 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.
 
-MaxText also supports `context_parallel_strategy=usp` ([USP](https://arxiv.org/abs/2405.07719), Ulysses over ring) on the same TPU Tokamax Splash path for training. This initial support is non-load-balanced. To use it, set `context_parallel_strategy=usp`, `ici_context_parallelism` for the ring size, and `ici_context_usp_ulysses_parallelism` for the Ulysses size (total context parallelism is their product). USP factors the context parallelism into a ring dimension on the `context` mesh axis and a Ulysses dimension on the `context_usp_ulysses` mesh axis: the Ulysses all-to-all exchanges sequence ownership for head ownership over the Ulysses axis at each fixed ring position, and the ring kernel then rotates K and V across the ring axis inside each head subset. The strategy is hybrid-only: both dimensions must be greater than one, and the single-dimension endpoints are the existing `ring` and `ulysses` strategies. It shares the Ulysses restrictions (explicit positive ICI-only sizes, `attention=flash` with Tokamax Splash, global causal attention, head counts divisible by the Ulysses size, no MQA, no load balancing, no dropout, no multi-token prediction, no dKV megacore) and additionally requires `max_target_length` divisible by the total context parallelism and by the ring size squared.
+MaxText also supports `context_parallel_strategy=usp` ([USP](https://arxiv.org/abs/2405.07719), Ulysses over ring) on the same TPU Tokamax Splash path for training. To use it, set `context_parallel_strategy=usp`, `ici_context_parallelism` for the ring size, and `ici_context_usp_ulysses_parallelism` for the Ulysses size (total context parallelism is their product). USP factors the context parallelism into a ring dimension on the `context` mesh axis and a Ulysses dimension on the `context_usp_ulysses` mesh axis: the Ulysses all-to-all exchanges sequence ownership for head ownership over the Ulysses axis at each fixed ring position, and the ring kernel then rotates K and V across the ring axis inside each head subset. The strategy is hybrid-only: both dimensions must be greater than one, and the single-dimension endpoints are the existing `ring` and `ulysses` strategies. It shares the Ulysses restrictions (explicit positive ICI-only sizes, `attention=flash` with Tokamax Splash, global causal attention, head counts divisible by the Ulysses size, no MQA, no dropout, no multi-token prediction, no dKV megacore) and additionally requires `max_target_length` divisible by the total context parallelism and by the ring size squared; load balancing requires an even ring size.
 
 ### CP Arithmetic Intensity
 

src/maxtext/configs/types.py

@@ -2968,8 +2968,10 @@ def _validate_usp_context_parallelism(self):
       raise ValueError("TPU USP attention does not support sparse indexer masks.")
     if self.attention_type != "global":
       raise ValueError("TPU USP attention is initially supported only for global causal attention.")
-    if self.context_parallel_load_balance:
-      raise ValueError("TPU USP attention does not support context_parallel_load_balance=True.")
+    if self.context_parallel_load_balance and usp_ring_size % 2 != 0:
+      raise ValueError(
+          "TPU USP attention with context_parallel_load_balance=True requires an even ici_context_parallelism."
+      )
     if self.use_ragged_attention:
       raise ValueError("TPU USP attention does not support ragged attention.")
     if self.attention_sink:

tests/unit/attention_test.py

@@ -1943,13 +1943,13 @@ def attention_loss(x, pos, seg):
       )
       self.assertLen(hlo_test_utils.collective_lines(hlo_text, "collective-permute"), 0)
 
-  def _usp_test_config(self, packing=False):
+  def _usp_test_config(self, packing=False, context_parallel_load_balance=False):
     return pyconfig.initialize(
         [sys.argv[0], get_test_config_path()],
         **self.config_arguments,
         attention="flash",
         context_parallel_strategy="usp",
-        context_parallel_load_balance=False,
+        context_parallel_load_balance=context_parallel_load_balance,
         ici_context_parallelism=2,
         ici_context_usp_ulysses_parallelism=2,
         use_tokamax_splash=True,
@@ -1959,14 +1959,16 @@ def _usp_test_config(self, packing=False):
     )
 
   @parameterized.named_parameters(
-      {"testcase_name": "usp_2x2", "packing": False},
-      {"testcase_name": "usp_2x2_packed", "packing": True},
+      {"testcase_name": "usp_2x2", "context_parallel_load_balance": False, "packing": False},
+      {"testcase_name": "usp_2x2_load_balance", "context_parallel_load_balance": True, "packing": False},
+      {"testcase_name": "usp_2x2_packed", "context_parallel_load_balance": False, "packing": True},
+      {"testcase_name": "usp_2x2_packed_load_balance", "context_parallel_load_balance": True, "packing": True},
   )
   @pytest.mark.tpu_only
-  def test_tpu_flash_attention_usp_context_parallel(self, packing):
+  def test_tpu_flash_attention_usp_context_parallel(self, context_parallel_load_balance, packing):
     """Test equivalence between dot_product and flash attention + USP context parallelism"""
 
-    cfg_cp = self._usp_test_config(packing=packing)
+    cfg_cp = self._usp_test_config(packing=packing, context_parallel_load_balance=context_parallel_load_balance)
     devices_array_cp = maxtext_utils.create_device_mesh(cfg_cp)
     mesh_cp = Mesh(devices_array_cp, cfg_cp.mesh_axes)
     if packing:
@@ -1999,18 +2001,20 @@ def test_tpu_flash_attention_usp_context_parallel(self, packing):
     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 + USP context parallelism are not close. "
-        f"packing={packing}.",
+        f"context_parallel_load_balance={context_parallel_load_balance}, packing={packing}.",
     )
 
   @parameterized.named_parameters(
-      {"testcase_name": "usp_2x2", "packing": False},
-      {"testcase_name": "usp_2x2_packed", "packing": True},
+      {"testcase_name": "usp_2x2", "context_parallel_load_balance": False, "packing": False},
+      {"testcase_name": "usp_2x2_load_balance", "context_parallel_load_balance": True, "packing": False},
+      {"testcase_name": "usp_2x2_packed", "context_parallel_load_balance": False, "packing": True},
+      {"testcase_name": "usp_2x2_packed_load_balance", "context_parallel_load_balance": True, "packing": True},
   )
   @pytest.mark.tpu_only
-  def test_tpu_flash_attention_usp_context_parallel_grad(self, packing):
+  def test_tpu_flash_attention_usp_context_parallel_grad(self, context_parallel_load_balance, packing):
     """Test input-gradient equivalence between dot_product and flash attention + USP context parallelism"""
 
-    cfg_cp = self._usp_test_config(packing=packing)
+    cfg_cp = self._usp_test_config(packing=packing, context_parallel_load_balance=context_parallel_load_balance)
     devices_array_cp = maxtext_utils.create_device_mesh(cfg_cp)
     mesh_cp = Mesh(devices_array_cp, cfg_cp.mesh_axes)
     if packing:
@@ -2032,11 +2036,19 @@ def generic_loss(lnx):
       return jnp.mean(output.astype(jnp.float32) ** 2)
 
     def usp_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)
+        usp_decoder_segment_ids = max_utils.reorder_sequence(decoder_segment_ids, cp_size=context_parallel_size)
+        usp_decoder_positions = max_utils.reorder_sequence(decoder_positions, cp_size=context_parallel_size)
+      else:
+        usp_decoder_segment_ids = decoder_segment_ids
+        usp_decoder_positions = decoder_positions
       output, _ = attention_as_mha_flash_cp(
           lnx,
           lnx,
-          decoder_segment_ids=decoder_segment_ids,
-          inputs_positions=decoder_positions,
+          decoder_segment_ids=usp_decoder_segment_ids,
+          inputs_positions=usp_decoder_positions,
           deterministic=True,
           model_mode=MODEL_MODE_TRAIN,
       )
@@ -2051,7 +2063,7 @@ def usp_loss(lnx):
     self.assertTrue(
         jax.numpy.allclose(generic_grad, usp_grad, rtol=1e-02, atol=1e-07, equal_nan=False),
         msg="Input gradients from generic dot product and flash attention + USP context parallelism are not close. "
-        f"packing={packing}.",
+        f"context_parallel_load_balance={context_parallel_load_balance}, packing={packing}.",
     )
 
   @pytest.mark.tpu_only

tests/unit/configs_value_test.py

@@ -488,6 +488,29 @@ def test_tpu_usp_config_validation_accepts_packing(self):
 
     self.assertTrue(config.packing)
 
+  def test_tpu_usp_config_validation_accepts_load_balance(self):
+    argv = [
+        "",
+        _BASE_CONFIG_PATH,
+        "run_name=test",
+        "attention=flash",
+        "use_tokamax_splash=True",
+        "use_jax_splash=False",
+        "context_parallel_strategy=usp",
+        "context_parallel_load_balance=True",
+        "ici_context_parallelism=2",
+        "ici_context_usp_ulysses_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.context_parallel_load_balance)
+
   def test_context_usp_ulysses_parallelism_requires_usp(self):
     argv = [
         "",
@@ -522,7 +545,11 @@ def test_tpu_usp_config_validation_rejects_unsupported_configs(self):
         "skip_jax_distributed_system=True",
     ]
     cases = [
-        (["context_parallel_load_balance=True"], ["context_parallel_load_balance=False"], "load_balance"),
+        (
+            ["context_parallel_load_balance=True", "ici_context_parallelism=3"],
+            ["context_parallel_load_balance=False", "ici_context_parallelism=2"],
+            "even ici_context_parallelism",
+        ),
         (["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"),