Merge pull request #3341 from ROCm:ici-parallelism-PartialRotaryEmbeddingTest

PiperOrigin-RevId: 881501786

Author: Google-ML-Automation

tests/unit/partial_rotary_embedding_test.py

@@ -33,7 +33,7 @@
 from maxtext.layers.embeddings import PartialRotaryEmbedding, RotaryEmbedding
 from maxtext.configs import pyconfig
 from maxtext.utils import maxtext_utils
-from tests.utils.test_helpers import get_test_config_path
+from tests.utils.test_helpers import get_test_config_path, get_decoupled_parallelism_overrides
 
 
 class PartialRotaryEmbeddingTest(unittest.TestCase):
@@ -42,10 +42,12 @@ class PartialRotaryEmbeddingTest(unittest.TestCase):
   def setUp(self):
     super().setUp()
     # build a simple config and mesh like other embedding tests
+    extra_args = get_decoupled_parallelism_overrides()
     self.cfg = pyconfig.initialize(
         [sys.argv[0], get_test_config_path()],
         run_name="test_embeddings",
         enable_checkpointing=False,
+        **extra_args,
     )
     devices_array = maxtext_utils.create_device_mesh(self.cfg)
     self.mesh = Mesh(devices_array, self.cfg.mesh_axes)