From 65d55f2452428ed634562a98815bb8add7bc8219 Mon Sep 17 00:00:00 2001 From: Tianyu Gu Date: Wed, 2 Sep 2026 07:15:26 +0000 Subject: [PATCH] Disable the profiler on an empty profiling window instead of raising `Profiler.__init__` raised ValueError when the first profile step was not before the last one. That window is derived from the trainer's initial and max steps plus the profiler options, and in RL setups the PeftTrainer is constructed with initial_step == max_step, so any profiler_options at all aborted the run at construction time, well after the models had loaded. An empty profiling window is a misconfiguration of profiling, not of the training run. Warn with the derived step numbers and disable profiling, matching how the profiler already behaves without options or on a non-primary process. --- tests/sft/profiler_test.py | 21 +++++++++++++++++---- tunix/sft/profiler.py | 18 +++++++++++++++--- 2 files changed, 32 insertions(+), 7 deletions(-) diff --git a/tests/sft/profiler_test.py b/tests/sft/profiler_test.py index 5ed0cee7c..1f1c61648 100644 --- a/tests/sft/profiler_test.py +++ b/tests/sft/profiler_test.py @@ -289,20 +289,33 @@ def test_multiple_profiler_instances( profiler_steps=0, ), ) - def test_invalid_step_numbers( - self, max_step, initial_step, skip_first_n_steps, profiler_steps + @mock.patch.object(jax, 'process_index', return_value=0) + @mock.patch.object(jax.profiler, 'start_trace') + def test_invalid_step_numbers_disable_profiling( + self, + mock_start_trace, + _, + max_step, + initial_step, + skip_first_n_steps, + profiler_steps, ): profiler_options = profiler.ProfilerOptions( log_dir=self.log_dir, skip_first_n_steps=skip_first_n_steps, profiler_steps=profiler_steps, ) - with self.assertRaises(ValueError): - profiler.Profiler( + with self.assertLogs(level='WARNING') as logs: + p = profiler.Profiler( initial_step=initial_step, max_step=max_step, profiler_options=profiler_options, ) + self.assertTrue(any('Profiler disabled' in line for line in logs.output)) + self.assertTrue(p._do_not_profile) + + p.maybe_activate(initial_step + skip_first_n_steps) + mock_start_trace.assert_not_called() def test_profiler_with_max_num_hosts_supported(self): def dummy_start_trace( diff --git a/tunix/sft/profiler.py b/tunix/sft/profiler.py index 616548dfd..77a862ed0 100644 --- a/tunix/sft/profiler.py +++ b/tunix/sft/profiler.py @@ -70,10 +70,22 @@ def __init__( ) # We use >= instead of > because last_profile_step is step number + 1. if self._first_profile_step >= self._last_profile_step: - raise ValueError( - f"First profile step {self._first_profile_step} cannot be greater" - f" than the last profile step {self._last_profile_step}." + # An empty profiling window is a misconfiguration of the profiler, not + # of the training run: warn and skip profiling rather than abort a run + # that would otherwise train fine. + logging.warning( + "Profiler disabled: first profile step %d is not before the last" + " profile step %d (initial_step=%d, max_step=%d," + " skip_first_n_steps=%d, profiler_steps=%d).", + self._first_profile_step, + self._last_profile_step, + initial_step, + max_step, + profiler_options.skip_first_n_steps, + profiler_options.profiler_steps, ) + self._do_not_profile = True + return self._started_by_this_instance = False def _check_if_max_num_hosts_supported(self) -> bool: