Update on "apply less heavy profiling"

While I try to mitigate #266 by inserting a barrier in the dump step, this PR simplifies profiling to only provide minimal stats, thus making it less heavy. [ghstack-poisoned]
pytorch · Apr 25, 2024 · 1c56dd4 · 1c56dd4
1 parent 2b58706
commit 1c56dd4
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/torchtitan/profiling.py b/torchtitan/profiling.py
@@ -53,7 +53,9 @@ def trace_handler(prof):
 
         warmup, active = WARMUP, 1
         wait = profile_freq - (active + warmup)
-        assert wait >= 0, "profile_freq must be greater than or equal to warmup + active"
+        assert (
+            wait >= 0
+        ), "profile_freq must be greater than or equal to warmup + active"
         with torch.profiler.profile(
             activities=[
                 torch.profiler.ProfilerActivity.CPU,