Skip to content

[observability] Profile Neuron devices with the PrivateUse1 activity #6

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

Closed
kvshbg-aws wants to merge 1 commit into from
Closed

[observability] Profile Neuron devices with the PrivateUse1 activity #6

Show file tree
Hide file tree
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
34 changes: 34 additions & 0 deletions tests/unit_tests/cpu/observability/test_profiler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -170,6 +170,40 @@ def test_build_torch_profiler_returns_active_handle(self):
with profiler:
self.assertIsNotNone(profiler.torch_profiler)

def test_neuron_adds_privateuse1_activity_and_stacks(self):
import tempfile

import torch

with (
tempfile.TemporaryDirectory() as tmpdir,
mock.patch("torch.cuda.is_available", return_value=False),
mock.patch("torch.xpu.is_available", return_value=False),
mock.patch("torchtitan.observability.profiler.device_type", "neuron"),
mock.patch("torch.profiler.profile") as profile,
):
profiler = Profiler(
Profiler.Config(
enable_profiling=True,
profile_freq=4,
profiler_warmup=1,
profiler_active=1,
),
global_step=0,
base_folder=tmpdir,
)
profiler.build_torch_profiler(
global_step=0, base_folder=tmpdir, leaf_folder=""
)
self.assertEqual(
profile.call_args.kwargs["activities"],
[
torch.profiler.ProfilerActivity.CPU,
torch.profiler.ProfilerActivity.PrivateUse1,
],
)
self.assertTrue(profile.call_args.kwargs["with_stack"])

def test_memory_snapshot_frequency_is_independent(self):
import tempfile

Expand Down
8 changes: 7 additions & 1 deletion torchtitan/observability/profiler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
from torchtitan.config import Configurable
from torchtitan.distributed.cuda_graph import get_cuda_graph_annotations
from torchtitan.observability import structured_logger as sl
from torchtitan.tools.utils import device_module
from torchtitan.tools.utils import device_module, device_type

# torch's export_chrome_trace gained cuda_graph_annotations when the offline joiner
# (torch.cuda._annotate_cuda_graph_trace) was removed. Older versions still export, just
Expand Down Expand Up @@ -357,10 +357,15 @@ def trace_handler(prof):

wait = profile_freq - (active + warmup)
activities = [torch.profiler.ProfilerActivity.CPU]
with_stack = False
if torch.cuda.is_available():
activities.append(torch.profiler.ProfilerActivity.CUDA)
elif torch.xpu.is_available():
activities.append(torch.profiler.ProfilerActivity.XPU)
elif device_type == "neuron":
activities.append(torch.profiler.ProfilerActivity.PrivateUse1)
# Neuron can link framework events to source lines through the stacks.
with_stack = True

torch_profiler = torch.profiler.profile(
activities=activities,
Expand All @@ -369,6 +374,7 @@ def trace_handler(prof):
),
on_trace_ready=trace_handler,
record_shapes=True,
with_stack=with_stack,
)
torch_profiler.__enter__()
torch_profiler.step_num = global_step
Expand Down
Loading