diff --git a/src/maxdiffusion/max_utils.py b/src/maxdiffusion/max_utils.py index 28b747894..6e272847c 100644 --- a/src/maxdiffusion/max_utils.py +++ b/src/maxdiffusion/max_utils.py @@ -85,7 +85,7 @@ def _jax_profiler_enabled(config): def _ml_diagnostics_profiler_enabled(config): - return "enable_ml_diagnostics" in config.get_keys() and config.enable_ml_diagnostics and jax.process_index() == 0 + return "enable_ml_diagnostics" in config.get_keys() and config.enable_ml_diagnostics def profiler_enabled(config): @@ -120,22 +120,27 @@ class Profiler: def __init__(self, config, session_name=None): self.config = config self.session_name = session_name + self.mld_xprof = None + self._active = None # "mld" | "jax" | None def start(self): - if _jax_profiler_enabled(self.config) and _ml_diagnostics_profiler_enabled(self.config): + use_mld = _ml_diagnostics_profiler_enabled(self.config) and xprof is not None + use_jax = _jax_profiler_enabled(self.config) + + if use_mld and use_jax: max_logging.log( - "Warning: Both ML Diagnostics profiler and JAX profiler are enabled. " - "This may cause increased overhead and duplicate profiling data. It " - "is recommended to enable only one profiler at a time for accurate " - "performance analysis." + "Both ML Diagnostics and JAX profilers are enabled. They share the same " + "underlying JAX tracer and cannot run concurrently. Using ML Diagnostics " + "and skipping the standalone JAX profiler." ) + use_jax = False - if _ml_diagnostics_profiler_enabled(self.config) and xprof is not None: + if use_mld: ensure_machinelearning_job_runs(self.config) self.mld_xprof = xprof() self.mld_xprof.start(self.session_name) - - if _jax_profiler_enabled(self.config): + self._active = "mld" + elif use_jax: log_dir = self.config.tensorboard_dir if log_dir.startswith("gs://"): log_dir = os.path.join("/tmp/profiler_traces", self.config.run_name) @@ -144,14 +149,13 @@ def start(self): os.makedirs(log_dir, exist_ok=True) max_logging.log(f"Starting profiler trace in: {log_dir}") jax.profiler.start_trace(log_dir) + self._active = "jax" def stop(self): - if _ml_diagnostics_profiler_enabled(self.config) and xprof is not None: - ensure_machinelearning_job_runs(self.config) + if self._active == "mld": if self.mld_xprof is not None: self.mld_xprof.stop() - - if _jax_profiler_enabled(self.config): + elif self._active == "jax": jax.profiler.stop_trace() trace_dir = self.config.tensorboard_dir @@ -171,6 +175,7 @@ def stop(self): blob = bucket.blob(blob_name) blob.upload_from_filename(local_file) max_logging.log(f"Uploaded {local_file} to gs://{bucket_name}/{blob_name}") + self._active = None def __enter__(self): self.start() diff --git a/src/maxdiffusion/tests/profiler_test.py b/src/maxdiffusion/tests/profiler_test.py index 7b7683b37..763be4ca4 100644 --- a/src/maxdiffusion/tests/profiler_test.py +++ b/src/maxdiffusion/tests/profiler_test.py @@ -63,6 +63,52 @@ def test_ml_diagnostics_profiler(self, mock_process_index, mock_xprof, mock_ml_r mock_xprof.return_value.stop.assert_called_once() + @patch("maxdiffusion.max_utils.machinelearning_run") + @patch("maxdiffusion.max_utils.xprof") + @patch("jax.process_index", return_value=1) + def test_ml_diagnostics_profiler_non_master_host( + self, mock_process_index, mock_xprof, mock_ml_run + ): + """Tests that ML Diagnostics profiler is also enabled on non-master hosts (process_index != 0).""" + config = MockConfig( + enable_ml_diagnostics=True, + profiler_gcs_path="gs://fake-bucket/profiler", + enable_ondemand_xprof=True, + run_name="test_run", + enable_profiler=False, + tensorboard_dir="/tmp/fake_tensorboard", + ) + + self.assertTrue(max_utils._ml_diagnostics_profiler_enabled(config)) + + with max_utils.Profiler(config, session_name="test_session"): + mock_xprof.return_value.start.assert_called_once_with("test_session") + + mock_xprof.return_value.stop.assert_called_once() + + @patch("jax.profiler.start_trace") + @patch("maxdiffusion.max_utils.machinelearning_run") + @patch("maxdiffusion.max_utils.xprof") + @patch("jax.process_index", return_value=0) + def test_both_profilers_enabled_prioritizes_mld( + self, mock_process_index, mock_xprof, mock_ml_run, mock_start_trace + ): + """Tests that when both ML Diagnostics and JAX profiler are enabled, ML Diagnostics is prioritized and JAX profiler is skipped.""" + config = MockConfig( + enable_ml_diagnostics=True, + profiler_gcs_path="gs://fake-bucket/profiler", + enable_ondemand_xprof=True, + run_name="test_run", + enable_profiler=True, + tensorboard_dir="/tmp/fake_tensorboard", + ) + + with max_utils.Profiler(config, session_name="test_session"): + mock_xprof.return_value.start.assert_called_once_with("test_session") + mock_start_trace.assert_not_called() + + mock_xprof.return_value.stop.assert_called_once() + @patch("jax.profiler.start_trace") @patch("jax.profiler.stop_trace") @patch("jax.process_index", return_value=0) diff --git a/src/maxdiffusion/train_utils.py b/src/maxdiffusion/train_utils.py index bbf44dfe6..e82174e92 100644 --- a/src/maxdiffusion/train_utils.py +++ b/src/maxdiffusion/train_utils.py @@ -61,6 +61,20 @@ def _validate_gcs_bucket_name(bucket_name, config_var): " Set either train_text_encoder, or cache_latents_text_encoder_outputs to False" ) + keys = config.get_keys() + if ( + "enable_profiler" in keys + and config.enable_profiler + and "enable_ml_diagnostics" in keys + and config.enable_ml_diagnostics + ): + max_logging.log( + "Both enable_profiler and enable_ml_diagnostics are set to True. They" + " share the underlying JAX tracer and cannot run concurrently." + " Prioritizing ML Diagnostics profiler and disabling standard JAX" + " profiler." + ) + def record_scalar_metrics(metrics, step_time_delta, per_device_tflops, lr): """Records scalar metrics to be written to tensorboard"""