Merge pull request #3387 from AI-Hypercomputer:xfgu-fixit

PiperOrigin-RevId: 882293460

Author: Google-ML-Automation

tests/unit/train_rl_test.py

@@ -28,10 +28,14 @@
 )
 
 
-def _get_mock_devices(num_devices):
-  mock_devices = [mock.MagicMock() for _ in range(num_devices)]
-  for i, d in enumerate(mock_devices):
-    d.id = i
+def _get_mock_devices(devices_per_slice, num_slices=1):
+  mock_devices = []
+  for slice_idx in range(num_slices):
+    for _ in range(devices_per_slice):
+      d = mock.MagicMock()
+      d.id = len(mock_devices)
+      d.slice_index = slice_idx
+      mock_devices.append(d)
   return mock_devices
 
 
@@ -93,6 +97,85 @@ def test_setup_configs_and_devices_pathways_fractional_split(self):
       self.assertEqual(trainer_devices, mock_devices[:2])
       self.assertEqual(sampler_devices, mock_devices[2:])
 
+  @pytest.mark.cpu_only
+  def test_setup_configs_and_devices_multislice_not_enough_slices(self):
+    """Test setup_configs_and_devices raises ValueError when not enough slices."""
+    mock_devices = _get_mock_devices(num_slices=2, devices_per_slice=4)
+    mock_config = SimpleNamespace(
+        num_trainer_slices=2,
+        num_samplers_slices=1,
+    )
+
+    def side_effect(argv, **kwargs):
+      res = SimpleNamespace(**vars(mock_config))
+      for k, v in kwargs.items():
+        setattr(res, k, v)
+      return res
+
+    with (
+        mock.patch.object(jax, "devices", return_value=mock_devices),
+        mock.patch(
+            "maxtext.trainers.post_train.rl.train_rl.pyconfig.initialize_pydantic",
+            side_effect=side_effect,
+        ),
+    ):
+      with self.assertRaisesRegex(ValueError, "Not enough slices for trainer and samplers"):
+        train_rl.setup_configs_and_devices(["dummy", "dummy"])
+
+  @pytest.mark.cpu_only
+  def test_setup_configs_and_devices_multislice_invalid_tp(self):
+    """Test setup_configs_and_devices raises ValueError for invalid TP."""
+    mock_devices = _get_mock_devices(num_slices=4, devices_per_slice=8)
+    mock_config = SimpleNamespace(
+        num_trainer_slices=2,
+        num_samplers_slices=2,
+        ici_tensor_parallelism=3,  # 8 is not divisible by 3
+        ici_fsdp_parallelism=-1,
+    )
+
+    def side_effect(argv, **kwargs):
+      res = SimpleNamespace(**vars(mock_config))
+      for k, v in kwargs.items():
+        setattr(res, k, v)
+      return res
+
+    with (
+        mock.patch.object(jax, "devices", return_value=mock_devices),
+        mock.patch(
+            "maxtext.trainers.post_train.rl.train_rl.pyconfig.initialize_pydantic",
+            side_effect=side_effect,
+        ),
+    ):
+      with self.assertRaisesRegex(ValueError, "must be divisible by tensor parallelism"):
+        train_rl.setup_configs_and_devices(["dummy", "dummy"])
+
+  @pytest.mark.cpu_only
+  def test_setup_configs_and_devices_multislice_invalid_tp_fsdp(self):
+    """Test setup_configs_and_devices raises ValueError for inconsistent TP and FSDP."""
+    mock_devices = _get_mock_devices(num_slices=4, devices_per_slice=8)
+    mock_config = SimpleNamespace(
+        num_trainer_slices=2,
+        num_samplers_slices=2,
+        ici_tensor_parallelism=4,
+        ici_fsdp_parallelism=3,  # 4 * 3 != 8
+    )
+
+    def side_effect(argv, **kwargs):
+      res = SimpleNamespace(**vars(mock_config))
+      for k, v in kwargs.items():
+        setattr(res, k, v)
+      return res
+
+    with (
+        mock.patch.object(jax, "devices", return_value=mock_devices),
+        mock.patch(
+            "maxtext.trainers.post_train.rl.train_rl.pyconfig.initialize_pydantic",
+            side_effect=side_effect,
+        ),
+    ):
+      with self.assertRaisesRegex(ValueError, "must equal devices_per_slice"):
+        train_rl.setup_configs_and_devices(["dummy", "dummy"])
+
   @pytest.mark.cpu_only
   def test_get_rollout_kwargs_no_dp(self):
     """Test case 1: sampler_config.rollout_data_parallelism=-1 -> verify result is calculated."""