Merge pull request AI-Hypercomputer#4252 from AI-Hypercomputer:sanbao/gpt

PiperOrigin-RevId: 940703754

Author: Google-ML-Automation

src/maxtext/common/goodput.py

@@ -66,31 +66,18 @@ def maybe_monitor_goodput(config):
         enable_gcp_goodput_metrics=config.enable_gcp_goodput_metrics,
         enable_gcp_step_deviation_metrics=config.enable_gcp_step_deviation_metrics,
     )
-    monitor_class = monitoring.GoodputMonitor
-
-    if config.elastic_enabled:
-      try:
-        from ml_goodput_measurement import monitoring_elastic  # pylint: disable=import-outside-toplevel
-
-        monitor_class = monitoring_elastic.ElasticGoodputMonitor
-      except ImportError:
-        max_logging.log("Elastic monitor failed!")
-
-    kwargs = {
-        "job_name": config.run_name,
-        "logger_name": f"goodput_{config.run_name}",
-        "tensorboard_dir": config.tensorboard_dir,
-        "upload_interval": config.goodput_upload_interval_seconds,
-        "monitoring_enabled": True,
-        "include_badput_breakdown": True,
-        "include_step_deviation": config.monitor_step_time_deviation,
-        "step_deviation_interval_seconds": config.step_deviation_interval_seconds,
-        "gcp_options": gcp_options,
-    }
-    if monitor_class == monitoring.GoodputMonitor:
-      kwargs["pathway_enabled"] = config.enable_pathways_goodput
-
-    goodput_monitor = monitor_class(**kwargs)
+    goodput_monitor = monitoring.GoodputMonitor(
+        job_name=config.run_name,
+        logger_name=f"goodput_{config.run_name}",
+        tensorboard_dir=config.tensorboard_dir,
+        upload_interval=config.goodput_upload_interval_seconds,
+        monitoring_enabled=True,
+        pathway_enabled=config.enable_pathways_goodput,
+        include_badput_breakdown=True,
+        include_step_deviation=config.monitor_step_time_deviation,
+        step_deviation_interval_seconds=config.step_deviation_interval_seconds,
+        gcp_options=gcp_options,
+    )
     goodput_monitor.start_goodput_uploader()
     max_logging.log("Started Goodput upload to Tensorboard & GCM in the background!")
     yield
@@ -134,23 +121,8 @@ def create_goodput_recorder(config):
     if config.enable_goodput_recording and jax.process_index() == 0:
       max_logging.log("[GOODPUT NO-OP] recorder skipped (decoupled stub).")
     return None
-
-  if not config.enable_goodput_recording:
-    return None
-
-  logger_name = f"goodput_{config.run_name}"
-
-  # Detect if we should use the elastic-aware recorder
-  if config.elastic_enabled:
-    try:
-      from ml_goodput_measurement import goodput_elastic  # pylint: disable=import-outside-toplevel
-      from maxtext.utils import elastic_utils  # pylint: disable=import-outside-toplevel
-
-      recorder = goodput_elastic.ElasticGoodputRecorder(config.run_name, logger_name, jax.process_index() == 0)
-      elastic_utils.record_slice_state(recorder)
-    except ImportError as e:
-      max_logging.log(f"Could not create elastic goodput recorder: {e}")
-    else:
-      return recorder
-
-  return goodput.GoodputRecorder(config.run_name, logger_name, jax.process_index() == 0)
+  if config.enable_goodput_recording:
+    logger_name = f"goodput_{config.run_name}"
+    recorder = goodput.GoodputRecorder(config.run_name, logger_name, jax.process_index() == 0)
+    return recorder
+  return None

src/maxtext/examples/chat_templates/gpt_oss_rl.json

@@ -0,0 +1,4 @@
+{
+  "SYSTEM_PROMPT": "You are given a problem. Think about the problem and provide your reasoning. Place it between {reasoning_start_token} and {reasoning_end_token}. Then, provide the final answer (i.e., just one numerical value) between {solution_start_token} and {solution_end_token}.",
+  "TEMPLATE": "<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI.\nKnowledge cutoff: 2024-06\nCurrent date: 2026-06-19\n\nReasoning: medium\n\n# Valid channels: analysis, commentary, final. Channel must be included for every message.\nCalls to these tools must go to the commentary channel: 'functions'.<|end|><|start|>user<|message|><start_of_turn>user\n{system_prompt}\n\n{question}<end_of_turn>\n<start_of_turn>model<|end|><|start|>assistant"
+}

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

@@ -226,9 +226,9 @@ def math_verify_pool(
 
   cpu_count = multiprocessing.cpu_count()
   if num_procs is None:
-    num_procs = min(_DEFAULT_MAX_PROCS, len(items), cpu_count)
+    num_procs = min(_DEFAULT_MAX_PROCS, cpu_count)
   else:
-    num_procs = max(1, min(num_procs, len(items), cpu_count))
+    num_procs = max(1, min(num_procs, cpu_count))
 
   cnt = 0
   pool = _get_pool(num_procs)

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

@@ -339,7 +339,7 @@ def _use_raw_prompt(x):
   dataset_size = int(trainer_config.num_batches * trainer_config.batch_size * trainer_config.train_fraction)
   train_dataset = train_dataset[:dataset_size]
   train_dataset = train_dataset.repeat(trainer_config.num_epoch)
-  train_dataset = train_dataset.to_iter_dataset().batch(trainer_config.batch_size)
+  train_dataset = train_dataset.to_iter_dataset().batch(trainer_config.batch_size, drop_remainder=True)
 
   if trainer_config.num_test_batches > 0:
     # eval_batch_size = -1 (default) → use trainer_config.batch_size (legacy
@@ -358,7 +358,7 @@ def _use_raw_prompt(x):
     test_dataset = test_dataset[
         trainer_config.test_batch_start_index : trainer_config.num_test_batches * eval_batch_size_for_eval
     ]
-    test_dataset = test_dataset.to_iter_dataset().batch(eval_batch_size_for_eval)
+    test_dataset = test_dataset.to_iter_dataset().batch(eval_batch_size_for_eval, drop_remainder=True)
 
   return train_dataset, test_dataset
 

src/maxtext/utils/elastic_utils.py

@@ -28,37 +28,13 @@
 pending_elastic_event_type = None
 
 
-def record_slice_state(recorder, active_slices_override: int | None = None) -> None:
-  """Records live slice counts and logs them to the GoodputRecorder."""
-  if (
-      recorder is None
-      or not hasattr(recorder, "record_elastic_slice_counts")
-      or not pathwaysutils.is_pathways_backend_used()
-      or elastic_manager is None
-  ):
-    return
-
-  available_slices = len(pathwaysutils.elastic.get_active_slice_indices())
-  active_slices = (
-      active_slices_override if active_slices_override is not None else len(elastic_manager.active_slice_indices)
-  )
-  total_slices = len(pathwaysutils.elastic.get_slice_to_devices(jax.devices()))
-
-  recorder.record_elastic_slice_counts(
-      available_slices=available_slices,
-      active_slices=active_slices,
-      total_slices=total_slices,
-  )
-
-
 def record_elastic_event_start(recorder, config) -> None:
   """Records start of an elastic scale up event."""
   global pending_elastic_event_type
   event_type = "elastic_scale_up" if is_scale_up_event(config) else "elastic_slice_down"
   pending_elastic_event_type = event_type
   if recorder:
-    recorder.record_elastic_wait_start_time(event_type=event_type)
-    record_slice_state(recorder, active_slices_override=0)
+    recorder.record_custom_badput_event_start_time(custom_badput_event_type=event_type)
 
 
 def record_elastic_wait_end_and_reinit_start(recorder) -> None:
@@ -69,18 +45,16 @@ def record_elastic_wait_end_and_reinit_start(recorder) -> None:
   event_type = pending_elastic_event_type
   pending_elastic_event_type = None
   if recorder:
-    recorder.record_elastic_wait_end_time(event_type=event_type)
-    recorder.record_elastic_reinit_start_time()
-    record_slice_state(recorder)
+    recorder.record_custom_badput_event_end_time(custom_badput_event_type=event_type)
+    recorder.record_custom_badput_event_start_time(custom_badput_event_type="elastic_reinitialization")
   pending_reinit_recorder = recorder
 
 
 def record_elastic_reinit_end() -> None:
   """Records end of elastic reinitialization event."""
   global pending_reinit_recorder
   if pending_reinit_recorder is not None:
-    pending_reinit_recorder.record_elastic_reinit_end_time()
-    record_slice_state(pending_reinit_recorder)
+    pending_reinit_recorder.record_custom_badput_event_end_time(custom_badput_event_type="elastic_reinitialization")
     pending_reinit_recorder = None
 
 

tests/post_training/unit/train_rl_test.py

@@ -359,6 +359,7 @@ def get_filtered_data_side_effect(dataset_name, model_tokenizer, template_config
         data_shuffle_seed=42,
         max_prefill_predict_length=10,
         batch_size=2,
+        eval_batch_size=1,
         num_batches=2,
         train_fraction=1.0,
         num_epoch=1,
@@ -422,7 +423,8 @@ def test_prepare_datasets_with_split(self, mock_load):
         data_template_path="maxtext/examples/chat_templates/gsm8k_rl.json",
         data_shuffle_seed=42,
         num_batches=1,
-        batch_size=5,
+        batch_size=2,
+        eval_batch_size=1,
         train_fraction=1.0,
         num_epoch=1,
         num_test_batches=1,

tests/unit/elastic_utils_test.py

@@ -70,12 +70,6 @@ def setUp(self):
 
     # Configure default behaviors if needed
     self.fake_pathwaysutils.is_pathways_backend_used.return_value = True
-    self.fake_pathwaysutils.elastic = Mock()
-    self.fake_pathwaysutils.elastic.get_active_slice_indices.return_value = [0, 1]
-    self.fake_pathwaysutils.elastic.get_slice_to_devices.return_value = {
-        0: [FakeDevice()],
-        1: [FakeDevice()],
-    }
     self.fake_jax.process_index.return_value = 0
 
     # Inject fakes into elastic_utils namespace
@@ -103,28 +97,9 @@ def tearDown(self):
     elastic_utils.pending_elastic_event_type = None
     super().tearDown()
 
-  def test_record_slice_state(self):
-    elastic_utils.elastic_manager = self.fake_manager
-    self.fake_manager.active_slice_indices = {0}
-    self.fake_manager.slice_to_devices = {0: [FakeDevice()], 1: [FakeDevice()]}
-    self.fake_pathwaysutils.elastic.get_active_slice_indices.return_value = {0, 1}
-
-    fake_recorder = Mock()
-    fake_recorder.record_elastic_slice_counts = Mock()
-
-    elastic_utils.record_slice_state(fake_recorder)
-
-    fake_recorder.record_elastic_slice_counts.assert_called_once_with(available_slices=2, active_slices=1, total_slices=2)
-
-    fake_recorder.record_elastic_slice_counts.reset_mock()
-    elastic_utils.record_slice_state(fake_recorder, active_slices_override=0)
-    fake_recorder.record_elastic_slice_counts.assert_called_once_with(available_slices=2, active_slices=0, total_slices=2)
-
   def test_elastic_enabled(self):
     config = FakeConfig()
     self.fake_pathwaysutils.is_pathways_backend_used.return_value = True
-    self.fake_pathwaysutils.elastic = Mock()
-    self.fake_pathwaysutils.elastic.get_active_slice_indices.return_value = [0, 1]
     config.elastic_enabled = True
     self.assertTrue(elastic_utils.elastic_enabled(config))
 
@@ -169,8 +144,6 @@ def test_live_devices_no_pathways(self):
   def test_live_devices_pathways(self):
     """Tests live_devices when pathways is used."""
     self.fake_pathwaysutils.is_pathways_backend_used.return_value = True
-    self.fake_pathwaysutils.elastic = Mock()
-    self.fake_pathwaysutils.elastic.get_active_slice_indices.return_value = [0, 1]
     device0 = FakeDevice(slice_index=0)
     device1 = FakeDevice(slice_index=1)
     self.fake_jax.devices.return_value = [device0, device1]
@@ -195,8 +168,6 @@ def test_live_devices_disabled(self):
   def test_elastic_retry_disabled(self):
     """Tests elastic_retry when disabled but pathways is used."""
     self.fake_pathwaysutils.is_pathways_backend_used.return_value = True
-    self.fake_pathwaysutils.elastic = Mock()
-    self.fake_pathwaysutils.elastic.get_active_slice_indices.return_value = [0, 1]
     config = FakeConfig()
     config.elastic_enabled = False
     msg = (
@@ -348,28 +319,28 @@ def test_record_elastic_event_start(self):
     """Tests recording an elastic slice down start."""
     elastic_utils.elastic_manager = self.fake_manager
     self.fake_manager.new_slice_event.is_set.return_value = False
-    self.fake_manager.slice_to_devices = {0: [FakeDevice()], 1: [FakeDevice()]}
     fake_recorder = Mock()
     config = FakeConfig()
 
     elastic_utils.record_elastic_event_start(fake_recorder, config)
 
-    fake_recorder.record_elastic_wait_start_time.assert_called_once_with(event_type="elastic_slice_down")
-    fake_recorder.record_elastic_slice_counts.assert_called_once()
+    fake_recorder.record_custom_badput_event_start_time.assert_called_once_with(
+        custom_badput_event_type="elastic_slice_down"
+    )
     self.assertEqual(elastic_utils.pending_elastic_event_type, "elastic_slice_down")
 
   def test_record_elastic_event_start_scale_up(self):
     """Tests recording an elastic slice scale up start."""
     elastic_utils.elastic_manager = self.fake_manager
     self.fake_manager.new_slice_event.is_set.return_value = True
-    self.fake_manager.slice_to_devices = {0: [FakeDevice()], 1: [FakeDevice()]}
     fake_recorder = Mock()
     config = FakeConfig()
 
     elastic_utils.record_elastic_event_start(fake_recorder, config)
 
-    fake_recorder.record_elastic_wait_start_time.assert_called_once_with(event_type="elastic_scale_up")
-    fake_recorder.record_elastic_slice_counts.assert_called_once()
+    fake_recorder.record_custom_badput_event_start_time.assert_called_once_with(
+        custom_badput_event_type="elastic_scale_up"
+    )
 
   def test_record_elastic_wait_end_and_reinit_start_noop_on_first_attempt(self):
     """Tests recording elastic event end and elastic reinit start."""
@@ -378,39 +349,36 @@ def test_record_elastic_wait_end_and_reinit_start_noop_on_first_attempt(self):
 
     elastic_utils.record_elastic_wait_end_and_reinit_start(fake_recorder)
 
-    fake_recorder.record_elastic_wait_end_time.assert_not_called()
-    fake_recorder.record_elastic_reinit_start_time.assert_not_called()
+    fake_recorder.record_custom_badput_event_end_time.assert_not_called()
+    fake_recorder.record_custom_badput_event_start_time.assert_not_called()
     self.assertIsNone(elastic_utils.pending_reinit_recorder)
 
   def test_record_elastic_wait_end_and_reinit_start(self):
     """Test recording end of slice down and start of reinit."""
     elastic_utils.pending_elastic_event_type = "elastic_slice_down"
-    elastic_utils.elastic_manager = self.fake_manager
-    self.fake_manager.active_slice_indices = {0}
-    self.fake_manager.slice_to_devices = {0: [FakeDevice()], 1: [FakeDevice()]}
-    elastic_utils.pending_elastic_event_type = "elastic_slice_down"
     fake_recorder = Mock()
 
     elastic_utils.record_elastic_wait_end_and_reinit_start(fake_recorder)
 
-    fake_recorder.record_elastic_wait_end_time.assert_called_once_with(event_type="elastic_slice_down")
-    fake_recorder.record_elastic_reinit_start_time.assert_called_once()
-    fake_recorder.record_elastic_slice_counts.assert_called_once()
+    fake_recorder.record_custom_badput_event_end_time.assert_called_once_with(
+        custom_badput_event_type="elastic_slice_down"
+    )
+    fake_recorder.record_custom_badput_event_start_time.assert_called_once_with(
+        custom_badput_event_type="elastic_reinitialization"
+    )
     self.assertIs(elastic_utils.pending_reinit_recorder, fake_recorder)
     self.assertIsNone(elastic_utils.pending_elastic_event_type)
 
   def test_record_elastic_reinit_end(self):
     """Tests recording end of elastic reinit."""
     fake_recorder = Mock()
     elastic_utils.pending_reinit_recorder = fake_recorder
-    elastic_utils.elastic_manager = self.fake_manager
-    self.fake_manager.active_slice_indices = {0}
-    self.fake_manager.slice_to_devices = {0: [FakeDevice()], 1: [FakeDevice()]}
 
     elastic_utils.record_elastic_reinit_end()
 
-    fake_recorder.record_elastic_reinit_end_time.assert_called_once()
-    fake_recorder.record_elastic_slice_counts.assert_called_once()
+    fake_recorder.record_custom_badput_event_end_time.assert_called_once_with(
+        custom_badput_event_type="elastic_reinitialization"
+    )
     self.assertIsNone(elastic_utils.pending_reinit_recorder)
 
   def test_record_elastic_reinit_end_on_cold_start(self):
@@ -433,8 +401,6 @@ def __setattr__(self, name, value):
 
     config = ReadOnlyConfig()
     self.fake_pathwaysutils.is_pathways_backend_used.return_value = True
-    self.fake_pathwaysutils.elastic = Mock()
-    self.fake_pathwaysutils.elastic.get_active_slice_indices.return_value = [0, 1]
 
     # Should not raise ValueError
     elastic_utils.ensure_elastic_manager_initialized(config)