feat(goodput): add goodput_job_name config and instrument RL & Distillation trainers

src/dependencies/requirements/base_requirements/tpu-post-train-requirements.txt

@@ -16,6 +16,7 @@ llguidance<1.8.0
 loguru
 lxml
 mistral_common
+ml-goodput-measurement>=0.2.3
 model-hosting-container-standards
 openai
 openai-harmony

src/maxtext/configs/post_train/rl.yml

@@ -105,6 +105,10 @@ debug: false
 # If true, Tunix-managed metrics measurement will be enabled. The metrics will be
 # uploaded to tensorboard.
 enable_tunix_perf_metrics: false
+enable_goodput_recording: false
+goodput_job_name: ""
+monitor_goodput: false
+enable_pathways_goodput: false
 
 # ====== Training ======
 batch_size: 1

src/maxtext/trainers/post_train/distillation/train_distill.py

@@ -56,8 +56,18 @@
 from maxtext.layers.learn_to_init_layer import apply_lti_model_update
 from maxtext.optimizers import optimizers
 from maxtext.trainers.post_train.distillation import distillation_utils, lti_utils
+from maxtext.trainers.post_train.hooks import GoodputTrainingHooks
 from maxtext.utils import max_logging
 from maxtext.utils import maxtext_utils
+from maxtext.common.goodput import (
+    GoodputEvent,
+    RECORD_JOB_END_TIME,
+    RECORD_JOB_START_TIME,
+    create_goodput_recorder,
+    maybe_monitor_goodput,
+    maybe_record_goodput,
+    record_goodput,
+)
 from maxtext.utils import model_creation_utils
 
 # Tunix Imports
@@ -682,121 +692,146 @@ def train_distill(
         "Distillation requires matching vocabularies."
     )
 
-  # Build Training Components (No hardware context required)
-  strategy, optimizer, train_config = build_training_components(
-      student_config, teacher_config, is_offline, offline_data_dir
-  )
+  goodput_recorder = create_goodput_recorder(student_config)
+  record_goodput(goodput_recorder, RECORD_JOB_START_TIME)
+  _job_completed_gracefully = False
+  try:
+    with maybe_monitor_goodput(student_config):
+      _train_distill_impl(student_config, teacher_config, is_offline, offline_data_dir, goodput_recorder)
+      _job_completed_gracefully = True
+  finally:
+    if _job_completed_gracefully:
+      record_goodput(goodput_recorder, RECORD_JOB_END_TIME)
+
+
+def _train_distill_impl(
+    student_config: pyconfig.HyperParameters,
+    teacher_config: pyconfig.HyperParameters,
+    is_offline: bool = False,
+    offline_data_dir: str | None = None,
+    goodput_recorder: Any = None,
+) -> None:
+  """Internal distillation training implementation."""
+  with maybe_record_goodput(goodput_recorder, GoodputEvent.TPU_INIT):
+    # 1. Setup Mesh
+    devices = jax.devices()
+    devices_array = maxtext_utils.create_device_mesh(student_config, devices)
+    mesh = jax.sharding.Mesh(devices_array, student_config.mesh_axes)
+
+    # 2. Load Models
+    max_logging.log("Applying logical axis rules for model initialization and training...")
+    with jax.set_mesh(mesh), logical_axis_rules(student_config.logical_axis_rules):
+      if is_offline:
+        max_logging.log("Offline Distillation: Skipping Teacher Model loading.")
+        teacher_model = None
+      else:
+        max_logging.log(f"Loading Teacher from {teacher_config.load_parameters_path}...")
+        _log_config_details(teacher_config, "Teacher")
+        teacher_model = get_maxtext_model(teacher_config, mesh)
+        teacher_model.eval()
+
+      # LTI phase needs the student initialization step to know about the teacher configuration
+      student_config.get_keys()["teacher_config"] = teacher_config
+
+      max_logging.log(f"Loading Student from {student_config.load_parameters_path}...")
+      _log_config_details(student_config, "Student")
+      student_model = get_maxtext_model(student_config, mesh)
+      student_params_to_update = getattr(student_config, "student_params_to_update", []) or []
+      student_param_update_templates = [re.compile(t) for t in student_params_to_update]
+
+      def student_freeze_param_fn(path) -> bool:
+        path_str = "/".join(str(p) for p in path)
+        return not any(regex.search(path_str) for regex in student_param_update_templates)
+
+      # Inject the teacher's frozen weights into the student model
+      if teacher_model:
+        lti_utils.prepare_student_weights(
+            student_model,
+            teacher_model,
+            teacher_weights_copy_map=getattr(student_config, "distill_weights_copy_map", {}),
+            student_weights_share_map=getattr(student_config, "distill_student_weights_share_map", {}),
+        )
 
-  # 1. Setup Mesh
-  devices = jax.devices()
-  devices_array = maxtext_utils.create_device_mesh(student_config, devices)
-  mesh = jax.sharding.Mesh(devices_array, student_config.mesh_axes)
+      student_model.train()
+      model_bundle = ModelBundle(teacher_model, student_model)
 
   # Hardware Execution (Safe Context)
-  max_logging.log("Applying logical axis rules for model initialization and training...")
   with jax.set_mesh(mesh), logical_axis_rules(student_config.logical_axis_rules):
-    # 2. Load Models
-    if is_offline:
-      max_logging.log("Offline Distillation: Skipping Teacher Model loading.")
-      teacher_model = None
-    else:
-      max_logging.log(f"Loading Teacher from {teacher_config.load_parameters_path}...")
-      _log_config_details(teacher_config, "Teacher")
-      teacher_model = get_maxtext_model(teacher_config, mesh)
-      teacher_model.eval()
-
-    # LTI phase needs the student initialization step to know about the teacher configuration
-    student_config.get_keys()["teacher_config"] = teacher_config
-
-    max_logging.log(f"Loading Student from {student_config.load_parameters_path}...")
-    _log_config_details(student_config, "Student")
-    student_model = get_maxtext_model(student_config, mesh)
-    student_params_to_update = getattr(student_config, "student_params_to_update", []) or []
-    student_param_update_templates = [re.compile(t) for t in student_params_to_update]
-
-    def student_freeze_param_fn(path) -> bool:
-      path_str = "/".join(str(p) for p in path)
-      return not any(regex.search(path_str) for regex in student_param_update_templates)
-
-    # Inject the teacher's frozen weights into the student model
-    if teacher_model:
-      lti_utils.prepare_student_weights(
-          student_model,
-          teacher_model,
-          teacher_weights_copy_map=getattr(student_config, "distill_weights_copy_map", {}),
-          student_weights_share_map=getattr(student_config, "distill_student_weights_share_map", {}),
+    with maybe_record_goodput(goodput_recorder, GoodputEvent.TRAINING_PREPARATION):
+      # Build Training Components
+      strategy, optimizer, train_config = build_training_components(
+          student_config, teacher_config, is_offline, offline_data_dir
       )
 
-    student_model.train()
-    model_bundle = ModelBundle(teacher_model, student_model)
-
-    # 3. Initialize Trainer
-    trainer = MaxTextDistillationTrainer(
-        model=model_bundle,
-        strategy=strategy,
-        optimizer=optimizer,
-        training_config=train_config,
-        student_config=student_config,
-        teacher_config=teacher_config,
-        is_offline=is_offline,
-        student_freeze_param_filter=student_freeze_param_fn if student_params_to_update else None,
-    )
-    trainer.is_managed_externally = True
-    trainer._has_aux = True  # pylint: disable=protected-access
-
-    if is_offline:
-      max_logging.log("Initializing Data Iterators via MaxText pipeline...")
-
-      # Point Grain to offline files instead of the raw text dataset
-      student_config.get_keys()["grain_train_files"] = offline_data_dir
-
-      student_config.get_keys()["eval_interval"] = 0
-      student_config.get_keys()["is_offline_distillation"] = True
-      student_config.get_keys()["dataset_shuffle_buffer_size"] = 0
-      student_config.get_keys()["dataset_shuffle_seed"] = 0
-
-    raw_train_iter, raw_eval_iter = input_pipeline_interface.create_data_iterator(student_config, mesh)
-
-    # 5. Input Pipeline Checkpointing & Restoration
-    # Replace the default CheckpointManager with a Grain-aware one, which enables iterator checkpointing for grain datasets.
-    raw_train_iter = trainer.setup_checkpoint_manager_and_restore(raw_train_iter, student_config)
-
-    # Sync the ModelBundle step counter with the restored training step so that
-    # loss weight schedules resume from the correct position after checkpoint restore.
-    model_bundle.training_step.set_value(jnp.array(trainer._train_steps, dtype=jnp.int32))  # pylint: disable=protected-access
-
-    # 6. Configure Input Mapping
-    def custom_gen_model_input_fn(batch):
-      inputs_dict = {
-          "input_tokens": batch.input_tokens,
-          "positions": batch.positions,
-          "attention_mask": batch.input_mask,
-          "decoder_segment_ids": batch.decoder_segment_ids,
-          "targets": batch.targets,  # Passed to strategy (labels_fn)
-          "targets_position": batch.targets_position,  # Passed to strategy (labels_fn)
-          "targets_segmentation": batch.targets_segmentation,  # Passed to strategy (labels_fn)
-          "cache": None,
-      }
-      # If we are in online mode then we exit
-      if getattr(batch, "top_k_logits", None) is None:
-        return inputs_dict
-
-      inputs_dict["teacher_output"] = distillation_utils.DistillationForwardOutput(
-          logits=batch.top_k_logits, out_projection_activations=None, top_k_indices=batch.top_k_indices
+      # 3. Initialize Trainer
+      trainer = MaxTextDistillationTrainer(
+          model=model_bundle,
+          strategy=strategy,
+          optimizer=optimizer,
+          training_config=train_config,
+          student_config=student_config,
+          teacher_config=teacher_config,
+          is_offline=is_offline,
+          student_freeze_param_filter=student_freeze_param_fn if student_params_to_update else None,
       )
-      return inputs_dict
+      trainer.is_managed_externally = True
+      trainer._has_aux = True  # pylint: disable=protected-access
+
+      if is_offline:
+        max_logging.log("Initializing Data Iterators via MaxText pipeline...")
+
+        # Point Grain to offline files instead of the raw text dataset
+        student_config.get_keys()["grain_train_files"] = offline_data_dir
+
+        student_config.get_keys()["eval_interval"] = 0
+        student_config.get_keys()["is_offline_distillation"] = True
+        student_config.get_keys()["dataset_shuffle_buffer_size"] = 0
+        student_config.get_keys()["dataset_shuffle_seed"] = 0
+
+      raw_train_iter, raw_eval_iter = input_pipeline_interface.create_data_iterator(student_config, mesh)
+
+      # 5. Input Pipeline Checkpointing & Restoration
+      # Replace the default CheckpointManager with a Grain-aware one, which enables
+      # iterator checkpointing for grain datasets.
+      raw_train_iter = trainer.setup_checkpoint_manager_and_restore(raw_train_iter, student_config)
+
+      # Sync the ModelBundle step counter with the restored training step so that
+      # loss weight schedules resume from the correct position after checkpoint restore.
+      model_bundle.training_step.set_value(jnp.array(trainer._train_steps, dtype=jnp.int32))  # pylint: disable=protected-access
+
+      # 6. Configure Input Mapping
+      def custom_gen_model_input_fn(batch):
+        inputs_dict = {
+            "input_tokens": batch.input_tokens,
+            "positions": batch.positions,
+            "attention_mask": batch.input_mask,
+            "decoder_segment_ids": batch.decoder_segment_ids,
+            "targets": batch.targets,  # Passed to strategy (labels_fn)
+            "targets_position": batch.targets_position,  # Passed to strategy (labels_fn)
+            "targets_segmentation": batch.targets_segmentation,  # Passed to strategy (labels_fn)
+            "cache": None,
+        }
+        # If we are in online mode then we exit
+        if getattr(batch, "top_k_logits", None) is None:
+          return inputs_dict
+
+        inputs_dict["teacher_output"] = distillation_utils.DistillationForwardOutput(
+            logits=batch.top_k_logits, out_projection_activations=None, top_k_indices=batch.top_k_indices
+        )
+        return inputs_dict
 
-    trainer = trainer.with_gen_model_input_fn(custom_gen_model_input_fn)
+      trainer = trainer.with_gen_model_input_fn(custom_gen_model_input_fn)
+      trainer.with_training_hooks(GoodputTrainingHooks(student_config, goodput_recorder))
 
-    # 7. Create Iterator Wrappers (Use Utils)
-    train_iter = distillation_utils.MaxTextToTunixIterator(raw_train_iter)
+      # 7. Create Iterator Wrappers (Use Utils)
+      train_iter = distillation_utils.MaxTextToTunixIterator(raw_train_iter)
 
-    eval_iter = None
-    if raw_eval_iter is not None:
-      max_logging.log("Evaluation iterator successfully initialized.")
-      eval_iter = distillation_utils.MaxTextToTunixIterator(raw_eval_iter)
-    elif student_config.eval_interval > 0:
-      max_logging.log("Warning: eval_interval > 0 but create_data_iterator returned None for eval_iter.")
+      eval_iter = None
+      if raw_eval_iter is not None:
+        max_logging.log("Evaluation iterator successfully initialized.")
+        eval_iter = distillation_utils.MaxTextToTunixIterator(raw_eval_iter)
+      elif student_config.eval_interval > 0:
+        max_logging.log("Warning: eval_interval > 0 but create_data_iterator returned None for eval_iter.")
 
     # 8. Train
     max_logging.log("Starting Distillation Training...")

src/maxtext/trainers/post_train/hooks.py

@@ -38,14 +38,54 @@
 from maxtext.utils import sharding
 
 
-class BaseTrainingHooks(TrainingHooks, abc.ABC):
+class GoodputTrainingHooks(TrainingHooks):
+  """Lightweight training hooks for recording Goodput step start times."""
+
+  def __init__(self, config, goodput_recorder):
+    self.config = config
+    self.goodput_recorder = goodput_recorder
+
+  @override
+  def on_train_start(self, train_ctx: peft_trainer.PeftTrainer):
+    """Called at the beginning of training."""
+    del train_ctx
+
+  @override
+  def on_train_end(self, train_ctx: peft_trainer.PeftTrainer):
+    """Called at the end of training."""
+    del train_ctx
+
+  @override
+  def on_train_step_start(self, train_ctx: peft_trainer.PeftTrainer):
+    """Called at the beginning of a training step."""
+    if self.config.enable_goodput_recording:
+      record_goodput(self.goodput_recorder, f"record_{GoodputEvent.STEP.value}_start_time", train_ctx.train_steps)
+
+  @override
+  def on_train_step_end(
+      self, train_ctx: peft_trainer.PeftTrainer, train_step: int, train_loss: float, step_time: float = 0.0
+  ):
+    """Called at the end of a training step."""
+    del train_ctx, train_step, train_loss, step_time
+
+  @override
+  def on_eval_step_start(self, train_ctx: peft_trainer.PeftTrainer):
+    """Called at the beginning of an evaluation step."""
+    del train_ctx
+
+  @override
+  def on_eval_step_end(self, train_ctx: peft_trainer.PeftTrainer, eval_loss: float):
+    """Called at the end of an evaluation step."""
+    del train_ctx, eval_loss
+
+
+class BaseTrainingHooks(GoodputTrainingHooks, abc.ABC):
   """Shared training hooks for post-training."""
 
   def __init__(self, config, mesh, learning_rate_schedule, goodput_recorder):
-    self.config = config
+    super().__init__(config, goodput_recorder)
     self.mesh = mesh
     self.metric_logger = MetricLogger(self.config, learning_rate_schedule)
-    self.goodput_recorder = goodput_recorder
     self.metadata = {}
     self.train_metadata = defaultdict(float)
     self.eval_metadata = defaultdict(float)
@@ -90,8 +130,7 @@ def on_train_end(self, train_ctx: peft_trainer.PeftTrainer):  # pylint: disable=
   @override
   def on_train_step_start(self, train_ctx: peft_trainer.PeftTrainer):
     """Called at the beginning of a training step."""
-    if self.config.enable_goodput_recording:
-      record_goodput(self.goodput_recorder, f"record_{GoodputEvent.STEP.value}_start_time", train_ctx.train_steps)
+    super().on_train_step_start(train_ctx)
 
     # Calculate the number of non-padded tokens in the batch
     self.train_metadata[train_ctx.train_steps] = {