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: