diff --git a/src/dependencies/requirements/base_requirements/tpu-post-train-requirements.txt b/src/dependencies/requirements/base_requirements/tpu-post-train-requirements.txt index 1369939346..066e64f61b 100644 --- a/src/dependencies/requirements/base_requirements/tpu-post-train-requirements.txt +++ b/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 diff --git a/src/maxtext/configs/post_train/rl.yml b/src/maxtext/configs/post_train/rl.yml index 1841cd5752..4485736612 100644 --- a/src/maxtext/configs/post_train/rl.yml +++ b/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 diff --git a/src/maxtext/trainers/post_train/distillation/train_distill.py b/src/maxtext/trainers/post_train/distillation/train_distill.py index 7e390d72f0..46ef43e169 100644 --- a/src/maxtext/trainers/post_train/distillation/train_distill.py +++ b/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...") diff --git a/src/maxtext/trainers/post_train/hooks.py b/src/maxtext/trainers/post_train/hooks.py index 8dd5f4d1ad..36dd60e0d2 100644 --- a/src/maxtext/trainers/post_train/hooks.py +++ b/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] = { diff --git a/src/maxtext/trainers/post_train/rl/hooks.py b/src/maxtext/trainers/post_train/rl/hooks.py index 161db6be09..09bfdfe040 100644 --- a/src/maxtext/trainers/post_train/rl/hooks.py +++ b/src/maxtext/trainers/post_train/rl/hooks.py @@ -18,6 +18,7 @@ from tunix.sft import hooks as _tunix_hooks +from maxtext.common.goodput import GoodputEvent, record_goodput from maxtext.trainers.post_train.rl.evaluate_rl import evaluate from maxtext.utils import max_logging @@ -46,12 +47,14 @@ def __init__( test_dataset: Any, eval_interval: int, reward_fns: Optional[list[Callable[..., Any]]] = None, + goodput_recorder: Any = None, ): self._rl_cluster = rl_cluster self._trainer_config = trainer_config self._test_dataset = test_dataset self._eval_interval = eval_interval self._reward_fns = reward_fns + self._goodput_recorder = goodput_recorder self._last_step_evaluated = -1 # The five lifecycle methods below are abstract in `tunix.sft.hooks.TrainingHooks`, @@ -64,7 +67,9 @@ def on_train_end(self, train_ctx): # noqa: ARG002 del train_ctx def on_train_step_start(self, train_ctx): # noqa: ARG002 - del train_ctx + if getattr(self._trainer_config, "enable_goodput_recording", False): + # Ensure step is monotonically increasing across RL iterations (0, 1, 2, ...) + record_goodput(self._goodput_recorder, f"record_{GoodputEvent.STEP.value}_start_time", train_ctx.train_steps) def on_eval_step_start(self, train_ctx): # noqa: ARG002 del train_ctx @@ -75,6 +80,8 @@ def on_eval_step_end(self, train_ctx, *args, **kwargs): # noqa: ARG002 def on_train_step_end(self, trainer, step, loss): # noqa: ARG002 """Fire `evaluate(...)` once per `eval_interval` outer steps.""" del trainer, loss + if self._eval_interval <= 0 or getattr(self._trainer_config, "num_test_batches", 0) <= 0: + return try: outer_step = int(self._rl_cluster.global_steps) except Exception: # pylint: disable=broad-exception-caught diff --git a/src/maxtext/trainers/post_train/rl/train_rl.py b/src/maxtext/trainers/post_train/rl/train_rl.py index 21a751a01f..5306d587b7 100644 --- a/src/maxtext/trainers/post_train/rl/train_rl.py +++ b/src/maxtext/trainers/post_train/rl/train_rl.py @@ -77,6 +77,15 @@ from maxtext.common.common_types import DecoderBlockType from maxtext.configs import pyconfig, types +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.globals import MAXTEXT_CONFIGS_DIR from maxtext.integration.vllm.maxtext_vllm_rollout import MaxTextVllmRollout from maxtext.trainers.post_train.rl.evaluate_rl import evaluate @@ -648,17 +657,27 @@ def rl_train(argv: Sequence[str], kwargs: dict): trainer_devices: JAX devices for the trainer. sampler_devices: JAX devices for the sampler. """ - _rl_train_impl(argv, kwargs) - - -def _rl_train_impl(argv: Sequence[str], kwargs: dict): - """rl_train execution body.""" trainer_config, sampler_config, trainer_devices, sampler_devices = model_creation_utils.setup_configs_and_devices( argv, kwargs, config_class=types.RLConfig, ) + goodput_recorder = create_goodput_recorder(trainer_config) + record_goodput(goodput_recorder, RECORD_JOB_START_TIME) + _job_completed_gracefully = False + try: + with maybe_monitor_goodput(trainer_config): + _rl_train_impl(trainer_config, sampler_config, trainer_devices, sampler_devices, goodput_recorder) + _job_completed_gracefully = True + finally: + if _job_completed_gracefully: + record_goodput(goodput_recorder, RECORD_JOB_END_TIME) + +def _rl_train_impl( + trainer_config: Any, sampler_config: Any, trainer_devices: Any, sampler_devices: Any, goodput_recorder: Any = None +): + """rl_train execution body.""" # 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 @@ -669,13 +688,16 @@ def _rl_train_impl(argv: Sequence[str], kwargs: dict): ) configure_tokenizer_chat_template(model_tokenizer, trainer_config) - reference_model, reference_mesh, actor_model, actor_mesh, rollout_mesh = model_creation_utils.create_models_and_meshes( - trainer_config, - sampler_config, - trainer_devices, - sampler_devices, - tokenizer_pad_id=model_tokenizer.pad_token_id, - ) + with maybe_record_goodput(goodput_recorder, GoodputEvent.TPU_INIT): + reference_model, reference_mesh, actor_model, actor_mesh, rollout_mesh = ( + model_creation_utils.create_models_and_meshes( + trainer_config, + sampler_config, + trainer_devices, + sampler_devices, + tokenizer_pad_id=model_tokenizer.pad_token_id, + ) + ) if not trainer_config.debug: # Apply filter to suppress noisy logs @@ -714,17 +736,18 @@ def _rl_train_impl(argv: Sequence[str], kwargs: dict): max_logging.log(f"Policy mesh shape: {actor_mesh.shape}") max_logging.log(f"Rollout_mesh shape: {rollout_mesh.shape}") - rl_cluster, rl_trainer, _, reward_fns = create_rl_components( - trainer_config, - sampler_config, - sampler_devices, - actor_model, - actor_mesh, - reference_model, - reference_mesh, - rollout_mesh, - model_tokenizer, - ) + with maybe_record_goodput(goodput_recorder, GoodputEvent.TRAINING_PREPARATION): + rl_cluster, rl_trainer, _, reward_fns = create_rl_components( + trainer_config, + sampler_config, + sampler_devices, + actor_model, + actor_mesh, + reference_model, + reference_mesh, + rollout_mesh, + model_tokenizer, + ) # Run evaluation before training if trainer_config.num_test_batches > 0: @@ -757,7 +780,7 @@ def _rl_train_impl(argv: Sequence[str], kwargs: dict): # Wire intermediate eval: fire greedy `evaluate(...)` every `eval_interval` # outer steps. No-op when eval_interval <= 0 or num_test_batches <= 0. - utils_rl.install_training_hooks(rl_cluster, trainer_config, test_dataset, reward_fns) + utils_rl.install_training_hooks(rl_cluster, trainer_config, test_dataset, reward_fns, goodput_recorder=goodput_recorder) max_logging.warning("Starting RL training...") rl_trainer.train(train_dataset) diff --git a/src/maxtext/trainers/post_train/rl/utils_rl.py b/src/maxtext/trainers/post_train/rl/utils_rl.py index 0517ee07d1..b66eca1c3a 100644 --- a/src/maxtext/trainers/post_train/rl/utils_rl.py +++ b/src/maxtext/trainers/post_train/rl/utils_rl.py @@ -866,20 +866,22 @@ def install_training_hooks( trainer_config: Any, test_dataset: Any, reward_fns: Optional[list[Callable[..., Any]]] = None, + goodput_recorder: Any = None, ) -> None: - """Install maxtext's `RLTrainingHooks` on the actor trainer. + """Attach `RLTrainingHooks` to `rl_cluster.actor_trainer` if enabled. When `reward_fns` is provided, intermediate eval logs the per-example `mean_reward` alongside the correctness metrics, mirroring the training-time reward stack. - No-op if `eval_interval <= 0` or `num_test_batches <= 0` or tunix's hooks - module is unavailable. + No-op if both intermediate eval (`eval_interval > 0` and `num_test_batches > 0`) + and `enable_goodput_recording` are disabled, or if tunix's hooks module is unavailable. """ - if trainer_config.num_test_batches <= 0: - return eval_interval = int(getattr(trainer_config, "eval_interval", 0)) - if eval_interval <= 0: + num_test_batches = int(getattr(trainer_config, "num_test_batches", 0)) + eval_enabled = num_test_batches > 0 and eval_interval > 0 + goodput_enabled = bool(getattr(trainer_config, "enable_goodput_recording", False)) + if not eval_enabled and not goodput_enabled: return try: # `hooks` hard-imports `tunix.sft.hooks`. If that's missing (stock-only @@ -894,7 +896,9 @@ def install_training_hooks( try: actor = rl_cluster.actor_trainer if getattr(actor, "training_hooks", None) is None: - actor.training_hooks = RLTrainingHooks(rl_cluster, trainer_config, test_dataset, eval_interval, reward_fns) + actor.training_hooks = RLTrainingHooks( + rl_cluster, trainer_config, test_dataset, eval_interval, reward_fns, goodput_recorder=goodput_recorder + ) max_logging.warning( f"[intermediate-eval] hook installed: evaluate(...) will fire every {eval_interval} outer steps." ) diff --git a/tests/post_training/unit/hooks_test.py b/tests/post_training/unit/hooks_test.py index e7796382af..0ec06d0d87 100644 --- a/tests/post_training/unit/hooks_test.py +++ b/tests/post_training/unit/hooks_test.py @@ -33,7 +33,7 @@ from maxtext.configs import pyconfig from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR -from maxtext.trainers.post_train.hooks import BaseDataHooks, BaseTrainingHooks +from maxtext.trainers.post_train.hooks import BaseDataHooks, BaseTrainingHooks, GoodputTrainingHooks from maxtext.common.metric_logger import MetricLogger from maxtext.utils import exceptions from maxtext.utils import maxtext_utils @@ -151,6 +151,20 @@ def test_on_train_end_asserts_if_on_train_start_not_called(self): with self.assertRaises(AssertionError): self.training_hooks.on_train_end(self.mock_train_ctx) + def test_goodput_training_hooks(self): + mock_recorder = MagicMock() + mock_cfg = MagicMock() + mock_cfg.enable_goodput_recording = True + hooks = GoodputTrainingHooks(mock_cfg, mock_recorder) + self.mock_train_ctx.train_steps = 3 + hooks.on_train_start(self.mock_train_ctx) + hooks.on_train_step_start(self.mock_train_ctx) + hooks.on_train_step_end(self.mock_train_ctx, train_step=4, train_loss=1.5) + hooks.on_eval_step_start(self.mock_train_ctx) + hooks.on_eval_step_end(self.mock_train_ctx, eval_loss=1.2) + hooks.on_train_end(self.mock_train_ctx) + mock_recorder.record_step_start_time.assert_called_once_with(3) + @patch("maxtext.trainers.post_train.hooks.create_data_iterator") def test_load_next_train_batch_stop_iteration_returns_none(self, mock_create_data_iterator): mock_train_iter = MagicMock() diff --git a/tests/post_training/unit/train_distill_test.py b/tests/post_training/unit/train_distill_test.py index 845014db32..9c1c5e6b86 100644 --- a/tests/post_training/unit/train_distill_test.py +++ b/tests/post_training/unit/train_distill_test.py @@ -1139,6 +1139,8 @@ def test_main_offline_mode_skips_teacher_loading( mock_student_cfg.per_device_batch_size = 1 mock_student_cfg.max_target_length = 16 mock_student_cfg.gradient_accumulation_steps = 1 + mock_student_cfg.enable_goodput_recording = False + mock_student_cfg.monitor_goodput = False mock_teacher_cfg = mock.Mock() mock_teacher_cfg.vocab_size = 32000 @@ -1249,6 +1251,8 @@ def test_main_online_mode_loads_teacher( mock_student_cfg.per_device_batch_size = 1 mock_student_cfg.max_target_length = 16 mock_student_cfg.gradient_accumulation_steps = 1 + mock_student_cfg.enable_goodput_recording = False + mock_student_cfg.monitor_goodput = False mock_teacher_cfg = mock.Mock() mock_teacher_cfg.vocab_size = 32000 @@ -1404,6 +1408,106 @@ def test_save_run_manifest_swallows_errors(self): # Must not raise — failures here should not kill training. train_distill._save_run_manifest(argv, config) # pylint: disable=protected-access + def _run_train_distill_with_goodput(self, enable_goodput_recording: bool): + """Helper to run train_distill with mocked components and goodput_recorder.""" + devices = jax.devices()[:1] + mock_recorder = mock.MagicMock() + + mock_student_cfg = mock.Mock() + mock_student_cfg.vocab_size = 32000 + mock_student_cfg.mesh_axes = ("data",) + mock_student_cfg.dataset_type = "grain" + mock_student_cfg.data_sharding = ("fsdp",) + mock_student_cfg.learning_rate = 1e-4 + mock_student_cfg.warmup_steps_fraction = 0.1 + mock_student_cfg.learning_rate_final_fraction = 0.1 + mock_student_cfg.steps = 10 + mock_student_cfg.checkpoint_period = 5 + mock_student_cfg.gradient_clipping_threshold = 0.0 + mock_student_cfg.eval_interval = -1 + mock_student_cfg.gradient_accumulation_steps = 1 + mock_student_cfg.global_batch_size = 8 + mock_student_cfg.distill_temperature = 1.0 + mock_student_cfg.distill_alpha = 0.5 + mock_student_cfg.distill_beta = 0.0 + mock_student_cfg.distill_layer_indices = None + mock_student_cfg.distill_feature_loss_type = "cosine" + mock_student_cfg.use_sft = False + mock_student_cfg.enable_dropout = False + mock_student_cfg.learn_to_init_mode = False + mock_student_cfg.distill_weights_copy_map = {} + mock_student_cfg.distill_student_weights_share_map = {} + mock_student_cfg.get_keys.return_value = {} + mock_student_cfg.student_params_to_update = [] + mock_student_cfg.distill_alpha_end = None + mock_student_cfg.distill_alpha_schedule = "constant" + mock_student_cfg.distill_temperature_end = None + mock_student_cfg.distill_temperature_schedule = "constant" + mock_student_cfg.distill_beta_end = None + mock_student_cfg.distill_beta_schedule = "constant" + mock_student_cfg.max_num_checkpoints_to_keep = 1 + mock_student_cfg.async_checkpointing = False + mock_student_cfg.profiler = "none" + mock_student_cfg.tensorboard_dir = "" + mock_student_cfg.checkpoint_dir = "" + mock_student_cfg.log_period = 10 + mock_student_cfg.save_checkpoint_on_completion = False + mock_student_cfg.logical_axis_rules = [] + mock_student_cfg.enable_goodput_recording = enable_goodput_recording + mock_student_cfg.monitor_goodput = False + mock_student_cfg.goodput_job_name = "test-distill-job" + mock_student_cfg.run_name = "test-distill-run" + + mock_teacher_cfg = mock.Mock() + mock_teacher_cfg.vocab_size = 32000 + + def fake_create_recorder(cfg): + return mock_recorder if cfg.enable_goodput_recording else None + + with ( + mock.patch.object(train_distill, "create_goodput_recorder", side_effect=fake_create_recorder), + mock.patch.object(train_distill.maxtext_utils, "create_device_mesh", return_value=np.array(devices)), + mock.patch.object(train_distill.tokenizer, "build_tokenizer", return_value=mock.Mock(pad_id=0)), + mock.patch.object(train_distill, "get_maxtext_model", return_value=mock.Mock()), + mock.patch.object(train_distill.input_pipeline_interface, "create_data_iterator", return_value=(None, None)), + mock.patch.object(train_distill, "MaxTextDistillationTrainer") as mock_trainer_cls, + ): + mock_trainer_instance = mock_trainer_cls.return_value + mock_trainer_instance._train_steps = 0 + mock_trainer_instance.with_gen_model_input_fn.return_value = mock_trainer_instance + + def fake_train(train_iter, eval_iter): # pylint: disable=unused-argument + hooks_arg = mock_trainer_instance.with_training_hooks.call_args[0][0] + self.assertIsInstance(hooks_arg, train_distill.GoodputTrainingHooks) + hooks_arg.on_train_start(mock.Mock(train_steps=0)) + hooks_arg.on_train_step_start(mock.Mock(train_steps=0)) + hooks_arg.on_train_step_end(mock.Mock(train_steps=1), train_step=1, train_loss=1.0) + hooks_arg.on_train_step_start(mock.Mock(train_steps=1)) + hooks_arg.on_train_step_end(mock.Mock(train_steps=2), train_step=2, train_loss=0.8) + hooks_arg.on_train_end(mock.Mock(train_steps=2)) + + mock_trainer_instance.train.side_effect = fake_train + + train_distill.train_distill(mock_student_cfg, mock_teacher_cfg, is_offline=True, offline_data_dir="gs://dummy") + + return mock_recorder + + def test_train_distill_goodput_enabled(self): + """Verifies train_distill with enable_goodput_recording=True records lifecycle and step events.""" + mock_recorder = self._run_train_distill_with_goodput(enable_goodput_recording=True) + mock_recorder.record_job_start_time.assert_called_once() + mock_recorder.record_tpu_init_start_time.assert_called_once() + mock_recorder.record_tpu_init_end_time.assert_called_once() + mock_recorder.record_training_preparation_start_time.assert_called_once() + mock_recorder.record_training_preparation_end_time.assert_called_once() + self.assertEqual(mock_recorder.record_step_start_time.call_args_list, [mock.call(0), mock.call(1)]) + mock_recorder.record_job_end_time.assert_called_once() + + def test_train_distill_goodput_disabled(self): + """Verifies train_distill with enable_goodput_recording=False performs zero recorder calls.""" + mock_recorder = self._run_train_distill_with_goodput(enable_goodput_recording=False) + self.assertEqual(len(mock_recorder.method_calls), 0) + if __name__ == "__main__": absltest.main() diff --git a/tests/post_training/unit/train_rl_test.py b/tests/post_training/unit/train_rl_test.py index c620b0ef57..4899c44fed 100644 --- a/tests/post_training/unit/train_rl_test.py +++ b/tests/post_training/unit/train_rl_test.py @@ -1020,5 +1020,131 @@ def apply_chat_template(self, conversation, tokenize=False): self.assertEqual(rendered, "Hello!") +class RLGoodputTest(unittest.TestCase): + """Unit tests for Goodput recording in RL training and hooks.""" + + def test_rl_training_hooks_goodput_with_eval_interval_zero(self): + """Verify RLTrainingHooks records step start time and does not raise ZeroDivisionError when eval_interval=0.""" + from maxtext.trainers.post_train.rl import hooks as rl_hooks # pylint: disable=import-outside-toplevel + + mock_recorder = mock.MagicMock() + trainer_config = SimpleNamespace( + enable_goodput_recording=True, + eval_interval=0, + num_test_batches=0, + num_eval_passes=1, + eval_corr_lst=False, + eval_make_lst=False, + ) + rl_cluster = SimpleNamespace(global_steps=0, actor_trainer=SimpleNamespace(training_hooks=None)) + + train_rl.utils_rl.install_training_hooks( + rl_cluster, trainer_config, test_dataset=None, goodput_recorder=mock_recorder + ) + hook = rl_cluster.actor_trainer.training_hooks + self.assertIsInstance(hook, rl_hooks.RLTrainingHooks) + + # Simulate step 0 and step 1 start/end + train_ctx = SimpleNamespace(train_steps=0) + hook.on_train_step_start(train_ctx) + hook.on_train_step_end(trainer=train_ctx, step=1, loss=0.5) # Must not raise ZeroDivisionError + + train_ctx.train_steps = 1 + hook.on_train_step_start(train_ctx) + hook.on_train_step_end(trainer=train_ctx, step=2, loss=0.4) + + self.assertEqual(mock_recorder.record_step_start_time.call_args_list, [mock.call(0), mock.call(1)]) + + def _run_rl_train_with_goodput(self, enable_goodput_recording: bool): + """Helper to run rl_train with mocked components and goodput_recorder.""" + mock_recorder = mock.MagicMock() + trainer_config = SimpleNamespace( + enable_goodput_recording=enable_goodput_recording, + monitor_goodput=False, + goodput_job_name="test-rl-job", + run_name="test-rl-run", + debug=SimpleNamespace(rl=False), + load_parameters_path="gs://dummy/ckpt", + tokenizer_path="dummy_tokenizer", + hf_access_token=None, + checkpoint_dir="/tmp/rl_ckpt", + tensorboard_dir="/tmp/rl_tb", + profiler="", + steps=1, + batch_size=1, + num_batches=1, + num_test_batches=0, + eval_interval=0, + num_eval_passes=1, + eval_corr_lst=False, + eval_make_lst=False, + train_fraction=1.0, + load_checkpoint_only_once=False, + save_checkpoint_on_completion=False, + enable_single_controller=False, + reward_functions_path="", + reward_functions="", + rl=SimpleNamespace(use_agentic_rollout=False), + ) + sampler_config = SimpleNamespace(debug=SimpleNamespace(rl=False)) + + def fake_create_recorder(cfg): + return mock_recorder if cfg.enable_goodput_recording else None + + mock_rl_cluster = SimpleNamespace( + global_steps=0, actor_trainer=SimpleNamespace(training_hooks=None), close=mock.MagicMock() + ) + mock_rl_trainer = mock.MagicMock() + + def fake_train(dataset): # pylint: disable=unused-argument + hook = mock_rl_cluster.actor_trainer.training_hooks + if hook is not None: + hook.on_train_step_start(SimpleNamespace(train_steps=0)) + hook.on_train_step_end(trainer=None, step=1, loss=1.0) + + mock_rl_trainer.train.side_effect = fake_train + + with ( + mock.patch.object(train_rl, "create_goodput_recorder", side_effect=fake_create_recorder), + mock.patch.object( + train_rl.model_creation_utils, + "setup_configs_and_devices", + return_value=(trainer_config, sampler_config, [mock.MagicMock()], [mock.MagicMock()]), + ), + mock.patch.object( + train_rl.model_creation_utils, + "create_models_and_meshes", + return_value=(mock.MagicMock(), mock.MagicMock(), mock.MagicMock(), mock.MagicMock(), mock.MagicMock()), + ), + mock.patch.object(train_rl.epath, "Path", return_value=mock.MagicMock()), + mock.patch.object(train_rl.AutoTokenizer, "from_pretrained", return_value=mock.MagicMock()), + mock.patch.object(train_rl, "configure_tokenizer_chat_template"), + mock.patch.object(train_rl, "prepare_datasets", return_value=(["batch"], [])), + mock.patch.object( + train_rl, "create_rl_components", return_value=(mock_rl_cluster, mock_rl_trainer, mock.MagicMock(), []) + ), + mock.patch.object(train_rl.jax, "effects_barrier"), + ): + train_rl.rl_train(["dummy"], {}) + + return mock_recorder + + def test_rl_train_goodput_enabled_records_lifecycle_and_steps(self): + """Verify rl_train with enable_goodput_recording=True records all lifecycle and step events.""" + mock_recorder = self._run_rl_train_with_goodput(enable_goodput_recording=True) + mock_recorder.record_job_start_time.assert_called_once() + mock_recorder.record_tpu_init_start_time.assert_called_once() + mock_recorder.record_tpu_init_end_time.assert_called_once() + mock_recorder.record_training_preparation_start_time.assert_called_once() + mock_recorder.record_training_preparation_end_time.assert_called_once() + mock_recorder.record_step_start_time.assert_called_once_with(0) + mock_recorder.record_job_end_time.assert_called_once() + + def test_rl_train_goodput_disabled_performs_zero_recorder_calls(self): + """Verify rl_train with enable_goodput_recording=False performs zero recorder calls.""" + mock_recorder = self._run_rl_train_with_goodput(enable_goodput_recording=False) + self.assertEqual(len(mock_recorder.method_calls), 0) + + if __name__ == "__main__": unittest.main()