Skip to content

MaxText/Tunix: Expose enable_continuous_profiling CLI flag in MaxText and Tunix profiler to allow saving >2GB profile as xplane.riegeli #2035

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

MaxText/Tunix: Expose enable_continuous_profiling CLI flag in MaxText and Tunix profiler to allow saving >2GB profile as xplane.riegeli #2035

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 14 additions & 7 deletions tunix/sft/profiler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,13 +25,15 @@
@dataclasses.dataclass(frozen=True)
class ProfilerOptions:
"""Options for configuring the JAX profiler."""

# Directory to write the profile to.
log_dir: str
# Number of steps to skip before profiling.
skip_first_n_steps: int
# Number of steps to profile.
profiler_steps: int
# Whether to set the profile options.
# Whether to set tracer level options (host_tracer_level,
# python_tracer_level).
set_profile_options: bool = True
# https://github.com/jax-ml/jax/blob/0b1b909dd66a113ee0d7e54e55d0efef480e2a8a/docs/profiling.md?plain=1#L285
host_tracer_level: int = 2 # set to 2 to capture HBM profiles.
Expand All @@ -40,6 +42,8 @@ class ProfilerOptions:
# Maximum number of hosts to profile. Only supported for Pathways workloads
# using pathwaysutils.
max_num_hosts: int | None = None
# Whether to enable continuous profiling via JAX advanced_configuration.
enable_continuous_profiling: bool = False


class Profiler:
Expand Down Expand Up @@ -127,19 +131,22 @@ def maybe_activate(self, step: int):
)
return
logging.info("Starting JAX profiler at step %d.", step)
profile_options = jax.profiler.ProfileOptions()
if self._profiler_options.set_profile_options:
profile_options = jax.profiler.ProfileOptions()
profile_options.host_tracer_level = (
self._profiler_options.host_tracer_level
)
profile_options.python_tracer_level = (
self._profiler_options.python_tracer_level
)
self._start_trace(
log_dir=self._output_path, profiler_options=profile_options
)
else:
self._start_trace(log_dir=self._output_path)
if self._profiler_options.enable_continuous_profiling:
profile_options.advanced_configuration = {
"enable_continuous_profiling": True
}
self._start_trace(
log_dir=self._output_path, profiler_options=profile_options
)

Profiler._is_active = True
self._started_by_this_instance = True

Expand Down
Loading