profile every step

AI-Hypercomputer · Nov 14, 2023 · 868ce45 · 868ce45
1 parent 00c27ea
commit 868ce45
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/MaxText/train.py b/MaxText/train.py
@@ -295,8 +295,12 @@ def train_loop(config, state=None):
 
     # Start profiling at end of first step to avoid compilation.
     # Move before for loop to include.
-    if step == 0:
-      max_utils.activate_profiler(config)
+    # if step == 0:
+    #  max_utils.activate_profiler(config)
+    if jax.process_index() == 0 and config.enable_profiler:
+      if step > 0:
+        jax.profiler.stop_trace()
+      jax.profiler.start_trace(os.path.join(config.tensorboard_dir, str(step)))
 
   max_utils.deactivate_profiler(config)
   writer.close()