From ab8062c6eb4c1e8c5cba428e81accbf809a15148 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Carlos=20Mochol=C3=AD?= <carlossmocholi@gmail.com>
Date: Mon, 10 Feb 2025 19:15:23 +0100
Subject: [PATCH] profile with modules and stack

---
 torchtitan/profiling.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/torchtitan/profiling.py b/torchtitan/profiling.py
index f85ae2b12..e72bfea9a 100644
--- 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