profile with modules and stack

pytorch · Feb 10, 2025 · ab8062c · ab8062c
1 parent 49c6d6f
commit ab8062c
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/torchtitan/profiling.py b/torchtitan/profiling.py
@@ -64,6 +64,8 @@ def trace_handler(prof):
             schedule=torch.profiler.schedule(wait=wait, warmup=warmup, active=active),
             on_trace_ready=trace_handler,
             record_shapes=True,
+            with_modules=True,
+            with_stack=True,
         ) as torch_profiler:
             torch_profiler.step_num = global_step
             yield torch_profiler