From d6c207e32f87a395b1e89a8a874549af3fa078bf Mon Sep 17 00:00:00 2001
From: Kavish Gandhi <kvshbg@amazon.com>
Date: Mon, 5 Oct 2026 21:57:02 +0000
Subject: [PATCH] [observability] Profile Neuron devices with the PrivateUse1
 activity

On Neuron, the profiler recorded only CPU activity. Add
ProfilerActivity.PrivateUse1 when device_type is "neuron", as the CUDA
and XPU branches do. Also set with_stack, so Neuron Explorer can link
framework events to source lines.
---
 .../cpu/observability/test_profiler.py        | 34 +++++++++++++++++++
 torchtitan/observability/profiler.py          |  8 ++++-
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/tests/unit_tests/cpu/observability/test_profiler.py b/tests/unit_tests/cpu/observability/test_profiler.py
index 78d91181dae..84df5a01205 100644
--- a/tests/unit_tests/cpu/observability/test_profiler.py
+++ b/tests/unit_tests/cpu/observability/test_profiler.py
@@ -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
 
diff --git a/torchtitan/observability/profiler.py b/torchtitan/observability/profiler.py
index 65dfa435570..5ffb1f315cb 100644
--- a/torchtitan/observability/profiler.py
+++ b/torchtitan/observability/profiler.py
@@ -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
@@ -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,
@@ -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