fix creating leaf folder

Summary:
the leaf folder wasn't being created so and no profiles were being written, so create it if it doesn't exist

Author: tushar00jain

torchtitan/tools/profiling.py

@@ -40,16 +40,14 @@ def maybe_enable_profiling(
 
         def trace_handler(prof):
             curr_trace_dir_name = "iteration_" + str(prof.step_num)
-            curr_trace_dir = os.path.join(trace_dir, curr_trace_dir_name)
+            curr_trace_dir = os.path.join(trace_dir, curr_trace_dir_name, leaf_folder)
             if not os.path.exists(curr_trace_dir):
                 os.makedirs(curr_trace_dir, exist_ok=True)
 
             logger.info(f"Dumping profiler traces at step {prof.step_num}")
             begin = time.monotonic()
 
-            output_file = os.path.join(
-                curr_trace_dir, leaf_folder, f"rank{rank}_trace.json"
-            )
+            output_file = os.path.join(curr_trace_dir, f"rank{rank}_trace.json")
             prof.export_chrome_trace(output_file)
             logger.info(
                 f"Finished dumping profiler traces in {time.monotonic() - begin:.2f} seconds"
@@ -123,13 +121,13 @@ def step(self, exit_ctx: bool = False):
                     # dump as iteration_0_exit if OOM at iter 1
                     curr_step = self.step_num - 1
                     dir_name = f"iteration_{curr_step}_exit"
-                curr_snapshot_dir = os.path.join(snapshot_dir, dir_name)
+                curr_snapshot_dir = os.path.join(snapshot_dir, dir_name, leaf_folder)
                 if not os.path.exists(curr_snapshot_dir):
                     os.makedirs(curr_snapshot_dir, exist_ok=True)
                 logger.info(f"Dumping memory snapshot at step {curr_step}")
                 begin = time.monotonic()
                 output_file = os.path.join(
-                    curr_snapshot_dir, leaf_folder, f"rank{rank}_memory_snapshot.pickle"
+                    curr_snapshot_dir, f"rank{rank}_memory_snapshot.pickle"
                 )
                 with open(output_file, "wb") as output:
                     pickle.dump(torch.cuda.memory._snapshot(), output)