Error out when num_vocab_tiling > 1 is configured for RL training

Author: shralex

src/maxtext/trainers/post_train/rl/train_rl.py

@@ -602,6 +602,12 @@ def _rl_train_impl(argv: Sequence[str], kwargs: dict):
       argv, kwargs
   )
 
+  if trainer_config.num_vocab_tiling > 1:
+    raise ValueError(
+        f"Vocab Tiling is not supported with RL. "
+        f"num_vocab_tiling was configured to {trainer_config.num_vocab_tiling}, but it must be 1 when running train_rl."
+    )
+
   # Create model tokenizer first so we can plumb its pad_id into the model
   # adapter (used to synthesize segment_ids that mask pad positions from
   # attention — without this the trainer attends to pad tokens and produces

tests/post_training/unit/train_rl_test.py

@@ -471,6 +471,17 @@ def test_prepare_datasets_without_split(self, mock_load):
     mock_load.assert_has_calls(expected_calls, any_order=True)
     assert mock_load.call_count == len(expected_calls)
 
+  @pytest.mark.cpu_only
+  @mock.patch("maxtext.trainers.post_train.rl.train_rl.model_creation_utils.setup_configs_and_devices")
+  def test_rl_train_invalid_vocab_tiling(self, mock_setup):
+    mock_config = SimpleNamespace(
+        num_vocab_tiling=2,
+    )
+    mock_setup.return_value = (mock_config, mock_config, [], [])
+
+    with self.assertRaisesRegex(ValueError, "Vocab Tiling is not supported with RL"):
+      train_rl._rl_train_impl([], {})
+
 
 if __name__ == "__main__":
   unittest.main()